diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 7a188669cb8..d6d6f1328ae 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -102,6 +102,7 @@ Changelog - Fix ``examples/vllm_serve`` serving shared experts uncalibrated: their ``gate_proj``/``up_proj`` quantizer keys were not merged into ``gate_up_proj`` on reload, so they matched no module and were dropped. - Fix Qwen3-VL MoE PTQ failing on ``transformers>=5.12`` with ``AttributeError: 'QuantQwen3VLMoeTextExperts' object has no attribute 'hidden_size'`` (NVBug 6518551). transformers 5.12 moved ``Qwen3VLMoeTextExperts`` onto the standard ``@use_experts_implementation`` fused layout (``hidden_size``/``expert_dim`` renamed to ``hidden_dim``/``intermediate_dim``, ``gate_up_proj`` transposed to ``(num_experts, 2*intermediate_dim, hidden_dim)``, two ``F.linear`` calls per expert), but the legacy ``_QuantQwen3VLMoeTextExperts`` wrapper stayed statically registered and shadowed on-the-fly detection. The new layout is now left to ``register_fused_experts_on_the_fly``, which claims it with the generic ``_QuantFusedExperts``; the legacy wrapper is still registered on ``transformers<5.12``, whose ``torch.bmm``-based forward the generic wrapper cannot intercept. - Fix ``examples/hf_ptq`` multi-node FSDP2 export (``--use_fsdp2``) failing with ``RuntimeError: Cannot set version_counter for inference tensor``. ``export_quantized`` wrapped its whole body in ``torch.inference_mode()``, so the full params gathered by ``get_model_state_dict(full_state_dict=True)`` were inference tensors and the subsequent ``state_dict()`` -> ``param.detach()`` could not set their version counter. The export context is now ``torch.no_grad()``, which still disables autograd but keeps the gathered params as normal tensors. +- Fix QLoRA export in ``examples/llm_qat/export.py`` failing with ``AssertionError: Model already has modelopt state!`` (NVBug 6542481). The QLoRA training output is an adapter-only checkpoint, so ``from_pretrained`` resolves the quantized base model from ``adapter_config.json`` and ``enable_huggingface_checkpointing`` already restores its ModelOpt state; the export then restored a second time. It now restores only when the loaded model is not already converted. Two further breakages on the same path are also fixed: ``_restore_qtensor_wrappers`` matched no modules because PEFT re-parents the quantized linear as ``.base_layer`` while ``q_tensor_state`` is keyed by the name it was saved with (the packed NVFP4 weight then reached ``F.linear`` and raised a shape error), and ``postprocess_state_dict`` silently dropped every ``base_layer.*`` key missing from a hand-maintained rename map — losing the NVFP4 ``weight_scale_2`` global scale and any linear ``bias`` (Qwen2-style q/k/v biases), and leaving ``base_layer`` in the exported AWQ ``pre_quant_scale`` key. The rename is now a generic ``.base_layer.`` strip. 0.45 (2026-07-02) ^^^^^^^^^^^^^^^^^ diff --git a/examples/llm_qat/export.py b/examples/llm_qat/export.py index f48e85c3ee4..afe2bd4d1cf 100644 --- a/examples/llm_qat/export.py +++ b/examples/llm_qat/export.py @@ -23,7 +23,7 @@ import modelopt.torch.opt as mto from modelopt.torch.export.convert_hf_config import convert_hf_quant_config_format from modelopt.torch.export.unified_export_hf import _export_transformers_checkpoint -from modelopt.torch.opt.conversion import restore_from_modelopt_state +from modelopt.torch.opt.conversion import ModeloptStateManager, restore_from_modelopt_state from modelopt.torch.quantization.utils import set_quantizer_state_dict from modelopt.torch.utils import print_rank_0 @@ -48,8 +48,11 @@ def get_model( # Load model model = AutoModelForCausalLM.from_pretrained(ckpt_path, device_map=device_map) - # Restore modelopt state for LoRA models. For QAT/QAD models from_pretrained call handles this - if hasattr(model, "peft_config"): + # Restore modelopt state for LoRA models. For QAT/QAD models from_pretrained call handles this. + # For QLoRA the base checkpoint is quantized, so from_pretrained already restored the state. + # Skipping is safe only because QATTrainer writes modelopt_state_train.pth at trainer init, + # from that same base state. + if hasattr(model, "peft_config") and not ModeloptStateManager.is_converted(model): modelopt_state = mto.load_modelopt_state(f"{ckpt_path}/modelopt_state_train.pth") restore_from_modelopt_state(model, modelopt_state) print_rank_0("Restored modelopt state") diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index ab2ef0d9029..8af87c817ec 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -991,16 +991,10 @@ def postprocess_state_dict( "weight_shape", ] - # For modelopt-trained LoRA models, we need to remove the base_layer prefix from the keys for deployment - if is_modelopt_qlora: - replacements.update( - { - "base_layer.weight": "weight", - "base_layer.input_scale": "input_scale", - "base_layer.weight_scale": "weight_scale", - } - ) - skip_keys.append("base_layer") + def _export_key(key: str) -> str: + # PEFT nests the quantized module under `base_layer`, which deployment does not expect. + # Strip it generically so new key types (bias, scales) do not need to be enumerated here. + return key.replace(".base_layer.", ".") if is_modelopt_qlora else key post_state_dict = {} @@ -1012,7 +1006,7 @@ def postprocess_state_dict( # Skip keys not related to quantizers if all(skip_key not in key for skip_key in skip_keys): - post_state_dict[key] = value + post_state_dict[_export_key(key)] = value continue # Apply replacements if the key matches any suffix in the replacements dict @@ -1033,7 +1027,7 @@ def postprocess_state_dict( logger.warning( "Large KV activations detected. Quantized KV cache may lead to higher accuracy drop." ) - post_state_dict[prefix + new_suffix] = value + post_state_dict[_export_key(prefix + new_suffix)] = value break # Squeeze scales with a leading dimension of 1 diff --git a/modelopt/torch/opt/plugins/transformers.py b/modelopt/torch/opt/plugins/transformers.py index f72715d410c..0052525b6e9 100644 --- a/modelopt/torch/opt/plugins/transformers.py +++ b/modelopt/torch/opt/plugins/transformers.py @@ -110,16 +110,34 @@ def _restore_qtensor_wrappers(model, model_path): q_tensor_state = mode_config.get("metadata", {}).get("q_tensor_state", {}) if not q_tensor_state: continue - for name, module in model.named_modules(): - if ( - isinstance(module, RealQuantLinear) - and name in q_tensor_state - and not isinstance(module.weight, QTensorWrapper) - ): - module._parameters["weight"] = QTensorWrapper( - qtensor=module.weight.data, - metadata=q_tensor_state[name]["metadata"], - ) + # PEFT nests the quantized linear as `.base_layer`, and either the saved keys or the + # live names may carry that suffix. Normalize both so the lookup works in either direction. + q_tensor_state = {k.removesuffix(".base_layer"): v for k, v in q_tensor_state.items()} + + pending = [ + (name, module) + for name, module in model.named_modules() + if isinstance(module, RealQuantLinear) and not isinstance(module.weight, QTensorWrapper) + ] + matched = 0 + for name, module in pending: + key = name.removesuffix(".base_layer") + if key not in q_tensor_state: + continue + module._parameters["weight"] = QTensorWrapper( + qtensor=module.weight.data, + metadata=q_tensor_state[key]["metadata"], + ) + matched += 1 + + # A total miss means some wrapper renamed the modules. Warn instead of letting it surface + # as an opaque shape error at dequantization. + if pending and not matched: + warnings.warn( + f"Found {len(q_tensor_state)} compressed weight(s) in {modelopt_state_path} but " + f"re-wrapped none of the {len(pending)} candidate module(s); their names may have " + "been remapped. The model will likely fail when the packed weights are used." + ) def _new_from_pretrained(cls, /, pretrained_model_name_or_path, *args, **kwargs): diff --git a/tests/examples/llm_qat/test_llm_qat.py b/tests/examples/llm_qat/test_llm_qat.py index f87117501f5..4cd67d7f905 100644 --- a/tests/examples/llm_qat/test_llm_qat.py +++ b/tests/examples/llm_qat/test_llm_qat.py @@ -14,8 +14,12 @@ # limitations under the License. +import json + import pytest +import torch from _test_utils.examples.run_command import run_example_command +from safetensors.torch import load_file # Mapping from backend name to accelerate config file BACKEND_CONFIGS = { @@ -86,6 +90,18 @@ def _run_train(config: str, extra_cmd_args: list[str], backend: str = "fsdp2", c setup_free_port=True, ) + +def _run_export(ckpt_dir: str, export_dir: str): + run_example_command( + [ + "python", "export.py", + "--pyt_ckpt_path", ckpt_dir, + "--export_path", export_dir, + ], + "llm_qat", + ) + + def test_dataset_utils_pretokenize(tiny_qwen3_path, tmp_path): """Test dataset_utils.py standalone CLI pre-tokenization.""" cache_dir = tmp_path / "dataset_cache" @@ -152,18 +168,43 @@ def test_qwen3_lora_qat_nvfp4(tiny_qwen3_path, tmp_path): ) # Step 2: LoRA QAT + lora_qat_output_dir = tmp_path / "lora_qat" _run_train( "configs/train/qat_nvfp4.yaml", [ "--model_name_or_path", str(ptq_output_dir), "--do_train", "True", "--lora", "True", - "--output_dir", str(tmp_path / "lora_qat"), + "--output_dir", str(lora_qat_output_dir), ], backend="fsdp2", cache_dir=cache_dir, ) + # Step 3: Export. This checkpoint is fake-quantized, so the calibrated amaxes rather than + # packed weights are what must survive the load. + export_dir = tmp_path / "lora_qat_export" + _run_export(str(lora_qat_output_dir), str(export_dir)) + + base_model_dir = export_dir / "base_model" + with open(base_model_dir / "hf_quant_config.json") as f: + assert json.load(f)["quantization"]["quant_algo"] == "NVFP4" + + base_weights = load_file(base_model_dir / "model.safetensors") + assert not any("base_layer" in k or k.endswith("_amax") for k in base_weights) + + # LoRA freezes the base model, so a direct PTQ export is a trusted oracle for every calibrated + # value. This catches scales that keep their key but were reset to defaults. + ptq_export_dir = tmp_path / "ptq_export" + _run_export(str(ptq_output_dir), str(ptq_export_dir)) + reference = load_file(ptq_export_dir / "model.safetensors") + + scales = [k for k in reference if k.endswith(("_scale", "_scale_2"))] + assert scales, "no NVFP4 scales in the reference PTQ export" + for key in scales: + assert key in base_weights, f"{key} missing from the LoRA-QAT export" + assert torch.equal(base_weights[key], reference[key]), f"{key} does not match PTQ export" + @pytest.mark.parametrize("backend", [ "fsdp2", @@ -219,14 +260,37 @@ def test_qwen3_qlora_nvfp4(tiny_qwen3_path, tmp_path): ) # Step 2: QLoRA training + qlora_output_dir = tmp_path / "qlora" _run_train( "configs/train/qlora_nvfp4.yaml", [ "--model_name_or_path", str(ptq_output_dir), "--do_train", "True", "--lora", "True", - "--output_dir", str(tmp_path / "qlora"), + "--output_dir", str(qlora_output_dir), ], backend="ddp", cache_dir=cache_dir, ) + + # Step 3: Export the QLoRA checkpoint for deployment + export_dir = tmp_path / "qlora_export" + _run_export(str(qlora_output_dir), str(export_dir)) + + # The base model is exported compressed; the adapters stay at the top level. + base_model_dir = export_dir / "base_model" + assert (export_dir / "adapter_model.safetensors").is_file() + assert (base_model_dir / "hf_quant_config.json").is_file() + + with open(base_model_dir / "hf_quant_config.json") as f: + assert json.load(f)["quantization"]["quant_algo"] == "NVFP4" + + # NVFP4 needs the packed weight and *both* scales to be dequantizable downstream. + base_weights = load_file(base_model_dir / "model.safetensors") + packed_weights = [k for k, v in base_weights.items() if k.endswith(".weight") and v.dtype == torch.uint8] + assert packed_weights, "no NVFP4-packed weights found in the exported base model" + for key in packed_weights: + prefix = key.removesuffix(".weight") + assert f"{prefix}.weight_scale" in base_weights + assert f"{prefix}.weight_scale_2" in base_weights + assert not any("base_layer" in k or "lora" in k for k in base_weights) diff --git a/tests/gpu/torch/export/test_export.py b/tests/gpu/torch/export/test_export.py index cac0a9a9aef..55137a64639 100644 --- a/tests/gpu/torch/export/test_export.py +++ b/tests/gpu/torch/export/test_export.py @@ -259,6 +259,41 @@ def test_postprocess_state_dict(state_dict, quantization, maxbound, expected_sta assert processed_state_dict == expected_state_dict +def test_postprocess_state_dict_qlora_strips_base_layer(): + """Every QLoRA `base_layer.*` tensor needed for deployment must survive the rename. + + Dropping the NVFP4 global scale or a bias yields an undeployable checkpoint. + """ + state_dict = { + "layer1.base_layer.weight": torch.ones(4, 2, dtype=torch.uint8), + "layer1.base_layer.weight_scale": torch.ones(4, 1), + "layer1.base_layer.weight_scale_2": torch.tensor([0.5]), + "layer1.base_layer.input_scale": torch.tensor([0.25]), + "layer1.base_layer.bias": torch.arange(4.0), + "layer1.base_layer.input_quantizer._pre_quant_scale": torch.ones(2), + # Quantizer internals must still be dropped. + "layer1.base_layer.weight_quantizer._amax": torch.tensor([1.0]), + "layer1.base_layer.input_quantizer._amax": torch.tensor([1.0]), + "layer1.base_layer.weight_quantizer._scale": torch.ones(4, 1), + "layer1.base_layer.weight_quantizer._double_scale": torch.tensor([0.5]), + } + + processed_state_dict = postprocess_state_dict( + state_dict, 448.0, QUANTIZATION_NONE, is_modelopt_qlora=True + ) + + assert set(processed_state_dict) == { + "layer1.weight", + "layer1.weight_scale", + "layer1.weight_scale_2", + "layer1.input_scale", + "layer1.bias", + "layer1.pre_quant_scale", + } + assert torch.equal(processed_state_dict["layer1.weight_scale_2"], torch.tensor([0.5])) + assert torch.equal(processed_state_dict["layer1.bias"], torch.arange(4.0)) + + @pytest.mark.parametrize( ("config", "expected"), [ diff --git a/tests/unit/torch/opt/plugins/test_hf_patching.py b/tests/unit/torch/opt/plugins/test_hf_patching.py index 8a44ad23c76..0476f44c56c 100644 --- a/tests/unit/torch/opt/plugins/test_hf_patching.py +++ b/tests/unit/torch/opt/plugins/test_hf_patching.py @@ -14,6 +14,8 @@ # limitations under the License. import pytest +import torch +import torch.nn as nn from _test_utils.torch.transformers_models import ( create_tiny_llama_dir, get_tiny_qwen3, @@ -23,6 +25,9 @@ import modelopt.torch.distill as mtd import modelopt.torch.opt as mto +import modelopt.torch.quantization as mtq +from modelopt.torch.opt.plugins.transformers import _restore_qtensor_wrappers +from modelopt.torch.quantization.qtensor import QTensorWrapper @pytest.mark.parametrize( @@ -54,3 +59,61 @@ def test_nested_model_save_restore(tmp_path, model_cls, teacher_model_type): tf_output_tester(model, model_test) # KD state is not saved and it should be empty assert not mto.ModeloptStateManager(model_test).has_state + + +class _LoraLike(nn.Module): + """Stand-in for peft's `lora.Linear`, which nests the original module under `base_layer`.""" + + def __init__(self, base_layer): + super().__init__() + self.base_layer = base_layer + + +def _compressed_model_and_state_dir(tmp_path, state_keyed_with_base_layer=False): + model = nn.Sequential() + model.fc = nn.Linear(64, 32) + mtq.quantize(model, mtq.NVFP4_DEFAULT_CFG, lambda m: m(torch.randn(2, 64))) + mtq.compress(model) + assert isinstance(model.fc.weight, QTensorWrapper) + + state = mto.modelopt_state(model) + if state_keyed_with_base_layer: + # Compressing after the adapters are attached saves the keys with the peft suffix. + for _, mode_config in state["modelopt_state_dict"]: + q_tensor_state = mode_config.get("metadata", {}).get("q_tensor_state", {}) + for key in list(q_tensor_state): + q_tensor_state[f"{key}.base_layer"] = q_tensor_state.pop(key) + torch.save(state, tmp_path / "modelopt_state.pth") + + # transformers>=5 assigns a plain Parameter holding the packed data, dropping the wrapper. + packed = model.fc.weight.data.clone() + del model.fc._parameters["weight"] + model.fc._parameters["weight"] = nn.Parameter(packed, requires_grad=False) + assert not isinstance(model.fc.weight, QTensorWrapper) + return model + + +@pytest.mark.parametrize("wrap_in_lora", [False, True]) +@pytest.mark.parametrize("state_keyed_with_base_layer", [False, True]) +def test_restore_qtensor_wrappers(tmp_path, wrap_in_lora, state_keyed_with_base_layer): + """Either side may carry the `.base_layer` suffix, so the lookup must work in both directions.""" + model = _compressed_model_and_state_dir(tmp_path, state_keyed_with_base_layer) + if wrap_in_lora: + model.fc = _LoraLike(model.fc) + + _restore_qtensor_wrappers(model, str(tmp_path)) + + linear = model.fc.base_layer if wrap_in_lora else model.fc + assert isinstance(linear.weight, QTensorWrapper) + assert linear.weight.metadata["shape"] == torch.Size([32, 64]) + + +def test_restore_qtensor_wrappers_warns_when_nothing_matches(tmp_path): + """A total miss must be loud -- it otherwise surfaces as an opaque shape error at dequant.""" + model = _compressed_model_and_state_dir(tmp_path) + model.fc = _LoraLike(_LoraLike(model.fc)) # a nesting the lookup does not know about + + with pytest.warns(UserWarning, match="re-wrapped none"): + _restore_qtensor_wrappers(model, str(tmp_path)) + + assert not isinstance(model.fc.base_layer.base_layer.weight, QTensorWrapper)