diff --git a/CHANGELOG.rst b/CHANGELOG.rst index e30bb9e245e..9a84dd3d39e 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -63,6 +63,7 @@ Changelog **Bug Fixes** +- Megatron-Core fake quantization of Gated DeltaNet (Qwen3.5 / Qwen3.6) with a Transformer Engine ``in_proj`` now keeps its beta / alpha gates in BF16 during evaluation, QAD, ``fold_weight`` and the vLLM fake-quant export, matching deployment. Other ``in_proj`` layers stay fully quantized and warn. - ``hf_quant_config.json`` now declares ``pre_quant_scale`` only for INT4 / W4A8 layers that have the tensor; loading INT4 / W4A8 fused-MoE checkpoints in TensorRT-LLM also needs a release whose MoE loader treats it as optional. AWQ, SmoothQuant and SVDQuant now warn when they skip quantized fused-MoE experts. - Hugging Face fused MoE experts stored transposed, ``(num_experts, in_dim, out_dim)``, without a quantized wrapper of their own (for example Aria's) are now left unquantized with a warning. They were quantized with gate and up split along the wrong axis, so their exported weights were wrong; the Hugging Face export now refuses such a model's MoE block instead. - Calibrating a Megatron-Core MLA model now gives ``linear_q_down_proj`` and ``linear_kv_down_proj`` one per-tensor input and weight range, as vLLM and TensorRT-LLM fuse them into one GEMM with one scale per tensor. Re-calibrate existing MLA checkpoints to get the shared ranges. diff --git a/modelopt/torch/export/plugins/vllm_fakequant_megatron.py b/modelopt/torch/export/plugins/vllm_fakequant_megatron.py index 64c4deb4ed4..019ee7bf902 100644 --- a/modelopt/torch/export/plugins/vllm_fakequant_megatron.py +++ b/modelopt/torch/export/plugins/vllm_fakequant_megatron.py @@ -33,6 +33,7 @@ from modelopt.torch.export.quant_format import QUANTIZATION_NONE from modelopt.torch.export.unified_export_megatron import GPTModelExporter from modelopt.torch.quantization.nn import GroupedQuantizer, SequentialQuantizer, TensorQuantizer +from modelopt.torch.quantization.utils.core_utils import apply_without_qdq from modelopt.torch.utils import get_unwrapped_name from modelopt.torch.utils.distributed import DistributedProcessGroup, is_master @@ -405,8 +406,12 @@ def _self_attention_scaling( self._state_dict.update({prefix + name: value for name, value in state.items()}) @staticmethod - def _fakequant_weight(weight, weight_quantizer, dtype): - """Return a CPU QDQ weight while preserving the quantizer's device.""" + def _fakequant_weight(weight, weight_quantizer, dtype, unquantized_rows=None): + """Return a CPU QDQ weight while preserving the quantizer's device. + + ``unquantized_rows`` skip the QDQ but keep the quantizer's scaling / rotation, as in the + fake-quant forward (GatedDeltaNet's beta / alpha gates). + """ weight = weight.to(dtype) # Fold the weight_quantizer into the weight by applying fake-quantization # (quantize then dequantize). The weight_quantizer amax is not exported; @@ -430,7 +435,17 @@ def _fakequant_weight(weight, weight_quantizer, dtype): if need_move: weight_quantizer.to(quant_device) try: - weight = weight_quantizer(weight.to(quant_device)).to(dtype) + weight = weight.to(quant_device) + qweight = weight_quantizer(weight) + if unquantized_rows is not None: + qweight = torch.cat( + [ + qweight[: unquantized_rows.start], + apply_without_qdq(weight_quantizer, weight[unquantized_rows]), + qweight[unquantized_rows.stop :], + ] + ) + weight = qweight.to(dtype) finally: if need_move: weight_quantizer.to(wq_dev) @@ -478,7 +493,10 @@ def _get_quantized_state( # Grouped slicing reads QDQ weights from the hooked module.state_dict(). if "weight" in name_to_value and module not in self._grouped_fq_modules: name_to_value["weight"] = self._fakequant_weight( - name_to_value["weight"], getattr(module, "weight_quantizer", None), dtype + name_to_value["weight"], + getattr(module, "weight_quantizer", None), + dtype, + getattr(module, "_modelopt_unquantized_output_rows", None), ) # Save activation ranges; weight quantizers are folded into the weights above. diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py index 84fd26cc4a3..5cdb008bd2c 100644 --- a/modelopt/torch/quantization/plugins/megatron.py +++ b/modelopt/torch/quantization/plugins/megatron.py @@ -1141,6 +1141,21 @@ def _setup(self): data_parallel_group, mcore_parallel.get_tensor_model_parallel_group(), ) + # The fused in_proj ends with the beta / alpha gates, which recipes and every export + # keep in BF16 (vLLM's in_proj_ba is too narrow for FP4 kernels), so leave them + # unquantized here too. + names = list(getattr(self, "in_proj_split_names", ())) + if not isinstance(self.in_proj, TELayerNormColumnParallelLinear): + warn_rank_0( + f"GatedDeltaNet in_proj is a {type(self.in_proj).__name__}, so fake quantization" + " also quantizes its beta / alpha gates, which exports keep in BF16." + ) + elif names[-2:] == ["beta", "alpha"]: + local_rows = sum(self.in_proj_split_sections) + gate_rows = sum(self.in_proj_split_sections[-2:]) + self.in_proj._modelopt_unquantized_output_rows = slice( + local_rows - gate_rows, local_rows + ) @contextmanager def _quantized_gdn_kernel(self): diff --git a/modelopt/torch/quantization/plugins/transformer_engine.py b/modelopt/torch/quantization/plugins/transformer_engine.py index 32c9a3cffac..e80eee5e73e 100644 --- a/modelopt/torch/quantization/plugins/transformer_engine.py +++ b/modelopt/torch/quantization/plugins/transformer_engine.py @@ -26,6 +26,7 @@ import transformer_engine.pytorch.module.linear as te_linear from modelopt.torch.quantization.utils import replace_function +from modelopt.torch.quantization.utils.core_utils import apply_without_qdq from ..nn import GroupedQuantizer, QuantModuleRegistry, SequentialQuantizer, TensorQuantizer from .custom import _ParallelLinear @@ -315,19 +316,41 @@ class _QuantLayerNormLinearFunc(torch.autograd.Function): @staticmethod def forward(ctx, inp, ln_weight, ln_bias, weight, *args, **kwargs): - input_quantizer, weight_quantizer = _QuantLayerNormLinearFunc.modelopt_quantizers + input_quantizer, weight_quantizer, unquantized_rows = ( + _QuantLayerNormLinearFunc.modelopt_quantizers + ) qweight = weight_quantizer(weight) + if unquantized_rows is not None: + # Deployment keeps these rows in BF16: they skip QDQ but keep the quantizers' scaling + # and rotation. Not in place: the fake-quant weight cache may hold qweight. + row_weight = apply_without_qdq(weight_quantizer, weight[unquantized_rows]) + qweight = torch.cat( + ( + qweight[: unquantized_rows.start], + row_weight.to(qweight.dtype), + qweight[unquantized_rows.stop :], + ) + ) qweight.requires_grad = weight.requires_grad if ctx is not None: # We need to recompute the quantized input for the backward pass, so we save the input_quantizer ctx.modelopt_input_quantizer = input_quantizer + ctx.modelopt_unquantized_rows = unquantized_rows original_gemm = te_layernorm_linear.general_gemm - def _patched_general_gemm(weight, input, *gemm_args, **gemm_kwargs): + def _patched_general_gemm(gemm_weight, input, *gemm_args, **gemm_kwargs): qinput = input_quantizer(input) - return original_gemm(weight, qinput, *gemm_args, **gemm_kwargs) + outputs = original_gemm(gemm_weight, qinput, *gemm_args, **gemm_kwargs) + if unquantized_rows is not None: + bias = gemm_kwargs.get("bias") + outputs[0][..., unquantized_rows] = torch.nn.functional.linear( + apply_without_qdq(input_quantizer, input).to(row_weight.dtype), + row_weight, + None if bias is None else bias[unquantized_rows], + ).to(outputs[0].dtype) + return outputs with replace_function( te_layernorm_linear, @@ -366,7 +389,16 @@ def _patched_general_gemm(a, b, *gemm_args, **gemm_kwargs): # but gemm is called with the unquantized input (a) # So lets first get the quantized input (qinput) and then call the gemm qinput = ctx.modelopt_input_quantizer(a) - return original_gemm(qinput, b, *gemm_args, **gemm_kwargs) + outputs = original_gemm(qinput, b, *gemm_args, **gemm_kwargs) + rows = ctx.modelopt_unquantized_rows + if rows is not None and outputs[0] is not None: + # The forward computed these rows from the unquantized input. + x = apply_without_qdq(ctx.modelopt_input_quantizer, a) + dy = b[..., rows].reshape(-1, rows.stop - rows.start) + outputs[0][rows] = (dy.t() @ x.reshape(dy.shape[0], -1).to(dy.dtype)).to( + outputs[0].dtype + ) + return outputs with replace_function( te_layernorm_linear, @@ -378,6 +410,7 @@ def _patched_general_gemm(a, b, *gemm_args, **gemm_kwargs): outputs = te_layernorm_linear._LayerNormLinear.backward(ctx, *grad_outputs) delattr(ctx, "modelopt_input_quantizer") + delattr(ctx, "modelopt_unquantized_rows") return outputs @@ -394,6 +427,18 @@ def _setup(self): ) self.fuse_wgrad_accumulation = False + def fold_weight(self, keep_attrs: bool = False): + """Fold the weight quantizer, keeping the rows the forward pass leaves unquantized.""" + rows = getattr(self, "_modelopt_unquantized_output_rows", None) + unquantized = None + if rows is not None: + with torch.no_grad(): + # Cloned: without transforms this returns a view of the weight that folding overwrites. + unquantized = apply_without_qdq(self.weight_quantizer, self.weight[rows]).clone() + super().fold_weight(keep_attrs) + if unquantized is not None: + self.weight.data[rows] = unquantized.to(self.weight.dtype) + def forward(self, *args, **kwargs): """Call ModelOpt patch for _LayerNormLinear functional.""" _assert_te_fp8_enabled() @@ -401,6 +446,7 @@ def forward(self, *args, **kwargs): _QuantLayerNormLinearFunc.modelopt_quantizers = ( self.input_quantizer, self.weight_quantizer, + getattr(self, "_modelopt_unquantized_output_rows", None), ) with replace_function( te_layernorm_linear, diff --git a/modelopt/torch/quantization/utils/core_utils.py b/modelopt/torch/quantization/utils/core_utils.py index b4120720b5c..612a737a691 100644 --- a/modelopt/torch/quantization/utils/core_utils.py +++ b/modelopt/torch/quantization/utils/core_utils.py @@ -1024,6 +1024,25 @@ def enable_quant(quantizer): quantizer._if_quant = original_if_quant +def apply_without_qdq(quantizer, inputs: torch.Tensor) -> torch.Tensor: + """Apply ``quantizer``'s pre-quant scale and rotation to ``inputs``, without quantizing them. + + For rows a quantized module computes in high precision, so they keep the transforms (e.g. + an AWQ / SmoothQuant scale that the stored weight compensates for) of the quantized rows. + """ + from ..nn import SequentialQuantizer # local import: ..nn imports this module + + members = list(quantizer) if isinstance(quantizer, SequentialQuantizer) else [quantizer] + for member in members: + disabled = member._disabled + member._disabled = True + try: + inputs = member(inputs) + finally: + member._disabled = disabled + return inputs + + @contextmanager def disable_calib(quantizer): """Temporarily disable calibration for a quantizer. diff --git a/tests/gpu_megatron/torch/export/test_vllm_fakequant_megatron_export.py b/tests/gpu_megatron/torch/export/test_vllm_fakequant_megatron_export.py index 5f0b947890d..a4f5ee3cce2 100644 --- a/tests/gpu_megatron/torch/export/test_vllm_fakequant_megatron_export.py +++ b/tests/gpu_megatron/torch/export/test_vllm_fakequant_megatron_export.py @@ -22,6 +22,7 @@ import pytest import torch +import transformers import yaml from _test_utils.torch.megatron.models import get_mcore_gpt_model, get_mcore_hybrid_model from _test_utils.torch.megatron.utils import initialize_for_megatron, run_mcore_inference @@ -29,6 +30,7 @@ create_tiny_gpt_oss_dir, create_tiny_llama_dir, create_tiny_nemotron_h_dir, + get_tiny_qwen3_5_vl_offline, ) from megatron.core.parallel_state import is_pipeline_last_stage from safetensors import safe_open @@ -45,6 +47,7 @@ gather_mcore_vllm_fq_quantizer_recipe, ) from modelopt.torch.quantization.nn import GroupedQuantizer, SequentialQuantizer, TensorQuantizer +from modelopt.torch.quantization.utils.core_utils import apply_without_qdq def _assert_exported_quantizers(export_dir, expected_names, amax=1.001, disabled_names=()): @@ -414,6 +417,80 @@ def test_mcore_vllm_export_mtp(request, tmp_path): workers.run(partial(_test_mcore_vllm_export_mtp, tmp_path)) +def _test_mcore_vllm_export_gdn(source, tmp_path, rank, size): + text_cfg = transformers.AutoConfig.from_pretrained(source).text_config + model = get_mcore_gpt_model( + tensor_model_parallel_size=size, + initialize_megatron=True, + num_layers=text_cfg.num_hidden_layers, + hidden_size=text_cfg.hidden_size, + num_attention_heads=text_cfg.num_attention_heads, + num_query_groups=text_cfg.num_key_value_heads, + ffn_hidden_size=text_cfg.intermediate_size, + max_sequence_length=text_cfg.max_position_embeddings, + vocab_size=text_cfg.vocab_size, + normalization="RMSNorm", + transformer_impl="modelopt", + kv_channels=text_cfg.head_dim, + experimental_attention_variant="gated_delta_net", + linear_attention_freq=len(text_cfg.layer_types), + ).cuda() + model.eval() + + def forward_loop(model): + with torch.no_grad(): + run_mcore_inference(model, torch.randint(0, 32, (1, 32), device="cuda")) + + model = mtq.quantize(model, mtq.FP8_DEFAULT_CFG, forward_loop) + gdn = model.decoder.layers[0].self_attention + weight_quantizer = gdn.in_proj.weight_quantizer + if find_spec("fast_hadamard_transform") is not None: + # A CUDA-only transform that the gate rows keep: the export must stage them onto CUDA. + weight_quantizer.set_from_attribute_config( + {"rotate": {"enable": True, "mode": "rotate_back"}} + ) + with torch.no_grad(): + weight = gdn.in_proj.weight.to(torch.bfloat16) + transformed = apply_without_qdq(weight_quantizer, weight).cpu() + weight = weight.cpu() + model.cpu() # CPU-offloaded weights + export_dir = tmp_path / "gdn_export" + export_mcore_gpt_to_hf_vllm_fq( + model, str(source), dtype=torch.bfloat16, export_dir=str(export_dir) + ) + + sections = dict(zip(gdn.in_proj_split_names, gdn.in_proj_split_sections)) + qkvz_rows = sum(sections[name] for name in ("query", "key", "value", "z")) + gate_rows = { + "in_proj_b": transformed[qkvz_rows : qkvz_rows + sections["beta"]], + "in_proj_a": transformed[qkvz_rows + sections["beta"] :], + } + prefix = "model.language_model.layers.0.linear_attn." + state = torch.load(export_dir / "quantizer_state.pth", weights_only=True, map_location="cpu") + recipe = yaml.safe_load((export_dir / "quant_recipe.yaml").read_text())["quantizers"] + weight_map = json.loads((export_dir / "model.safetensors.index.json").read_text())["weight_map"] + # The beta / alpha gates are written unquantized and get no quantizers; qkv / z keep theirs. + for name, expected in gate_rows.items(): + key = f"{prefix}{name}.weight" + with safe_open(export_dir / weight_map[key], framework="pt") as f: + assert torch.equal(f.get_tensor(key), expected) + assert not any(f"{prefix}{name}." in key for key in state.keys() | recipe.keys()) + key = f"{prefix}in_proj_qkv.weight" + qkv_rows = sum(sections[name] for name in ("query", "key", "value")) + with safe_open(export_dir / weight_map[key], framework="pt") as f: + assert not torch.equal(f.get_tensor(key), weight[:qkv_rows]) + for name in ("in_proj_qkv", "in_proj_z"): + assert f"{prefix}{name}.input_quantizer" in recipe + + +def test_mcore_vllm_export_gdn(dist_workers_size_1, tmp_path): + """GatedDeltaNet's beta / alpha gates export unquantized, with no quantizers of their own.""" + pytest.importorskip("fla") # Megatron-Core GatedDeltaNet needs the optional fla + source = tmp_path / "tiny_qwen3_5_vl" + get_tiny_qwen3_5_vl_offline().save_pretrained(source) + dist_workers_size_1.run(partial(_test_mcore_vllm_export_gdn, source, tmp_path)) + + def _assert_unsupported_settings(model, source, export_dir, rank, size): attribute_cfgs = [ ( diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py index 0756f2d0d40..59705346a07 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py @@ -27,6 +27,7 @@ pytest.importorskip("fla") # Megatron-Core GatedDeltaNet and the state QDQ kernel need fla GatedDeltaNet = pytest.importorskip("megatron.core.ssm.gated_delta_net").GatedDeltaNet +from modelopt.torch.export.plugins.vllm_fakequant_megatron import VllmFqGPTModelExporter from modelopt.torch.quantization.plugins.gated_delta_net import _state_qdq_chunk_gated_delta_rule from modelopt.torch.quantization.plugins.megatron import _QuantGatedDeltaNet @@ -38,7 +39,7 @@ SEED = 1234 -def _make_model(tp_size): +def _make_model(tp_size, num_heads=1): model = ( get_mcore_gpt_model( tensor_model_parallel_size=tp_size, @@ -48,8 +49,8 @@ def _make_model(tp_size): vocab_size=32, max_sequence_length=128, experimental_attention_variant="gated_delta_net", - linear_num_key_heads=1, - linear_num_value_heads=1, + linear_num_key_heads=num_heads, + linear_num_value_heads=num_heads, linear_key_head_dim=32, linear_value_head_dim=32, ) @@ -173,3 +174,77 @@ def test_gdn_qat_and_sharded_restore(compiled_gdn_workers, tmp_path): """Train through QDQ after a Megatron distributed-checkpoint round trip.""" workers, sites = compiled_gdn_workers workers.run(_test_gdn_qat_helper, tmp_path, sites) + + +def _gate_forward_backward(in_proj, x, gates, grad): + """Gate rows of in_proj(x), and the input / gate-weight gradients of a loss on them alone.""" + x = x.detach().requires_grad_() + in_proj.weight.grad = None + output = in_proj(x)[0] + (output[..., gates] * grad).sum().backward() + return output.detach(), x.grad, in_proj.weight.grad[gates].clone() + + +def _test_in_proj_gates_unquantized_helper(rank, size): + initialize_for_megatron( + tensor_model_parallel_size=size, pipeline_model_parallel_size=1, seed=SEED + ) + model = _make_model(size, num_heads=2) + gdn = next(m for m in model.modules() if isinstance(m, GatedDeltaNet)) + in_proj = gdn.in_proj + num_rows = in_proj.weight.shape[0] + gates = slice(num_rows - 2 * 2 // size, num_rows) # beta + alpha: one row per value head each + weight = in_proj.weight.detach().clone() + x = torch.randn(8, 2, 64, device="cuda", dtype=weight.dtype) + grad = torch.randn(8, 2, gates.stop - gates.start, device="cuda", dtype=weight.dtype) + reference = _gate_forward_backward(in_proj, x, gates, grad) + + fp8 = {"num_bits": (4, 3), "axis": None} + config = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*in_proj.weight_quantizer", "cfg": fp8}, + {"quantizer_name": "*in_proj.input_quantizer", "cfg": fp8}, + ], + "algorithm": "max", + } + mtq.quantize(model, config, get_forward(model)) + + # The gate rows, and the gradients through them alone, are the BF16 ones; the others are not. + output, x_grad, gate_weight_grad = _gate_forward_backward(in_proj, x, gates, grad) + torch.testing.assert_close(output[..., gates], reference[0][..., gates]) + assert not torch.allclose(output[..., : gates.start], reference[0][..., : gates.start]) + torch.testing.assert_close(x_grad, reference[1]) + torch.testing.assert_close(gate_weight_grad, reference[2]) + + exporter = object.__new__(VllmFqGPTModelExporter) + exporter._quantizer_recipe_markers, exporter.exclude_modules = [], [] + exporter._grouped_fq_modules, exporter._packing_experts = set(), False + exported = exporter._get_quantized_state(in_proj, weight.dtype)[0]["weight"] + assert torch.equal(exported[gates], weight[gates].cpu()) + assert not torch.equal(exported[: gates.start], weight[: gates.start].cpu()) + + # An AWQ / SmoothQuant input scale, which the stored weight compensates for, still applies. + # Powers of two, so that scaling the input and dividing the weight round nothing. + scale = 2.0 ** torch.randint(-1, 2, (64,), device="cuda").to(weight.dtype) + in_proj.input_quantizer.pre_quant_scale = scale + with torch.no_grad(): + in_proj.weight.div_(scale) + scaled_weight = in_proj.weight.detach().clone() + with torch.no_grad(): + output = in_proj(x)[0] + torch.testing.assert_close(output[..., gates], reference[0][..., gates]) + + mtq.fold_weight(model) + assert torch.equal(in_proj.weight[gates], scaled_weight[gates]) + assert not torch.equal(in_proj.weight[: gates.start], scaled_weight[: gates.start]) + + +def test_in_proj_keeps_gates_unquantized(dist_workers_size_1): + """The beta / alpha rows of the fused in_proj stay BF16 in fake quant, as in every export.""" + dist_workers_size_1.run(_test_in_proj_gates_unquantized_helper) + + +def test_in_proj_keeps_gates_unquantized_tp2(dist_workers_size_2): + """Same with each rank holding its own gate rows (TP <= 2, the number of value heads).""" + dist_workers_size_2.run(_test_in_proj_gates_unquantized_helper)