From f46445fa1502328fcae8dc326cdaf732b10a71be Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Mon, 20 Apr 2026 20:52:28 +0000 Subject: [PATCH 01/25] Update [ghstack-poisoned] --- test/prototype/gptq/test_gptqv2.py | 21 ++++ torchao/prototype/gptq/api.py | 157 ++++++++++++++++++++++++- torchao/prototype/gptq/gptq_example.py | 16 ++- 3 files changed, 188 insertions(+), 6 deletions(-) diff --git a/test/prototype/gptq/test_gptqv2.py b/test/prototype/gptq/test_gptqv2.py index 366fe6d242..0bae6ec4da 100644 --- a/test/prototype/gptq/test_gptqv2.py +++ b/test/prototype/gptq/test_gptqv2.py @@ -15,6 +15,9 @@ gptq_quantize, ) from torchao.prototype.gptq.observer import GPTQObserverTensor +from torchao.prototype.mx_formats.inference_workflow import ( + NVFP4DynamicActivationNVFP4WeightConfig, +) from torchao.quantization import Int4WeightOnlyConfig, Int8WeightOnlyConfig, quantize_ from torchao.quantization.granularity import PerRow from torchao.utils import _is_mslk_available @@ -435,6 +438,13 @@ def test_gptq_quantize_function(self, base_config): pytest.param( Int8WeightOnlyConfig(granularity=PerRow(), version=2), id="int8" ), + pytest.param( + NVFP4DynamicActivationNVFP4WeightConfig( + use_dynamic_per_tensor_scale=True, + use_triton_kernel=True, + ), + id="nvfp4", + ), ], ) def test_gptq_quantize_better_than_naive(self, base_config): @@ -506,6 +516,13 @@ def test_gptq_quantize_better_than_naive(self, base_config): pytest.param( Int8WeightOnlyConfig(granularity=PerRow(), version=2), id="int8" ), + pytest.param( + NVFP4DynamicActivationNVFP4WeightConfig( + use_dynamic_per_tensor_scale=True, + use_triton_kernel=True, + ), + id="nvfp4", + ), ], ) def test_gptq_sqnr(self, base_config): @@ -556,6 +573,10 @@ def test_gptq_sqnr(self, base_config): assert sqnr_gptq > 25, f"GPTQ SQNR: {sqnr_gptq} is too low" elif isinstance(base_config, Int8WeightOnlyConfig): assert sqnr_gptq > 30, f"GPTQ SQNR: {sqnr_gptq} is too low" + elif isinstance(base_config, NVFP4DynamicActivationNVFP4WeightConfig): + assert sqnr_gptq > 15, f"GPTQ SQNR: {sqnr_gptq} is too low" + else: + raise AssertionError("unsupported") assert sqnr_gptq > sqnr_rtn, ( f"GPTQ SQNR: {sqnr_gptq} is not better than RTN SQNR: {sqnr_rtn}" ) diff --git a/torchao/prototype/gptq/api.py b/torchao/prototype/gptq/api.py index 0059bd1d51..ac6f7cdc6e 100644 --- a/torchao/prototype/gptq/api.py +++ b/torchao/prototype/gptq/api.py @@ -19,6 +19,25 @@ pack_int4 = None from torchao.core.config import AOBaseConfig +from torchao.prototype.mx_formats.constants import F4_E2M1_MAX +from torchao.prototype.mx_formats.inference_workflow import ( + NVFP4DynamicActivationNVFP4WeightConfig, +) +from torchao.prototype.mx_formats.kernels import ( + f4_unpacked_to_f32, + f32_to_f4_unpacked, + pack_uint4, +) +from torchao.prototype.mx_formats.nvfp4_tensor import ( + NVFP4Tensor, + QuantizeTensorToNVFP4Kwargs, + nvfp4_quantize, + per_tensor_amax_to_scale, +) +from torchao.prototype.mx_formats.utils import ( + hp_data_dims_to_swizzled_scale_dims_nvfp4, + to_blocked, +) from torchao.quantization import Int4Tensor, Int8Tensor from torchao.quantization.granularity import PerRow from torchao.quantization.quant_api import ( @@ -33,6 +52,7 @@ CONFIG_TO_TORCHAO_BASE_TENSOR = { Int4WeightOnlyConfig: Int4Tensor, Int8WeightOnlyConfig: Int8Tensor, + NVFP4DynamicActivationNVFP4WeightConfig: NVFP4Tensor, } @@ -61,7 +81,11 @@ class GPTQConfig(AOBaseConfig): """ step: str = "observe" # "observe" or "convert" - base_config: Union[Int4WeightOnlyConfig, Int8WeightOnlyConfig] = None + base_config: Union[ + Int4WeightOnlyConfig, + Int8WeightOnlyConfig, + NVFP4DynamicActivationNVFP4WeightConfig, + ] = None percdamp: float = 0.01 gptq_quantize_block_size: int = 256 @@ -179,6 +203,59 @@ def _int4_row_dequantize_zp( return torch.cat(dequant_chunks, dim=-1) +def _nvfp4_with_precalculated_scales_qdq( + data_hp: torch.Tensor, + per_tensor_scale: torch.Tensor, + block_scale: torch.Tensor, +) -> torch.Tensor: + """ + Same as torchao.prototype.mx_formats.nvfp4_tensor.nvfp4_quantize, but with + per_tensor_scale and block_scale precalculated and the end result dequantized. + """ + assert per_tensor_scale.dtype is torch.float32 + assert block_scale.dtype is torch.float8_e4m3fn + # this function only works for data_hp.shape == (N, k_slice) + # and block_scale.shape == (N,) + assert len(block_scale.shape) == 1 + + scaled_block_scales_fp32 = block_scale.to(torch.float32) + reciprocal_scale = (1.0 / per_tensor_scale) / scaled_block_scales_fp32 + data_scaled = data_hp * reciprocal_scale.unsqueeze(-1) + data_scaled = torch.clamp(data_scaled, -F4_E2M1_MAX, F4_E2M1_MAX) + data_lp = f32_to_f4_unpacked(data_scaled) + data_lp_hp = f4_unpacked_to_f32(data_lp) + data_lp_hp_unscaled = data_lp_hp / reciprocal_scale.unsqueeze(-1) + return data_lp_hp_unscaled + + +def _nvfp4_with_precalculated_scales_q( + data_hp: torch.Tensor, + per_tensor_scale: torch.Tensor, + block_scale: torch.Tensor, +) -> torch.Tensor: + """ + Same as torchao.prototype.mx_formats.nvfp4_tensor.nvfp4_quantize, but with + per_tensor_scale and block_scale precalculated. + """ + assert per_tensor_scale.dtype is torch.float32 + assert block_scale.dtype is torch.float8_e4m3fn + + # TODO(future): figure out what to reuse vs leave copy-pasted vs + # nvfp4_tensor.py + scaled_block_scales_fp32 = block_scale.to(torch.float32) + reciprocal_scale = (1.0 / per_tensor_scale) / scaled_block_scales_fp32 + N, K = data_hp.shape + # reshape to 3d to properly broadcast for scaling + data_hp = data_hp.view(N, K // 16, 16) + data_scaled = data_hp * reciprocal_scale.unsqueeze(-1) + data_scaled = torch.clamp(data_scaled, -F4_E2M1_MAX, F4_E2M1_MAX) + data_lp = f32_to_f4_unpacked(data_scaled) + data_lp_packed = pack_uint4(data_lp) + # reshape back to 2d + data_lp_packed = data_lp_packed.view(N, K // 2) + return data_lp_packed + + def gptq_quantize(H: torch.Tensor, W_t: torch.Tensor, config: GPTQConfig): """ This function implements the GPTQ algorithm described in this paper: https://arxiv.org/abs/2210.17323 (Algorithm 1) @@ -241,6 +318,15 @@ def gptq_quantize(H: torch.Tensor, W_t: torch.Tensor, config: GPTQConfig): block_size = get_block_size(W_t.shape, base_config.granularity) block_size = list(block_size) group_size = block_size[-1] + elif isinstance(base_config, NVFP4DynamicActivationNVFP4WeightConfig): + group_size = 16 + block_size = [1, group_size] + # for per-tensor nvfp4, we need to calculate the global scale over the + # entire tensor before we enter the GPTQ loop + tensor_amax = torch.max(torch.abs(W_t)) + nvfp4_global_scale = per_tensor_amax_to_scale(tensor_amax) + else: + raise AssertionError("unsupported") assert group_size > 0 @@ -337,6 +423,27 @@ def gptq_quantize(H: torch.Tensor, W_t: torch.Tensor, config: GPTQConfig): ], base_config.granularity, ) + elif isinstance(base_config, NVFP4DynamicActivationNVFP4WeightConfig): + tensor_slice = B_cur[ + :, + G_k_start - B_cur_k_start : G_k_end - B_cur_k_start, + ].contiguous() + # quantize this slice using pre-calculated global scale, to + # get the blockwise dynamic scales, they will be frozen + # after this point + scale, _data_lp = nvfp4_quantize( + tensor_slice, + per_tensor_scale=nvfp4_global_scale, + ) + group_qparams.append(scale) + # TODO(future PR): simpler version of `nvfp4_quantize` which + # just calculates the scale, since we are throwing away the + # quantized packed data here. For now, just call the full + # one. + del _data_lp + + else: + raise AssertionError("unsupported") # Quantize each column and propagate errors to subsequent columns for k in range(G_k_start - B_cur_k_start, G_k_end - B_cur_k_start): @@ -354,6 +461,12 @@ def gptq_quantize(H: torch.Tensor, W_t: torch.Tensor, config: GPTQConfig): scale=quantized_tensor.scale, ) dq = q.dequantize(output_dtype=torch.float) + elif isinstance(base_config, NVFP4DynamicActivationNVFP4WeightConfig): + dq = _nvfp4_with_precalculated_scales_qdq( + w_t, + nvfp4_global_scale, + scale.squeeze(-1), + ) err1 = (w_t - dq) / Hinv_cur[k, k] B_cur[:, k:] -= err1.matmul(Hinv_cur[k, k:].unsqueeze(0)) @@ -385,6 +498,48 @@ def gptq_quantize(H: torch.Tensor, W_t: torch.Tensor, config: GPTQConfig): result = Int8Tensor.from_hp( W_t, granularity=base_config.granularity, scale=quantized_tensor.scale ) + else: + N, K = W_t.shape + # TODO(future PR): clean up the line below. Context: the current nvfp4 + # code follows the int4 code - we save the blockwise scales to an array + # and concat it. This leads to the necessity of t().contiguous().t() to + # get the scales back into the right layout for W_t. This is not + # intuitive, likely better to initialize the scales holder ahead of time + # and write the scales directly to their final place. + combined_scale = ( + torch.cat(group_qparams, dim=0).reshape(K // group_size, N).t().contiguous() + ) + qdata = _nvfp4_with_precalculated_scales_q( + W_t, + nvfp4_global_scale, + combined_scale, + ) + + act_quant_kwargs = QuantizeTensorToNVFP4Kwargs( + use_dynamic_per_tensor_scale=base_config.use_dynamic_per_tensor_scale, + use_triton_kernel=base_config.use_triton_kernel, + is_swizzled_scales=True, + ) + + # swizzle the block scales + combined_scale_swizzled = to_blocked(combined_scale).flatten() + scale_N, scale_K = hp_data_dims_to_swizzled_scale_dims_nvfp4(N, K) + combined_scale_swizzled = combined_scale_swizzled.view(scale_N, scale_K) + + result = NVFP4Tensor( + qdata, + combined_scale_swizzled, + block_size=group_size, + orig_dtype=W_t.dtype, + per_tensor_scale=nvfp4_global_scale, + # TODO(future): get act_per_tensor_scale from calibration data? + # for now, set it to None here to calculate it dynamically at + # runtime + act_per_tensor_scale=None, + is_swizzled_scales=True, + use_triton_kernel=base_config.use_triton_kernel, + act_quant_kwargs=act_quant_kwargs, + ) return result diff --git a/torchao/prototype/gptq/gptq_example.py b/torchao/prototype/gptq/gptq_example.py index 57ab7ed974..a8415828c2 100644 --- a/torchao/prototype/gptq/gptq_example.py +++ b/torchao/prototype/gptq/gptq_example.py @@ -230,7 +230,8 @@ def main(): # Generate output directory name from args model_name = args.model_id.split("/")[-1] # Get last part of model ID - output_dir = f"{model_name}_{args.quantization}" + # TODO(before land): make output dir configurable + output_dir = f"/home/dev/tmp/20260420_{model_name}_{args.quantization}" if args.quantization != "none": output_dir += f"_gs{args.group_size}" @@ -358,11 +359,16 @@ def main(): "--model_args", f"pretrained={output_dir}", "--tasks", - "leaderboard_bbh", - "--num_fewshot", - "3", + # "leaderboard_bbh", + # "gsm8k", + "wikitext", + # "--num_fewshot", + # "3", "--batch_size", - "auto", + # "auto", + "1", + # "--limit", + # "20", ] print(f"Running command: {' '.join(lm_eval_cmd)}") From 3c92c1a36ffe165c0e6c39f38808754f01494b54 Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Mon, 20 Apr 2026 20:52:31 +0000 Subject: [PATCH 02/25] Update [ghstack-poisoned] --- torchao/prototype/gptq/api.py | 1 + torchao/prototype/gptq/gptq_example.py | 62 +++++++++++++++++++++++--- 2 files changed, 57 insertions(+), 6 deletions(-) diff --git a/torchao/prototype/gptq/api.py b/torchao/prototype/gptq/api.py index ac6f7cdc6e..c2076e1d2d 100644 --- a/torchao/prototype/gptq/api.py +++ b/torchao/prototype/gptq/api.py @@ -120,6 +120,7 @@ def _gptq_config_transform( ) return module elif config.step == "convert": + print('gptq convert') # Quantization phase: tensor should be an GPTQObserverTensor if not isinstance(tensor, GPTQObserverTensor): raise ValueError( diff --git a/torchao/prototype/gptq/gptq_example.py b/torchao/prototype/gptq/gptq_example.py index a8415828c2..289ec98ebb 100644 --- a/torchao/prototype/gptq/gptq_example.py +++ b/torchao/prototype/gptq/gptq_example.py @@ -17,6 +17,9 @@ from transformers import AutoModelForCausalLM, AutoTokenizer from torchao.prototype.gptq import GPTQConfig +from torchao.prototype.mx_formats.inference_workflow import ( + NVFP4DynamicActivationNVFP4WeightConfig, +) from torchao.quantization import Int4WeightOnlyConfig, Int8WeightOnlyConfig, quantize_ from torchao.quantization.granularity import PerRow @@ -98,6 +101,8 @@ def prepare_dataset( hf_dataset_id = dataset_map.get(dataset_id, dataset_id) # Load dataset and preprocess + # TODO(before land): clean up below + dataset_split = "train" train_dataset_raw = load_dataset(hf_dataset_id, split=dataset_split, streaming=True) train_dataset_raw = train_dataset_raw.shuffle(seed=seed, buffer_size=1_000) @@ -149,7 +154,8 @@ def parse_args(): parser.add_argument( "--num-calibration-samples", type=int, - default=128, + # default=128, + default=2, help="Number of calibration samples to use", ) parser.add_argument( @@ -177,6 +183,9 @@ def parse_args(): "int8-rtn", "int8-gptq-sequential", "int8-gptq-nonsequential", + "nvfp4-rtn", + "nvfp4-gptq-sequential", + "nvfp4-gptq-nonsequential", ], help="Quantization method to use", ) @@ -241,6 +250,8 @@ def main(): "int4-gptq-nonsequential", "int8-gptq-sequential", "int8-gptq-nonsequential", + "nvfp4-gptq-sequential", + "nvfp4-gptq-nonsequential", ]: output_dir += f"_{args.dataset_id}_n{args.num_calibration_samples}" output_dir += f"_damp{args.percdamp}_bs{args.gptq_block_size}" @@ -250,29 +261,51 @@ def main(): # Handle different quantization methods quantization_start_time = time.time() + def skip_lm_head(module, fqn): + # TODO(before land): remove o_proj from below, it's just for debugging + return isinstance(module, torch.nn.Linear) and "lm_head" not in fqn and "o_proj" in fqn + + if args.quantization == "int4-rtn": print("Applying Int4 RTN (Round-To-Nearest) quantization...") config = Int4WeightOnlyConfig(group_size=args.group_size) - quantize_(model, config, filter_fn=None) + quantize_(model, config, filter_fn=skip_lm_head) elif args.quantization == "int8-rtn": print("Applying Int8 RTN (Round-To-Nearest) quantization...") config = Int8WeightOnlyConfig(version=2, granularity=PerRow()) - quantize_(model, config, filter_fn=None) + quantize_(model, config, filter_fn=skip_lm_head) + + elif args.quantization == "nvfp4-rtn": + print("Applying NVFP4 RTN (Round-To-Nearest) quantization...") + + config = NVFP4DynamicActivationNVFP4WeightConfig( + use_dynamic_per_tensor_scale=True, + use_triton_kernel=True, + ) + quantize_(model, config, filter_fn=skip_lm_head) elif args.quantization in [ "int4-gptq-sequential", "int4-gptq-nonsequential", "int8-gptq-sequential", "int8-gptq-nonsequential", + "nvfp4-gptq-sequential", + "nvfp4-gptq-nonsequential", ]: # Determine base config based on quantization type if "int4" in args.quantization: base_config = Int4WeightOnlyConfig(group_size=args.group_size) quant_type = "Int4" - else: # int8 + elif "int8" in args.quantization: base_config = Int8WeightOnlyConfig(granularity=PerRow(), version=2) quant_type = "Int8" + else: # nvfp4 + base_config = NVFP4DynamicActivationNVFP4WeightConfig( + use_dynamic_per_tensor_scale=True, + use_triton_kernel=True, + ) + quant_type = "NVFP4" # First application: wrap weights with GPTQObserverTensor (observe step) print( @@ -284,7 +317,7 @@ def main(): percdamp=args.percdamp, gptq_quantize_block_size=args.gptq_block_size, ) - quantize_(model, observe_config, filter_fn=None) + quantize_(model, observe_config, filter_fn=skip_lm_head) # Prepare calibration dataset print( @@ -316,7 +349,7 @@ def main(): for seq in tqdm(dataset, desc="Calibrating"): model(seq.to(input_device)) # Apply quantization - quantize_(model, convert_config, filter_fn=None) + quantize_(model, convert_config, filter_fn=skip_lm_head) else: # sequential print(f"Applying {quant_type} GPTQ quantization (sequential)...") sequential_quantize(model, dataset, convert_config) @@ -334,6 +367,23 @@ def main(): # Save model to generated output directory print(f"Saving model to {output_dir}...") tokenizer.save_pretrained(output_dir) + print(model) + if model.config.tie_word_embeddings: + model.config.tie_word_embeddings = False + model._tied_weights_keys = {} + model.lm_head.weight = torch.nn.Parameter( + model.lm_head.weight.clone(), requires_grad=False + ) + + # monkey patch huggingface to skip weight tie checks + # TODO(before land): debug why i'm hitting this error here, + # as nvfp4 works in other hf models just fine + # import transformers.modeling_utils + # transformers.modeling_utils.remove_tied_weights_from_state_dict = lambda state_dict, model: state_dict + + import transformers + assert transformers.__version__ == '4.57.6', "unsupported" + model.save_pretrained(output_dir, safe_serialization=False) print("DONE!") From a669b9e8b5b331dec619d5995e2d87674a637450 Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Tue, 21 Apr 2026 00:28:42 +0000 Subject: [PATCH 03/25] Update [ghstack-poisoned] --- test/prototype/gptq/test_gptqv2.py | 18 +++++++++++++++++- torchao/prototype/gptq/gptq_example.py | 16 +++++----------- 2 files changed, 22 insertions(+), 12 deletions(-) diff --git a/test/prototype/gptq/test_gptqv2.py b/test/prototype/gptq/test_gptqv2.py index 0bae6ec4da..a67804590c 100644 --- a/test/prototype/gptq/test_gptqv2.py +++ b/test/prototype/gptq/test_gptqv2.py @@ -20,7 +20,10 @@ ) from torchao.quantization import Int4WeightOnlyConfig, Int8WeightOnlyConfig, quantize_ from torchao.quantization.granularity import PerRow -from torchao.utils import _is_mslk_available +from torchao.utils import ( + _is_mslk_available, + is_sm_at_least_100, +) def _calculate_hessian(inputs, device=None): @@ -449,6 +452,13 @@ def test_gptq_quantize_function(self, base_config): ) def test_gptq_quantize_better_than_naive(self, base_config): """Test that GPTQ produces lower error than naive quantization.""" + + if ( + isinstance(base_config, NVFP4DynamicActivationNVFP4WeightConfig) + and not is_sm_at_least_100() + ): + pytest.skip("CUDA capability >= 10.0 required for nvfp4") + torch.manual_seed(43) # Create weight and realistic Hessian from actual activations @@ -526,6 +536,12 @@ def test_gptq_quantize_better_than_naive(self, base_config): ], ) def test_gptq_sqnr(self, base_config): + if ( + isinstance(base_config, NVFP4DynamicActivationNVFP4WeightConfig) + and not is_sm_at_least_100() + ): + pytest.skip("CUDA capability >= 10.0 required for nvfp4") + torch.manual_seed(43) model = ToyLinearModel(m=512, n=2048, k=1024).cuda().to(torch.bfloat16) diff --git a/torchao/prototype/gptq/gptq_example.py b/torchao/prototype/gptq/gptq_example.py index a8415828c2..57ab7ed974 100644 --- a/torchao/prototype/gptq/gptq_example.py +++ b/torchao/prototype/gptq/gptq_example.py @@ -230,8 +230,7 @@ def main(): # Generate output directory name from args model_name = args.model_id.split("/")[-1] # Get last part of model ID - # TODO(before land): make output dir configurable - output_dir = f"/home/dev/tmp/20260420_{model_name}_{args.quantization}" + output_dir = f"{model_name}_{args.quantization}" if args.quantization != "none": output_dir += f"_gs{args.group_size}" @@ -359,16 +358,11 @@ def main(): "--model_args", f"pretrained={output_dir}", "--tasks", - # "leaderboard_bbh", - # "gsm8k", - "wikitext", - # "--num_fewshot", - # "3", + "leaderboard_bbh", + "--num_fewshot", + "3", "--batch_size", - # "auto", - "1", - # "--limit", - # "20", + "auto", ] print(f"Running command: {' '.join(lm_eval_cmd)}") From 53bd8d0214d30d1d2c287cd5660f3011afd71226 Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Tue, 21 Apr 2026 01:00:51 +0000 Subject: [PATCH 04/25] Update [ghstack-poisoned] --- torchao/prototype/gptq/api.py | 10 +++++++- torchao/prototype/gptq/gptq_example.py | 33 +++++++++++--------------- 2 files changed, 23 insertions(+), 20 deletions(-) diff --git a/torchao/prototype/gptq/api.py b/torchao/prototype/gptq/api.py index b7011e68f6..6d5d0f2d80 100644 --- a/torchao/prototype/gptq/api.py +++ b/torchao/prototype/gptq/api.py @@ -99,6 +99,11 @@ def __post_init__(self): ) +# simple progress counter for GPTQ convert +# TODO(future): make this cleaner, will require a refactor +gptq_convert_layer_counter = 0 + + @register_quantize_module_handler(GPTQConfig) def _gptq_config_transform( module: torch.nn.Module, config: GPTQConfig, *, parameter_name="weight" @@ -120,7 +125,10 @@ def _gptq_config_transform( ) return module elif config.step == "convert": - print("gptq convert") + global gptq_convert_layer_counter + print(f"gptq convert {gptq_convert_layer_counter}") + gptq_convert_layer_counter += 1 + # Quantization phase: tensor should be an GPTQObserverTensor if not isinstance(tensor, GPTQObserverTensor): raise ValueError( diff --git a/torchao/prototype/gptq/gptq_example.py b/torchao/prototype/gptq/gptq_example.py index 041ef572b7..38473d11f9 100644 --- a/torchao/prototype/gptq/gptq_example.py +++ b/torchao/prototype/gptq/gptq_example.py @@ -154,8 +154,7 @@ def parse_args(): parser.add_argument( "--num-calibration-samples", type=int, - # default=128, - default=2, + default=128, help="Number of calibration samples to use", ) parser.add_argument( @@ -265,9 +264,8 @@ def main(): def skip_lm_head(module, fqn): # TODO(before land): remove o_proj from below, it's just for debugging return ( - isinstance(module, torch.nn.Linear) - and "lm_head" not in fqn - and "o_proj" in fqn + isinstance(module, torch.nn.Linear) and "lm_head" not in fqn + # and "o_proj" in fqn ) if args.quantization == "int4-rtn": @@ -352,6 +350,13 @@ def skip_lm_head(module, fqn): # Run calibration for seq in tqdm(dataset, desc="Calibrating"): model(seq.to(input_device)) + # Print total # of GPTQ modules + num_gptq_weights = 0 + for name, param in model.named_parameters(): + # TODO(before land): better check + if "GPTQ" in str(type(param)): + num_gptq_weights += 1 + print(f"Total GPTQ weights to convert: {num_gptq_weights}") # Apply quantization quantize_(model, convert_config, filter_fn=skip_lm_head) else: # sequential @@ -372,22 +377,12 @@ def skip_lm_head(module, fqn): print(f"Saving model to {output_dir}...") tokenizer.save_pretrained(output_dir) print(model) - if model.config.tie_word_embeddings: - model.config.tie_word_embeddings = False - model._tied_weights_keys = {} - model.lm_head.weight = torch.nn.Parameter( - model.lm_head.weight.clone(), requires_grad=False - ) - - # monkey patch huggingface to skip weight tie checks - # TODO(before land): debug why i'm hitting this error here, - # as nvfp4 works in other hf models just fine - # import transformers.modeling_utils - # transformers.modeling_utils.remove_tied_weights_from_state_dict = lambda state_dict, model: state_dict - import transformers + # transformers 5.0.0 have a lot of errors with nvfp4 subclasses + # TODO(before land): debug this further + import transformers - assert transformers.__version__ == "4.57.6", "unsupported" + assert transformers.__version__ == "4.57.6", "unsupported" model.save_pretrained(output_dir, safe_serialization=False) From 4c86363542dc359f3a3c684c2931145e50109d25 Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Tue, 21 Apr 2026 10:21:32 +0000 Subject: [PATCH 05/25] Update [ghstack-poisoned] --- torchao/prototype/gptq/api.py | 1 + 1 file changed, 1 insertion(+) diff --git a/torchao/prototype/gptq/api.py b/torchao/prototype/gptq/api.py index ac6f7cdc6e..1a83aee684 100644 --- a/torchao/prototype/gptq/api.py +++ b/torchao/prototype/gptq/api.py @@ -319,6 +319,7 @@ def gptq_quantize(H: torch.Tensor, W_t: torch.Tensor, config: GPTQConfig): block_size = list(block_size) group_size = block_size[-1] elif isinstance(base_config, NVFP4DynamicActivationNVFP4WeightConfig): + assert base_config.use_dynamic_per_tensor_scale, "unsupported" group_size = 16 block_size = [1, group_size] # for per-tensor nvfp4, we need to calculate the global scale over the From 9b7dc74acae46d3f017792490981f6646affa52a Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Tue, 21 Apr 2026 13:55:17 +0000 Subject: [PATCH 06/25] Update [ghstack-poisoned] --- torchao/prototype/gptq/gptq_example.py | 133 ++++++++++++++---- ...vfp4_llama3_2_1b_nonsequential_wikitext.sh | 22 +++ 2 files changed, 125 insertions(+), 30 deletions(-) create mode 100755 torchao/prototype/gptq/gptq_nvfp4_llama3_2_1b_nonsequential_wikitext.sh diff --git a/torchao/prototype/gptq/gptq_example.py b/torchao/prototype/gptq/gptq_example.py index 38473d11f9..08c71cd1d0 100644 --- a/torchao/prototype/gptq/gptq_example.py +++ b/torchao/prototype/gptq/gptq_example.py @@ -16,7 +16,12 @@ from tqdm import tqdm from transformers import AutoModelForCausalLM, AutoTokenizer +from packaging.version import Version + +import transformers + from torchao.prototype.gptq import GPTQConfig +from torchao.prototype.gptq.observer import GPTQObserverTensor from torchao.prototype.mx_formats.inference_workflow import ( NVFP4DynamicActivationNVFP4WeightConfig, ) @@ -96,14 +101,19 @@ def prepare_dataset( dataset_map = { "hellaswag": "Rowan/hellaswag", "ultrachat200k": "HuggingFaceH4/ultrachat_200k", + "c4": ("allenai/c4", "en"), } - hf_dataset_id = dataset_map.get(dataset_id, dataset_id) + dataset_entry = dataset_map.get(dataset_id, dataset_id) + if isinstance(dataset_entry, tuple): + hf_dataset_id, hf_config_name = dataset_entry + else: + hf_dataset_id, hf_config_name = dataset_entry, None # Load dataset and preprocess - # TODO(before land): clean up below - dataset_split = "train" - train_dataset_raw = load_dataset(hf_dataset_id, split=dataset_split, streaming=True) + train_dataset_raw = load_dataset( + hf_dataset_id, hf_config_name, split=dataset_split, streaming=True + ) train_dataset_raw = train_dataset_raw.shuffle(seed=seed, buffer_size=1_000) def preprocess_hellaswag(example): @@ -167,9 +177,15 @@ def parse_args(): "--dataset-id", type=str, default="ultrachat200k", - choices=["hellaswag", "ultrachat200k"], + choices=["hellaswag", "ultrachat200k", "c4"], help="Dataset for calibration (hellaswag or ultrachat200k)", ) + parser.add_argument( + "--dataset-split", + type=str, + default="train_sft", + help="Dataset split to use for calibration", + ) parser.add_argument( "--quantization", type=str, @@ -206,12 +222,62 @@ def parse_args(): default=1024, help="Block size for GPTQ quantization", ) + parser.add_argument( + "--lm-eval-tasks", + type=str, + default="leaderboard_bbh", + help="Comma-separated tasks for lm_eval", + ) + parser.add_argument( + "--num-fewshot", + type=int, + default=3, + help="Number of few-shot examples for lm_eval (0 to disable)", + ) + parser.add_argument( + "--lm-eval-batch-size", + type=str, + default="auto", + help="Batch size for lm_eval (default: auto)", + ) + parser.add_argument( + "--lm-eval-limit", + type=int, + default=None, + help="Limit number of examples per task for lm_eval (default: no limit)", + ) + parser.add_argument( + "--output-dir-prefix", + type=str, + required=True, + help="Prefix for the output directory (e.g. /home/user/tmp/20260420)", + ) + parser.add_argument( + "--skip-lm-eval", + action="store_true", + default=False, + help="Skip running lm_eval after quantization, useful for quickly iterating on lm_eval arguments", + ) + parser.add_argument( + "--o-proj-only", + action="store_true", + default=False, + help="Only quantize `o_proj` layers, useful for faster GPTQ runs for debugging", + ) return parser.parse_args() def main(): args = parse_args() + # lm_eval batch_size="auto" with nvfp4 gptq causes the error in + # MSLK nvfp4 triton kernel, likely an unsupported shape: + # https://gist.github.com/vkuzo/b71ca46365dee017d1602e9638d91603 + # TODO(future): debug and fix this. For now, the workaround is + # for the user to manually specify lm_eval batch_size. + if "nvfp4" in args.quantization: + assert args.lm_eval_batch_size != "auto", "unsupported" + # Map dtype string to torch dtype dtype_map = { "float32": torch.float32, @@ -238,9 +304,7 @@ def main(): # Generate output directory name from args model_name = args.model_id.split("/")[-1] # Get last part of model ID - output_dir = f"{model_name}_{args.quantization}" - # TODO(before land): make output dir configurable - output_dir = f"/home/dev/tmp/20260420_{model_name}_{args.quantization}" + output_dir = f"{args.output_dir_prefix}_{model_name}_{args.quantization}" if args.quantization != "none": output_dir += f"_gs{args.group_size}" @@ -262,21 +326,28 @@ def main(): quantization_start_time = time.time() def skip_lm_head(module, fqn): - # TODO(before land): remove o_proj from below, it's just for debugging + return isinstance(module, torch.nn.Linear) and "lm_head" not in fqn + + def skip_lm_head_o_proj(module, fqn): return ( - isinstance(module, torch.nn.Linear) and "lm_head" not in fqn - # and "o_proj" in fqn + isinstance(module, torch.nn.Linear) + and "lm_head" not in fqn + and "o_proj" in fqn ) + filter_fn_to_use = skip_lm_head + if args.o_proj_only: + filter_fn_to_use = skip_lm_head_o_proj + if args.quantization == "int4-rtn": print("Applying Int4 RTN (Round-To-Nearest) quantization...") config = Int4WeightOnlyConfig(group_size=args.group_size) - quantize_(model, config, filter_fn=skip_lm_head) + quantize_(model, config, filter_fn=filter_fn_to_use) elif args.quantization == "int8-rtn": print("Applying Int8 RTN (Round-To-Nearest) quantization...") config = Int8WeightOnlyConfig(version=2, granularity=PerRow()) - quantize_(model, config, filter_fn=skip_lm_head) + quantize_(model, config, filter_fn=filter_fn_to_use) elif args.quantization == "nvfp4-rtn": print("Applying NVFP4 RTN (Round-To-Nearest) quantization...") @@ -285,7 +356,7 @@ def skip_lm_head(module, fqn): use_dynamic_per_tensor_scale=True, use_triton_kernel=True, ) - quantize_(model, config, filter_fn=skip_lm_head) + quantize_(model, config, filter_fn=filter_fn_to_use) elif args.quantization in [ "int4-gptq-sequential", @@ -319,7 +390,7 @@ def skip_lm_head(module, fqn): percdamp=args.percdamp, gptq_quantize_block_size=args.gptq_block_size, ) - quantize_(model, observe_config, filter_fn=skip_lm_head) + quantize_(model, observe_config, filter_fn=filter_fn_to_use) # Prepare calibration dataset print( @@ -330,7 +401,7 @@ def skip_lm_head(module, fqn): max_seq_length, args.num_calibration_samples, dataset_id=args.dataset_id, - dataset_split="train_sft", + dataset_split=args.dataset_split, seed=42, ) @@ -353,14 +424,14 @@ def skip_lm_head(module, fqn): # Print total # of GPTQ modules num_gptq_weights = 0 for name, param in model.named_parameters(): - # TODO(before land): better check - if "GPTQ" in str(type(param)): + if isinstance(param, GPTQObserverTensor): num_gptq_weights += 1 print(f"Total GPTQ weights to convert: {num_gptq_weights}") # Apply quantization - quantize_(model, convert_config, filter_fn=skip_lm_head) + quantize_(model, convert_config, filter_fn=filter_fn_to_use) else: # sequential print(f"Applying {quant_type} GPTQ quantization (sequential)...") + assert filter_fn_to_use == skip_lm_head, "unsupported" sequential_quantize(model, dataset, convert_config) quantization_end_time = time.time() @@ -380,9 +451,9 @@ def skip_lm_head(module, fqn): # transformers 5.0.0 have a lot of errors with nvfp4 subclasses # TODO(before land): debug this further - import transformers - - assert transformers.__version__ == "4.57.6", "unsupported" + assert Version(transformers.__version__) < Version("5.0.0"), ( + f"transformers {transformers.__version__} is not supported, need < 5.0.0" + ) model.save_pretrained(output_dir, safe_serialization=False) @@ -409,18 +480,20 @@ def skip_lm_head(module, fqn): "--model_args", f"pretrained={output_dir}", "--tasks", - # "leaderboard_bbh", - "wikitext", - # "--num_fewshot", - # "3", + args.lm_eval_tasks, "--batch_size", - # "auto", - "1", - # "--limit", - # "20", + args.lm_eval_batch_size, ] + if args.num_fewshot > 0: + lm_eval_cmd += ["--num_fewshot", str(args.num_fewshot)] + if args.lm_eval_limit is not None: + lm_eval_cmd += ["--limit", str(args.lm_eval_limit)] + print(f"Running command: {' '.join(lm_eval_cmd)}") + if args.skip_lm_eval: + print("Terminating early due to skip_lm_eval=True") + return try: subprocess.run(lm_eval_cmd, check=True) except subprocess.CalledProcessError as e: diff --git a/torchao/prototype/gptq/gptq_nvfp4_llama3_2_1b_nonsequential_wikitext.sh b/torchao/prototype/gptq/gptq_nvfp4_llama3_2_1b_nonsequential_wikitext.sh new file mode 100755 index 0000000000..e60d17ab7a --- /dev/null +++ b/torchao/prototype/gptq/gptq_nvfp4_llama3_2_1b_nonsequential_wikitext.sh @@ -0,0 +1,22 @@ +#!/bin/bash + +# +# A quick smoke test for non-sequential GPTQ on `unsloth/Llama-3.2-1B` +# + +COMMON_ARGS="--output-dir-prefix /home/dev/tmp/20260421 --model-id unsloth/Llama-3.2-1B --lm-eval-tasks wikitext --num-fewshot 0 --lm-eval-batch-size 16" + +# baseline (bf16) +echo -e "\n\nbaseline (bf16)\n\n" +python -u torchao/prototype/gptq/gptq_example.py $COMMON_ARGS --quantization none +echo -e "done" + +# nvfp4-rtn +echo -e "\n\nnvfp4-rtn\n\n" +python -u torchao/prototype/gptq/gptq_example.py $COMMON_ARGS --quantization nvfp4-rtn +echo -e "done" + +# nvfp4-gptq-nonsequential +echo -e "\n\nnvfp4-gptq-nonsequential\n\n" +python -u torchao/prototype/gptq/gptq_example.py $COMMON_ARGS --quantization nvfp4-gptq-nonsequential --dataset-id c4 --dataset-split train +echo -e "done" From d69b32a6714efbe917ce849353597861d5d18c8a Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Tue, 21 Apr 2026 14:03:33 +0000 Subject: [PATCH 07/25] Update [ghstack-poisoned] --- torchao/prototype/gptq/gptq_example.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/torchao/prototype/gptq/gptq_example.py b/torchao/prototype/gptq/gptq_example.py index 08c71cd1d0..d24097ce96 100644 --- a/torchao/prototype/gptq/gptq_example.py +++ b/torchao/prototype/gptq/gptq_example.py @@ -12,14 +12,12 @@ from typing import Any, List, Optional import torch +import transformers from datasets import load_dataset from tqdm import tqdm from transformers import AutoModelForCausalLM, AutoTokenizer from packaging.version import Version - -import transformers - from torchao.prototype.gptq import GPTQConfig from torchao.prototype.gptq.observer import GPTQObserverTensor from torchao.prototype.mx_formats.inference_workflow import ( From 294c9ccfdf495ed7cb365cff7b97d858eb4903ad Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Tue, 21 Apr 2026 18:31:12 +0000 Subject: [PATCH 08/25] Update [ghstack-poisoned] --- benchmarks/benchmark_gptq.py | 76 +++++++++++++++++++++++++++++++++++ torchao/prototype/gptq/api.py | 33 ++++++++++++++- 2 files changed, 107 insertions(+), 2 deletions(-) create mode 100644 benchmarks/benchmark_gptq.py diff --git a/benchmarks/benchmark_gptq.py b/benchmarks/benchmark_gptq.py new file mode 100644 index 0000000000..142626286d --- /dev/null +++ b/benchmarks/benchmark_gptq.py @@ -0,0 +1,76 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +import time + +import fire +import torch + +from torchao.prototype.gptq import GPTQConfig, gptq_quantize +from torchao.prototype.mx_formats.inference_workflow import ( + NVFP4DynamicActivationNVFP4WeightConfig, +) + + +def run( + K: int = 2048, + N: int = 4096, + profile_fname: str = None, +): + print(f"K={K}, N={N}") + + A = torch.randn(K, K, dtype=torch.float32, device="cuda") + H = A.t() @ A + + W_t = torch.randn(N, K, dtype=torch.bfloat16, device="cuda") + + config = GPTQConfig( + step="convert", + base_config=NVFP4DynamicActivationNVFP4WeightConfig( + use_dynamic_per_tensor_scale=True, + use_triton_kernel=True, + ), + ) + + # Warmup + print("Warmup...") + gptq_quantize(H.clone(), W_t.clone(), config) + torch.cuda.synchronize() + + num_runs = 5 + if profile_fname is not None: + print("Profiling run...") + with torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ], + record_shapes=True, + with_stack=True, + ) as prof: + torch.cuda.synchronize() + start = time.time() + gptq_quantize(H.clone(), W_t.clone(), config) + torch.cuda.synchronize() + elapsed = time.time() - start + print(f"gptq_quantize time: {elapsed:.3f}s") + prof.export_chrome_trace(profile_fname) + print(f"Saved: {profile_fname}") + else: + print(f"Timed run ({num_runs} iterations)...") + times = [] + for _ in range(num_runs): + torch.cuda.synchronize() + start = time.time() + gptq_quantize(H.clone(), W_t.clone(), config) + torch.cuda.synchronize() + times.append(time.time() - start) + avg = sum(times) / len(times) + print(f"gptq_quantize avg time: {avg:.3f}s") + + +if __name__ == "__main__": + fire.Fire(run) diff --git a/torchao/prototype/gptq/api.py b/torchao/prototype/gptq/api.py index 14881d974e..d3d7329c0a 100644 --- a/torchao/prototype/gptq/api.py +++ b/torchao/prototype/gptq/api.py @@ -4,7 +4,9 @@ # This source code is licensed under the BSD 3-Clause license found in the # LICENSE file in the root directory of this source tree. +import os import types +import warnings from dataclasses import dataclass from functools import partial from typing import Union @@ -265,6 +267,33 @@ def _nvfp4_with_precalculated_scales_q( return data_lp_packed +# Set to True to torch.compile the NVFP4 quantize/dequantize functions +# inside gptq_quantize. Gives ~3x speedup. +_use_torch_compile = True + +if _use_torch_compile: + _nvfp4_qdq_fn = torch.compile(_nvfp4_with_precalculated_scales_qdq) + _nvfp4_q_fn = torch.compile(_nvfp4_with_precalculated_scales_q) + + # Triton's default f32 division uses approximate reciprocal which + # introduces ~1 ULP error per division. In GPTQ's error propagation + # loop this compounds across columns. IEEE-compliant division rounding + # eliminates the drift. + import torch._inductor.config as _inductor_config + + if os.environ.get("TORCHINDUCTOR_EMULATE_DIVISION_ROUNDING") == "0": + warnings.warn( + "TORCHINDUCTOR_EMULATE_DIVISION_ROUNDING=0 may cause numerical " + "drift in GPTQ with torch.compile. " + "Consider unsetting it or setting it to 1." + ) + else: + _inductor_config.eager_numerics.division_rounding = True +else: + _nvfp4_qdq_fn = _nvfp4_with_precalculated_scales_qdq + _nvfp4_q_fn = _nvfp4_with_precalculated_scales_q + + def gptq_quantize(H: torch.Tensor, W_t: torch.Tensor, config: GPTQConfig): """ This function implements the GPTQ algorithm described in this paper: https://arxiv.org/abs/2210.17323 (Algorithm 1) @@ -472,7 +501,7 @@ def gptq_quantize(H: torch.Tensor, W_t: torch.Tensor, config: GPTQConfig): ) dq = q.dequantize(output_dtype=torch.float) elif isinstance(base_config, NVFP4DynamicActivationNVFP4WeightConfig): - dq = _nvfp4_with_precalculated_scales_qdq( + dq = _nvfp4_qdq_fn( w_t, nvfp4_global_scale, scale.squeeze(-1), @@ -519,7 +548,7 @@ def gptq_quantize(H: torch.Tensor, W_t: torch.Tensor, config: GPTQConfig): combined_scale = ( torch.cat(group_qparams, dim=0).reshape(K // group_size, N).t().contiguous() ) - qdata = _nvfp4_with_precalculated_scales_q( + qdata = _nvfp4_q_fn( W_t, nvfp4_global_scale, combined_scale, From 65fae6264a1accca9a1e44b7baf7d3b019fddc4d Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Tue, 21 Apr 2026 18:52:00 +0000 Subject: [PATCH 09/25] Update [ghstack-poisoned] --- test/prototype/gptq/test_gptqv2.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/test/prototype/gptq/test_gptqv2.py b/test/prototype/gptq/test_gptqv2.py index a67804590c..78af648e14 100644 --- a/test/prototype/gptq/test_gptqv2.py +++ b/test/prototype/gptq/test_gptqv2.py @@ -10,6 +10,13 @@ import torch import torch.nn.functional as F +from torchao.utils import torch_version_at_least + +pytestmark = pytest.mark.skipif( + not torch_version_at_least("2.11.0"), + reason="GPTQ prototype requires PyTorch 2.11+", +) + from torchao.prototype.gptq import ( GPTQConfig, gptq_quantize, From 5ee2ad2b6029eef0dee6f7317bf2a1eacd3560a7 Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Wed, 22 Apr 2026 11:21:39 +0000 Subject: [PATCH 10/25] Update [ghstack-poisoned] --- torchao/prototype/gptq/api.py | 33 +++++++++++++++++++++------------ 1 file changed, 21 insertions(+), 12 deletions(-) diff --git a/torchao/prototype/gptq/api.py b/torchao/prototype/gptq/api.py index d3d7329c0a..6869353273 100644 --- a/torchao/prototype/gptq/api.py +++ b/torchao/prototype/gptq/api.py @@ -14,6 +14,8 @@ import torch import torch.nn as nn +from torchao.utils import torch_version_at_least + try: from mslk.quantize.shuffle import int4_row_quantize_zp, pack_int4 except: @@ -275,20 +277,27 @@ def _nvfp4_with_precalculated_scales_q( _nvfp4_qdq_fn = torch.compile(_nvfp4_with_precalculated_scales_qdq) _nvfp4_q_fn = torch.compile(_nvfp4_with_precalculated_scales_q) - # Triton's default f32 division uses approximate reciprocal which - # introduces ~1 ULP error per division. In GPTQ's error propagation - # loop this compounds across columns. IEEE-compliant division rounding - # eliminates the drift. - import torch._inductor.config as _inductor_config - - if os.environ.get("TORCHINDUCTOR_EMULATE_DIVISION_ROUNDING") == "0": + if torch_version_at_least("2.11.0"): + # Triton's default f32 division uses approximate reciprocal which + # introduces ~1 ULP error per division. In GPTQ's error propagation + # loop this compounds across columns. IEEE-compliant division rounding + # eliminates the drift. + import torch._inductor.config as _inductor_config + + if os.environ.get("TORCHINDUCTOR_EMULATE_DIVISION_ROUNDING") == "0": + warnings.warn( + "TORCHINDUCTOR_EMULATE_DIVISION_ROUNDING=0 may cause numerical " + "drift in GPTQ with torch.compile. " + "Consider unsetting it or setting it to 1." + ) + else: + _inductor_config.eager_numerics.division_rounding = True + else: warnings.warn( - "TORCHINDUCTOR_EMULATE_DIVISION_ROUNDING=0 may cause numerical " - "drift in GPTQ with torch.compile. " - "Consider unsetting it or setting it to 1." + "PyTorch < 2.11.0 detected. Upgrade to PyTorch 2.11.0+ for " + "better GPTQ numerics with torch.compile (IEEE-compliant " + "division rounding)." ) - else: - _inductor_config.eager_numerics.division_rounding = True else: _nvfp4_qdq_fn = _nvfp4_with_precalculated_scales_qdq _nvfp4_q_fn = _nvfp4_with_precalculated_scales_q From 2adda75ff0a6f84596bf1e88782b060fed8e5821 Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Wed, 22 Apr 2026 11:21:39 +0000 Subject: [PATCH 11/25] Update [ghstack-poisoned] --- torchao/prototype/gptq/gptq_example.py | 28 ++++++++++++++++++-------- 1 file changed, 20 insertions(+), 8 deletions(-) diff --git a/torchao/prototype/gptq/gptq_example.py b/torchao/prototype/gptq/gptq_example.py index d24097ce96..016f0ff174 100644 --- a/torchao/prototype/gptq/gptq_example.py +++ b/torchao/prototype/gptq/gptq_example.py @@ -12,12 +12,11 @@ from typing import Any, List, Optional import torch -import transformers from datasets import load_dataset from tqdm import tqdm -from transformers import AutoModelForCausalLM, AutoTokenizer +from transformers import AutoModelForCausalLM, AutoTokenizer, TorchAoConfig +from transformers.quantizers.quantizer_torchao import TorchAoHfQuantizer -from packaging.version import Version from torchao.prototype.gptq import GPTQConfig from torchao.prototype.gptq.observer import GPTQObserverTensor from torchao.prototype.mx_formats.inference_workflow import ( @@ -447,11 +446,24 @@ def skip_lm_head_o_proj(module, fqn): tokenizer.save_pretrained(output_dir) print(model) - # transformers 5.0.0 have a lot of errors with nvfp4 subclasses - # TODO(before land): debug this further - assert Version(transformers.__version__) < Version("5.0.0"), ( - f"transformers {transformers.__version__} is not supported, need < 5.0.0" - ) + if "nvfp4" in args.quantization: + import inspect + + source = inspect.getsource(TorchAoHfQuantizer.get_weight_conversions) + if "_weight_per_tensor_scale" not in source: + raise RuntimeError( + "Your version of `transformers` does not support NVFP4 serialization. " + "Please install a version that includes " + "https://github.com/huggingface/transformers/pull/45573" + ) + + if args.quantization != "none": + # Attach hf_quantizer so save_pretrained uses the flatten path for tensor + # subclasses (e.g. NVFP4Tensor) that don't have a valid storage pointer. + ao_config = base_config if "gptq" in args.quantization else config + torchao_config = TorchAoConfig(quant_type=ao_config) + model.config.quantization_config = torchao_config + model.hf_quantizer = TorchAoHfQuantizer(torchao_config) model.save_pretrained(output_dir, safe_serialization=False) From 5fe657469f1cff71df891fc9ebbdb4b823d4e722 Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Wed, 22 Apr 2026 18:01:24 +0000 Subject: [PATCH 12/25] Update [ghstack-poisoned] --- .../mx_formats/test_inference_workflow.py | 69 +++++++++++++++++++ torchao/prototype/mx_formats/mx_tensor.py | 4 +- torchao/prototype/mx_formats/nvfp4_tensor.py | 11 +++ 3 files changed, 82 insertions(+), 2 deletions(-) diff --git a/test/prototype/mx_formats/test_inference_workflow.py b/test/prototype/mx_formats/test_inference_workflow.py index cdc2b31aa5..23bc6b5a8b 100644 --- a/test/prototype/mx_formats/test_inference_workflow.py +++ b/test/prototype/mx_formats/test_inference_workflow.py @@ -17,6 +17,7 @@ NVFP4DynamicActivationNVFP4WeightConfig, NVFP4WeightOnlyConfig, ) +from torchao.prototype.mx_formats.nvfp4_tensor import NVFP4Tensor from torchao.quantization import quantize_ from torchao.quantization.quantize_.common import KernelPreference from torchao.quantization.utils import compute_error @@ -426,3 +427,71 @@ def test_nvfp4_static_vs_dynamic_quantization(): assert torch.equal(y_dynamic, y_static), ( "Expect dynamic and static quant result to be equal" ) + + +class GroupedMMModel(nn.Module): + """A toy model whose only op in forward is torch._grouped_mm.""" + + def __init__(self, E, K, N, device="cuda", dtype=torch.bfloat16): + super().__init__() + self.weight = nn.Parameter(torch.randn(E, N, K, device=device, dtype=dtype)) + + def forward(self, x, offs): + return torch._grouped_mm(x, self.weight.transpose(-2, -1), offs=offs) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif( + not torch_version_at_least("2.8.0"), reason="torch >= 2.8.0 required" +) +@pytest.mark.skipif( + not is_sm_at_least_100(), reason="CUDA capability >= 10.0 required for NVFP4" +) +@torch.no_grad() +@skip_if_rocm("ROCm float4 gemm require gfx950") +def test_grouped_mm_nvfp4(): + """Smoke test: torch._grouped_mm forward pass with bfloat16 inputs.""" + E, K, N = 4, 128, 256 + m_per_group = [32, 96, 16, 112] + total_m = sum(m_per_group) + + device = "cuda" + dtype = torch.bfloat16 + + model_ref = GroupedMMModel(E, K, N, device=device, dtype=dtype) + model = copy.deepcopy(model_ref) + + x = torch.randn(total_m, K, device=device, dtype=dtype) + offs = torch.tensor( + [sum(m_per_group[: i + 1]) for i in range(E)], + device=device, + dtype=torch.int32, + ) + + # reference path + y_ref = model_ref(x, offs) + + # quantized path + quantize_( + model, + NVFP4DynamicActivationNVFP4WeightConfig( + use_triton_kernel=False, + ), + filter_fn=lambda mod, *args: isinstance(mod, GroupedMMModel) + and hasattr(mod, "weight"), + ) + assert isinstance(model.weight, NVFP4Tensor), ( + f"Expected NVFP4Tensor weight, got {type(model.weight)}" + ) + w_sqnr = compute_error(model_ref.weight, model.weight.dequantize()) + assert w_sqnr > 18.0 + + ww = model.weight + wwt = ww.transpose(-2, -1) + assert tuple(wwt.shape) == (ww.shape[0], ww.shape[2], ww.shape[1]) + + # For now, this is emulated. In the near future we'll hook up + # a real nvfp4 grouped gemm. + y = model(x, offs) + y_sqnr = compute_error(y_ref, y) + assert y_sqnr > 18.0 diff --git a/torchao/prototype/mx_formats/mx_tensor.py b/torchao/prototype/mx_formats/mx_tensor.py index 0f0f45a5fd..c6ca81d08a 100644 --- a/torchao/prototype/mx_formats/mx_tensor.py +++ b/torchao/prototype/mx_formats/mx_tensor.py @@ -486,7 +486,7 @@ def tensor_size_hp_to_fp4x2(orig_size, is_contiguous): else: assert len(orig_size) == 3, "unsupported" # only supporting dim0, dim1, dim2 and dim0, dim2, dim1 orders - new_size = [new_size[0], new_size[2] // 2, new_size[1]] + new_size = [new_size[0], new_size[1] // 2, new_size[2]] return new_size @@ -500,7 +500,7 @@ def tensor_size_fp4x2_to_hp(orig_size, is_contiguous): else: assert len(orig_size) == 3, "unsupported" # only supporting dim0, dim1, dim2 and dim0, dim2, dim1 orders - new_size = [new_size[0], new_size[2] * 2, new_size[1]] + new_size = [new_size[0], new_size[1] * 2, new_size[2]] return new_size diff --git a/torchao/prototype/mx_formats/nvfp4_tensor.py b/torchao/prototype/mx_formats/nvfp4_tensor.py index 766692bacf..d7c4a3b671 100644 --- a/torchao/prototype/mx_formats/nvfp4_tensor.py +++ b/torchao/prototype/mx_formats/nvfp4_tensor.py @@ -656,6 +656,17 @@ def nvfp4_addmm(func, types, args, kwargs): return _addmm_nvfp4_dispatch(input_tensor, weight_tensor, func, bias=bias) +@implements([aten._grouped_mm.default]) +def nvfp4_grouped_mm(func, types, args, kwargs): + mat_a, mat_b = args[0], args[1] + # temporary: dequantize nvfp4 tensor to call original grouped_mm + # TODO(future PR): enable nvfp4 with per-expert outer scale (current code + # has a single outer scale across all experts). + # TODO(future PR): hook up real nvfp4 grouped_mm + mat_b_dq = mat_b.dequantize() + return func(mat_a, mat_b_dq, *args[2:], **kwargs) + + def per_tensor_amax_to_scale(amax: torch.Tensor) -> torch.Tensor: """Convert per-tensor amax to per-tensor scale for NVFP4 quantization. From 5292f2f1ba3047e3c9a22808456853af57d7b2d3 Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Wed, 22 Apr 2026 19:59:55 +0000 Subject: [PATCH 13/25] Update [ghstack-poisoned] --- .../prototype/mx_formats/test_nvfp4_tensor.py | 60 +++++++++++++++++++ torchao/prototype/mx_formats/nvfp4_tensor.py | 35 ++++++++++- 2 files changed, 92 insertions(+), 3 deletions(-) diff --git a/test/prototype/mx_formats/test_nvfp4_tensor.py b/test/prototype/mx_formats/test_nvfp4_tensor.py index c90a9cc68a..afcb94f0ad 100644 --- a/test/prototype/mx_formats/test_nvfp4_tensor.py +++ b/test/prototype/mx_formats/test_nvfp4_tensor.py @@ -712,3 +712,63 @@ def test_nvfp4_matmul_optional_per_tensor_scale(shapes, a_has_scale, use_triton_ sqnr = compute_error(C_ref, C_nvfp4) SQNR_THRESHOLD = 16.0 assert sqnr >= SQNR_THRESHOLD, f"SQNR {sqnr:.2f} < {SQNR_THRESHOLD}, {a_has_scale=}" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif( + not torch_version_at_least("2.8.0"), reason="torch.compile requires PyTorch 2.8+" +) +def test_nvfp4_per_expert_scale(): + # per-tensor scale reference + E, K, N = 2, 64, 128 + x0 = torch.randn(N, K, dtype=torch.bfloat16, device="cuda") + x1 = torch.randn(N, K, dtype=torch.bfloat16, device="cuda") * 2 + tensor_amax_x0 = torch.max(torch.abs(x0)) + per_tensor_scale_x0 = per_tensor_amax_to_scale(tensor_amax_x0) + tensor_amax_x1 = torch.max(torch.abs(x1)) + per_tensor_scale_x1 = per_tensor_amax_to_scale(tensor_amax_x1) + x0_nvfp4 = NVFP4Tensor.to_nvfp4( + x0, per_tensor_scale=per_tensor_scale_x0, is_swizzled_scales=False + ) + x1_nvfp4 = NVFP4Tensor.to_nvfp4( + x1, per_tensor_scale=per_tensor_scale_x1, is_swizzled_scales=False + ) + + xc = torch.cat([x0, x1], dim=0).view(E, N, K) + scalec = torch.cat( + [ + per_tensor_scale_x0.view( + 1, + ), + per_tensor_scale_x1.view( + 1, + ), + ], + dim=0, + ).view(E, 1, 1) + + xc_nvfp4 = NVFP4Tensor.to_nvfp4( + xc, per_tensor_scale=scalec, is_swizzled_scales=False + ) + + # internals must match + torch.testing.assert_close( + torch.cat([x0_nvfp4.qdata, x1_nvfp4.qdata], dim=0).view(E, N, K // 2), + xc_nvfp4.qdata, + atol=0, + rtol=0, + ) + + torch.testing.assert_close( + torch.cat([x0_nvfp4.scale, x1_nvfp4.scale], dim=0).view(E, N, K // 16), + xc_nvfp4.scale, + atol=0, + rtol=0, + ) + + x0_dq = x0_nvfp4.dequantize() + x1_dq = x1_nvfp4.dequantize() + xc_dq = xc_nvfp4.dequantize() + + xc_dq_ref = torch.cat([x0_dq, x1_dq], dim=0).view(E, N, K) + torch.testing.assert_close(xc_dq_ref, xc_dq, atol=0, rtol=0) diff --git a/torchao/prototype/mx_formats/nvfp4_tensor.py b/torchao/prototype/mx_formats/nvfp4_tensor.py index d7c4a3b671..af132d96a9 100644 --- a/torchao/prototype/mx_formats/nvfp4_tensor.py +++ b/torchao/prototype/mx_formats/nvfp4_tensor.py @@ -104,6 +104,11 @@ def __new__( requires_grad=False, ) + if per_tensor_scale is not None: + # 0: per-tensor scale, with shape (,) + # 3: per-expert scale, with shape (E, 1, 1) + assert len(per_tensor_scale.shape) in (0, 3), "unsupported" + self.qdata = qdata self.scale = scale self.block_size = block_size @@ -224,7 +229,7 @@ def dequantize(self, output_dtype: Optional[torch.dtype] = None) -> torch.Tensor return result def get_hp_scales(self) -> torch.Tensor: - """Get the scales of the NVFP4Tensor in original dtype. + """Get the scales of the NVFP4Tensor in float32. Returns: torch.Tensor: Scales of the NVFP4Tensor @@ -243,9 +248,9 @@ def get_hp_scales(self) -> torch.Tensor: ) return ( - scale_e4m3.to(self.orig_dtype) + scale_e4m3.to(torch.float) if self.per_tensor_scale is None - else self.per_tensor_scale * scale_e4m3.to(self.orig_dtype) + else self.per_tensor_scale * scale_e4m3.to(torch.float) ) @classmethod @@ -283,6 +288,13 @@ def _same_metadata(cls, self: "NVFP4Tensor", src: "NVFP4Tensor") -> bool: implements = NVFP4Tensor.implements +def _assert_no_per_expert_scale(t): + if t.per_tensor_scale is not None: + assert len(t.per_tensor_scale.shape) == 0, ( + "per-expert scale not supported in this codepath" + ) + + # TODO(future PR): move this to AOBaseTensor (will require debugging/fixing CI) @implements([aten._to_copy.default]) def nvfp4_to_copy(func, types, args, kwargs): @@ -365,6 +377,7 @@ def nvfp4_slice(func, types, args, kwargs): assert len(x.shape) == 2, ( f"only rank 2 is supported for slice, got rank {len(x.shape)}" ) + _assert_no_per_expert_scale(x) sliced_data, sliced_scale = _swizzle_aware_slice(x, dim, start, end, step) @@ -388,6 +401,7 @@ def nvfp4_slice(func, types, args, kwargs): def nvfp4_t(func, types, args, kwargs): # For now, only transpose(input, 0, 1) is supported. old = args[0] + _assert_no_per_expert_scale(old) new = NVFP4Tensor( old.qdata.t(), old.scale.t(), @@ -405,6 +419,7 @@ def nvfp4_t(func, types, args, kwargs): @implements([aten.transpose.int]) def nvfp4_transpose(func, types, args, kwargs): old, dim0, dim1 = args + _assert_no_per_expert_scale(old) assert len(old.shape) == 3, f"unsupported rank {len(old.shape)}" valid_3d_dims = ((1, 2), (2, 1), (-1, -2), (-2, -1)) assert (dim0, dim1) in valid_3d_dims, f"transpose unsupported for {dim0=} {dim1=}" @@ -427,6 +442,7 @@ def nvfp4_transpose(func, types, args, kwargs): @implements([aten.view.default]) def nvfp4_view_op(func, types, args, kwargs): data = args[0].qdata + _assert_no_per_expert_scale(args[0]) new_size = args[1] new_size = tensor_size_hp_to_fp4x2(new_size, data.is_contiguous()) new_data = func(data, new_size, *args[2:], **kwargs) @@ -448,6 +464,7 @@ def nvfp4_select(func, types, args, kwargs): old, dim, index = args assert dim == 0, f"NVFP4Tensor aten.select.int with {dim=} is not yet supported" assert len(old.qdata.shape) == len(old.scale.shape), "unsupported" + _assert_no_per_expert_scale(old) new = old.__class__( old.qdata[index], old.scale[index], @@ -476,6 +493,8 @@ def _addmm_nvfp4_dispatch( assert a.block_size == 16, f"NVFP4 requires block_size=16, got {a.block_size}" assert b.block_size == 16, f"NVFP4 requires block_size=16, got {b.block_size}" assert len(a.shape) == 2 and len(b.shape) == 2 + _assert_no_per_expert_scale(a) + _assert_no_per_expert_scale(b) M, K = a.shape[0], a.shape[1] N = b.shape[1] @@ -564,6 +583,7 @@ def nvfp4_linear(func, types, args, kwargs): if not isinstance(weight_tensor, NVFP4Tensor): raise NotImplementedError("NVFP4Tensor: weight must be NVFP4Tensor") + _assert_no_per_expert_scale(weight_tensor) if weight_tensor.act_quant_kwargs is None: # weight_only quant @@ -597,6 +617,7 @@ def nvfp4_mm(func, types, args, kwargs): if not isinstance(weight_tensor, NVFP4Tensor): raise NotImplementedError("NVFP4Tensor: weight must be NVFP4Tensor") + _assert_no_per_expert_scale(weight_tensor) if weight_tensor.act_quant_kwargs is None: weight_dequant = weight_tensor.dequantize(weight_tensor.orig_dtype) @@ -629,6 +650,7 @@ def nvfp4_addmm(func, types, args, kwargs): if not isinstance(weight_tensor, NVFP4Tensor): raise NotImplementedError("NVFP4Tensor: weight must be NVFP4Tensor") + _assert_no_per_expert_scale(weight_tensor) if weight_tensor.act_quant_kwargs is None: weight_dequant = weight_tensor.dequantize(weight_tensor.orig_dtype) @@ -740,6 +762,13 @@ def nvfp4_quantize( # This will likely be calibrated but # we want the per_tensor_scale ~= amax of the block_scale_fp32 block_scale_fp32 = block_scale.to(torch.float32) + + # if the per_tensor_scale is multidimensional, we are in the special + # case that handles the 3d weight tensor with per-expert outer scales. + # Since the code below is 2d, we convert this scale to 2d. + if len(per_tensor_scale.shape) == 3: + per_tensor_scale = per_tensor_scale.squeeze(-1) + # Quantize the blockwise scales w/ the per_tensor_scale scaled_block_scales = block_scale_fp32 / per_tensor_scale scaled_block_scales_fp8 = torch.clamp( From f6792167d217531b73d836811e904fdabed4e858 Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Wed, 22 Apr 2026 21:08:34 +0000 Subject: [PATCH 14/25] Update [ghstack-poisoned] --- .../mx_formats/test_inference_workflow.py | 23 ++++---- .../mx_formats/inference_workflow.py | 15 +++++- torchao/prototype/mx_formats/nvfp4_tensor.py | 53 ++++++++++++++++--- 3 files changed, 68 insertions(+), 23 deletions(-) diff --git a/test/prototype/mx_formats/test_inference_workflow.py b/test/prototype/mx_formats/test_inference_workflow.py index 23bc6b5a8b..9091a60594 100644 --- a/test/prototype/mx_formats/test_inference_workflow.py +++ b/test/prototype/mx_formats/test_inference_workflow.py @@ -268,17 +268,6 @@ def test_narrow_similar_to_vllm(self): ) self._test_narrow_similar_to_vllm(config) - @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") - @pytest.mark.skipif( - not torch_version_at_least("2.8.0"), - reason="torch.compile requires PyTorch 2.8+", - ) - def test_nvfp4_quantize_3d_param_similar_to_vllm(self): - config = NVFP4WeightOnlyConfig( - use_dynamic_per_tensor_scale=False, - ) - self._test_quantize_3d_param_similar_to_vllm(config) - @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @pytest.mark.skipif( @@ -459,9 +448,15 @@ def test_grouped_mm_nvfp4(): dtype = torch.bfloat16 model_ref = GroupedMMModel(E, K, N, device=device, dtype=dtype) + + # make the future nvfp4 weight scales interesting + model_ref.weight[0, :, :] *= 10.0 + model_ref.weight[-1, :, :] *= 1e-3 + model = copy.deepcopy(model_ref) x = torch.randn(total_m, K, device=device, dtype=dtype) + offs = torch.tensor( [sum(m_per_group[: i + 1]) for i in range(E)], device=device, @@ -483,6 +478,8 @@ def test_grouped_mm_nvfp4(): assert isinstance(model.weight, NVFP4Tensor), ( f"Expected NVFP4Tensor weight, got {type(model.weight)}" ) + assert model.weight.per_tensor_scale.shape == (E, 1, 1) + # breakpoint() w_sqnr = compute_error(model_ref.weight, model.weight.dequantize()) assert w_sqnr > 18.0 @@ -490,8 +487,6 @@ def test_grouped_mm_nvfp4(): wwt = ww.transpose(-2, -1) assert tuple(wwt.shape) == (ww.shape[0], ww.shape[2], ww.shape[1]) - # For now, this is emulated. In the near future we'll hook up - # a real nvfp4 grouped gemm. y = model(x, offs) y_sqnr = compute_error(y_ref, y) - assert y_sqnr > 18.0 + assert y_sqnr > 15.0 diff --git a/torchao/prototype/mx_formats/inference_workflow.py b/torchao/prototype/mx_formats/inference_workflow.py index d8d6e27842..2c0db1af16 100644 --- a/torchao/prototype/mx_formats/inference_workflow.py +++ b/torchao/prototype/mx_formats/inference_workflow.py @@ -264,6 +264,7 @@ def _nvfp4_inference_linear_transform( return NVFP4ObservedLinear.from_float(module) elif step == QuantizationStep.CONVERT or step == "convert": + assert len(weight.shape) == 2, "3D weights not yet supported here" if not isinstance(module, NVFP4ObservedLinear): return module @@ -314,8 +315,17 @@ def _nvfp4_inference_linear_transform( per_tensor_scale = None if config.use_dynamic_per_tensor_scale: - tensor_amax = torch.max(torch.abs(weight)) - per_tensor_scale = per_tensor_amax_to_scale(tensor_amax) + if len(weight.shape) == 2: + tensor_amax = torch.max(torch.abs(weight)) + per_tensor_scale = per_tensor_amax_to_scale(tensor_amax) + else: + assert len(weight.shape) == 3, f"unsupported {weight.shape=}" + tensor_amax = torch.amax(torch.abs(weight), dim=(1, 2)) + per_tensor_scale = per_tensor_amax_to_scale(tensor_amax) + # 1D -> 3D + per_tensor_scale = per_tensor_scale.view( + per_tensor_scale.shape[0], 1, 1 + ) act_quant_kwargs = QuantizeTensorToNVFP4Kwargs( use_dynamic_per_tensor_scale=config.use_dynamic_per_tensor_scale, @@ -383,6 +393,7 @@ def _nvfp4_weight_only_linear_transform( """Quantization handler for NVFP4WeightOnlyConfig""" weight = module.weight + assert len(weight.shape) == 2, "3D weights not yet supported in this workflow" if weight.shape[-2] % 16 != 0 or weight.shape[-1] % 16 != 0: raise RuntimeError( f"NVFP4 only supports weight shape with last 2 dims divisible by 16, got {weight.shape}" diff --git a/torchao/prototype/mx_formats/nvfp4_tensor.py b/torchao/prototype/mx_formats/nvfp4_tensor.py index af132d96a9..8b18b90220 100644 --- a/torchao/prototype/mx_formats/nvfp4_tensor.py +++ b/torchao/prototype/mx_formats/nvfp4_tensor.py @@ -246,6 +246,7 @@ def get_hp_scales(self) -> torch.Tensor: scale_e4m3 = from_blocked( scale_e4m3, math.prod(leading_dims) * M, K // self.block_size ) + scale_e4m3 = scale_e4m3.view(*leading_dims, M, K // self.block_size) return ( scale_e4m3.to(torch.float) @@ -419,7 +420,6 @@ def nvfp4_t(func, types, args, kwargs): @implements([aten.transpose.int]) def nvfp4_transpose(func, types, args, kwargs): old, dim0, dim1 = args - _assert_no_per_expert_scale(old) assert len(old.shape) == 3, f"unsupported rank {len(old.shape)}" valid_3d_dims = ((1, 2), (2, 1), (-1, -2), (-2, -1)) assert (dim0, dim1) in valid_3d_dims, f"transpose unsupported for {dim0=} {dim1=}" @@ -680,13 +680,52 @@ def nvfp4_addmm(func, types, args, kwargs): @implements([aten._grouped_mm.default]) def nvfp4_grouped_mm(func, types, args, kwargs): + # TODO(before land): wrap MSLK kernel with a custom op, guard for MSLK + # not installed, etc + from mslk.quantize.triton.fp4_quantize import nvfp4_quantize_stacked + from torch.nn.functional import ScalingType, SwizzleType, scaled_grouped_mm + mat_a, mat_b = args[0], args[1] - # temporary: dequantize nvfp4 tensor to call original grouped_mm - # TODO(future PR): enable nvfp4 with per-expert outer scale (current code - # has a single outer scale across all experts). - # TODO(future PR): hook up real nvfp4 grouped_mm - mat_b_dq = mat_b.dequantize() - return func(mat_a, mat_b_dq, *args[2:], **kwargs) + offs = args[2] if len(args) > 2 else kwargs.get("offs", None) + assert offs is not None, "offs is required for nvfp4 grouped_mm" + + # mat_b is a transposed NVFP4Tensor: the model stores weight as (E, N, K) + # and calls weight.transpose(-2, -1). Undo the transpose to get the + # original NVFP4Tensor so we can extract qdata/scale in the right layout. + assert isinstance(mat_b, NVFP4Tensor) + is_transposed = mat_b.qdata.stride(-2) < mat_b.qdata.stride(-1) + if is_transposed: + mat_b = mat_b.transpose(-2, -1) + + E = offs.shape[0] + m_sizes = torch.diff(offs, prepend=offs.new_zeros(1)).to(torch.int64) + + # For now, quantize activation with per-expert global scales of 1.0 + # TODO(future PR): compute real per-token-group global scales from activation amax, + # will need some slow eager code or a triton kernel + a_global_scale = torch.ones(E, device=mat_a.device, dtype=torch.float32) + mat_a_qdata, mat_a_scale = nvfp4_quantize_stacked(m_sizes, mat_a, a_global_scale) + + # Weight: extract qdata and already-swizzled scales + mat_b_qdata = mat_b.qdata.view(torch.float4_e2m1fn_x2) + # Flatten 3D scale (E, padded_N, padded_K//16) -> 2D (E, padded_N * padded_K//16) + # as required by _scaled_grouped_mm for 2D-3D grouped GEMM + mat_b_scale = mat_b.scale.view(torch.float8_e4m3fn).flatten(1) + b_global_scale = mat_b.per_tensor_scale.view(E) + + # _scaled_grouped_mm expects wq.transpose(-2, -1) + return scaled_grouped_mm( + mat_a_qdata, + mat_b_qdata.transpose(-2, -1), + scale_a=[mat_a_scale, a_global_scale], + scale_recipe_a=[ScalingType.BlockWise1x16, ScalingType.TensorWise], + scale_b=[mat_b_scale, b_global_scale], + scale_recipe_b=[ScalingType.BlockWise1x16, ScalingType.TensorWise], + swizzle_a=SwizzleType.SWIZZLE_32_4_4, + swizzle_b=SwizzleType.SWIZZLE_32_4_4, + offs=offs, + output_dtype=mat_a.dtype, + ) def per_tensor_amax_to_scale(amax: torch.Tensor) -> torch.Tensor: From be9dc1b3998438643cfc8b0db832813db15933b5 Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Thu, 23 Apr 2026 09:27:57 +0000 Subject: [PATCH 15/25] Update [ghstack-poisoned] --- torchao/prototype/mx_formats/nvfp4_tensor.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/torchao/prototype/mx_formats/nvfp4_tensor.py b/torchao/prototype/mx_formats/nvfp4_tensor.py index af132d96a9..52b5876a64 100644 --- a/torchao/prototype/mx_formats/nvfp4_tensor.py +++ b/torchao/prototype/mx_formats/nvfp4_tensor.py @@ -56,8 +56,8 @@ class NVFP4Tensor(TorchAOBaseTensor): Attributes: qdata: Packed FP4 data (2 values per byte) scale: Blockwise scales in float8_e4m3fn format (may be swizzled) - per_tensor_scale: Optional global per-tensor scale in float32 format - act_per_tensor_scale: Optional global per-tensor scale in float32 format, for activation + per_tensor_scale: Optional global per-tensor or per-expert scale in float32 format + act_per_tensor_scale: Optional global per-tensor or per-token-group scale in float32 format, for activation block_size (int): Block size for quantization (fixed at 16) orig_dtype (torch.dtype): Original tensor dtype before quantization is_swizzled_scales (bool): Whether scales are stored in swizzled (blocked) format From 83283cfffdb0840f0d459a55882a8ce7fb78fd26 Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Thu, 23 Apr 2026 10:55:00 +0000 Subject: [PATCH 16/25] Update [ghstack-poisoned] --- .../mx_formats/test_inference_workflow.py | 3 +- torchao/prototype/mx_formats/kernels.py | 84 +++++++++++++++++++ torchao/prototype/mx_formats/nvfp4_tensor.py | 33 ++++---- 3 files changed, 101 insertions(+), 19 deletions(-) diff --git a/test/prototype/mx_formats/test_inference_workflow.py b/test/prototype/mx_formats/test_inference_workflow.py index 9091a60594..de088e36ea 100644 --- a/test/prototype/mx_formats/test_inference_workflow.py +++ b/test/prototype/mx_formats/test_inference_workflow.py @@ -441,7 +441,7 @@ def forward(self, x, offs): def test_grouped_mm_nvfp4(): """Smoke test: torch._grouped_mm forward pass with bfloat16 inputs.""" E, K, N = 4, 128, 256 - m_per_group = [32, 96, 16, 112] + m_per_group = [1, 3, 4, 16] total_m = sum(m_per_group) device = "cuda" @@ -479,7 +479,6 @@ def test_grouped_mm_nvfp4(): f"Expected NVFP4Tensor weight, got {type(model.weight)}" ) assert model.weight.per_tensor_scale.shape == (E, 1, 1) - # breakpoint() w_sqnr = compute_error(model_ref.weight, model.weight.dequantize()) assert w_sqnr > 18.0 diff --git a/torchao/prototype/mx_formats/kernels.py b/torchao/prototype/mx_formats/kernels.py index 7f891d6e39..ce614d4daf 100644 --- a/torchao/prototype/mx_formats/kernels.py +++ b/torchao/prototype/mx_formats/kernels.py @@ -1249,3 +1249,87 @@ def _(x, global_scale=None): scales = scales.view(*orig_leading_dims, -1, padded_cols) xq = xq.view(*orig_leading_dims, -1, N // 2) return scales, xq + + +def mslk_calculate_group_max(x: torch.Tensor, m_sizes: torch.Tensor) -> torch.Tensor: + """Compute per-expert activation global scale (encoding convention). + + Args: + x: [M, K] concatenated activation tensor (bf16/fp16). + m_sizes: [E] int64 tensor of rows per expert. + + Returns: + Per-expert global scale in encoding convention (448 * FP4_MAX / amax), + shape [E], dtype float32. + """ + return _mslk_calculate_group_max_custom_op(x, m_sizes) + + +@torch.library.custom_op("ao::mslk_calculate_group_max", mutates_args=()) +def _mslk_calculate_group_max_custom_op( + x: torch.Tensor, m_sizes: torch.Tensor +) -> torch.Tensor: + assert _mslk_available, ( + "mslk is required for calculate_group_max. " + "Install from https://github.com/meta-pytorch/MSLK" + ) + from mslk.quantize.triton.fp4_quantize import calculate_group_max + + global_scale, _ = calculate_group_max(x, m_sizes) + return global_scale + + +@_mslk_calculate_group_max_custom_op.register_fake +def _(x, m_sizes): + E = m_sizes.shape[0] + return x.new_empty(E, dtype=torch.float32) + + +def mslk_quantize_nvfp4_stacked( + m_sizes: torch.Tensor, + x: torch.Tensor, + global_scale: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Quantize concatenated MoE activations to NVFP4 with per-expert global scales. + + Args: + m_sizes: [E] int64 tensor of rows per expert. + x: [M, K] concatenated activation tensor (bf16/fp16). + global_scale: [E] fp32 per-expert global scales in encoding convention. + + Returns: + Tuple of (xq, scale): + xq: [M, K//2] float4_e2m1fn_x2 packed FP4 data. + scale: Padded+swizzled float8_e4m3fn block scales. + """ + return _mslk_quantize_nvfp4_stacked_custom_op(m_sizes, x, global_scale) + + +@torch.library.custom_op("ao::mslk_quantize_nvfp4_stacked", mutates_args=()) +def _mslk_quantize_nvfp4_stacked_custom_op( + m_sizes: torch.Tensor, + x: torch.Tensor, + global_scale: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + assert _mslk_available, ( + "mslk is required for nvfp4_quantize_stacked. " + "Install from https://github.com/meta-pytorch/MSLK" + ) + from mslk.quantize.triton.fp4_quantize import nvfp4_quantize_stacked + + xq, scale = nvfp4_quantize_stacked(m_sizes, x, global_scale) + return xq, scale + + +@_mslk_quantize_nvfp4_stacked_custom_op.register_fake +def _(m_sizes, x, global_scale): + M, K = x.shape[0], x.shape[1] + num_segments = m_sizes.shape[0] + # Upper-bound on padded total rows (each segment can add at most 127 padding rows) + padded_total_M_ub = M + num_segments * 127 + num_scales_per_row = K // 16 + n_col_blocks = triton.cdiv(num_scales_per_row, 4) + padded_cols = n_col_blocks * 4 + xq = x.new_empty(M, K // 2, dtype=torch.float4_e2m1fn_x2) + scale = x.new_empty(padded_total_M_ub, padded_cols, dtype=torch.float8_e4m3fn) + return xq, scale diff --git a/torchao/prototype/mx_formats/nvfp4_tensor.py b/torchao/prototype/mx_formats/nvfp4_tensor.py index cf8ac3865f..034ea377c0 100644 --- a/torchao/prototype/mx_formats/nvfp4_tensor.py +++ b/torchao/prototype/mx_formats/nvfp4_tensor.py @@ -15,7 +15,9 @@ from torchao.prototype.mx_formats.kernels import ( f4_unpacked_to_f32, f32_to_f4_unpacked, + mslk_calculate_group_max, mslk_quantize_nvfp4, + mslk_quantize_nvfp4_stacked, pack_uint4, unpack_uint4, ) @@ -680,9 +682,6 @@ def nvfp4_addmm(func, types, args, kwargs): @implements([aten._grouped_mm.default]) def nvfp4_grouped_mm(func, types, args, kwargs): - # TODO(before land): wrap MSLK kernel with a custom op, guard for MSLK - # not installed, etc - from mslk.quantize.triton.fp4_quantize import nvfp4_quantize_stacked from torch.nn.functional import ScalingType, SwizzleType, scaled_grouped_mm mat_a, mat_b = args[0], args[1] @@ -694,32 +693,32 @@ def nvfp4_grouped_mm(func, types, args, kwargs): # original NVFP4Tensor so we can extract qdata/scale in the right layout. assert isinstance(mat_b, NVFP4Tensor) is_transposed = mat_b.qdata.stride(-2) < mat_b.qdata.stride(-1) - if is_transposed: - mat_b = mat_b.transpose(-2, -1) + assert is_transposed, "unsupported" E = offs.shape[0] m_sizes = torch.diff(offs, prepend=offs.new_zeros(1)).to(torch.int64) - # For now, quantize activation with per-expert global scales of 1.0 - # TODO(future PR): compute real per-token-group global scales from activation amax, - # will need some slow eager code or a triton kernel - a_global_scale = torch.ones(E, device=mat_a.device, dtype=torch.float32) - mat_a_qdata, mat_a_scale = nvfp4_quantize_stacked(m_sizes, mat_a, a_global_scale) + # mslk_calculate_group_max returns encoding scale (448 * FP4_MAX / amax) + # mslk_quantize_nvfp4_stacked needs encoding scale, scaled_grouped_mm + # needs decoding scale (1 / encoding_scale) + a_global_scale_enc = mslk_calculate_group_max(mat_a, m_sizes) + mat_a_qdata, mat_a_scale = mslk_quantize_nvfp4_stacked( + m_sizes, mat_a, a_global_scale_enc + ) + a_global_scale_dec = 1.0 / a_global_scale_enc - # Weight: extract qdata and already-swizzled scales - mat_b_qdata = mat_b.qdata.view(torch.float4_e2m1fn_x2) # Flatten 3D scale (E, padded_N, padded_K//16) -> 2D (E, padded_N * padded_K//16) # as required by _scaled_grouped_mm for 2D-3D grouped GEMM - mat_b_scale = mat_b.scale.view(torch.float8_e4m3fn).flatten(1) + mat_b_t_scale = mat_b.scale.transpose(-2, -1).flatten(1) + # [E, 1, 1] -> E b_global_scale = mat_b.per_tensor_scale.view(E) - # _scaled_grouped_mm expects wq.transpose(-2, -1) return scaled_grouped_mm( mat_a_qdata, - mat_b_qdata.transpose(-2, -1), - scale_a=[mat_a_scale, a_global_scale], + mat_b.qdata.view(torch.float4_e2m1fn_x2), + scale_a=[mat_a_scale, a_global_scale_dec], scale_recipe_a=[ScalingType.BlockWise1x16, ScalingType.TensorWise], - scale_b=[mat_b_scale, b_global_scale], + scale_b=[mat_b_t_scale, b_global_scale], scale_recipe_b=[ScalingType.BlockWise1x16, ScalingType.TensorWise], swizzle_a=SwizzleType.SWIZZLE_32_4_4, swizzle_b=SwizzleType.SWIZZLE_32_4_4, From ed9e39f9447790600fcd4f83db0e891f77781b8f Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Thu, 23 Apr 2026 13:03:12 +0000 Subject: [PATCH 17/25] Update [ghstack-poisoned] --- scripts/prototype/test_nvfp4_moe.py | 69 +++++++++++++++++++++++++++++ 1 file changed, 69 insertions(+) create mode 100644 scripts/prototype/test_nvfp4_moe.py diff --git a/scripts/prototype/test_nvfp4_moe.py b/scripts/prototype/test_nvfp4_moe.py new file mode 100644 index 0000000000..c28faab738 --- /dev/null +++ b/scripts/prototype/test_nvfp4_moe.py @@ -0,0 +1,69 @@ +"""Minimal OLMoE inference example with optional NVFP4 quantization for expert weights.""" + +from contextlib import nullcontext + +import fire +import torch +from transformers import AutoModelForCausalLM, AutoTokenizer +from transformers.models.olmoe.modeling_olmoe import OlmoeExperts + + +def main(recipe: str = "bf16"): + print(f"{recipe=}") + model_id = "allenai/OLMoE-1B-7B-0924" + + tokenizer = AutoTokenizer.from_pretrained(model_id) + model = AutoModelForCausalLM.from_pretrained( + model_id, + dtype=torch.bfloat16, + device_map="cuda", + experts_implementation="grouped_mm", + ) + print(model) + + if recipe == "nvfp4": + from torchao.prototype.mx_formats.inference_workflow import ( + NVFP4DynamicActivationNVFP4WeightConfig, + _nvfp4_inference_linear_transform, + ) + from torchao.prototype.mx_formats.nvfp4_tensor import NVFP4Tensor + + # Quantize expert weights (gate_up_proj and down_proj) to NVFP4 + # TODO(future PR): make quantize_ work instead of having to work around + # such as below, need to allow parameter name other than `weight` + config = NVFP4DynamicActivationNVFP4WeightConfig(use_triton_kernel=False) + for name, mod in model.named_modules(): + if isinstance(mod, OlmoeExperts): + for pname in ("gate_up_proj", "down_proj"): + _nvfp4_inference_linear_transform(mod, config, parameter_name=pname) + + # Verify quantization worked + for name, mod in model.named_modules(): + if isinstance(mod, OlmoeExperts): + for pname in ("gate_up_proj", "down_proj"): + param = getattr(mod, pname) + assert isinstance(param, NVFP4Tensor), ( + f"{name}.{pname} is {type(param).__name__}, expected NVFP4Tensor" + ) + print(f"{name}: gate_up_proj and down_proj are NVFP4Tensor") + + # generate() switches to batched_mm for decoding, which doesn't support + # NVFP4Tensor (needs aten.index.Tensor). Override to keep grouped_mm. + # TODO(future PR): implement bmm for nvfp4 and remove this workaround + model._optimize_model_for_decode = nullcontext + elif recipe == "bf16": + pass + else: + raise ValueError(f"Unknown recipe: {recipe}") + + prompt = "The capital of France is" + inputs = tokenizer(prompt, return_tensors="pt").to("cuda") + + with torch.no_grad(): + output = model.generate(**inputs, max_new_tokens=50) + + print(tokenizer.decode(output[0], skip_special_tokens=True)) + + +if __name__ == "__main__": + fire.Fire(main) From 238667049da90560bcf93756537bf85421346e0f Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Thu, 23 Apr 2026 13:16:36 +0000 Subject: [PATCH 18/25] Update [ghstack-poisoned] --- scripts/prototype/test_nvfp4_moe.py | 30 +++++++++++++++-------------- 1 file changed, 16 insertions(+), 14 deletions(-) diff --git a/scripts/prototype/test_nvfp4_moe.py b/scripts/prototype/test_nvfp4_moe.py index c28faab738..3b273a1a59 100644 --- a/scripts/prototype/test_nvfp4_moe.py +++ b/scripts/prototype/test_nvfp4_moe.py @@ -7,6 +7,12 @@ from transformers import AutoModelForCausalLM, AutoTokenizer from transformers.models.olmoe.modeling_olmoe import OlmoeExperts +from torchao.prototype.mx_formats.inference_workflow import ( + NVFP4DynamicActivationNVFP4WeightConfig, +) +from torchao.prototype.mx_formats.nvfp4_tensor import NVFP4Tensor +from torchao.quantization import FqnToConfig, quantize_ + def main(recipe: str = "bf16"): print(f"{recipe=}") @@ -22,20 +28,17 @@ def main(recipe: str = "bf16"): print(model) if recipe == "nvfp4": - from torchao.prototype.mx_formats.inference_workflow import ( - NVFP4DynamicActivationNVFP4WeightConfig, - _nvfp4_inference_linear_transform, - ) - from torchao.prototype.mx_formats.nvfp4_tensor import NVFP4Tensor - - # Quantize expert weights (gate_up_proj and down_proj) to NVFP4 - # TODO(future PR): make quantize_ work instead of having to work around - # such as below, need to allow parameter name other than `weight` config = NVFP4DynamicActivationNVFP4WeightConfig(use_triton_kernel=False) - for name, mod in model.named_modules(): - if isinstance(mod, OlmoeExperts): - for pname in ("gate_up_proj", "down_proj"): - _nvfp4_inference_linear_transform(mod, config, parameter_name=pname) + quantize_( + model, + FqnToConfig( + { + r"re:.*\.experts\.gate_up_proj": config, + r"re:.*\.experts\.down_proj": config, + } + ), + filter_fn=None, + ) # Verify quantization worked for name, mod in model.named_modules(): @@ -45,7 +48,6 @@ def main(recipe: str = "bf16"): assert isinstance(param, NVFP4Tensor), ( f"{name}.{pname} is {type(param).__name__}, expected NVFP4Tensor" ) - print(f"{name}: gate_up_proj and down_proj are NVFP4Tensor") # generate() switches to batched_mm for decoding, which doesn't support # NVFP4Tensor (needs aten.index.Tensor). Override to keep grouped_mm. From c1da8494df515bbc7e84fe4a41b3fd48f9657ecc Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Thu, 23 Apr 2026 13:57:06 +0000 Subject: [PATCH 19/25] Update [ghstack-poisoned] --- .../mx_formats/test_inference_workflow.py | 57 +++++++++++++++++++ torchao/prototype/mx_formats/nvfp4_tensor.py | 29 +++++++++- 2 files changed, 84 insertions(+), 2 deletions(-) diff --git a/test/prototype/mx_formats/test_inference_workflow.py b/test/prototype/mx_formats/test_inference_workflow.py index de088e36ea..c326351785 100644 --- a/test/prototype/mx_formats/test_inference_workflow.py +++ b/test/prototype/mx_formats/test_inference_workflow.py @@ -489,3 +489,60 @@ def test_grouped_mm_nvfp4(): y = model(x, offs) y_sqnr = compute_error(y_ref, y) assert y_sqnr > 15.0 + + +class BatchedMMModel(nn.Module): + """A toy model whose only op in forward is torch.bmm.""" + + def __init__(self, E, K, N, device="cuda", dtype=torch.bfloat16): + super().__init__() + self.weight = nn.Parameter(torch.randn(E, N, K, device=device, dtype=dtype)) + + def forward(self, x): + return torch.bmm(x, self.weight.transpose(-2, -1)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif( + not torch_version_at_least("2.8.0"), reason="torch >= 2.8.0 required" +) +@pytest.mark.skipif( + not is_sm_at_least_100(), reason="CUDA capability >= 10.0 required for NVFP4" +) +@torch.no_grad() +@skip_if_rocm("ROCm float4 gemm require gfx950") +def test_bmm_nvfp4(): + """Smoke test: torch.bmm forward pass with NVFP4 quantized weights.""" + E, K, N = 4, 128, 256 + M = 16 + + device = "cuda" + dtype = torch.bfloat16 + + model_ref = BatchedMMModel(E, K, N, device=device, dtype=dtype) + + model_ref.weight[0, :, :] *= 10.0 + model_ref.weight[-1, :, :] *= 1e-3 + + model = copy.deepcopy(model_ref) + + x = torch.randn(E, M, K, device=device, dtype=dtype) + + y_ref = model_ref(x) + + quantize_( + model, + NVFP4DynamicActivationNVFP4WeightConfig( + use_triton_kernel=False, + ), + filter_fn=lambda mod, *args: isinstance(mod, BatchedMMModel) + and hasattr(mod, "weight"), + ) + assert isinstance(model.weight, NVFP4Tensor), ( + f"Expected NVFP4Tensor weight, got {type(model.weight)}" + ) + + y = model(x) + y_sqnr = compute_error(y_ref, y) + assert y_sqnr > 15.0 + print(y_sqnr) diff --git a/torchao/prototype/mx_formats/nvfp4_tensor.py b/torchao/prototype/mx_formats/nvfp4_tensor.py index 034ea377c0..9830bd11d7 100644 --- a/torchao/prototype/mx_formats/nvfp4_tensor.py +++ b/torchao/prototype/mx_formats/nvfp4_tensor.py @@ -466,13 +466,16 @@ def nvfp4_select(func, types, args, kwargs): old, dim, index = args assert dim == 0, f"NVFP4Tensor aten.select.int with {dim=} is not yet supported" assert len(old.qdata.shape) == len(old.scale.shape), "unsupported" - _assert_no_per_expert_scale(old) + if old.per_tensor_scale is not None and len(old.per_tensor_scale.shape) == 3: + new_per_tensor_scale = old.per_tensor_scale[index].squeeze() + else: + new_per_tensor_scale = old.per_tensor_scale new = old.__class__( old.qdata[index], old.scale[index], old.block_size, old.orig_dtype, - old.per_tensor_scale, + new_per_tensor_scale, old.act_per_tensor_scale, old.is_swizzled_scales, old.use_triton_kernel, @@ -646,6 +649,28 @@ def nvfp4_mm(func, types, args, kwargs): return _addmm_nvfp4_dispatch(input_tensor, weight_tensor, func) +@implements([aten.bmm.default]) +def nvfp4_bmm(func, types, args, kwargs): + input_tensor, weight_tensor = args[0], args[1] + # For now, implement nvfp4 bmm with a for loop over 2d gemms. Correct + # numerics, bad performance. + # TODO(future): hook up a kernel once we have one. + res = [] + for e_idx in range(input_tensor.shape[0]): + # Note: `i_e` will be quantized to nvfp4 in the + # override of `torch.mm` + i_e = input_tensor[e_idx] + w_e = weight_tensor[e_idx] + # Note: unsqueeze is for the `cat` op to write directly + # to the correct output shape + o_e = torch.mm(i_e, w_e).unsqueeze(0) + res.append(o_e) + out = torch.cat(res, dim=0) + return out + + # return func(input_tensor, weight_tensor) + + @implements([aten.addmm.default]) def nvfp4_addmm(func, types, args, kwargs): bias, input_tensor, weight_tensor = args[0], args[1], args[2] From cdcd2b315709dbaa266d98f5a3e734e7841e2668 Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Thu, 23 Apr 2026 13:58:25 +0000 Subject: [PATCH 20/25] Update [ghstack-poisoned] --- test/prototype/mx_formats/test_inference_workflow.py | 1 - torchao/prototype/mx_formats/nvfp4_tensor.py | 2 -- 2 files changed, 3 deletions(-) diff --git a/test/prototype/mx_formats/test_inference_workflow.py b/test/prototype/mx_formats/test_inference_workflow.py index c326351785..93ddf9d662 100644 --- a/test/prototype/mx_formats/test_inference_workflow.py +++ b/test/prototype/mx_formats/test_inference_workflow.py @@ -545,4 +545,3 @@ def test_bmm_nvfp4(): y = model(x) y_sqnr = compute_error(y_ref, y) assert y_sqnr > 15.0 - print(y_sqnr) diff --git a/torchao/prototype/mx_formats/nvfp4_tensor.py b/torchao/prototype/mx_formats/nvfp4_tensor.py index 9830bd11d7..6bbcde79b4 100644 --- a/torchao/prototype/mx_formats/nvfp4_tensor.py +++ b/torchao/prototype/mx_formats/nvfp4_tensor.py @@ -668,8 +668,6 @@ def nvfp4_bmm(func, types, args, kwargs): out = torch.cat(res, dim=0) return out - # return func(input_tensor, weight_tensor) - @implements([aten.addmm.default]) def nvfp4_addmm(func, types, args, kwargs): From 196d4398f90ffde0670cc4dbb30792a82cec74f3 Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Thu, 23 Apr 2026 15:52:00 +0000 Subject: [PATCH 21/25] Update [ghstack-poisoned] --- scripts/prototype/test_nvfp4_moe.py | 50 ++++++++++++++++++++++++++++- 1 file changed, 49 insertions(+), 1 deletion(-) diff --git a/scripts/prototype/test_nvfp4_moe.py b/scripts/prototype/test_nvfp4_moe.py index 3b273a1a59..cc39ceef41 100644 --- a/scripts/prototype/test_nvfp4_moe.py +++ b/scripts/prototype/test_nvfp4_moe.py @@ -1,11 +1,16 @@ """Minimal OLMoE inference example with optional NVFP4 quantization for expert weights.""" +import gc +import inspect +import subprocess +import tempfile from contextlib import nullcontext import fire import torch from transformers import AutoModelForCausalLM, AutoTokenizer from transformers.models.olmoe.modeling_olmoe import OlmoeExperts +from transformers.quantizers.quantizer_torchao import TorchAoHfQuantizer from torchao.prototype.mx_formats.inference_workflow import ( NVFP4DynamicActivationNVFP4WeightConfig, @@ -14,7 +19,7 @@ from torchao.quantization import FqnToConfig, quantize_ -def main(recipe: str = "bf16"): +def main(recipe: str = "bf16", run_lm_eval: bool = False): print(f"{recipe=}") model_id = "allenai/OLMoE-1B-7B-0924" @@ -66,6 +71,49 @@ def main(recipe: str = "bf16"): print(tokenizer.decode(output[0], skip_special_tokens=True)) + if run_lm_eval: + if recipe == "nvfp4": + source = inspect.getsource(TorchAoHfQuantizer.get_weight_conversions) + if "gate_up_proj" not in source: + raise RuntimeError( + "Your version of `transformers` does not support NVFP4 MoE serialization. " + "Please install a version that includes " + "https://github.com/huggingface/transformers/pull/45609" + ) + + with tempfile.TemporaryDirectory() as output_dir: + print(f"\nSaving model to {output_dir}...") + + if recipe != "bf16": + from transformers import TorchAoConfig + + torchao_config = TorchAoConfig(quant_type=config) + model.config.quantization_config = torchao_config + model.hf_quantizer = TorchAoHfQuantizer(torchao_config) + + model.save_pretrained(output_dir, safe_serialization=False) + tokenizer.save_pretrained(output_dir) + + del model + gc.collect() + torch.cuda.empty_cache() + + lm_eval_cmd = [ + "lm_eval", + "--model", + "hf", + "--model_args", + f"pretrained={output_dir}", + "--tasks", + "wikitext", + "--num_fewshot", + "0", + "--batch_size", + "16", + ] + print(f"Running: {' '.join(lm_eval_cmd)}") + subprocess.run(lm_eval_cmd, check=True) + if __name__ == "__main__": fire.Fire(main) From 5a0db1645d29b9bc69c2b65882093ec5cff56138 Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Thu, 23 Apr 2026 17:21:53 +0000 Subject: [PATCH 22/25] Update [ghstack-poisoned] --- scripts/prototype/test_nvfp4_moe.py | 89 +++++++++++++++++++++++++++-- 1 file changed, 84 insertions(+), 5 deletions(-) diff --git a/scripts/prototype/test_nvfp4_moe.py b/scripts/prototype/test_nvfp4_moe.py index cc39ceef41..56320fd032 100644 --- a/scripts/prototype/test_nvfp4_moe.py +++ b/scripts/prototype/test_nvfp4_moe.py @@ -12,6 +12,7 @@ from transformers.models.olmoe.modeling_olmoe import OlmoeExperts from transformers.quantizers.quantizer_torchao import TorchAoHfQuantizer +from torchao.prototype.gptq.gptq_example import prepare_dataset from torchao.prototype.mx_formats.inference_workflow import ( NVFP4DynamicActivationNVFP4WeightConfig, ) @@ -19,7 +20,60 @@ from torchao.quantization import FqnToConfig, quantize_ -def main(recipe: str = "bf16", run_lm_eval: bool = False): +def install_expert_counters(model): + """Install forward pre-hooks on OlmoeExperts to count per-expert token routing. + + Returns a dict mapping module FQN to a (num_experts,) int64 tensor of counts, + and a list of hook handles for removal. + """ + expert_counts = {} + handles = [] + + for name, mod in model.named_modules(): + if not isinstance(mod, OlmoeExperts): + continue + counts = torch.zeros(mod.num_experts, dtype=torch.int64) + expert_counts[name] = counts + + def make_hook(c, n): + def hook(module, args, kwargs): + top_k_index = args[1] + c.add_(top_k_index.flatten().bincount(minlength=n).cpu()) + + return hook + + handles.append( + mod.register_forward_pre_hook( + make_hook(counts, mod.num_experts), with_kwargs=True + ) + ) + + return expert_counts, handles + + +def print_expert_counts(expert_counts): + print("\n=== Per-expert token counts ===") + for name, counts in expert_counts.items(): + total = counts.sum().item() + print(f"{name}: total={total}, per_expert={counts.tolist()}") + + print("\n=== Global expert utilization summary ===") + all_counts = torch.cat([c for c in expert_counts.values()]) + n = len(all_counts) + for threshold in range(129): + if threshold % 10 != 0: + continue + num = int((all_counts <= threshold).sum().item()) + print(f"experts with <= {threshold} tokens: {num}/{n} ({num / n * 100:.1f}%)") + + +def main( + recipe: str = "bf16", + run_lm_eval: bool = False, + calibrate_on_c4: bool = False, + num_calibration_samples: int = 128, + max_sequence_length: int = 2048, +): print(f"{recipe=}") model_id = "allenai/OLMoE-1B-7B-0924" @@ -54,15 +108,40 @@ def main(recipe: str = "bf16", run_lm_eval: bool = False): f"{name}.{pname} is {type(param).__name__}, expected NVFP4Tensor" ) - # generate() switches to batched_mm for decoding, which doesn't support - # NVFP4Tensor (needs aten.index.Tensor). Override to keep grouped_mm. - # TODO(future PR): implement bmm for nvfp4 and remove this workaround - model._optimize_model_for_decode = nullcontext elif recipe == "bf16": pass else: raise ValueError(f"Unknown recipe: {recipe}") + # generate() switches to batched_mm for decoding, which doesn't support + # NVFP4Tensor (needs aten.index.Tensor). Override to keep grouped_mm. + # TODO(future PR): implement bmm for nvfp4 and remove this workaround + model._optimize_model_for_decode = nullcontext + + if calibrate_on_c4: + assert recipe == "bf16", ( + "calibrate_on_c4 is only supported with recipe=bf16 for now" + ) + + expert_counts, hooks = install_expert_counters(model) + + dataset = prepare_dataset( + tokenizer, + max_sequence_length, + num_calibration_samples=num_calibration_samples, + dataset_id="c4", + dataset_split="train", + ) + print(f"Running calibration on {len(dataset)} C4 samples...") + with torch.no_grad(): + for seq in dataset: + model(seq.to("cuda")) + print("Calibration complete.") + + print_expert_counts(expert_counts) + for h in hooks: + h.remove() + prompt = "The capital of France is" inputs = tokenizer(prompt, return_tensors="pt").to("cuda") From 932677ba34c61ec6e02a65b2d0d0eb8e51b94170 Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Thu, 23 Apr 2026 20:05:39 +0000 Subject: [PATCH 23/25] Update [ghstack-poisoned] --- test/prototype/gptq/test_gptqv2.py | 18 +++++++++++++----- torchao/prototype/gptq/api.py | 4 ++-- torchao/prototype/gptq/gptq_example.py | 2 +- torchao/prototype/gptq/observer.py | 21 +++++++++++---------- 4 files changed, 27 insertions(+), 18 deletions(-) diff --git a/test/prototype/gptq/test_gptqv2.py b/test/prototype/gptq/test_gptqv2.py index 78af648e14..fa06c415d2 100644 --- a/test/prototype/gptq/test_gptqv2.py +++ b/test/prototype/gptq/test_gptqv2.py @@ -102,8 +102,10 @@ def test_observer_tensor_creation(self): # Check hp_data is stored correctly torch.testing.assert_close(observer.hp_data, weight) - # Check hessian is initialized as None - assert observer.hessian is None + # Check hessian is initialized as zeros + assert torch.equal( + observer.hessian, torch.zeros(64, 64, dtype=torch.float32, device="cuda") + ) # Check total_batches is initialized as 0 assert observer.total_batches == 0 @@ -120,7 +122,9 @@ def test_observer_tensor_attributes(self): # Test hessian attribute assert hasattr(observer, "hessian") - assert observer.hessian is None + assert torch.equal( + observer.hessian, torch.zeros(32, 32, dtype=torch.float32, device="cuda") + ) # Test total_batches attribute assert hasattr(observer, "total_batches") @@ -193,6 +197,7 @@ def test_multiple_observations(self): # Check total_batches matches total samples assert observer_weight.total_batches == total_samples + @pytest.mark.skip(reason="bmm math is incorrect, will fix in next PR") @pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA available") def test_bmm_operation_with_observer(self): """Test torch.bmm with GPTQObserverTensor updates Hessian correctly.""" @@ -250,8 +255,11 @@ def test_observer_config_transform(self, base_config): # Check hp_data matches original weight torch.testing.assert_close(linear.weight.hp_data, original_weight) - # Check hessian is None initially - assert linear.weight.hessian is None + # Check hessian is initialized as zeros + assert torch.equal( + linear.weight.hessian, + torch.zeros(64, 64, dtype=torch.float32, device="cuda"), + ) assert linear.weight.total_batches == 0 # Perform a forward pass diff --git a/torchao/prototype/gptq/api.py b/torchao/prototype/gptq/api.py index 6869353273..b680472c87 100644 --- a/torchao/prototype/gptq/api.py +++ b/torchao/prototype/gptq/api.py @@ -141,10 +141,10 @@ def _gptq_config_transform( ) # Validate that observations were recorded - if tensor.hessian is None: + if tensor.total_batches == 0: raise ValueError( f"No observations recorded for {parameter_name}. " - f"Hessian is None. Did you run forward passes during the observe step?" + f"total_batches is 0. Did you run forward passes during the observe step?" ) # Use pre-computed Hessian directly diff --git a/torchao/prototype/gptq/gptq_example.py b/torchao/prototype/gptq/gptq_example.py index 016f0ff174..521df63072 100644 --- a/torchao/prototype/gptq/gptq_example.py +++ b/torchao/prototype/gptq/gptq_example.py @@ -450,7 +450,7 @@ def skip_lm_head_o_proj(module, fqn): import inspect source = inspect.getsource(TorchAoHfQuantizer.get_weight_conversions) - if "_weight_per_tensor_scale" not in source: + if "per_tensor_scale" not in source: raise RuntimeError( "Your version of `transformers` does not support NVFP4 serialization. " "Please install a version that includes " diff --git a/torchao/prototype/gptq/observer.py b/torchao/prototype/gptq/observer.py index e6307c2c66..cbfb382da5 100644 --- a/torchao/prototype/gptq/observer.py +++ b/torchao/prototype/gptq/observer.py @@ -29,6 +29,17 @@ def __init__(self, hp_data: torch.Tensor, total_batches: int, hessian=None): self.hessian = hessian self.total_batches = total_batches + # initialize hessian + assert self.hp_data.is_contiguous() + if self.hessian is None: + feature_dim = self.hp_data.shape[-1] + self.hessian = torch.zeros( + feature_dim, + feature_dim, + dtype=torch.float32, + device=self.hp_data.device, + ) + def update(self, input: torch.Tensor): """Incrementally update Hessian matrix from input activations.""" # Move input to same device as hp_data and convert to float @@ -39,16 +50,6 @@ def update(self, input: torch.Tensor): n = 1 if len(shape) == 2 else shape[0] x = x.reshape(-1, shape[-1]) - # Lazily initialize Hessian on first call - if self.hessian is None: - feature_dim = x.shape[-1] - self.hessian = torch.zeros( - feature_dim, - feature_dim, - dtype=torch.float32, - device=self.hp_data.device, - ) - # Apply running average formula if self.total_batches > 0: self.hessian *= self.total_batches / (self.total_batches + n) From 0c74af8de0cbc3485965dd6f1091be678be45d5a Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Fri, 24 Apr 2026 12:52:36 +0000 Subject: [PATCH 24/25] Update [ghstack-poisoned] --- test/prototype/gptq/test_gptqv2.py | 84 +++++++++++------------ torchao/prototype/gptq/api.py | 2 +- torchao/prototype/gptq/observer.py | 106 +++++++++++++++++++++-------- 3 files changed, 117 insertions(+), 75 deletions(-) diff --git a/test/prototype/gptq/test_gptqv2.py b/test/prototype/gptq/test_gptqv2.py index fa06c415d2..fdd962fe4c 100644 --- a/test/prototype/gptq/test_gptqv2.py +++ b/test/prototype/gptq/test_gptqv2.py @@ -108,31 +108,7 @@ def test_observer_tensor_creation(self): ) # Check total_batches is initialized as 0 - assert observer.total_batches == 0 - - @pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA available") - def test_observer_tensor_attributes(self): - """Test GPTQObserverTensor attributes are correctly set.""" - weight = torch.randn(16, 32, dtype=torch.bfloat16, device="cuda") - observer = GPTQObserverTensor.from_hp(weight) - - # Test hp_data attribute - assert hasattr(observer, "hp_data") - assert isinstance(observer.hp_data, torch.Tensor) - - # Test hessian attribute - assert hasattr(observer, "hessian") - assert torch.equal( - observer.hessian, torch.zeros(32, 32, dtype=torch.float32, device="cuda") - ) - - # Test total_batches attribute - assert hasattr(observer, "total_batches") - assert observer.total_batches == 0 - - # Test update method exists - assert hasattr(observer, "update") - assert callable(observer.update) + assert (observer.total_batches == 0).all() @pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA available") def test_linear_operation_with_observer(self): @@ -161,7 +137,7 @@ def test_linear_operation_with_observer(self): # Check that Hessian was initialized and updated assert observer_weight.hessian is not None assert observer_weight.hessian.shape == (in_features, in_features) - assert observer_weight.total_batches == 1 + assert (observer_weight.total_batches == 1).all() # Verify output is correct expected_output = F.linear(input_tensor, weight) @@ -195,36 +171,47 @@ def test_multiple_observations(self): assert observer_weight.hessian.shape == (in_features, in_features) # Check total_batches matches total samples - assert observer_weight.total_batches == total_samples + assert (observer_weight.total_batches == total_samples).all() - @pytest.mark.skip(reason="bmm math is incorrect, will fix in next PR") @pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA available") def test_bmm_operation_with_observer(self): """Test torch.bmm with GPTQObserverTensor updates Hessian correctly.""" - batch = 4 + num_experts = 4 m = 8 n = 16 k = 12 + num_passes = 4 # Create input and weight tensors - input_tensor = torch.randn(batch, m, k, dtype=torch.float32, device="cuda") - weight = torch.randn(batch, k, n, dtype=torch.float32, device="cuda") - observer_weight = GPTQObserverTensor.from_hp(weight) + weight = torch.randn(num_experts, n, k, dtype=torch.float32, device="cuda") - # Perform bmm operation - output = torch.bmm(input_tensor, observer_weight) + inputs = [ + torch.randn(num_experts, m, k, dtype=torch.float32, device="cuda") + for _ in range(num_passes) + ] - # Check output shape - assert output.shape == (batch, m, n) + # 3D path: single observer with bmm + observer_3d = GPTQObserverTensor.from_hp(weight) + for x in inputs: + torch.bmm(x, observer_3d.transpose(-2, -1)) - # Check Hessian was initialized and updated - assert observer_weight.hessian is not None - # For bmm with batch dimension, the Hessian is computed on the last dimension - assert observer_weight.total_batches == batch - - # Verify output is correct - expected_output = torch.bmm(input_tensor, weight) - torch.testing.assert_close(output, expected_output) + # 2D path: per-expert observers with F.linear + observers_2d = [ + GPTQObserverTensor.from_hp(weight[e]) for e in range(num_experts) + ] + for x in inputs: + for e in range(num_experts): + F.linear(x[e], observers_2d[e]) + + # Verify per-expert hessians match bitwise to calculating each expert's + # hessian individually + for e in range(num_experts): + assert torch.equal(observer_3d.hessian[e], observers_2d[e].hessian), ( + f"Expert {e} hessian mismatch" + ) + assert torch.equal( + observer_3d.total_batches[e : e + 1], observers_2d[e].total_batches + ), f"Expert {e} total_batches mismatch" @pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA available") @pytest.mark.parametrize( @@ -260,7 +247,7 @@ def test_observer_config_transform(self, base_config): linear.weight.hessian, torch.zeros(64, 64, dtype=torch.float32, device="cuda"), ) - assert linear.weight.total_batches == 0 + assert (linear.weight.total_batches == 0).all() # Perform a forward pass input_tensor = torch.randn(4, 64, dtype=torch.float32, device="cuda") @@ -268,7 +255,7 @@ def test_observer_config_transform(self, base_config): # Check Hessian was initialized after forward pass assert linear.weight.hessian is not None - assert linear.weight.total_batches == 1 + assert (linear.weight.total_batches == 1).all() # Check output shape assert output.shape == (4, 32) @@ -384,6 +371,9 @@ def test_unified_config_two_phase(self, base_config): ) def test_gptq_quantize_function(self, base_config): """Test gptq_quantize function with synthetic Hessian and weights.""" + if isinstance(base_config, Int4WeightOnlyConfig) and is_sm_at_least_100(): + pytest.skip("int4 kernels do not work on sm100") + torch.manual_seed(42) # Create synthetic weight matrix @@ -556,6 +546,8 @@ def test_gptq_sqnr(self, base_config): and not is_sm_at_least_100() ): pytest.skip("CUDA capability >= 10.0 required for nvfp4") + if isinstance(base_config, Int4WeightOnlyConfig) and is_sm_at_least_100(): + pytest.skip("int4 kernels do not work on sm100") torch.manual_seed(43) diff --git a/torchao/prototype/gptq/api.py b/torchao/prototype/gptq/api.py index b680472c87..c1b7d2a185 100644 --- a/torchao/prototype/gptq/api.py +++ b/torchao/prototype/gptq/api.py @@ -141,7 +141,7 @@ def _gptq_config_transform( ) # Validate that observations were recorded - if tensor.total_batches == 0: + if (tensor.total_batches == 0).any(): raise ValueError( f"No observations recorded for {parameter_name}. " f"total_batches is 0. Did you run forward passes during the observe step?" diff --git a/torchao/prototype/gptq/observer.py b/torchao/prototype/gptq/observer.py index cbfb382da5..01f1930824 100644 --- a/torchao/prototype/gptq/observer.py +++ b/torchao/prototype/gptq/observer.py @@ -11,11 +11,11 @@ class GPTQObserverTensor(TorchAOBaseTensor): - tensor_data_names = ["hp_data"] + tensor_data_names = ["hp_data", "total_batches"] optional_tensor_data_names = ["hessian"] - tensor_attribute_names = ["total_batches"] + tensor_attribute_names = [] - def __new__(cls, hp_data: torch.Tensor, total_batches: int, hessian=None): + def __new__(cls, hp_data: torch.Tensor, total_batches, hessian=None): shape = hp_data.shape kwargs = {} kwargs["device"] = hp_data.device @@ -23,42 +23,80 @@ def __new__(cls, hp_data: torch.Tensor, total_batches: int, hessian=None): kwargs["requires_grad"] = False return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs) # type: ignore[attr-defined] - def __init__(self, hp_data: torch.Tensor, total_batches: int, hessian=None): + def __init__(self, hp_data: torch.Tensor, total_batches, hessian=None): super().__init__() self.hp_data = hp_data self.hessian = hessian - self.total_batches = total_batches + if isinstance(total_batches, torch.Tensor): + self.total_batches = total_batches + elif len(self.hp_data.shape) == 3: + self.total_batches = torch.zeros( + self.hp_data.shape[0], dtype=torch.int64, device=self.hp_data.device + ) + else: + self.total_batches = torch.zeros( + 1, dtype=torch.int64, device=self.hp_data.device + ) # initialize hessian - assert self.hp_data.is_contiguous() if self.hessian is None: + assert self.hp_data.is_contiguous() feature_dim = self.hp_data.shape[-1] - self.hessian = torch.zeros( - feature_dim, - feature_dim, - dtype=torch.float32, - device=self.hp_data.device, - ) - - def update(self, input: torch.Tensor): - """Incrementally update Hessian matrix from input activations.""" - # Move input to same device as hp_data and convert to float - x = input.float().to(self.hp_data.device) + if len(self.hp_data.shape) == 2: + self.hessian = torch.zeros( + feature_dim, + feature_dim, + dtype=torch.float32, + device=self.hp_data.device, + ) + else: + assert len(self.hp_data.shape) == 3, "unsupported" + expert_dim = self.hp_data.shape[0] + self.hessian = torch.zeros( + expert_dim, + feature_dim, + feature_dim, + dtype=torch.float32, + device=self.hp_data.device, + ) + + @staticmethod + def _update_single_hessian( + x: torch.Tensor, hessian: torch.Tensor, total_batches: torch.Tensor + ): + """Update a single 2D Hessian and total_batches in-place.""" shape = x.shape - - # Calculate batch size n = 1 if len(shape) == 2 else shape[0] x = x.reshape(-1, shape[-1]) - # Apply running average formula - if self.total_batches > 0: - self.hessian *= self.total_batches / (self.total_batches + n) + # cast to Python int64 for optimal type promotion semantics + # Note: there is definitely a better way to get ^, saving for + # a follow-up PR. For now, this preserves numerics. + tb = total_batches.item() + if tb > 0: + hessian *= tb / (tb + n) - self.total_batches += n + total_batches += n + # cast to Python int64 for optimal type promotion semantics + # Note: there is definitely a better way to get ^, saving for + # a follow-up PR. For now, this preserves numerics. + tb = total_batches.item() - # Update Hessian: x = ((2 / total_batches) ** (1 / 2)) * x.t() - x = ((2 / self.total_batches) ** (1 / 2)) * x.t() - self.hessian += x.matmul(x.t()) + x = ((2 / tb) ** (1 / 2)) * x.t() + hessian += x.matmul(x.t()) + + def update_2d(self, input: torch.Tensor): + x = input.float().to(self.hp_data.device) + self._update_single_hessian(x, self.hessian, self.total_batches[0:1]) + + def update_3d(self, input: torch.Tensor): + x = input.float().to(self.hp_data.device) + # TODO(future PR): optimize if this is too slow + for e_idx in range(self.hessian.shape[0]): + x_cur = x[e_idx] + h_cur = self.hessian[e_idx] + total_batches = self.total_batches[e_idx : e_idx + 1] + self._update_single_hessian(x_cur, h_cur, total_batches) @classmethod def from_hp(cls, hp_tensor): @@ -79,7 +117,7 @@ def _(func, types, args, kwargs): args[2] if len(args) > 2 else None, ) if isinstance(weight_tensor, GPTQObserverTensor): - weight_tensor.update(input_tensor.detach()) + weight_tensor.update_2d(input_tensor.detach()) return F.linear(input_tensor, weight_tensor.hp_data, bias) else: raise ValueError( @@ -87,11 +125,23 @@ def _(func, types, args, kwargs): ) +@implements(aten.transpose.int) +def _(func, types, args, kwargs): + self, dim0, dim1 = args[0], args[1], args[2] + assert {dim0, dim1} == {-2, -1} or {dim0, dim1} == { + self.hp_data.ndim - 2, + self.hp_data.ndim - 1, + }, f"only transpose of last two dims is supported, got dims {dim0}, {dim1}" + new_data = func(self.hp_data, dim0, dim1) + new_hessian = func(self.hessian, dim0, dim1) + return GPTQObserverTensor(new_data, self.total_batches, new_hessian) + + @implements(aten.bmm.default) def _(func, types, args, kwargs): input_tensor, weight_tensor = ( args[0], args[1], ) - weight_tensor.update(input_tensor.detach()) + weight_tensor.update_3d(input_tensor.detach()) return func(input_tensor, weight_tensor.hp_data) From dd7c1eeb0e6627f0b90033da10a4c89b04358893 Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Fri, 24 Apr 2026 13:25:44 +0000 Subject: [PATCH 25/25] Update [ghstack-poisoned] --- test/prototype/gptq/test_gptqv2.py | 79 ++++++++++++++++++++++++++++++ torchao/prototype/gptq/observer.py | 30 ++++++++++++ 2 files changed, 109 insertions(+) diff --git a/test/prototype/gptq/test_gptqv2.py b/test/prototype/gptq/test_gptqv2.py index fdd962fe4c..8d7d507cc2 100644 --- a/test/prototype/gptq/test_gptqv2.py +++ b/test/prototype/gptq/test_gptqv2.py @@ -213,6 +213,85 @@ def test_bmm_operation_with_observer(self): observer_3d.total_batches[e : e + 1], observers_2d[e].total_batches ), f"Expert {e} total_batches mismatch" + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA available") + @pytest.mark.skipif( + not is_sm_at_least_100(), + reason="CUDA capability >= 10.0 required for _grouped_mm", + ) + def test_grouped_mm_operation_with_observer(self): + """Test torch._grouped_mm with GPTQObserverTensor updates per-expert Hessians correctly.""" + num_experts = 4 + n = 16 + k = 12 + + weight = torch.randn(num_experts, n, k, dtype=torch.float32, device="cuda") + + # 4 different per-expert token distributions. Several of these have + # experts that see 0 tokens, which exercises the empty-slice skip path. + m_per_group_list = [ + [1, 3, 4, 16], # all experts active + [0, 3, 4, 13], # expert 0 sees 0 tokens + [5, 5, 0, 5], # expert 2 sees 0 tokens + [2, 0, 6, 4], # expert 1 sees 0 tokens + ] + + offs_list = [ + torch.tensor( + [sum(m_per_group[: i + 1]) for i in range(num_experts)], + device="cuda", + dtype=torch.int32, + ) + for m_per_group in m_per_group_list + ] + + inputs = [ + torch.randn(sum(m_per_group), k, dtype=torch.float32, device="cuda") + for m_per_group in m_per_group_list + ] + + # 3D path: single observer with _grouped_mm + observer_3d = GPTQObserverTensor.from_hp(weight) + for x, offs in zip(inputs, offs_list): + torch._grouped_mm(x, observer_3d.transpose(-2, -1), offs=offs) + + # 2D path: per-expert observers with F.linear + observers_2d = [ + GPTQObserverTensor.from_hp(weight[e]) for e in range(num_experts) + ] + for x, offs in zip(inputs, offs_list): + prev_end = 0 + for e in range(num_experts): + end = offs[e].item() + if end > prev_end: + F.linear(x[prev_end:end], observers_2d[e]) + prev_end = end + + # Verify per-expert hessians match bitwise to calculating each expert's + # hessian individually + for e in range(num_experts): + assert torch.equal(observer_3d.hessian[e], observers_2d[e].hessian), ( + f"Expert {e} hessian mismatch" + ) + assert torch.equal( + observer_3d.total_batches[e : e + 1], observers_2d[e].total_batches + ), f"Expert {e} total_batches mismatch" + + # Verify total_batches matches an independent count derived directly + # from the offsets: each non-empty forward pass contributes 1 per + # active expert (each expert's 2D slice has len(shape) == 2, so n=1). + expected_total_batches = torch.tensor( + [ + sum(1 for m_per_group in m_per_group_list if m_per_group[e] > 0) + for e in range(num_experts) + ], + dtype=torch.int64, + device="cuda", + ) + assert torch.equal(observer_3d.total_batches, expected_total_batches), ( + f"total_batches {observer_3d.total_batches.tolist()} " + f"does not match expected {expected_total_batches.tolist()}" + ) + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA available") @pytest.mark.parametrize( "base_config", diff --git a/torchao/prototype/gptq/observer.py b/torchao/prototype/gptq/observer.py index 01f1930824..df5832d051 100644 --- a/torchao/prototype/gptq/observer.py +++ b/torchao/prototype/gptq/observer.py @@ -30,6 +30,9 @@ def __init__(self, hp_data: torch.Tensor, total_batches, hessian=None): if isinstance(total_batches, torch.Tensor): self.total_batches = total_batches elif len(self.hp_data.shape) == 3: + # TODO(future PR): audit whether we need to change this + # from `total_batches` (current) to something like `total_tokens`, + # to ensure that each token is weighted equally in the 3d case. self.total_batches = torch.zeros( self.hp_data.shape[0], dtype=torch.int64, device=self.hp_data.device ) @@ -98,6 +101,23 @@ def update_3d(self, input: torch.Tensor): total_batches = self.total_batches[e_idx : e_idx + 1] self._update_single_hessian(x_cur, h_cur, total_batches) + def update_3d_with_offs(self, input: torch.Tensor, offs: torch.Tensor): + x = input.float().to(self.hp_data.device) + # offs is cumulative end indices; expert e gets rows [prev_end : offs[e]] + # Pull offs to CPU once to avoid a GPU->CPU sync per expert. + # TODO(future PR): optimize if this is too slow + offs_cpu = offs.tolist() + prev_end = 0 + for e_idx in range(self.hessian.shape[0]): + end = offs_cpu[e_idx] + if end == prev_end: + continue + x_cur = x[prev_end:end] + h_cur = self.hessian[e_idx] + total_batches = self.total_batches[e_idx : e_idx + 1] + self._update_single_hessian(x_cur, h_cur, total_batches) + prev_end = end + @classmethod def from_hp(cls, hp_tensor): return GPTQObserverTensor(hp_tensor, 0, None) @@ -145,3 +165,13 @@ def _(func, types, args, kwargs): ) weight_tensor.update_3d(input_tensor.detach()) return func(input_tensor, weight_tensor.hp_data) + + +@implements([aten._grouped_mm.default]) +def _(func, types, args, kwargs): + mat_a, mat_b = args[0], args[1] + offs = args[2] if len(args) > 2 else kwargs.get("offs", None) + assert offs is not None, "offs is required for grouped_mm" + assert isinstance(mat_b, GPTQObserverTensor) + mat_b.update_3d_with_offs(mat_a.detach(), offs) + return func(mat_a, mat_b.hp_data, offs)