From 0513575c2a22d5f8ebd584448ead7ab352bacddd Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Tue, 13 Jan 2026 17:01:34 -0800 Subject: [PATCH 1/2] works E2E with Llama3 8B WS=1 and WS=2 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- .../_torch/auto_deploy/config/default.yaml | 2 + .../auto_deploy/custom_ops/torch_quant.py | 58 ++++++++++++ .../auto_deploy/models/quant_config_reader.py | 4 +- .../transform/library/quantization.py | 90 +++++++++++++++++++ .../auto_deploy/transform/library/sharding.py | 45 ++++++++++ .../_torch/auto_deploy/utils/node_utils.py | 1 + 6 files changed, 199 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/auto_deploy/config/default.yaml b/tensorrt_llm/_torch/auto_deploy/config/default.yaml index 53446b0e8a85..9ce9c5bf42fe 100644 --- a/tensorrt_llm/_torch/auto_deploy/config/default.yaml +++ b/tensorrt_llm/_torch/auto_deploy/config/default.yaml @@ -63,6 +63,8 @@ transforms: stage: pattern_matcher quantize_nvfp4_linear_from_config: stage: pattern_matcher + quantize_hf_fp8_linear_from_config: + stage: pattern_matcher quantize_fp8_bmm_from_config: stage: pattern_matcher quantize_fp8_from_graph: diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/torch_quant.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/torch_quant.py index c1c02556866b..1fddea9dd1bf 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/torch_quant.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/torch_quant.py @@ -319,3 +319,61 @@ def _fake( N_half = weight_quantized.shape[-2] N = N_half * 2 return torch.empty((*input.shape[:-1], N), dtype=input.dtype, device=input.device) + + +@torch.library.custom_op("auto_deploy::torch_fake_quant_hf_fp8_linear", mutates_args=()) +def torch_fake_quant_hf_fp8_linear( + input: torch.Tensor, # [..., K] + weight_quantized: torch.Tensor, # [N, K] float8_e4m3fn + bias: Optional[torch.Tensor], # [N] or None + input_scale: List[torch.Tensor], # unused for HF FP8 (input quantized on the fly) + weight_scale: List[torch.Tensor], # [weight_scale_inv] + input_zp: List[torch.Tensor], # unused + weight_zp: List[torch.Tensor], # unused +) -> torch.Tensor: + """HuggingFace FineGrainedFP8 linear operation. + - weight_scale[0] = weight_scale_inv (per-block weight scale) + - input_scale, input_zp, weight_zp are unused + - block_size is inferred from weight and weight_scale_inv shapes + """ + from transformers.integrations.finegrained_fp8 import act_quant, w8a8_block_fp8_matmul_triton + + weight_scale_inv = weight_scale[0] + + # Infer block_size from weight and weight_scale_inv shapes + # weight shape: [N, K], weight_scale_inv shape: [N/block_n, K/block_k] + N, K = weight_quantized.shape + scale_n, scale_k = weight_scale_inv.shape + block_n = N // scale_n + block_k = K // scale_k + block_size = [block_n, block_k] + + qinput, scale = act_quant(input, block_size[1]) + output = w8a8_block_fp8_matmul_triton( + qinput, + weight_quantized, + scale, + weight_scale_inv, + block_size, + output_dtype=input.dtype, + ) + + if bias is not None: + output = output + bias + + return output.to(dtype=input.dtype) + + +@torch_fake_quant_hf_fp8_linear.register_fake +def _torch_fake_quant_hf_fp8_linear_fake( + input: torch.Tensor, + weight_quantized: torch.Tensor, + bias: Optional[torch.Tensor], + input_scale: List[torch.Tensor], + weight_scale: List[torch.Tensor], + input_zp: List[torch.Tensor], + weight_zp: List[torch.Tensor], +) -> torch.Tensor: + """Fake implementation for torch.export tracing.""" + out_features = weight_quantized.shape[0] + return torch.empty((*input.shape[:-1], out_features), dtype=input.dtype, device=input.device) diff --git a/tensorrt_llm/_torch/auto_deploy/models/quant_config_reader.py b/tensorrt_llm/_torch/auto_deploy/models/quant_config_reader.py index 232fe14f4ea6..81c0a97a50af 100644 --- a/tensorrt_llm/_torch/auto_deploy/models/quant_config_reader.py +++ b/tensorrt_llm/_torch/auto_deploy/models/quant_config_reader.py @@ -146,6 +146,8 @@ class HFQuantConfigReader(QuantConfigReader): Quantization reader that process transformers.quantizers.HFQuantizer functionality """ + _SUPPORTED_QUANT_METHODS = ("mxfp4", "fp8") + def __init__(self): super().__init__() self._hf_quantizer = None @@ -180,7 +182,7 @@ def from_file(cls, ckpt_dir: str) -> Optional[Tuple["HFQuantConfigReader", Dict[ # TODO(Fridah-nv):this class is only verified with GPT-OSS MXFP4, other hf quantizers # should have similar workflow and will be added to the pipeline quant_method = str(qconf.get("quant_method", "")).lower() - if quant_method != "mxfp4": + if quant_method not in cls._SUPPORTED_QUANT_METHODS: return None reader = cls() diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/quantization.py b/tensorrt_llm/_torch/auto_deploy/transform/library/quantization.py index 21d1ccd23485..dbad41f4077e 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/quantization.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/quantization.py @@ -633,3 +633,93 @@ def _apply( return gm, TransformInfo( skipped=False, num_matches=cnt, is_clean=cnt == 0, has_valid_shapes=True ) + + +@TransformRegistry.register("quantize_hf_fp8_linear_from_config") +class HFFineGrainedFP8LinearQuantization(Quantization): + """Quantization transform for HuggingFace FineGrainedFP8 (block-wise FP8) models. + + This transform replaces linear ops with the HF FineGrainedFP8 quantized op. + The HF FP8 format uses per-block weight scales (weight_scale_inv) and + dynamic input quantization. + + Config format (from HF config.json): + "quantization_config": { + "quant_method": "fp8", + "weight_block_size": [128, 128], + "modules_to_not_convert": ["lm_head"] + } + """ + + algo_name = "fp8" + + def target_op(self): + return torch.ops.auto_deploy.torch_fake_quant_hf_fp8_linear.default + + def quantize_weight(self, w: torch.Tensor) -> torch.Tensor: + return torch.empty_like(w, dtype=torch.float8_e4m3fn, device=w.device) + + def scale_names(self) -> List[str]: + return ["weight_scale_inv"] + + def default_scales(self, original_weight_shape: Tuple) -> Dict[str, torch.Tensor]: + # Default block size is 128x128 for HF FP8 + N, K = original_weight_shape + block_n, block_k = 128, 128 + scale_shape = (N // block_n, K // block_k) + return {"weight_scale_inv": torch.ones(scale_shape, dtype=torch.bfloat16)} + + def build_custom_args_for_linear(self, scales: Dict[str, Node]) -> Tuple: + return ([], [scales["weight_scale_inv"]], [], []) + + def load_hook(self, state_dict, prefix, *args, weight_name: str): + """Load hook to handle HF FineGrainedFP8 checkpoint format. + + HF FP8 checkpoints store: + - weight: float8_e4m3fn tensor + - weight_scale_inv: per-block scale tensor + """ + if weight_name not in state_dict: + return + + weight = state_dict[weight_name] + if weight.dtype == torch.float8_e4m3fn: + scale_inv_name = weight_name + "_scale_inv" + if scale_inv_name in state_dict: + # Rename to match our buffer name + mod_prefix = weight_name.rsplit(".", 1)[0] + state_dict[mod_prefix + ".weight_scale_inv"] = state_dict[scale_inv_name] + + def _apply( + self, + gm: GraphModule, + cm: CachedSequenceInterface, + factory: ModelFactory, + shared_config: SharedConfig, + ) -> Tuple[GraphModule, TransformInfo]: + qcfg = factory.get_quant_config() + if not qcfg: + return gm, TransformInfo( + skipped=True, num_matches=0, is_clean=True, has_valid_shapes=True + ) + + quant_method = str(qcfg.get("quant_method", "")).lower() + if quant_method != self.algo_name: + return gm, TransformInfo( + skipped=True, num_matches=0, is_clean=True, has_valid_shapes=True + ) + + excluded = qcfg.get("modules_to_not_convert", []) + + cnt = 0 + for n in gm.graph.nodes: + if not is_linear_op(n): + continue + if should_skip_quantization(n, excluded): + continue + self._insert_quantized_linear(gm, n, is_quantized_graph=False) + cnt += 1 + + return gm, TransformInfo( + skipped=False, num_matches=cnt, is_clean=False, has_valid_shapes=(cnt == 0) + ) diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py b/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py index deb82a43e661..4aadb77ea6d1 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py @@ -436,6 +436,47 @@ def shard_load_hook( return +class HFFineGrainedFP8WeightShardingInfo(QuantizationShardingMixin, WeightShardingInfo): + """Tensor-parallel sharding for HuggingFace FineGrainedFP8 quantized linears. + + HF FP8 uses per-block weight scales (weight_scale_inv) with shape [N/block_n, K/block_k]. + When sharding the weight along a dimension, we also need to shard the scale tensor. + """ + + def scale_names(self) -> List[str]: + return ["weight_scale_inv"] + + def shard_scales( + self, + dim: int, + rank: int, + world_size: int, + weight_shape: torch.Size, + *, + weight_scale_inv: torch.Tensor, + ) -> Dict[str, torch.Tensor]: + # weight_scale_inv has shape [N/block_n, K/block_k] + # When we shard weight along dim, we need to shard scale along the same dim + sharded_scale = torch.tensor_split(weight_scale_inv, world_size, dim=dim)[rank] + return {"weight_scale_inv": sharded_scale} + + def shard_load_hook( + self, + state_dict, + prefix, + *args, + weight_name: str, + weight_shape: torch.Size, + dim: int, + rank: int, + world_size: int, + ) -> None: + scale_key = weight_name + "_scale_inv" + if scale_key in state_dict: + scale = state_dict[scale_key] + state_dict[scale_key] = torch.tensor_split(scale, world_size, dim=dim)[rank] + + def _shard_fp4_weight_scale(weight_scale, sharded_uint8_weight_shape, dim, rank, world_size): assert weight_scale.dim() == 1 weight_shape_original = list(sharded_uint8_weight_shape) @@ -1051,6 +1092,10 @@ def _validate_sharded_shapes( lambda n: is_op(n, torch.ops.auto_deploy.torch_fake_quant_nvfp4_linear), FP4WeightShardingInfo, ), + ( + lambda n: is_op(n, torch.ops.auto_deploy.torch_fake_quant_hf_fp8_linear), + HFFineGrainedFP8WeightShardingInfo, + ), ] diff --git a/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py b/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py index 795bfa4f2db0..be03dbadf110 100644 --- a/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py +++ b/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py @@ -362,6 +362,7 @@ def is_fake_quantized_linear_op(node: Node) -> bool: quantized_linear_op = { torch.ops.auto_deploy.torch_fake_quant_fp8_linear, torch.ops.auto_deploy.torch_fake_quant_nvfp4_linear, + torch.ops.auto_deploy.torch_fake_quant_hf_fp8_linear, } return is_op(node, quantized_linear_op) From 5351cf430377f0e1fe976dee4655937d57a1fb4a Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Wed, 14 Jan 2026 15:12:19 -0800 Subject: [PATCH 2/2] unit tests Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- .../_utils_test/_model_test_utils.py | 40 ++++++++++++ .../library/test_tp_sharding.py | 45 +++++++++++++- .../unit/singlegpu/custom_ops/test_quant.py | 62 +++++++++++++++++++ .../library/test_quantization.py | 45 ++++++++++++++ 4 files changed, 191 insertions(+), 1 deletion(-) diff --git a/tests/unittest/_torch/auto_deploy/_utils_test/_model_test_utils.py b/tests/unittest/_torch/auto_deploy/_utils_test/_model_test_utils.py index a71a09b46521..6f6840419ebf 100644 --- a/tests/unittest/_torch/auto_deploy/_utils_test/_model_test_utils.py +++ b/tests/unittest/_torch/auto_deploy/_utils_test/_model_test_utils.py @@ -275,6 +275,46 @@ def forward(self, x): ) +class FakeHFFP8Linear(nn.Linear): + """Fake HuggingFace FineGrainedFP8 linear layer for testing. + + Mimics the behavior of transformers.integrations.finegrained_fp8.FP8Linear + with per-block quantization (block_size = [128, 128] by default). + """ + + def __init__(self, in_features, out_features, bias=True, block_size=None): + super().__init__(in_features, out_features, bias) + device = self.weight.device + + if block_size is None: + block_n = min(128, out_features) + block_k = min(128, in_features) + block_size = [block_n, block_k] + self.block_size = block_size + + N, K = self.weight.shape + block_n, block_k = block_size + + weight_reshaped = self.weight.detach().view(N // block_n, block_n, K // block_k, block_k) + amax = weight_reshaped.abs().amax(dim=(1, 3)).to(torch.float32) # [N/block_n, K/block_k] + + eps = torch.finfo(torch.float32).tiny + weight_scale_inv = torch.clamp(amax / FP8_MAX, min=eps).to(device) + + weight_fp8 = ( + self.weight.detach().float() + / weight_scale_inv.repeat_interleave(block_n, dim=0).repeat_interleave(block_k, dim=1) + ).to(torch.float8_e4m3fn) + + self.weight = nn.Parameter(weight_fp8) + self.register_buffer("weight_scale_inv", weight_scale_inv) + + def forward(self, x): + return torch.ops.auto_deploy.torch_fake_quant_hf_fp8_linear( + x, self.weight, self.bias, [], [self.weight_scale_inv], [], [] + ) + + def generate_dynamic_shapes(max_batch_size, max_seq_len): dynamic_shapes = ( { diff --git a/tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_tp_sharding.py b/tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_tp_sharding.py index b4f82edcfaea..67dde0dde857 100644 --- a/tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_tp_sharding.py +++ b/tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_tp_sharding.py @@ -9,12 +9,13 @@ import torch.nn.functional as F from _dist_test_utils import get_device_counts from _graph_test_helpers import run_sharding_pattern_detection_test, run_test_transformed_gm -from _model_test_utils import FakeFP8Linear +from _model_test_utils import FakeFP8Linear, FakeHFFP8Linear import tensorrt_llm._torch.auto_deploy.distributed.common as dist_common from tensorrt_llm._torch.auto_deploy.export import torch_export_to_gm from tensorrt_llm._torch.auto_deploy.transform.library.sharding import ( FP8WeightShardingInfo, + HFFineGrainedFP8WeightShardingInfo, LayerType, ShardingTransformConfig, SplitDimension, @@ -125,6 +126,23 @@ def forward(self, x): return self.linear2(y) +class HFFP8MLP(nn.Module): + """MLP using HuggingFace FineGrainedFP8 quantization for testing.""" + + def __init__(self, in_features, out_features, bias=False): + super().__init__() + self.in_features = in_features + self.out_features = out_features + # Use larger features divisible by block size (128) + hidden_features = max(4 * in_features, 128) + self.linear1 = FakeHFFP8Linear(in_features, hidden_features, bias=bias) + self.linear2 = FakeHFFP8Linear(hidden_features, out_features, bias=bias) + + def forward(self, x): + y = F.relu(self.linear1(x)) + return self.linear2(y) + + def _run_sharding_execution_job( model_cls: nn.Module, dist_op_expected: str, @@ -150,6 +168,9 @@ def _run_sharding_execution_job( ).to(device="cuda", dtype=torch.float16) elif model_cls == FP8MLP: model = model_cls(num_features, num_features, bias=bias).to("cuda") + elif model_cls == HFFP8MLP: + # HFFP8MLP needs features divisible by 128 (block size) + model = model_cls(128, 128, bias=bias).to("cuda") else: model = model_cls(num_features, num_features, bias=bias).to( device="cuda", dtype=torch.float16 @@ -344,6 +365,26 @@ def _run_pattern_detection_job( min_local_shape=1, ) ) + elif model_cls == HFFP8MLP: + for node in gm.graph.nodes: + if is_op(node, torch.ops.auto_deploy.torch_fake_quant_hf_fp8_linear): + # linear1 should be sharded on dim=0, add_dist=False, min_local_shape=1 + # linear2 should be sharded on dim=1, add_dist=True, min_local_shape=1 + if "linear1" in node.args[1].name: + dim = SplitDimension.COLUMN + dist_op = None + else: + dim = SplitDimension.ROW + dist_op = "all_reduce" + expected_transformations.append( + HFFineGrainedFP8WeightShardingInfo( + target_node=node.name, + split_dim=dim, + config=config, + dist_op=dist_op, + min_local_shape=1, + ) + ) # get detected transformations optimizer = InferenceOptimizer( @@ -376,6 +417,7 @@ def _run_pattern_detection_job( ( (MLP, "torch_dist_all_reduce"), (FP8MLP, "torch_dist_all_reduce"), + (HFFP8MLP, "torch_dist_all_reduce"), (nn.Linear, "torch_dist_all_gather"), (GQA_Block, "torch_dist_all_reduce"), ), @@ -401,6 +443,7 @@ def test_sharding( ( (MLP, "torch_dist_all_reduce"), (FP8MLP, "torch_dist_all_reduce"), + (HFFP8MLP, "torch_dist_all_reduce"), (nn.Linear, "torch_dist_all_gather"), (GQA_Block, "torch_dist_all_reduce"), ), diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_quant.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_quant.py index bfe8d75f1a40..5b1588194939 100644 --- a/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_quant.py +++ b/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_quant.py @@ -309,3 +309,65 @@ def test_fake_quant_int4_linear_matches_fp_reference(bias_opt, input_dtype): ).to(out_int4.dtype) cos = F.cosine_similarity(out_fp32.reshape(-1), out_int4.reshape(-1), dim=0) assert cos > 0.98 + + +@pytest.mark.parametrize("M", [3, 12]) +@pytest.mark.parametrize("N", [128, 256]) # Must be divisible by block_size +@pytest.mark.parametrize("K", [128, 256]) # Must be divisible by block_size +@pytest.mark.parametrize("bias", [True, False]) +@pytest.mark.skipif(not fp8_compatible(), reason="Requires fp8 support") +def test_hf_fp8_linear(M, N, K, bias): + """Test HuggingFace FineGrainedFP8 linear custom op. + + This tests the torch_fake_quant_hf_fp8_linear op which implements + per-block FP8 quantization matching HuggingFace's FineGrainedFP8 format. + """ + block_size = [128, 128] + block_n, block_k = block_size + + assert N % block_n == 0, f"N={N} must be divisible by block_n={block_n}" + assert K % block_k == 0, f"K={K} must be divisible by block_k={block_k}" + + input_tensor = torch.rand(M, K, device="cuda", dtype=torch.float16) + weight = torch.rand(N, K, device="cuda", dtype=torch.float16) + bias_tensor = torch.rand(N, device="cuda", dtype=torch.float16) if bias else None + + FP8_MAX = torch.finfo(torch.float8_e4m3fn).max + eps = torch.finfo(torch.float32).tiny + + # Reshape weight to blocks: [N, K] -> [N/block_n, block_n, K/block_k, block_k] + weight_reshaped = weight.detach().float().view(N // block_n, block_n, K // block_k, block_k) + # Compute per-block amax: [N/block_n, K/block_k] + amax = weight_reshaped.abs().amax(dim=(1, 3)).to(torch.float32) + weight_scale_inv = torch.clamp(amax / FP8_MAX, min=eps).to("cuda") + + # Quantize weight to FP8 using per-block scales + # Expand scale to match weight shape for element-wise division + scale_expanded = weight_scale_inv.repeat_interleave(block_n, dim=0).repeat_interleave( + block_k, dim=1 + ) + weight_fp8 = (weight.float() / scale_expanded).to(torch.float8_e4m3fn) + + output_hf_fp8 = torch.ops.auto_deploy.torch_fake_quant_hf_fp8_linear( + input_tensor, + weight_fp8, + bias_tensor, + [], # input_scale - unused for HF FP8 + [weight_scale_inv], # weight_scale + [], # input_zp - unused + [], # weight_zp - unused + ) + + weight_dequant = (weight_fp8.float() * scale_expanded).to(input_tensor.dtype) + output_ref = torch.nn.functional.linear(input_tensor, weight_dequant, bias_tensor) + + assert output_hf_fp8.shape == output_ref.shape, ( + f"Shape mismatch: {output_hf_fp8.shape} vs {output_ref.shape}" + ) + + torch.testing.assert_close(output_hf_fp8, output_ref, rtol=0.1, atol=0.1) + + cos = F.cosine_similarity( + output_hf_fp8.reshape(-1).float(), output_ref.reshape(-1).float(), dim=0 + ) + assert cos > 0.95, f"Cosine similarity too low: {cos}" diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_quantization.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_quantization.py index 38caf1c5dc95..d046f2700611 100644 --- a/tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_quantization.py +++ b/tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_quantization.py @@ -113,6 +113,51 @@ def test_quantization(quant_config, atol, rtol, num_p_og): torch_export_to_gm(gm_transformed, args=(x,)) +@pytest.mark.parametrize("bias", [True, False]) +@pytest.mark.skipif(not fp8_compatible(), reason="Requires fp8 support") +def test_hf_fp8_quantization(bias): + """Test HuggingFace FineGrained FP8 quantization transform. + + This tests the quantize_hf_fp8_linear_from_config transform which converts + linear layers to use the torch_fake_quant_hf_fp8_linear op with per-block + quantization matching HuggingFace's FineGrainedFP8 format. + """ + model = MLP(128, 256, 128).to(torch.float16).to("cuda") + x = torch.randn(3, 128, dtype=torch.float16).to("cuda") + + quant_config = {"quant_method": "fp8"} + QUANT_OP = torch.ops.auto_deploy.torch_fake_quant_hf_fp8_linear + + gm = torch_export_to_gm(model, args=(x,), clone=True) + gm_transformed = InferenceOptimizer( + DummyFactory(quant_config), + { + "quantize_hf_fp8_linear_from_config": { + "stage": "pattern_matcher", + }, + }, + )(None, gm) + gm_transformed.to("cuda") + run_test_transformed_gm( + model, + x, + gm_transformed, + lambda gm: any(is_op(n, QUANT_OP) for n in gm.graph.nodes), + lambda num_p_og: num_p_og, + 0.1, # atol + 0.1, # rtol + True, # test_load_hook + False, # strict_loading + None, # dynamic_shapes + None, # check_num_matches + False, # skip_output_assert + quant_config, + ) + + assert not torch.allclose(model(x), gm_transformed(x)) + torch_export_to_gm(gm_transformed, args=(x,)) + + @pytest.mark.parametrize( "quant_config,atol,rtol,num_p_og,model_class", [