diff --git a/CHANGELOG.rst b/CHANGELOG.rst index b51831f402f..eff4cea75bf 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -17,7 +17,7 @@ Changelog - **Deduplicate the modules shared at source** in the quantized export step: ``_export_quantized_weight`` and ``_export_fused_experts`` now alias bit-identical packed ``weight`` / ``weight_scale`` / ``weight_scale_2`` buffers across modules sharing a source weight ``data_ptr()`` so the downstream ``postprocess_state_dict`` dedup catches them (~42% storage reduction on ``nvfp4_experts_only`` for tied 26B MoE checkpoints). - New ``sync_tied_input_amax`` helper max-merges per-side ``input_quantizer.amax`` across tied modules before export so single-backbone consumers that load one ``input_scale`` per parameter don't clip either side. - The exported state_dict is also **reordered (decoder keys win instead of encoder)** so canonical-side keys per HF's ``_tied_weights_keys`` declaration win the data_ptr dedup; gated to the DiffusionGemma model class in ``_reorder_canonical_first``, no-op for every other model. - - New DiffusionGemma model-specific recipe under ``modelopt_recipes/huggingface/diffusion_gemma/ptq/`` (``nvfp4_experts_only.yaml`` + its ``disabled_quantizers.yaml`` unit) adds the ``*self_conditioning*`` exclude on top of the standard default, leaving the shared ``default_disabled_quantizers`` unit clean for non-diffusion models — pattern matches the existing ``phi4mm`` / ``nemotron_vl`` model-specific recipes. + - New DiffusionGemma model-specific recipe under ``modelopt_recipes/huggingface/diffusion_gemma/ptq/`` (``nvfp4_experts_only.yaml`` + its ``disabled_quantizers.yaml`` unit) adds the ``*self_conditioning*`` exclude on top of the standard default, leaving the shared ``default_disabled_quantizers`` unit clean for non-diffusion models — pattern matches the existing ``nemotron_vl`` model-specific recipes. - ``hf_ptq.py`` also unwraps ``ModelOutput`` dataclasses from ``.generate()`` so the preview decode works on diffusion models. Non-tied models see no behavioral change. - Add Torch-TensorRT FP8 deployment example for HuggingFace ViT (``examples/torch_trt/``): ``torch_tensorrt_ptq.py`` covers ``mtq.quantize`` → ``torch_tensorrt.compile(ir="dynamo")``, and ``torch_tensorrt_accuracy.py`` reports the compiled model's ImageNet-1k top-1/top-5 accuracy via the ``onnx_ptq`` ``evaluate`` harness (the unquantized baseline is Torch-TensorRT-compiled too, for an apples-to-apples comparison). Ships a ViT-tuned FP8 PTQ recipe under ``modelopt_recipes/huggingface/vit/ptq/`` (``fp8.yaml``) composed from the shared ``modelopt_recipes/configs/`` units: it quantizes the encoder Linears, patch-embed ``nn.Conv2d``, ``classifier``, and per-block LayerNorm inputs plus the attention Q/K/V BMMs and softmax. Verified on ``google/vit-base-patch16-224`` (ImageNet-1k 50k validation): FP8 stays within 0.13 pp Top-1 of the FP16 baseline. - Add **AutoQuantize recipe** support: ``mtq.auto_quantize`` can be driven declaratively from a YAML recipe (``RecipeType.AUTO_QUANTIZE`` / ``AutoQuantizeConfig``) specifying candidate formats, the ``effective_bits`` target, cost model (incl. ``active_moe`` and ``excluded_module_name_patterns``), scoring method, and disabled layers. Adds an ``effective_bits`` cost-model override on ``QuantizeConfig`` / ``QuantizerAttributeConfig`` (block-scale-accurate NVFP4 = 4.5 via ``configs/numerics/nvfp4``). Shipped recipes live under ``modelopt_recipes/general/auto_quantize/`` and model-specific ones under ``modelopt_recipes/huggingface//auto_quantize/``. @@ -60,6 +60,7 @@ Changelog - Remove the ``examples/diffusers/eval`` image-quality evaluation example (ImageReward / CLIP-IQA / CLIP metrics) and its references in ``examples/diffusers/README.md``. The example was deprecated in 0.45 and is no longer maintained. - Remove the deprecated ``examples/llm_autodeploy`` example (deprecated in 0.45). Use TensorRT-LLM's `AutoDeploy `_ directly together with ModelOpt PTQ in ``examples/llm_ptq``. - Remove the deprecated ``examples/llm_qad`` Megatron-LM QAD example (deprecated in 0.45). Use the `megatron_bridge QAD example `_ instead, which provides a simpler Python-based interface and better model coverage. +- Dropped **Phi-4-multimodal** PTQ support in ``examples/hf_ptq`` (NVBug 6563509). Its bundled remote code predates Transformers v5 and no longer loads on any version ModelOpt supports (``transformers>=4.57``): it requires ``transformers<4.52`` because it reaches ``prepare_inputs_for_generation`` through ``peft``, which needs ``PreTrainedModel`` to still inherit ``GenerationMixin``, and it declares ``_tied_weights_keys`` as a list, which Transformers 5.x rejects. **Phi-3-vision** is dropped alongside it: it is the older, superseded model in the same family, so with its successor unsupportable there is no reason to keep carrying the predecessor. (Phi-3-vision shares the list-valued ``_tied_weights_keys`` defect and so is likewise broken on Transformers 5.x, though it does not hit the ``peft`` blocker.) The support-matrix row, the ``phi4mm`` model type, the multimodal-detection heuristics that only ever matched these two (``vision_lora`` / ``audio_processor`` / ``embd_layer.image_embd_layer``), the ``Phi3Image`` / ``PhiImage`` embedding-export exclusions, and the ``modelopt_recipes/huggingface/phi4mm/`` recipes have been removed. Text-only Phi-3/Phi-4 and Phi-3.5-MoE are unaffected. **Deprecations** @@ -73,6 +74,7 @@ Changelog **Bug Fixes** +- Fix NemotronH dense MLP quantization with the ``nvfp4_mlp_only`` and ``nvfp4_omlp_only`` recipe families. NemotronH registers these projections as ``mixer.up_proj`` and ``mixer.down_proj``, which were missed by the previous ``*mlp*`` selector and produced checkpoints with a null ``quant_algo`` (NVBug 6571812). - Fix ``ShapeInferenceError`` during ONNX INT8 + FP16 quantization (``--high_precision_dtype fp16``) of weakly-typed models (e.g. TensorFlow exports) that carry stale rank-0 ``graph.output`` shapes or ops such as ``TopK`` that ONNX's static shape inference cannot resolve. ``clear_stale_value_info`` now reconciles stale output shapes via symbolic shape inference (keeping every output's shape field populated), and AutoCast runs ONNX shape inference in strict mode and falls back to schema-based standalone type inference when it fails, so unresolved ops no longer leave tensors untyped. - Fused MoE expert auto-detection (``register_fused_experts_on_the_fly``) no longer requires an ``act_fn`` attribute. Some fused-expert modules (e.g. ``MiniMaxM3VLExperts``) apply a custom gated activation between the two ``F.linear`` calls instead of exposing ``act_fn``; they were silently skipped, leaving routed experts unquantized (an experts-only recipe matched nothing) and failing HF export with ``NotImplementedError``. ``_QuantFusedExperts`` is activation-agnostic (it only intercepts the two ``F.linear`` calls), so the requirement was unnecessary. This enables NVFP4/FP8 quantization and export for MiniMax-M2 / MiniMax-M3. - Fix unified HF export emitting transformers' *in-memory* (post-``conversion_mapping``) tensor names instead of the original model-hub names, breaking the unified-checkpoint contract (observed on MiniMax-M3: exported ``model.language_model.*`` / ``mlp.experts.*.gate_proj`` instead of hub ``language_model.model.*`` / ``block_sparse_moe.experts.*.w{1,2,3}``). transformers' own save-side ``revert_weight_conversion`` is disabled by ModelOpt because it raises ``RuntimeError`` on 0-d scalar scale tensors, so a new quant-aware reverse conversion (``modelopt/torch/export/quant_aware_conversion.py``) derives rename/split rules from the model's conversion mapping via transformers' ``reverse_transform()`` and carries each weight's companion scale tensors (``weight_scale``, ``weight_scale_2``, ``input_scale``, ``weight_scale_inv``, ``bias``) through the renames and un-fusions, so quantized exports round-trip to the hub names. Any mapping op that cannot be reversed quant-aware yet (e.g. still-stacked fused experts) falls back to the previous in-memory names instead of aborting the export. @@ -83,6 +85,9 @@ 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 HF checkpoint export failing with ``AttributeError: 'list' object has no attribute 'keys'`` for models whose modeling code still declares tied weights in the ``transformers<5`` list format (NVBug 6518665, observed on ``stepfun-ai/Step-3.7-Flash``). transformers 5.0 changed ``_tied_weights_keys`` to a ``{target: source}`` dict and ``save_pretrained`` calls ``.keys()`` on every submodule's declaration without a type check, so such models — common among ``trust_remote_code`` checkpoints — load fine but die at the end of PTQ, after calibration. ModelOpt's ``save_pretrained`` patch now normalizes a list-style declaration to the equivalent dict for the duration of the save (each entry mapped to itself, which is what the legacy list meant) and restores the original attribute afterwards. +- Fix unified HF export of multimodal models whose vision tower carries its own ``PrefixChange`` conversion (``LlavaForConditionalGeneration`` on ``transformers>=5.12`` — NVBug 6525511). transformers collects conversion mappings recursively and scopes each sub-model's transforms to that sub-module via ``scope_prefix``, matching only keys under that prefix. ModelOpt's quant-aware reverse conversion read the raw patterns and ignored ``scope_prefix``, so the vision tower's "add a ``vision_model.`` prefix" rule was applied to *every* key in the state dict: an exported llava-1.5-13b checkpoint had all 758 tensors moved under a bogus top-level ``vision_model.`` namespace (``vision_model.language_model.*``, ``vision_model.lm_head.*``), and vLLM rejected it with ``ValueError: There is no module or parameter named 'vision_model' in LlavaForConditionalGeneration``. Reverse rename rules now carry their scope and are applied only to keys under it, matching transformers' own ``WeightTransform._scoped_match`` semantics. ``Gemma3ForConditionalGeneration`` was affected identically and is fixed by the same change. +- 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) ^^^^^^^^^^^^^^^^^ @@ -101,7 +106,7 @@ Changelog - The PTQ example scripts ``examples/llm_ptq/hf_ptq.py``, ``examples/llm_ptq/multinode_ptq.py`` and ``examples/megatron_bridge/quantize.py`` now derive their ``--qformat`` / ``--kv_cache_qformat`` (``--quant_cfg`` / ``--kv_cache_quant`` for Megatron-Bridge) CLI vocabularies by discovering the YAML presets under ``modelopt_recipes/configs/ptq/presets/{model,kv}/`` rather than carrying hardcoded ``QUANT_CFG_CHOICES`` / ``KV_QUANT_CFG_CHOICES`` tables. The discovery helper, alias table and ready-built ``QUANT_CFG_CHOICES`` / ``KV_QUANT_CFG_CHOICES`` mappings now live in ``modelopt.recipe.presets`` and are shared by all three scripts. Presets are loaded eagerly into a plain dict at import. Adding a new preset YAML makes it available on the CLI of all three with no script change — note this means each script now accepts every preset under those directories, not just a previously curated subset. All previously-supported short names (``int8_sq``, ``nvfp4_awq``, ``fp8_pb_wo``, ``nvfp4_mse``, ``w4a8_awq``, ``nvfp4_local_hessian``, ``fp8_pc_pt``, ``int8_wo``) keep working via a small deprecation alias table; new formats should be exposed as preset YAMLs (or, longer term, as full ``--recipe`` recipes). - Add ``configs/ptq/presets/kv/fp8_cast.yaml`` and ``configs/ptq/presets/kv/nvfp4_cast.yaml``, promoting ``fp8_cast`` / ``nvfp4_cast`` to first-class KV presets composed from the existing ``kv_fp8_cast`` / ``kv_nvfp4_cast`` unit fragments. The previous runtime ``use_constant_amax`` post-edit in ``hf_ptq.py`` is removed; ``use_constant_amax: true`` now lives in the YAML and is therefore authoritative. **Custom (out-of-tree) recipes that target a cast KV format must set ``use_constant_amax: true`` themselves on the ``[kv]_bmm_quantizer`` config** — in-tree recipes already do via the ``kv_*_cast`` units. - Add FP8 KV-cache cast variants for the partial-NVFP4 and weight-only general PTQ recipes: ``general/ptq/nvfp4_mlp_only-kv_fp8_cast``, ``general/ptq/nvfp4_experts_only-kv_fp8_cast``, ``general/ptq/nvfp4_omlp_only-kv_fp8_cast``, and ``general/ptq/nvfp4_weight_only-kv_fp8_cast``. These compose the same model-quant configs as their ``-kv_fp8`` siblings with the ``kv_fp8_cast`` unit (constant-amax FP8 KV cache, no KV calibration forward pass). -- Add Nemotron-3-Super-120B-A12B PTQ recipes ``modelopt_recipes/huggingface/models/nvidia/Nemotron-3-Super-120B-A12B/ptq/nvfp4-mse.yaml`` (MSE-mixed) and ``nvfp4-max-calib.yaml`` (max-calib mixed): NVFP4 W4A4 routed experts + FP8 per-tensor shared experts / Mamba in/out_proj + FP8 KV cache. +- Add Nemotron-3-Super-120B-A12B PTQ recipes ``modelopt_recipes/huggingface/models/nvidia/Nemotron-3-Super-120B-A12B-BF16/ptq/nvfp4-mse.yaml`` (MSE-mixed) and ``nvfp4-max-calib.yaml`` (max-calib mixed): NVFP4 W4A4 routed experts + FP8 per-tensor shared experts / Mamba in/out_proj + FP8 KV cache. - Group layerwise calibration options under a nested ``LayerwiseConfig`` and add two knobs: ``get_qdq_activations_from_prev_layer`` (correct GPTQ-Hessian vs max-calib activation semantics — defaults to True for GPTQ, False for max/mse/local_hessian) and ``save_every`` (gate per-window ``next_inputs.pt`` activation-cache writes). Legacy bool ``layerwise`` and flat ``layerwise_checkpoint_dir`` keys still work; the bool form emits a ``DeprecationWarning``. - Add two layerwise-calibration memory optimizations: ``calib_mutates_weights`` (set False for amax-only algorithms — max/mse/local_hessian — to skip the per-layer weight checkpoint blob and in-memory writeback, persisting only quantizer state), and meta-device skip-layer placeholders (already-calibrated layers emit zero-filled ``meta`` tensors instead of real-device buffers, eliminating their activation memory — models with real-device inter-layer ops on the hidden state are unsupported). - Add ``examples/alpamayo`` showing FP8, NVFP4, and AutoQuantize (mixed-precision) quantization of the Alpamayo (formerly Alpamayo-R1) ~10B vision-language-action model, with a joint VLM + diffusion calibration loop and both fake-quant and ``--real-quant`` packed-checkpoint export. See `examples/alpamayo/README.md `_ for details. diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index a3967565ae0..8d6be9c3845 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -118,7 +118,6 @@ Please reference our [framework scripts](#framework-scripts) and our [docs](http | Whisper9 | ✅ | ❌ | ❌ | ❌ | - | | Nemotron-3 | ✅ | ❌ | ❌ | ❌ | ✅ | | Llava (VLM)11 | ✅ | ✅12 | ✅ | ✅ | - | -| Phi-3-vision, Phi-4-multimodal (VLM)11 | ✅ | ✅12 | ✅ | ✅ | ✅ | | Qwen2, 2.5-VL (VLM)11 | ✅ | ✅12 | ✅ | ✅ | ✅ | | Gemma 3 (VLM)11 | ✅ | - | - | - | - | | Nemotron VL (VLM)11,13 | ✅ | - | - | - | ✅ | diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index d1031dd6084..b8eef827d95 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -168,12 +168,6 @@ def _is_multimodal_config(config): """Check if a config indicates a multimodal model (config-only version of is_multimodal_model).""" return ( hasattr(config, "vision_config") # Standard vision config (e.g., Qwen2.5-VL) - or getattr(config, "model_type", "") == "phi4mm" # Phi-4 multimodal - or hasattr(config, "vision_lora") # Vision LoRA configurations - or hasattr(config, "audio_processor") # Audio processing capabilities - or ( - hasattr(config, "embd_layer") and hasattr(config.embd_layer, "image_embd_layer") - ) # Image embedding layers or getattr(config, "is_encoder_decoder", False) # Encoder-decoder VL models or any( # Architecture-based detection for custom VL models (e.g., Nemotron-Parse) "conditionalgeneration" in arch.lower() for arch in getattr(config, "architectures", []) diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 436c8867227..23265a47e70 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -699,9 +699,6 @@ def load_model(args: argparse.Namespace): # Left padding usually provides better calibration result. tokenizer.padding_side = "left" - if model_type == "phi4mm": - warnings.warn("Please set the default input_mode to InputMode.LANGUAGE before quantizing.") - return ( full_model, language_model, 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/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md index 2fae7fe3545..030edaf8fd5 100644 --- a/examples/megatron_bridge/README.md +++ b/examples/megatron_bridge/README.md @@ -365,9 +365,6 @@ torchrun --nproc_per_node 1 prune_minitron.py --help > [!NOTE] > Multi-token-prediction (MTP) heads (e.g. Qwen3.5) are not pruned yet — they are dropped for the prune run and the saved checkpoint has no MTP. Autoregressive inference is unaffected; for speculative decoding, run a short MTP SFT on the pruned model. -> [!NOTE] -> If pruning a Nemotron model and you want to save the pruned model back in HF format, please downgrade to `transformers<5` via `python -m pip install "transformers<5"` before pruning. - ### Vision-Language Models (VLMs) For a vision-language model (e.g. Qwen3.5-VL, Gemma3-VL), `prune_minitron.py` automatically prunes only the **language model** and leaves the vision tower intact, then saves the full VLM back. All the pruning modes above (parameter count, active parameter count, memory footprint, and manual `export_config`) work unchanged, with two VLM-specific caveats: diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index 16dd37d5f8d..b35369c40f3 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -480,5 +480,7 @@ def _restore_student_hook(model_chunks): args = get_args() try: main(args) + except BaseException: + dist.abort() # peers may be stuck in a collective this rank will never reach finally: dist.cleanup() diff --git a/examples/megatron_bridge/export_distilled_megatron_to_hf.py b/examples/megatron_bridge/export_distilled_megatron_to_hf.py index 1ba76ea7ae4..d5e95ff9d18 100644 --- a/examples/megatron_bridge/export_distilled_megatron_to_hf.py +++ b/examples/megatron_bridge/export_distilled_megatron_to_hf.py @@ -288,5 +288,7 @@ def main(args: argparse.Namespace): args = get_args() try: main(args) + except BaseException: + dist.abort() # peers may be stuck in a collective this rank will never reach finally: dist.cleanup() diff --git a/examples/megatron_bridge/export_quantized_megatron_to_hf.py b/examples/megatron_bridge/export_quantized_megatron_to_hf.py index 17db5e6da34..e4e3703d8a5 100644 --- a/examples/megatron_bridge/export_quantized_megatron_to_hf.py +++ b/examples/megatron_bridge/export_quantized_megatron_to_hf.py @@ -164,5 +164,7 @@ def main(args: argparse.Namespace): args = get_args() try: main(args) + except BaseException: + dist.abort() # peers may be stuck in a collective this rank will never reach finally: dist.cleanup() diff --git a/examples/megatron_bridge/prune_minitron.py b/examples/megatron_bridge/prune_minitron.py index 6f2b7c3829f..7a9b4b8bd56 100644 --- a/examples/megatron_bridge/prune_minitron.py +++ b/examples/megatron_bridge/prune_minitron.py @@ -105,6 +105,24 @@ def _hf_config_has_mtp(hf_cfg) -> bool: ) +# HF names the shared expert size with or without the ``moe_`` prefix depending on the model +# (e.g. Qwen3.5-MoE uses ``shared_expert_intermediate_size``). +_SHARED_EXPERT_SIZE_FIELDS = ( + "moe_shared_expert_intermediate_size", + "shared_expert_intermediate_size", +) + + +def _is_deepseek_style_moe(text_cfg) -> bool: + """Whether the shared expert is sized as ``n_shared_experts * moe_intermediate_size``. + + Such configs can only represent a shared expert size that is a multiple of the routed one. + """ + return getattr(text_cfg, "n_shared_experts", None) is not None and not any( + hasattr(text_cfg, field) for field in _SHARED_EXPERT_SIZE_FIELDS + ) + + def get_args() -> argparse.Namespace: parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) parser.add_argument("--hf_model_name_or_path", type=str, required=True) @@ -400,7 +418,9 @@ def main(args: argparse.Namespace): "num_layers_in_last_pipeline_stage": args.num_layers_in_last_pipeline_stage, "pipeline_dtype": torch.bfloat16, "seq_length": args.seq_length, - "mtp_num_layers": 0, # MTP is not supported during calibration + # MTP is not supported during calibration; drop it + "mtp_num_layers": 0, + "mtp_hybrid_override_pattern": None, }, init_model_parallel=True, moe_grouped_gemm=not args.no_moe_grouped_gemm, @@ -542,6 +562,17 @@ def score_func(m): pruning_config["max_width_pruning"] = args.max_width_pruning pruning_config["max_depth_pruning"] = args.max_depth_pruning pruning_config["hparams_to_skip"] = args.hparams_to_skip + # DeepSeek-style MoE configs size the shared expert as n_shared_experts * moe_intermediate_size, + # so only candidates whose shared size is a multiple of the routed one can be saved to HF. + src_hf_cfg = bridge.hf_pretrained.config + if _is_deepseek_style_moe(getattr(src_hf_cfg, "text_config", src_hf_cfg)): + warn_rank_0( + "DeepSeek-style MoE config detected: restricting the search to candidates whose " + "moe_shared_expert_intermediate_size is a multiple of moe_ffn_hidden_size." + ) + pruning_config["candidate_filter"] = lambda cfg: ( + cfg["moe_shared_expert_intermediate_size"] % cfg["moe_ffn_hidden_size"] == 0 + ) pruning_config["top_k"] = args.top_k # memory_mb constraint requires batch_size and seq_length pruning_config["batch_size"] = args.inference_batch_size @@ -588,10 +619,11 @@ def score_func(m): else: print_rank_0(f"Saving pruned model to {args.output_hf_path} in HF checkpoint format") - # [WAR] Save the pruned HF model by hand until Megatron-Bridge natively supports it. - # TODO: Replace this whole block with ``AutoBridge.from_auto_config(...).save_hf_weights(...)`` - # once the Megatron-Bridge fix ships (nemo:26.08). - bridge.hf_pretrained.save_artifacts(args.output_hf_path) + # Build the pruned HF config field-by-field from the pruned Megatron config, then stream weights. + # Rank 0 only: a late write from another rank would leave config.json stale. + if dist.is_master(): + bridge.hf_pretrained.save_artifacts(args.output_hf_path) + dist.barrier() hf_cfg = AutoConfig.from_pretrained( args.output_hf_path, trust_remote_code=args.trust_remote_code ) @@ -610,12 +642,7 @@ def score_func(m): text_cfg.mamba_head_dim = mcore_cfg.mamba_head_dim if hasattr(text_cfg, "moe_intermediate_size"): text_cfg.moe_intermediate_size = mcore_cfg.moe_ffn_hidden_size - # HF names this field with or without the ``moe_`` prefix depending on the model - # (e.g. Qwen3.5-MoE uses ``shared_expert_intermediate_size``). - for shared_expert_field in ( - "moe_shared_expert_intermediate_size", - "shared_expert_intermediate_size", - ): + for shared_expert_field in _SHARED_EXPERT_SIZE_FIELDS: if hasattr(text_cfg, shared_expert_field): setattr( text_cfg, shared_expert_field, mcore_cfg.moe_shared_expert_intermediate_size @@ -624,7 +651,16 @@ def score_func(m): text_cfg.num_experts = mcore_cfg.num_moe_experts if hasattr(text_cfg, "n_routed_experts"): text_cfg.n_routed_experts = mcore_cfg.num_moe_experts - if hasattr(text_cfg, "n_shared_experts"): + # n_shared_experts is a fixed count; only DeepSeek-style configs record the pruned shared + # expert size through it. candidate_filter keeps the search divisible, so only a manual + # --prune_export_config can violate this. + if _is_deepseek_style_moe(text_cfg): + if mcore_cfg.moe_shared_expert_intermediate_size % mcore_cfg.moe_ffn_hidden_size: + raise ValueError( + f"{mcore_cfg.moe_shared_expert_intermediate_size=} must be a multiple of " + f"{mcore_cfg.moe_ffn_hidden_size=} for this config, which stores the shared " + "expert size as n_shared_experts * moe_intermediate_size. " + ) text_cfg.n_shared_experts = ( mcore_cfg.moe_shared_expert_intermediate_size // mcore_cfg.moe_ffn_hidden_size ) @@ -658,8 +694,11 @@ def score_func(m): "distillation cannot recover this vision-path change -- consider full VLM " "training/distillation instead of LM-only to recover vision quality." ) - if isinstance(provider, _HYBRID_PROVIDER_TYPES) and hasattr( - text_cfg, "hybrid_override_pattern" + # Only older remote-code configs need this; native configs carry the cadence in layer_types. + if ( + isinstance(provider, _HYBRID_PROVIDER_TYPES) + and not hasattr(text_cfg, "layer_types") + and hasattr(text_cfg, "hybrid_override_pattern") ): # MCore's pattern can carry an MTP suffix (``/...``) and PP boundaries (``|``) which we need to remove text_cfg.hybrid_override_pattern = "".join( @@ -671,15 +710,36 @@ def score_func(m): if hasattr(text_cfg, field): setattr(text_cfg, field, 0) - # Save dummy pruned HF model to get the correct bridge for saving pruned weights - dummy_model_cls = AutoModelForImageTextToText if is_vlm else AutoModelForCausalLM - dummy_model_cls.from_config( - hf_cfg, trust_remote_code=args.trust_remote_code - ).save_pretrained(args.output_hf_path, trust_remote_code=args.trust_remote_code) - pruned_bridge = AutoBridge.from_hf_pretrained( - args.output_hf_path, trust_remote_code=args.trust_remote_code + # Config-only bridge (hf_keys=None) keeps the embedding task when transformers' saved key + # differs from the bridge mapping (NemotronH's backbone.embedding vs ...embeddings). + use_config_only_export = ( + hasattr(AutoBridge, "from_hf_config") + and isinstance(provider, _HYBRID_PROVIDER_TYPES) + and not is_vlm ) - pruned_bridge.save_hf_weights(model, args.output_hf_path) + if use_config_only_export: + pruned_bridge = AutoBridge.from_hf_config(hf_cfg) + # save_hf_pretrained reads trust_remote_code off the bridge to fetch source artifacts; + # from_hf_config can't infer it since AutoConfig consumes the kwarg. + pruned_bridge.trust_remote_code = args.trust_remote_code + pruned_bridge.save_hf_pretrained( + model, args.output_hf_path, source_path=args.hf_model_name_or_path + ) + else: + if isinstance(provider, _HYBRID_PROVIDER_TYPES) and not is_vlm: + warn_rank_0( + "Megatron-Bridge lacks config-only HF export; falling back to the dummy-model " + "path, which cannot round-trip a pruned native NemotronH config. Use " + "transformers<5 or a newer Megatron-Bridge if the save fails." + ) + dummy_model_cls = AutoModelForImageTextToText if is_vlm else AutoModelForCausalLM + dummy_model_cls.from_config( + hf_cfg, trust_remote_code=args.trust_remote_code + ).save_pretrained(args.output_hf_path, trust_remote_code=args.trust_remote_code) + pruned_bridge = AutoBridge.from_hf_pretrained( + args.output_hf_path, trust_remote_code=args.trust_remote_code + ) + pruned_bridge.save_hf_weights(model, args.output_hf_path) copy_hf_ckpt_remote_code(args.hf_model_name_or_path, args.output_hf_path) print_rank_0(f"Saved pruned model to {args.output_hf_path} in HF checkpoint format") @@ -704,5 +764,7 @@ def score_func(m): args = get_args() try: main(args) + except BaseException: + dist.abort() # peers may be stuck in a collective this rank will never reach finally: dist.cleanup() diff --git a/examples/megatron_bridge/quantize.py b/examples/megatron_bridge/quantize.py index 6355e60e435..1e4ee79f574 100644 --- a/examples/megatron_bridge/quantize.py +++ b/examples/megatron_bridge/quantize.py @@ -454,5 +454,7 @@ def forward_loop(_model=None): args = get_args() try: main(args) + except BaseException: + dist.abort() # peers may be stuck in a collective this rank will never reach finally: dist.cleanup() diff --git a/examples/speculative_decoding/eagle_utils.py b/examples/speculative_decoding/eagle_utils.py index b12b9da1a52..68c6db45235 100644 --- a/examples/speculative_decoding/eagle_utils.py +++ b/examples/speculative_decoding/eagle_utils.py @@ -141,6 +141,9 @@ def make_speculative_data_module( train_len=train_len, local_image_path=data_args.vlm_img_dir, return_labels=True, + answer_only_loss=answer_only_loss, + shift_labels=shift_labels, + chat_template=chat_template, ) else: diff --git a/examples/vllm_serve/Dockerfile b/examples/vllm_serve/Dockerfile index 7213c6fc430..3406c62ba71 100644 --- a/examples/vllm_serve/Dockerfile +++ b/examples/vllm_serve/Dockerfile @@ -1,4 +1,4 @@ -FROM vllm/vllm-openai:v0.20.0 +FROM vllm/vllm-openai:v0.26.0 # Set environment variables ENV PIP_NO_CACHE_DIR=off \ @@ -25,12 +25,16 @@ RUN cd Model-Optimizer && \ # Llama4 requires this RUN pip install flash-attn==2.7.4.post1 --no-build-isolation -# Pre-compile CUDA extensions to avoid compilation time during runtime +# Pre-compile CUDA extensions into a world-accessible directory so the vllm +# user can use the cache at runtime. +ENV TORCH_EXTENSIONS_DIR=/workspace/torch_extensions RUN python3 -c "import modelopt.torch.quantization.extensions as ext; ext.precompile()" || true -# Allow users to run without root +# Allow the non-root vllm user to access the workspace RUN chmod -R 777 /workspace +USER vllm + # Override the ENTRYPOINT from the base image to allow flexible usage ENTRYPOINT [] diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index 858243686d0..75bcf37089d 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -4,7 +4,7 @@ This is a simple example to demonstrate calibrating and serving ModelOpt fakequa Compared with realquant, fakequant is 2-5x slower, but doesn't require dedicated kernel support and facilitates research. -The general fakequant example is tested with vLLM 0.9.0 and 0.19.1. The compact +The general fakequant example is tested with vLLM 0.9.0, 0.19.1, and 0.26.0. The compact NVFP4 attention worker documented below requires vLLM 0.15.0 or newer. ## Prepare environment diff --git a/examples/vllm_serve/vllm_ptq_utils.py b/examples/vllm_serve/vllm_ptq_utils.py index 88b31d54a70..709d6532fb3 100644 --- a/examples/vllm_serve/vllm_ptq_utils.py +++ b/examples/vllm_serve/vllm_ptq_utils.py @@ -14,6 +14,7 @@ # limitations under the License. import dataclasses +import warnings from collections.abc import Callable from typing import Any @@ -66,6 +67,7 @@ def calibrate_loop(model: Any) -> None: NewRequestData, req_id=req_id, prompt_token_ids=input_ids_list, + prefill_token_ids=input_ids_list, mm_kwargs=[], mm_hashes=[], mm_positions=[], @@ -95,10 +97,39 @@ def calibrate_loop(model: Any) -> None: structured_output_request_ids={}, grammar_bitmask=None, ) - output = self.execute_model(scheduler_output) - if hasattr(self, "sample_tokens"): - if output is None: # TODO: make this default when vllm <= 0.11 is outdated - self.sample_tokens(None) + try: + output = self.execute_model(scheduler_output) + if hasattr(self, "sample_tokens"): + if output is None: # TODO: make this default when vllm <= 0.11 is outdated + self.sample_tokens(None) + finally: + # finish_requests runs before add_requests inside execute_model, so + # req IDs aren't registered yet at that point — call it directly after. + # Wrap in try/except so a cleanup error never masks the original exception. + try: + if hasattr(self.model_runner, "finish_requests"): + cleanup_output = _create_new_data_cls( + type(scheduler_output), + scheduled_new_reqs=[], + scheduled_cached_reqs=scheduler_output.scheduled_cached_reqs, + num_scheduled_tokens={}, + total_num_scheduled_tokens=0, + scheduled_spec_decode_tokens={}, + scheduled_encoder_inputs={}, + num_common_prefix_blocks=scheduler_output.num_common_prefix_blocks, + finished_req_ids=set(num_scheduled_tokens.keys()), + free_encoder_mm_hashes=[], + kv_connector_metadata=None, + structured_output_request_ids={}, + grammar_bitmask=None, + ) + self.model_runner.finish_requests(cleanup_output) + else: + warnings.warn( + "model_runner.finish_requests not found; request state may leak during calibration." + ) + except Exception: + warnings.warn("Failed to clean up request state after calibration batch.") return calibrate_loop diff --git a/modelopt/torch/export/layer_utils.py b/modelopt/torch/export/layer_utils.py index d5f1fb2330d..de136fcd378 100755 --- a/modelopt/torch/export/layer_utils.py +++ b/modelopt/torch/export/layer_utils.py @@ -222,12 +222,7 @@ def is_conv(module: nn.Module) -> bool: def is_embedding(module: nn.Module) -> bool: """Returns whether the module is an embedding layer.""" module_type_name = type(module).__name__ - return ( - "Embedding" in module_type_name - and "Rotary" not in module_type_name - and "PhiImage" not in module_type_name - and "Phi3Image" not in module_type_name - ) + return "Embedding" in module_type_name and "Rotary" not in module_type_name def build_embedding_config(module: nn.Module, normalization_constant: float = 1) -> EmbeddingConfig: diff --git a/modelopt/torch/export/model_utils.py b/modelopt/torch/export/model_utils.py index 307ea9aac51..1729dbfffcf 100755 --- a/modelopt/torch/export/model_utils.py +++ b/modelopt/torch/export/model_utils.py @@ -44,7 +44,6 @@ "phi3small": "phi3small", "phi3": "phi3", "PhiMoEForCausalLM": "phi3", - "Phi4MMForCausalLM": "phi4mm", "phi": "phi", "TLGv4ForCausalLM": "phi", "MixtralForCausalLM": "llama", @@ -88,10 +87,6 @@ def is_multimodal_model(model): This function detects various multimodal model architectures by checking for: - Standard vision configurations (vision_config) - Language model attributes (language_model) - - Specific multimodal model types (phi4mm) - - Vision LoRA configurations - - Audio processing capabilities - - Image embedding layers - Nemotron-Parse conditional generation models Args: @@ -104,10 +99,6 @@ def is_multimodal_model(model): >>> model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct") >>> is_multimodal_model(model) True - - >>> model = AutoModelForCausalLM.from_pretrained("microsoft/Phi-4-multimodal-instruct") - >>> is_multimodal_model(model) - True """ config = model.config @@ -118,12 +109,6 @@ def is_multimodal_model(model): return ( hasattr(config, "vision_config") # Standard vision config (e.g., Qwen2.5-VL) or hasattr(model, "language_model") # Language model attribute (e.g., LLaVA) - or getattr(config, "model_type", "") == "phi4mm" # Phi-4 multimodal - or hasattr(config, "vision_lora") # Vision LoRA configurations - or hasattr(config, "audio_processor") # Audio processing capabilities - or ( - hasattr(config, "embd_layer") and hasattr(config.embd_layer, "image_embd_layer") - ) # Image embedding layers or is_nemotron_parse # Nemotron-Parse conditional generation model ) diff --git a/modelopt/torch/export/plugins/hf_spec_export.py b/modelopt/torch/export/plugins/hf_spec_export.py index caa93db3634..255b1d9ab04 100644 --- a/modelopt/torch/export/plugins/hf_spec_export.py +++ b/modelopt/torch/export/plugins/hf_spec_export.py @@ -29,6 +29,23 @@ ALL_SPEC_MODES = ["eagle", "dflash"] + +def _get_rope_theta(config, default=None): + """Get RoPE theta from either legacy or Transformers 5 config fields.""" + rope_theta = getattr(config, "rope_theta", None) + if rope_theta is not None: + return rope_theta + + # Transformers 5 stores this under rope_parameters (and exposes the same + # data through rope_scaling for backwards compatibility). + for attr in ("rope_parameters", "rope_scaling"): + rope_config = getattr(config, attr, None) + if isinstance(rope_config, dict) and rope_config.get("rope_theta") is not None: + return rope_config["rope_theta"] + + return default + + LLAMA_EAGLE_SINGLE_LAYER = { "required": { "layers.0.self_attn.q_proj", @@ -376,14 +393,10 @@ def _export_config(self): "initializer_range": getattr(base_config, "initializer_range", 0.02), "attention_bias": getattr(draft_config, "attention_bias", False), "attention_dropout": getattr(draft_config, "attention_dropout", 0.0), - # Inherit the target's rope_theta: DFlash injects the target's KV into every - # draft layer, so the draft's RoPE base must match the target's. (The draft - # arch config carries no rope_theta of its own.) - "rope_theta": ( - getattr(base_config, "rope_theta", None) - if getattr(base_config, "rope_theta", None) is not None - else getattr(draft_config, "rope_theta", 1000000.0) - ), + # Inherit the target's RoPE base: DFlash injects target KV into every draft + # layer, so their RoPE bases must match. Transformers 5 stores rope_theta + # in rope_parameters rather than a top-level config attribute. + "rope_theta": _get_rope_theta(base_config, _get_rope_theta(draft_config, 1000000.0)), # YaRN long-context scaling is injected below (see the rope_scaling block). "rope_scaling": None, "tie_word_embeddings": False, diff --git a/modelopt/torch/export/quant_aware_conversion.py b/modelopt/torch/export/quant_aware_conversion.py index ece6e335ca8..2e32f123869 100644 --- a/modelopt/torch/export/quant_aware_conversion.py +++ b/modelopt/torch/export/quant_aware_conversion.py @@ -94,10 +94,18 @@ class QuantConversionUnsupportedError(Exception): @dataclass(frozen=True) class RenameRule: - """Reverse of a ``WeightRenaming``: ``re.sub(pattern, repl, key)`` on every key.""" + """Reverse of a ``WeightRenaming``: ``re.sub(pattern, repl, key)`` on every key. + + ``scope_prefixes`` mirrors transformers' ``WeightTransform._scoped_match``: a rule + collected from a *sub-model* carries the sub-module path it was scoped to, and its + patterns are written relative to that sub-model's own root. Such a rule must only be + applied to keys under one of these prefixes, with the prefix stripped before the + match and re-attached after. Empty tuple means the rule is unscoped (whole-model). + """ pattern: str repl: str + scope_prefixes: tuple[str, ...] = () @dataclass(frozen=True) @@ -157,6 +165,35 @@ def _apply_split_rule(state_dict: dict[str, torch.Tensor], rule: SplitRule) -> N state_dict[target_key] = _split_leaf_tensor(leaf, tensor, n, idx, rule.dim) +def _compile_rename_rules(rename_rules: list[RenameRule]): + """Pre-compile rename rules into ``(compiled_pattern, repl, scope_prefixes)`` triples.""" + return [(re.compile(r.pattern), r.repl, r.scope_prefixes) for r in rename_rules] + + +def _sub_scoped(pattern: re.Pattern, repl: str, key: str, scope_prefixes: tuple[str, ...]) -> str: + """Apply one rename rule to ``key``, honoring the rule's sub-model scope. + + Mirrors transformers' ``WeightTransform._scoped_match``: for a scoped rule, the first + matching prefix is stripped, the pattern is applied to the remaining suffix, and the + prefix is re-attached. A scoped rule that matches no prefix never applies -- this is + what keeps a sub-model's rule (e.g. the vision tower's "add ``vision_model.``" prefix + change) from rewriting sibling namespaces of the parent multimodal model. + """ + if not scope_prefixes: + return pattern.sub(repl, key) + for prefix in scope_prefixes: + if key.startswith(prefix): + return prefix + pattern.sub(repl, key[len(prefix) :]) + return key + + +def _apply_rename_rules(key: str, compiled) -> str: + """Apply all compiled rename rules to ``key``, in order.""" + for pattern, repl, scope_prefixes in compiled: + key = _sub_scoped(pattern, repl, key, scope_prefixes) + return key + + def apply_reverse_rules( state_dict: dict[str, torch.Tensor], split_rules: list[SplitRule], @@ -171,12 +208,10 @@ def apply_reverse_rules( for rule in split_rules: _apply_split_rule(out, rule) - compiled = [(re.compile(r.pattern), r.repl) for r in rename_rules] + compiled = _compile_rename_rules(rename_rules) renamed: dict[str, torch.Tensor] = {} for key, value in out.items(): - new_key = key - for pattern, repl in compiled: - new_key = pattern.sub(repl, new_key) + new_key = _apply_rename_rules(key, compiled) if new_key in renamed: raise QuantConversionUnsupportedError(f"rename collision on '{new_key}'") renamed[new_key] = value @@ -218,7 +253,7 @@ def build_reverse_name_mapper(model): _, rename_rules, _ = _build_reverse_rules(model) if not rename_rules: return None - compiled = [(re.compile(r.pattern), r.repl) for r in rename_rules] + compiled = _compile_rename_rules(rename_rules) # The rename patterns are anchored on full weight keys and use ``.`` (any char) as a # path separator, so a trailing glob wildcard in an exclude pattern would be consumed # (e.g. ``...mlp.shared_experts.`` -> ``...`` would eat the ``*``). Append a sentinel @@ -227,9 +262,7 @@ def build_reverse_name_mapper(model): _sentinel = ".\x00modelopt_name_sentinel" def _apply(text: str) -> str: - for pattern, repl in compiled: - text = pattern.sub(repl, text) - return text + return _apply_rename_rules(text, compiled) def _map(name: str) -> str: base, suffix = name, "" @@ -287,6 +320,32 @@ def _assert_experts_pre_expanded( ) +def _scope_prefixes(rev) -> tuple[str, ...]: + """Candidate key prefixes a scoped sub-model transform may apply under. + + transformers tags a conversion collected from a sub-model with ``scope_prefix`` (the + sub-module path) and ``base_model_prefix``, then matches keys against + ``base_model_prefix.scope_prefix.`` first and ``scope_prefix.`` second (see + ``WeightTransform._scoped_match``). Returned in that same priority order, each with a + trailing dot. Empty tuple when the transform is unscoped (owned by the root model), + in which case its patterns already address the full key space. + """ + scope = getattr(rev, "scope_prefix", None) + if scope is None: + return () + scope_dot = f"{scope}." if scope != "" else "" + base = getattr(rev, "base_model_prefix", None) or "" + base_dot = f"{base}." if base != "" else "" + # Deduplicate while preserving order. An empty candidate is kept: it only arises for + # ``scope_prefix == ""`` and, matching transformers, acts as the always-matching + # fallback that applies the pattern to the whole key. + seen: list[str] = [] + for c in (base_dot + scope_dot, scope_dot): + if c not in seen: + seen.append(c) + return tuple(seen) + + def _drop_shadowed_prefix_renames(model, rules: list[RenameRule]) -> list[RenameRule]: """Drop child-model reverse renames when the child namespace already exists. @@ -294,6 +353,13 @@ def _drop_shadowed_prefix_renames(model, rules: list[RenameRule]) -> list[Rename ``model.language_model.*`` reverse can also reach its parent VLM. In that model, ``model.language_model`` is already registered and applying the rule globally would capture both that namespace and siblings such as ``model.visual``. + + Only rules that are genuinely confined are skipped. A rule is confined when every + one of its ``scope_prefixes`` is non-empty, because :func:`_sub_scoped` then applies + it solely under that prefix and it cannot reach a sibling namespace. An *empty* + candidate (which ``scope_prefix == ""`` produces) matches every key, so such a rule + has the reach of an unscoped one and must still face this heuristic -- otherwise a + root-scoped nested-text rename would bypass the guard added for NVBug 6525534. """ named_modules = getattr(model, "named_modules", None) if not callable(named_modules): @@ -303,6 +369,11 @@ def _drop_shadowed_prefix_renames(model, rules: list[RenameRule]) -> list[Rename probe_suffix = ".\x00modelopt_namespace_probe" kept: list[RenameRule] = [] for rule in rules: + # `all(...)` matters: an empty candidate matches every key, so a rule carrying one + # is not actually confined and still needs the check below. + if rule.scope_prefixes and all(rule.scope_prefixes): + kept.append(rule) + continue pattern = re.compile(rule.pattern) shadowed = False for module_name in module_names: @@ -375,9 +446,24 @@ def _build_reverse_rules(model) -> tuple[list[SplitRule], list[RenameRule], list for conv in conversions: rev = conv.reverse_transform() # hub<-in-memory; reversed name patterns + ops if isinstance(rev, WeightRenaming): + scope_prefixes = _scope_prefixes(rev) for pattern, repl in zip(_as_list(rev.source_patterns), _as_list(rev.target_patterns)): - weight_renamings.append(RenameRule(pattern=pattern, repl=repl)) + weight_renamings.append( + RenameRule(pattern=pattern, repl=repl, scope_prefixes=scope_prefixes) + ) elif isinstance(rev, WeightConverter): + # Converter-derived rules (expert leaf renames, dense split) are matched by + # module suffix, not by an anchored pattern, so they carry no scope and would + # reach identically-named modules in sibling namespaces. No current model + # scopes a WeightConverter -- transformers only scopes WeightRenaming / + # PrefixChange -- so rather than emit rules we cannot scope, refuse the + # conversion and let the caller fall back to in-memory names. That is a + # warning plus unchanged names, instead of a silently mis-named checkpoint. + if _scope_prefixes(rev): + raise QuantConversionUnsupportedError( + f"scoped WeightConverter (scope_prefix=" + f"{getattr(rev, 'scope_prefix', None)!r}) cannot be reversed scope-aware" + ) ops = list(rev.operations) if any(isinstance(op, SplitModulelist) for op in ops): # Expert converter: ModelOpt already un-stacked/un-fused experts to diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index ab2ef0d9029..cc894d0ffd5 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -959,6 +959,14 @@ def from_quantized_weight( raise NotImplementedError(f"quantization format {quantization} not supported") +def _strip_base_layer(key: str, is_modelopt_qlora: bool) -> str: + """Drop the `base_layer` component PEFT inserts, which deployment does not expect. + + Stripping generically means new key types (bias, scales) need no enumeration here. + """ + return key.replace(".base_layer.", ".") if is_modelopt_qlora else key + + def postprocess_state_dict( state_dict: dict, maxbound: float, @@ -991,16 +999,8 @@ 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: + return _strip_base_layer(key, is_modelopt_qlora) post_state_dict = {} @@ -1012,7 +1012,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 +1033,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..a291b5abf36 100644 --- a/modelopt/torch/opt/plugins/transformers.py +++ b/modelopt/torch/opt/plugins/transformers.py @@ -55,6 +55,10 @@ "ModelOptTrainerArguments", ] +# transformers 5.0 changed `_tied_weights_keys` from a list to a {target: source} dict and added +# the `load_config` parameter to `_load_state_dict_into_zero3_model`. +_TRANSFORMERS_GE_5_0 = Version(transformers.__version__) >= Version("5.0") + def is_liger_available(): try: @@ -110,16 +114,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): @@ -147,20 +169,54 @@ def _new_from_config(cls, /, config, **kwargs): return model +@contextmanager +def _legacy_tied_weights_keys_as_dict(model: nn.Module): + """Temporarily normalize legacy list-style ``_tied_weights_keys`` to dict-style. + + transformers 5.0 changed ``_tied_weights_keys`` from ``list[str]`` to a + ``{target: source}`` dict, and ``save_pretrained`` calls ``.keys()`` on the attribute of + every submodule (``_get_tied_weight_keys``) without a type check. Modeling code still on + the 4.x list format - common for ``trust_remote_code`` checkpoints, e.g. + ``stepfun-ai/Step-3.7-Flash`` - therefore loads fine but dies on save with + ``AttributeError: 'list' object has no attribute 'keys'``. + + Mapping each entry to itself preserves the legacy semantics: the list entries were exactly + the dedup patterns ``_get_tied_weight_keys`` is expected to return. The original attribute + is restored on exit so the shim stays invisible to the rest of the model's lifetime. + """ + if not _TRANSFORMERS_GE_5_0: + yield + return + + patched = [] + try: + for module in model.modules(): + tied = getattr(module, "_tied_weights_keys", None) + if isinstance(tied, (list, tuple, set)): + # The attribute is usually a class attribute; remember whether this instance + # had its own so the restore does not leave a shadowing copy behind. + patched.append((module, tied, "_tied_weights_keys" in module.__dict__)) + module._tied_weights_keys = {key: key for key in tied} + yield + finally: + for module, tied, had_own_attr in patched: + if had_own_attr: + module._tied_weights_keys = tied + else: + del module._tied_weights_keys + + def _save_pretrained_with_checks(self, save_directory, *args, **kwargs): if getattr(self, "_tp_size", None) is not None and ModeloptStateManager.is_converted(self): raise NotImplementedError( "ModelOpt does not support saving tensor parallel sharded Huggingface transformer models yet. " ) - return _new_save_pretrained(self, save_directory, *args, **kwargs) + with _legacy_tied_weights_keys_as_dict(self): + return _new_save_pretrained(self, save_directory, *args, **kwargs) # [Fix for huggingface bug] deepspeed zero3 training backend only loads params into the model from # state_dict, but not buffers. So lets explicitly load the buffers into the model from state_dict. -# The `load_config` parameter was added to `_load_state_dict_into_zero3_model` in transformers 5.0. -_TRANSFORMERS_GE_5_0 = Version(transformers.__version__) >= Version("5.0") - - def _load_params_and_buffers_into_zero3_model(model_to_load, state_dict, load_config=None): buffer_names = [name for name, _ in model_to_load.named_buffers()] buffer_state_dict = {k: v for k, v in state_dict.items() if k in buffer_names} diff --git a/modelopt/torch/prune/plugins/mcore_minitron.py b/modelopt/torch/prune/plugins/mcore_minitron.py index 28684fa3a4a..ad12fc7ebff 100644 --- a/modelopt/torch/prune/plugins/mcore_minitron.py +++ b/modelopt/torch/prune/plugins/mcore_minitron.py @@ -47,6 +47,7 @@ gather_from_tensor_model_parallel_region, reduce_from_tensor_model_parallel_region, ) +from megatron.core.transformer.multi_latent_attention import MLASelfAttention from pydantic import create_model from rich.console import Console from rich.markup import escape as rich_escape @@ -255,6 +256,11 @@ class MCoreMinitronSearcher(BaseSearcher): - `max_depth_pruning`: Maximum fraction per depth hyperparameter to prune (default: 0.20). Only top (1 - max_depth_pruning) choices will be considered. - `hparams_to_skip`: List of hparams to skip during the search (default: None). + - `candidate_filter`: Callable rejecting candidate configs the caller cannot use, e.g. ones + their checkpoint format cannot represent (default: None). Receives every supported hparam, + with non-searched ones filled in from the model config (unset ones are omitted, so a + filter using them raises `KeyError`). Like `score_func`, it is assumed + unchanged when resuming from a `checkpoint`, since rejected candidates are not cached. - `top_k`: Number of candidates to consider for score_func validation (default: 10). - `seq_length`: Sequence length for KV-cache memory estimate (default: 4096). Only used with the ``memory_mb`` constraint. @@ -280,6 +286,7 @@ def default_search_config(self) -> SearchConfig: "max_width_pruning": 0.40, "max_depth_pruning": 0.20, "hparams_to_skip": None, + "candidate_filter": None, "top_k": 10, # Memory footprint config (only used with memory_mb constraint) "seq_length": 4096, @@ -521,6 +528,7 @@ def search_best_arch_by_metrics(self) -> dict: max_width_pruning = self.config["max_width_pruning"] max_depth_pruning = self.config["max_depth_pruning"] hparams_to_skip = self.config["hparams_to_skip"] + candidate_filter = self.config["candidate_filter"] top_k = self.config["top_k"] constraints_str = ", ".join(f"{self._fmt_metric(v, k)} {k}" for k, v in max_metrics.items()) print_rank_0(f"\nSearching for the best pruned architecture under {constraints_str}...") @@ -549,12 +557,24 @@ def search_best_arch_by_metrics(self) -> dict: max_depth_pruning, hparams_to_skip, ) + # Only place to reject invalid hparam *combinations*; unsearched ones come from config. + base_config = { + hp: getattr(self.model.config, hp) + for hp in SUPPORTED_HPARAMS + if getattr(self.model.config, hp, None) is not None + } selected = [] + num_filtered = 0 for ss_config in tqdm( search_space_configs, desc="Finding all candidates fitting the constraints...", disable=not dist.is_master(), ): + if candidate_filter is not None and not candidate_filter( + {**base_config, **ss_config} + ): + num_filtered += 1 + continue candidate_metrics = self._compute_candidate_metrics(ss_config, max_num_layers) if all(candidate_metrics[k] <= max_metrics[k] for k in active_metric_keys): selected.append( @@ -562,7 +582,11 @@ def search_best_arch_by_metrics(self) -> dict: ss_config, {k: candidate_metrics[k] for k in active_metric_keys}, None ) ) - assert len(selected) > 0, "No subnets found fitting the constraints!" + if num_filtered: + print_rank_0(f"Rejected {num_filtered} candidates via candidate_filter.") + assert len(selected) > 0, "No subnets found fitting the constraints!" + ( + f" candidate_filter rejected all {num_filtered} candidates." if num_filtered else "" + ) print_rank_0(f"Found {len(selected)} candidates fitting the constraints!") self.all_candidates_per_constraint[constraints_cache_key] = sorted( selected, key=lambda x: x.metrics[primary_key], reverse=True @@ -952,6 +976,25 @@ def restore(self) -> RestoreEntrypoint: return restore_mcore_minitron +def _fused_ln_linears_over_hidden_size(module: nn.Module): + """Yield fused-layernorm linears whose layernorm is over ``hidden_size``. + + MLA's Q/KV up-projections are ``TELayerNormColumnParallelLinear`` too, but their layernorm is + over the latent rank and MCore unpacks their output as ``(out, bias)``. Setting + ``return_layernorm_output`` there makes it ``((out, ln_out), bias)`` and breaks the forward. + """ + up_projs = { + id(sub) + for m in module.modules() + if isinstance(m, MLASelfAttention) + for attr in ("linear_q_up_proj", "linear_kv_up_proj") + if (sub := getattr(m, attr, None)) is not None + } + for m in module.modules(): + if isinstance(m, TELayerNormColumnParallelLinear) and id(m) not in up_projs: + yield m + + class ImportanceEstimatorRegistry: """Register importance estimators and forward hooks for all supported modules in the model. @@ -1034,9 +1077,8 @@ def cleanup(self) -> None: self._hooks.clear() # Unpatch return_layernorm_output on fused TELayerNormColumnParallelLinear modules - for m in self.model.modules(): - if isinstance(m, TELayerNormColumnParallelLinear): - m.return_layernorm_output = False + for m in _fused_ln_linears_over_hidden_size(self.model): + m.return_layernorm_output = False def get_layer_scores(self) -> dict[int, torch.Tensor]: """Get the layer scores (1-indexed) from the model. @@ -1166,9 +1208,8 @@ def _estimate_hidden_size_importance(mod): # Layernorms are fused into TELayerNormColumnParallelLinear. We temporarily # patch return_layernorm_output=True so TE's fused kernel returns the layernorm output. # For MoE layers, pre_mlp_layernorm is a separate TENorm — use a regular forward hook. - for m in module.modules(): - if isinstance(m, TELayerNormColumnParallelLinear): - m.return_layernorm_output = True + for m in _fused_ln_linears_over_hidden_size(module): + m.return_layernorm_output = True for layer in module.decoder.layers: if isinstance(layer, _DynamicTransformerLayer): diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py index 752dd801a6e..d272fa400fe 100644 --- a/modelopt/torch/quantization/plugins/megatron.py +++ b/modelopt/torch/quantization/plugins/megatron.py @@ -133,6 +133,18 @@ def quant_module_get_extra_state(self) -> dict: QuantModule's extra_state with QuantModule.get_extra_state() which avoids the need to store the full module name. """ + # ``GPTModel.sharded_state_dict`` pops ``output_layer._extra_state`` and asserts it carries no + # data ("Expected output layer extra state to be empty", mcore models/gpt/gpt_model.py), so an + # output_layer with nothing quantized must contribute none. Scoped to output_layer: for every + # other module this quantizer_state is the only record that its quantizers were disabled (e.g. + # by auto_quantize or disable_quantizer), and dropping it would restore them enabled. + if ( + getattr(self, "_modelopt_output_layer", False) + and not isinstance(self, RealQuantLinear) + and not any(isinstance(m, TensorQuantizer) and m.is_enabled for m in self.modules()) + ): + return {} + extra_state = {} quantizer_state = {} @@ -248,6 +260,48 @@ def _incompatible_method(self, *args, **kwargs): return _incompatible_method +def _resolve_output_layer_untied(model: torch.nn.Module) -> bool | None: + """Whether ``output_layer`` weights are untied from the input embeddings, or None if unknown.""" + # named_modules() yields the root first, so the root's own flag wins when it has one. + for name, module in model.named_modules(): + # Skip subtrees that do not own the language model's output_layer: the vision tower (never + # quantized here) and a distillation teacher, which may be tied differently from the + # student it is wrapped with. + if "vision_model" in name or "_teacher_model" in name: + continue + shared = getattr(module, "share_embeddings_and_output_weights", None) + if shared is not None: + return not bool(shared) + return None + + +def _output_layer_untied(config) -> bool: + """Whether ``output_layer`` is untied, for use from ``sharded_state_dict``. + + Prefers the flag recorded by ``megatron_replace_quant_module_hook`` (the only source available + under Megatron-Bridge, which has no global args store), then Megatron-LM's + ``--untie-embeddings-and-output-weights``. + """ + untied = getattr(config, "modelopt_output_layer_untied", None) + if untied is not None: + return untied + try: + from megatron.training import get_args as _mlm_get_args + + return bool(getattr(_mlm_get_args(), "untie_embeddings_and_output_weights", False)) + except (ImportError, AssertionError) as e: + # ImportError: no megatron.training. AssertionError: get_args() before initialize_megatron. + # Warn once per config rather than on every save and every load. + if not getattr(config, "_modelopt_warned_output_layer_untied", False): + warn_rank_0( + f"Failed to get Megatron arg untie_embeddings_and_output_weights: {e}. " + "Treating output_layer as tied; its quantizer state will not be saved or " + "restored. If output_layer is in fact untied, it will be exported unquantized." + ) + config._modelopt_warned_output_layer_untied = True + return False + + def megatron_replace_quant_module_hook(model: torch.nn.Module): """Configure Megatron-Core model quantization support. @@ -260,6 +314,7 @@ def megatron_replace_quant_module_hook(model: torch.nn.Module): typing-matching the QuantModuleRegistry. 3. For Attention modules, we configure them to use core_attention path for KV cache quantization. """ + untied = _resolve_output_layer_untied(model) def _configure_attention_for_kv_cache_quant(module: Attention): """Configure Attention module for KV cache quantization compatibility.""" @@ -287,11 +342,8 @@ def _configure_attention_for_kv_cache_quant(module: Attention): def _register_extra_state_callbacks(model: torch.nn.Module): for name, module in model.named_modules(): if type(module) in QuantModuleRegistry: - # Skip output_layer w/o enabled weight_quantizer - if name.endswith("output_layer") and not getattr( - getattr(module, "weight_quantizer", None), "is_enabled", False - ): - continue + if name.endswith("output_layer"): + module._modelopt_output_layer = True register_modelopt_extra_state_callbacks( module, quant_module_get_extra_state, @@ -307,6 +359,10 @@ def _register_extra_state_callbacks(model: torch.nn.Module): if "vision_model" not in name: # We only enable hetereogenous_dist_checkpoint for language model, vision model is not quantized module.config.hetereogenous_dist_checkpoint = True + # Read back via ``self.config`` in sharded_state_dict; output_layer shares its + # parent MegatronModule's config object. The teacher may be tied differently. + if untied is not None and "_teacher_model" not in name: + module.config.modelopt_output_layer_untied = untied _register_extra_state_callbacks(module) @@ -374,16 +430,7 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): # output_layer.input_quantizer._amax but TP-only does not. This lead to # state_dict mismatch. if prefix.endswith("output_layer."): - try: - from megatron.training import get_args as _mlm_get_args - - _untied = bool( - getattr(_mlm_get_args(), "untie_embeddings_and_output_weights", False) - ) - except Exception as e: - warn_rank_0(f"Failed to get Megatron arg untie_embeddings_and_output_weights: {e}") - _untied = False - if not _untied: + if not _output_layer_untied(self.config): return super().sharded_state_dict(prefix, sharded_offsets, metadata) quantizer_state_dict = {} diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py index 00411eaa4ee..243db16a2bc 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py @@ -47,6 +47,19 @@ from modelopt.torch.kernels.quantization.attention.bmm2_qdq import fake_quant_v_onwrite +@functools.cache +def _flash_attention_kv_cache_layout() -> str: + """Return the installed vLLM backend's K/V packing contract.""" + cache_shape = FlashAttentionBackend.get_kv_cache_shape(3, 16, 1, 16) + if cache_shape == (2, 3, 16, 1, 16): + return "kv-first" + if cache_shape == (3, 2, 16, 1, 16): + return "blocks-first" + if cache_shape == (3, 1, 16, 32): + return "packed" + raise RuntimeError(f"Unsupported vLLM FlashAttention KV cache shape {cache_shape}") + + def _target_sparse_ratio_for_phase(target_sparse_ratio, phase: str) -> float: """Return target sparsity for a phase, defaulting old checkpoint metadata.""" if isinstance(target_sparse_ratio, float | int): @@ -514,7 +527,13 @@ def native_forward(): if resolved is None: return native_forward() - key_cache, value_cache = kv_cache.unbind(0) + cache_layout = _flash_attention_kv_cache_layout() + if cache_layout == "kv-first": + key_cache, value_cache = kv_cache.unbind(0) + elif cache_layout == "blocks-first": + key_cache, value_cache = kv_cache.unbind(1) + else: + key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1) is_decode_only = attn_metadata.max_query_len <= 1 common_kw = { "layer": layer, diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py index 7679ff0020b..e0d63bde136 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash.py +++ b/modelopt/torch/speculative/plugins/hf_dflash.py @@ -72,9 +72,11 @@ """ import logging +from typing import Any import torch import torch.nn.functional as F +import transformers from transformers import PreTrainedModel from transformers.models.qwen3.configuration_qwen3 import Qwen3Config as _Qwen3Config from transformers.trainer_pt_utils import LabelSmoother @@ -100,6 +102,54 @@ __all__ = ["HFDFlashModel"] +_QWEN3_VL_MROPE_WORKAROUND_VERSION = "5.3.0" +_MULTIMODAL_FORWARD_KWARGS = frozenset( + { + "pixel_values", + "pixel_values_videos", + "image_grid_thw", + "video_grid_thw", + "mm_token_type_ids", + "image_sizes", + "images", + "videos", + } +) + + +def _multimodal_forward_kwargs(model_kwargs: dict) -> dict: + """Return collator fields accepted by Hugging Face multimodal forwards.""" + return { + name: value + for name, value in model_kwargs.items() + if name in _MULTIMODAL_FORWARD_KWARGS and value is not None + } + + +def _expand_qwen3_video_grid_thw(video_grid_thw: torch.Tensor) -> torch.Tensor: + """Return the per-frame video grid representation used by Qwen3-VL RoPE. + + Qwen3-VL's video processor emits one ``[T, H, W]`` row per source video, but + its rendered prompt contains a separate visual-token group for every temporal + frame. Transformers 5.3's ``get_rope_index`` consumes one grid row per + rendered group, while the vision encoder still requires the original one-row- + per-video representation. This helper is therefore used *only* for mRoPE + position construction; callers must keep the original tensor for the model + forward. + """ + if video_grid_thw.ndim != 2 or video_grid_thw.shape[-1] != 3: + raise ValueError( + "Qwen3-VL video_grid_thw must have shape [num_videos, 3], got " + f"{tuple(video_grid_thw.shape)}." + ) + if torch.any(video_grid_thw[:, 0] <= 0): + raise ValueError("Qwen3-VL video_grid_thw temporal lengths must be positive.") + + expanded_grid_thw = torch.repeat_interleave(video_grid_thw, video_grid_thw[:, 0], dim=0) + expanded_grid_thw[:, 0] = 1 + return expanded_grid_thw + + def _dpace_position_weights( confidences: torch.Tensor, alpha: float, valid_mask: torch.Tensor | None = None ) -> torch.Tensor: @@ -183,6 +233,125 @@ def _base_llm_config(self): or self.config ) + def _qwen3_vl_position_ids( + self, + input_ids, + attention_mask, + position_ids, + past_key_values, + inputs_embeds, + model_kwargs, + ): + """Precompute Qwen3-VL mRoPE positions for Transformers 5.3.0 batches. + + The video encoder consumes one grid row per source video, whereas mRoPE + consumes one row per rendered temporal-frame group. Calling the + top-level model with the original video grid makes the two contracts + conflict. Construct the mRoPE positions with a frame-expanded copy, + then pass the original grid to the vision encoder in ``forward``. + + Transformers 5.4.0 performs this frame expansion in ``get_rope_index`` + itself; only 5.3.0 needs the external workaround. See + https://github.com/huggingface/transformers/blob/v5.4.0/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py + + Prefer ``get_rope_index`` over ``compute_3d_position_ids``. The latter + writes ``rope_deltas`` into the base model even though DFlash training + never supplies a cache. Keeping this calculation side-effect free is + important when the frozen target is reused for consecutive training + batches or validation. + """ + model_type = str(getattr(self.config, "model_type", "")) + if ( + position_ids is not None + or not model_type.startswith("qwen3_vl") + # Cached decoding uses the base model's rope_deltas path. DFlash + # training has no cache and is the only path that needs the + # frame-expanded construction below. + or past_key_values is not None + ): + return position_ids + + image_grid_thw = model_kwargs.get("image_grid_thw") + video_grid_thw = model_kwargs.get("video_grid_thw") + if not isinstance(image_grid_thw, torch.Tensor) and not isinstance( + video_grid_thw, torch.Tensor + ): + return position_ids + + if transformers.__version__ != _QWEN3_VL_MROPE_WORKAROUND_VERSION: + if transformers.__version__.startswith("5.3."): + raise RuntimeError( + "Qwen3-VL DFlash mRoPE supports Transformers 5.3.0 or >=5.4.0; " + f"got {transformers.__version__}. A 5.3.x patch release may already " + "expand video_grid_thw internally." + ) + return position_ids + + mm_token_type_ids = model_kwargs.get("mm_token_type_ids") + backbone = getattr(self, "model", None) + # Probed dynamically: which one exists depends on the Transformers version. + get_rope_index: Any = getattr(backbone, "get_rope_index", None) + compute_position_ids: Any = getattr(backbone, "compute_3d_position_ids", None) + if ( + not isinstance(mm_token_type_ids, torch.Tensor) + or input_ids is None + or (not callable(get_rope_index) and not callable(compute_position_ids)) + ): + raise ValueError( + "Qwen3-VL DFlash training requires input_ids, mm_token_type_ids, and " + "a Qwen3-VL model with get_rope_index or compute_3d_position_ids. " + "Use the Qwen3-VL AutoProcessor without dropping mm_token_type_ids." + ) + + if mm_token_type_ids.shape != input_ids.shape: + raise ValueError( + "Qwen3-VL mm_token_type_ids must have the same shape as input_ids, got " + f"{tuple(mm_token_type_ids.shape)} and {tuple(input_ids.shape)}." + ) + + rope_video_grid_thw = video_grid_thw + if isinstance(video_grid_thw, torch.Tensor) and video_grid_thw.numel() > 0: + video_token_mask = mm_token_type_ids == 2 + if isinstance(attention_mask, torch.Tensor): + video_token_mask = video_token_mask & attention_mask.bool() + video_group_starts = video_token_mask.clone() + video_group_starts[:, 1:] &= ~video_token_mask[:, :-1] + expected_video_groups = int(video_grid_thw[:, 0].sum()) + actual_video_groups = int(video_group_starts.sum()) + if actual_video_groups != expected_video_groups: + raise ValueError( + "Qwen3-VL video frame groups do not match video_grid_thw: " + f"expected {expected_video_groups}, found {actual_video_groups}." + ) + rope_video_grid_thw = _expand_qwen3_video_grid_thw(video_grid_thw) + + rope_kwargs = { + "input_ids": input_ids, + "image_grid_thw": image_grid_thw, + "video_grid_thw": rope_video_grid_thw, + "attention_mask": attention_mask, + "mm_token_type_ids": mm_token_type_ids, + } + if callable(get_rope_index): + position_ids, _ = get_rope_index(**rope_kwargs) + else: + position_ids = compute_position_ids( + **rope_kwargs, + inputs_embeds=inputs_embeds, + past_key_values=past_key_values, + ) + + expected_shape = (3, *input_ids.shape) + valid_position_ids = isinstance(position_ids, torch.Tensor) and ( + tuple(position_ids.shape) == expected_shape + ) + if not valid_position_ids: + raise RuntimeError( + "Qwen3-VL produced invalid mRoPE position ids: expected shape " + f"{expected_shape}, got {getattr(position_ids, 'shape', None)}." + ) + return position_ids + def _find_base_model_parts(self): """Locate base model submodules (backbone, embeddings, lm_head) by probing known paths. @@ -592,6 +761,16 @@ def forward( - Label alignment: position k predicts token at anchor+k - Optional loss decay weighting """ + if self.training: + position_ids = self._qwen3_vl_position_ids( + input_ids, + attention_mask, + position_ids, + past_key_values, + inputs_embeds, + kwargs, + ) + if not self.training: if self.dflash_offline: raise RuntimeError( @@ -638,12 +817,37 @@ def forward( ) target_hidden = base_outputs.target_hidden else: - # TODO: For co-training the base model, remove no_grad and eval() switch. + # Multimodal models need the top-level conditional-generation forward so their + # image/video features are inserted before the language model runs. Keep the + # long-standing narrow call for text-only models. + base_forward_kwargs = _multimodal_forward_kwargs(kwargs) + use_top_level_forward = bool(base_forward_kwargs) with torch.no_grad(): - raw_outputs = super().forward( - input_ids=input_ids, - attention_mask=attention_mask, - output_hidden_states=True, + if use_top_level_forward: + raw_outputs = super().forward( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=False, + output_attentions=output_attentions, + output_hidden_states=True, + cache_position=cache_position, + return_dict=True, + **base_forward_kwargs, + ) + else: + raw_outputs = super().forward( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + ) + + if not getattr(raw_outputs, "hidden_states", None): + raise RuntimeError( + "The base model did not return hidden states required for DFlash training. " + "Ensure its top-level multimodal forward supports output_hidden_states=True." ) offset = 1 selected = [raw_outputs.hidden_states[lid + offset] for lid in self.target_layer_ids] @@ -652,16 +856,15 @@ def forward( target_hidden=target_hidden, logits=raw_outputs.logits ) - # 2. Build loss mask. - # When labels are provided (answer_only_loss), they already encode both - # assistant masking and padding (-100 for both). When labels are not - # provided, fall back to attention_mask for padding only. + # 2. Build loss mask. Labels carry optional answer-only masking, but do + # not in general mark padded tokens with -100 (the VLM collator creates + # them from padded input_ids). Always intersect with attention_mask so + # anchor sampling and loss never include the padded tail. + loss_mask = torch.ones(bsz, seq_len, device=device) if labels is not None: - loss_mask = (labels != LabelSmoother.ignore_index).float() - elif attention_mask is not None: - loss_mask = attention_mask.float() - else: - loss_mask = torch.ones(bsz, seq_len, device=device) + loss_mask = loss_mask * (labels != LabelSmoother.ignore_index).float() + if attention_mask is not None: + loss_mask = loss_mask * attention_mask.float() # In offline training, assistant mask is dumped and passed as kwarg. if kwargs.get("loss_mask") is not None: @@ -674,8 +877,16 @@ def forward( n_blocks = anchor_positions.shape[1] if n_blocks == 0 or not block_keep_mask.any(): - # Zero loss that still flows through dflash_module for DDP gradient sync - dummy = self.dflash_module.fc.weight.sum() * 0.0 + # Keep all trainable draft parameters in the graph so DDP can reduce a rank + # that receives an all-masked answer-only batch. + dummy = sum( + ( + parameter.reshape(-1)[0] * 0.0 + for parameter in self.dflash_module.parameters() + if parameter.requires_grad + ), + torch.zeros((), device=device), + ) return ModelOutput(loss=dummy, logits=base_outputs.logits, train_acc=[[0.0]]) # 4. Build draft inputs diff --git a/modelopt/torch/speculative/utils.py b/modelopt/torch/speculative/utils.py index 6a3c19b993d..8c5418bb6b8 100644 --- a/modelopt/torch/speculative/utils.py +++ b/modelopt/torch/speculative/utils.py @@ -610,7 +610,13 @@ def load_vlm_or_llm( return FakeBaseModel.from_source(model_name_or_path, trust_remote_code=trust_remote_code) if _is_vlm: - model_cls = transformers.AutoModelForVision2Seq + # Transformers 5 renamed AutoModelForVision2Seq to + # AutoModelForImageTextToText. Prefer the pre-5 name so this loader + # continues to support the Transformers 4 environments used by older + # speculative-decoding jobs. + model_cls = getattr(transformers, "AutoModelForVision2Seq", None) + if model_cls is None: + model_cls = transformers.AutoModelForImageTextToText else: model_cls = transformers.AutoModelForCausalLM diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 12287865b7f..245ca81de79 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -18,7 +18,9 @@ import functools import io import os +import sys import time +import traceback from collections.abc import Callable from contextlib import suppress from datetime import timedelta @@ -212,13 +214,39 @@ def setup(timeout: timedelta | None = None): def cleanup(): - """Cleans up the distributed environment.""" + """Cleans up the distributed environment. + + The barrier is skipped when unwinding from an error, since peers may be blocked in a collective + this rank will never reach. ``SystemExit`` is treated as a clean exit (every rank reaches it). + That is not sufficient on its own -- ``destroy_process_group`` below blocks for the same reason + -- so error paths must call :func:`abort` before reaching this ``finally``. + """ if is_initialized(): - with suppress(Exception): - barrier() + exc = sys.exc_info()[1] + if exc is None or isinstance(exc, SystemExit): + with suppress(Exception): + barrier() torch.distributed.destroy_process_group() +def abort(exit_code: int = 1) -> None: + """Print the active exception and exit this rank immediately. + + Call from an ``except`` block in a distributed entrypoint. Both a barrier and + ``destroy_process_group`` stall when peers are blocked in a collective this rank will never + reach, and the traceback only prints once the enclosing ``finally`` returns -- so the run looks + hung rather than failed. Exiting lets the launcher (e.g. torchrun) terminate the peers. + ``SystemExit`` is re-raised instead, since every rank reaches an intentional exit. + """ + exc = sys.exc_info()[1] + if isinstance(exc, SystemExit): + raise exc + traceback.print_exc() + sys.stdout.flush() + sys.stderr.flush() + os._exit(exit_code) + + def is_fsdp2_model(model) -> bool: """Return True if any submodule of ``model`` has been wrapped with FSDP2 ``fully_shard``.""" return any(isinstance(m, FSDPModule) for m in model.modules()) diff --git a/modelopt/torch/utils/plugins/transformers_dataset.py b/modelopt/torch/utils/plugins/transformers_dataset.py index c27a3d09aea..97ae4ea2d14 100644 --- a/modelopt/torch/utils/plugins/transformers_dataset.py +++ b/modelopt/torch/utils/plugins/transformers_dataset.py @@ -325,6 +325,7 @@ def __init__( chat_template: str | None = None, add_generation_prompt: bool = False, answer_only_loss: bool = False, + shift_labels: bool = True, local_image_path: str = "", return_labels: bool = False, ): @@ -340,10 +341,97 @@ def __init__( chat_template=chat_template, add_generation_prompt=add_generation_prompt, answer_only_loss=answer_only_loss, + shift_labels=shift_labels, return_labels=return_labels, ) + def _verify_generation_tags(self): + """Accept VLM templates whose assistant spans have stable chat markers. + + Cosmos/Qwen ChatML templates do not necessarily use Hugging Face's + ``{% generation %}`` tags. For those templates we derive the same + assistant-only loss mask from the tokenized assistant boundaries. + """ + if self._assistant_marker_specs(): + return + super()._verify_generation_tags() + + def _assistant_marker_specs(self): + """Return tokenized assistant start/end boundaries for supported templates.""" + if hasattr(self, "_cached_assistant_marker_specs"): + return self._cached_assistant_marker_specs + + template = self.tokenizer.chat_template or "" + specs = [] + if "<|im_start|>" in template and "<|im_end|>" in template: + specs.append( + ( + self.tokenizer("<|im_start|>assistant\n", add_special_tokens=False)[ + "input_ids" + ], + [ + self.tokenizer("<|im_end|>\n", add_special_tokens=False)["input_ids"], + self.tokenizer("<|im_end|>", add_special_tokens=False)["input_ids"], + ], + ) + ) + self._cached_assistant_marker_specs = [ + (start, [end for end in ends if end]) for start, ends in specs if start and any(ends) + ] + return self._cached_assistant_marker_specs + + @staticmethod + def _find_subsequence(values, pattern, start=0, stop=None): + stop = len(values) if stop is None else stop + if not pattern or start >= stop: + return -1 + for index in range(start, stop - len(pattern) + 1): + if values[index : index + len(pattern)] == pattern: + return index + return -1 + + def _build_assistant_masks(self, tokenized_messages): + """Build assistant-content masks from ChatML boundaries.""" + input_ids = tokenized_messages["input_ids"] + attention_mask = tokenized_messages.get("attention_mask") + assistant_masks = torch.zeros_like(input_ids) + + for row_index, row in enumerate(input_ids): + tokens = row.tolist() + if isinstance(attention_mask, torch.Tensor): + active = attention_mask[row_index].nonzero(as_tuple=False).flatten() + if active.numel() == 0: + continue + sequence_start, sequence_end = int(active[0]), int(active[-1]) + 1 + else: + sequence_start, sequence_end = 0, len(tokens) + + for start_marker, end_markers in self._assistant_marker_specs(): + search_from = sequence_start + while search_from < sequence_end: + start = self._find_subsequence(tokens, start_marker, search_from, sequence_end) + if start == -1: + break + content_start = start + len(start_marker) + end_positions = [ + position + for marker in end_markers + if ( + position := self._find_subsequence( + tokens, marker, content_start, sequence_end + ) + ) + != -1 + ] + content_end = min(end_positions) if end_positions else sequence_end + if content_start < content_end: + assistant_masks[row_index, content_start:content_end] = 1 + search_from = max(content_start + 1, content_end + 1) + + return assistant_masks + def _process_multimodal_sample(self, examples): + derive_masks_from_markers = self.answer_only_loss and bool(self._assistant_marker_specs()) tokenized_messages = self.processor.apply_chat_template( examples, tokenize=True, @@ -353,9 +441,36 @@ def _process_multimodal_sample(self, examples): truncation=True, max_length=self.train_len, add_generation_prompt=self.add_generation_prompt, - return_assistant_tokens_mask=self.answer_only_loss, + return_assistant_tokens_mask=self.answer_only_loss and not derive_masks_from_markers, ) + if derive_masks_from_markers: + tokenized_messages["assistant_masks"] = self._build_assistant_masks(tokenized_messages) + + if self.return_labels: + input_ids = tokenized_messages["input_ids"] + labels = input_ids.new_full(input_ids.shape, IGNORE_TOKEN_ID) + if self.shift_labels: + labels[..., :-1] = input_ids[..., 1:] + else: + # DFlash predicts the token at the current position rather + # than the next autoregressive token. + labels[:] = input_ids + + if self.answer_only_loss: + if "assistant_masks" not in tokenized_messages: + raise ValueError( + "answer_only_loss requires assistant_masks from the VLM chat template." + ) + assistant_mask = tokenized_messages["assistant_masks"] + if not isinstance(assistant_mask, torch.Tensor) or not assistant_mask.any(): + labels[:] = IGNORE_TOKEN_ID + elif self.shift_labels: + labels[..., :-1][assistant_mask[..., 1:] == 0] = IGNORE_TOKEN_ID + else: + labels[assistant_mask == 0] = IGNORE_TOKEN_ID + tokenized_messages["labels"] = labels + return tokenized_messages def __call__(self, examples): @@ -385,6 +500,21 @@ def __call__(self, examples): msg["content"] = [{"type": "text", "text": msg["content"]}] for ctn in msg["content"]: + # Some JSONL producers use a fixed multimodal-part schema + # (text/image/video/fps on every part) so Arrow can load + # heterogeneous image and video datasets together. Drop + # the inactive placeholders before handing a part to the + # processor, which expects only fields relevant to its type. + content_type = ctn.get("type") + if content_type != "text" and ctn.get("text") == "": + del ctn["text"] + if content_type != "image" and ctn.get("image") == "": + del ctn["image"] + if content_type != "video": + if ctn.get("video") == "": + del ctn["video"] + if ctn.get("fps") == 0: + del ctn["fps"] if ctn["type"] == "image" and "image" in ctn: ctn["image"] = os.path.abspath( os.path.join(self.local_image_path, ctn["image"]) diff --git a/modelopt_recipes/configs/ptq/presets/model/nvfp4_mlp_only.yaml b/modelopt_recipes/configs/ptq/presets/model/nvfp4_mlp_only.yaml index c5d36fd9236..c8c269c1286 100644 --- a/modelopt_recipes/configs/ptq/presets/model/nvfp4_mlp_only.yaml +++ b/modelopt_recipes/configs/ptq/presets/model/nvfp4_mlp_only.yaml @@ -21,6 +21,7 @@ imports: block_sparse_moe_nvfp4: configs/ptq/units/block_sparse_moe_nvfp4 experts_nvfp4: configs/ptq/units/experts_nvfp4 default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + mixer_mlp_nvfp4: configs/ptq/units/mixer_mlp_nvfp4 nvfp4: configs/numerics/nvfp4 algorithm: max @@ -32,6 +33,7 @@ quant_cfg: - quantizer_name: '*mlp*input_quantizer' cfg: $import: nvfp4 + - $import: mixer_mlp_nvfp4 - $import: block_sparse_moe_nvfp4 - $import: experts_nvfp4 - $import: default_disabled_quantizers diff --git a/modelopt_recipes/configs/ptq/presets/model/nvfp4_omlp_only.yaml b/modelopt_recipes/configs/ptq/presets/model/nvfp4_omlp_only.yaml index 82bf401ea9f..cc908ced186 100644 --- a/modelopt_recipes/configs/ptq/presets/model/nvfp4_omlp_only.yaml +++ b/modelopt_recipes/configs/ptq/presets/model/nvfp4_omlp_only.yaml @@ -20,6 +20,7 @@ imports: base_disable_all: configs/ptq/units/base_disable_all block_sparse_moe_nvfp4: configs/ptq/units/block_sparse_moe_nvfp4 default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + mixer_mlp_nvfp4: configs/ptq/units/mixer_mlp_nvfp4 nvfp4: configs/numerics/nvfp4 algorithm: max @@ -37,5 +38,6 @@ quant_cfg: - quantizer_name: '*mlp*input_quantizer' cfg: $import: nvfp4 + - $import: mixer_mlp_nvfp4 - $import: block_sparse_moe_nvfp4 - $import: default_disabled_quantizers diff --git a/modelopt_recipes/configs/ptq/units/README.md b/modelopt_recipes/configs/ptq/units/README.md index cd738f62626..db37b9222ca 100644 --- a/modelopt_recipes/configs/ptq/units/README.md +++ b/modelopt_recipes/configs/ptq/units/README.md @@ -30,4 +30,5 @@ recipes (under `general/` or `models/`) or presets (under `presets/`). | `w4a4_nvfp4_nvfp4.yaml` | NVFP4 weight + activation quantizer entries (W4A4); supported on Blackwell+ GPUs | | `block_sparse_moe_nvfp4.yaml` | NVFP4 W4A4 on `*block_sparse_moe*` weight/input quantizers | | `experts_nvfp4.yaml` | NVFP4 W4A4 on `*.experts.*` weight/input quantizers | +| `mixer_mlp_nvfp4.yaml` | NVFP4 W4A4 on dense `*.mixer.{up,down}_proj` weight/input quantizers | | `attention_qkv_fp8.yaml` | FP8 E4M3 on attention q/k/v bmm and softmax quantizers | diff --git a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml index 5e48dc73b7e..3aadadd289c 100644 --- a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml +++ b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml @@ -38,6 +38,8 @@ enable: false - quantizer_name: '*router*' enable: false + - quantizer_name: 'mtp.*' + enable: false - quantizer_name: 'output.*' enable: false # Multimodal vision branch: keep the vision encoder (SigLIP / ViT) and any diff --git a/modelopt_recipes/huggingface/phi4mm/ptq/disabled_quantizers.yaml b/modelopt_recipes/configs/ptq/units/mixer_mlp_nvfp4.yaml similarity index 53% rename from modelopt_recipes/huggingface/phi4mm/ptq/disabled_quantizers.yaml rename to modelopt_recipes/configs/ptq/units/mixer_mlp_nvfp4.yaml index 1c6089f087f..f38b19d67e0 100644 --- a/modelopt_recipes/huggingface/phi4mm/ptq/disabled_quantizers.yaml +++ b/modelopt_recipes/configs/ptq/units/mixer_mlp_nvfp4.yaml @@ -13,22 +13,22 @@ # See the License for the specific language governing permissions and # limitations under the License. -# QuantizerCfgList snippet of disabled quantizers for Phi-4-Multimodal. -# Splices in the standard `default_disabled_quantizers` exclusions and appends -# Phi-4-MM-specific ones so that only the language model is quantized; -# speech/audio/image/vision branches are skipped. Recipes that import this -# should NOT also import `default_disabled_quantizers`. +# QuantizerCfgList snippet that enables dynamic NVFP4 on dense MLP projections +# registered directly under a ``mixer`` module. # modelopt-schema: modelopt.torch.quantization.config.QuantizerCfgListConfig imports: - default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + nvfp4: configs/numerics/nvfp4 --- - - $import: default_disabled_quantizers - - quantizer_name: '*speech*' - enable: false - - quantizer_name: '*audio*' - enable: false - - quantizer_name: '*image*' - enable: false - - quantizer_name: '*vision*' - enable: false + - quantizer_name: '*.mixer.up_proj.weight_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*.mixer.up_proj.input_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*.mixer.down_proj.weight_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*.mixer.down_proj.input_quantizer' + cfg: + $import: nvfp4 diff --git a/modelopt_recipes/general/ptq/nvfp4_mlp_only-kv_fp8.yaml b/modelopt_recipes/general/ptq/nvfp4_mlp_only-kv_fp8.yaml index a4cf71a1dbd..4fd2e0a7558 100644 --- a/modelopt_recipes/general/ptq/nvfp4_mlp_only-kv_fp8.yaml +++ b/modelopt_recipes/general/ptq/nvfp4_mlp_only-kv_fp8.yaml @@ -18,6 +18,7 @@ imports: base_disable_all: configs/ptq/units/base_disable_all default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + mixer_mlp_nvfp4: configs/ptq/units/mixer_mlp_nvfp4 nvfp4: configs/numerics/nvfp4 kv_fp8: configs/ptq/units/kv_fp8 @@ -36,6 +37,7 @@ quantize: - quantizer_name: '*mlp*input_quantizer' cfg: $import: nvfp4 + - $import: mixer_mlp_nvfp4 - quantizer_name: '*block_sparse_moe*weight_quantizer' cfg: $import: nvfp4 diff --git a/modelopt_recipes/general/ptq/nvfp4_mlp_only-kv_fp8_cast.yaml b/modelopt_recipes/general/ptq/nvfp4_mlp_only-kv_fp8_cast.yaml index 225ecf7f086..a12951bda65 100644 --- a/modelopt_recipes/general/ptq/nvfp4_mlp_only-kv_fp8_cast.yaml +++ b/modelopt_recipes/general/ptq/nvfp4_mlp_only-kv_fp8_cast.yaml @@ -18,6 +18,7 @@ imports: base_disable_all: configs/ptq/units/base_disable_all default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + mixer_mlp_nvfp4: configs/ptq/units/mixer_mlp_nvfp4 nvfp4: configs/numerics/nvfp4 kv_fp8_cast: configs/ptq/units/kv_fp8_cast @@ -36,6 +37,7 @@ quantize: - quantizer_name: '*mlp*input_quantizer' cfg: $import: nvfp4 + - $import: mixer_mlp_nvfp4 - quantizer_name: '*block_sparse_moe*weight_quantizer' cfg: $import: nvfp4 diff --git a/modelopt_recipes/general/ptq/nvfp4_mlp_only-novit-kv_fp8.yaml b/modelopt_recipes/general/ptq/nvfp4_mlp_only-novit-kv_fp8.yaml index 7bbf21393f8..2d80d7a5701 100644 --- a/modelopt_recipes/general/ptq/nvfp4_mlp_only-novit-kv_fp8.yaml +++ b/modelopt_recipes/general/ptq/nvfp4_mlp_only-novit-kv_fp8.yaml @@ -33,6 +33,7 @@ imports: base_disable_all: configs/ptq/units/base_disable_all default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + mixer_mlp_nvfp4: configs/ptq/units/mixer_mlp_nvfp4 nvfp4: configs/numerics/nvfp4 kv_fp8: configs/ptq/units/kv_fp8 @@ -52,6 +53,7 @@ quantize: - quantizer_name: '*mlp*input_quantizer' cfg: $import: nvfp4 + - $import: mixer_mlp_nvfp4 - quantizer_name: '*block_sparse_moe*weight_quantizer' cfg: $import: nvfp4 diff --git a/modelopt_recipes/general/ptq/nvfp4_mlp_only_mse-kv_fp8_cast.yaml b/modelopt_recipes/general/ptq/nvfp4_mlp_only_mse-kv_fp8_cast.yaml index 71c354ee1b1..18fed45d266 100644 --- a/modelopt_recipes/general/ptq/nvfp4_mlp_only_mse-kv_fp8_cast.yaml +++ b/modelopt_recipes/general/ptq/nvfp4_mlp_only_mse-kv_fp8_cast.yaml @@ -43,6 +43,18 @@ quantize: - quantizer_name: '*mlp*input_quantizer' cfg: $import: nvfp4 + - quantizer_name: '*.mixer.up_proj.weight_quantizer' + cfg: + $import: nvfp4_static + - quantizer_name: '*.mixer.up_proj.input_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*.mixer.down_proj.weight_quantizer' + cfg: + $import: nvfp4_static + - quantizer_name: '*.mixer.down_proj.input_quantizer' + cfg: + $import: nvfp4 - quantizer_name: '*block_sparse_moe*weight_quantizer' cfg: $import: nvfp4_static diff --git a/modelopt_recipes/general/ptq/nvfp4_omlp_only-kv_fp8.yaml b/modelopt_recipes/general/ptq/nvfp4_omlp_only-kv_fp8.yaml index 5348e8c7123..41541e4b2e4 100644 --- a/modelopt_recipes/general/ptq/nvfp4_omlp_only-kv_fp8.yaml +++ b/modelopt_recipes/general/ptq/nvfp4_omlp_only-kv_fp8.yaml @@ -18,6 +18,7 @@ imports: base_disable_all: configs/ptq/units/base_disable_all default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + mixer_mlp_nvfp4: configs/ptq/units/mixer_mlp_nvfp4 nvfp4: configs/numerics/nvfp4 kv_fp8: configs/ptq/units/kv_fp8 @@ -36,6 +37,7 @@ quantize: - quantizer_name: '*mlp*input_quantizer' cfg: $import: nvfp4 + - $import: mixer_mlp_nvfp4 - quantizer_name: '*block_sparse_moe*weight_quantizer' cfg: $import: nvfp4 diff --git a/modelopt_recipes/general/ptq/nvfp4_omlp_only-kv_fp8_cast.yaml b/modelopt_recipes/general/ptq/nvfp4_omlp_only-kv_fp8_cast.yaml index ba9e1e1c27a..14da2d92d4a 100644 --- a/modelopt_recipes/general/ptq/nvfp4_omlp_only-kv_fp8_cast.yaml +++ b/modelopt_recipes/general/ptq/nvfp4_omlp_only-kv_fp8_cast.yaml @@ -18,6 +18,7 @@ imports: base_disable_all: configs/ptq/units/base_disable_all default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + mixer_mlp_nvfp4: configs/ptq/units/mixer_mlp_nvfp4 nvfp4: configs/numerics/nvfp4 kv_fp8_cast: configs/ptq/units/kv_fp8_cast @@ -36,6 +37,7 @@ quantize: - quantizer_name: '*mlp*input_quantizer' cfg: $import: nvfp4 + - $import: mixer_mlp_nvfp4 - quantizer_name: '*block_sparse_moe*weight_quantizer' cfg: $import: nvfp4 diff --git a/modelopt_recipes/huggingface/models/nvidia/Nemotron-3-Nano-4B/ptq/nvfp4_w4a16.yaml b/modelopt_recipes/huggingface/models/nvidia/Nemotron-3-Nano-4B-BF16/ptq/nvfp4_w4a16.yaml similarity index 100% rename from modelopt_recipes/huggingface/models/nvidia/Nemotron-3-Nano-4B/ptq/nvfp4_w4a16.yaml rename to modelopt_recipes/huggingface/models/nvidia/Nemotron-3-Nano-4B-BF16/ptq/nvfp4_w4a16.yaml diff --git a/modelopt_recipes/huggingface/models/nvidia/Nemotron-3-Super-120B-A12B/ptq/nvfp4-max-calib.yaml b/modelopt_recipes/huggingface/models/nvidia/Nemotron-3-Super-120B-A12B-BF16/ptq/nvfp4-max-calib.yaml similarity index 100% rename from modelopt_recipes/huggingface/models/nvidia/Nemotron-3-Super-120B-A12B/ptq/nvfp4-max-calib.yaml rename to modelopt_recipes/huggingface/models/nvidia/Nemotron-3-Super-120B-A12B-BF16/ptq/nvfp4-max-calib.yaml diff --git a/modelopt_recipes/huggingface/models/nvidia/Nemotron-3-Super-120B-A12B/ptq/nvfp4-mse.yaml b/modelopt_recipes/huggingface/models/nvidia/Nemotron-3-Super-120B-A12B-BF16/ptq/nvfp4-mse.yaml similarity index 100% rename from modelopt_recipes/huggingface/models/nvidia/Nemotron-3-Super-120B-A12B/ptq/nvfp4-mse.yaml rename to modelopt_recipes/huggingface/models/nvidia/Nemotron-3-Super-120B-A12B-BF16/ptq/nvfp4-mse.yaml diff --git a/modelopt_recipes/huggingface/models/nvidia/Nemotron-3-Ultra-550B-A55B/ptq/nvfp4-4o6.yaml b/modelopt_recipes/huggingface/models/nvidia/Nemotron-3-Ultra-550B-A55B-BF16/ptq/nvfp4-4o6.yaml similarity index 100% rename from modelopt_recipes/huggingface/models/nvidia/Nemotron-3-Ultra-550B-A55B/ptq/nvfp4-4o6.yaml rename to modelopt_recipes/huggingface/models/nvidia/Nemotron-3-Ultra-550B-A55B-BF16/ptq/nvfp4-4o6.yaml diff --git a/modelopt_recipes/huggingface/models/nvidia/Nemotron-3.5-Lightning-30B-A3B-BF16/ptq/w4a16_nvfp4_4o6.yaml b/modelopt_recipes/huggingface/models/nvidia/Nemotron-3.5-Lightning-30B-A3B-BF16/ptq/w4a16_nvfp4_4o6.yaml new file mode 100644 index 00000000000..ab6933007b8 --- /dev/null +++ b/modelopt_recipes/huggingface/models/nvidia/Nemotron-3.5-Lightning-30B-A3B-BF16/ptq/w4a16_nvfp4_4o6.yaml @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Quantization config for nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4: +# NVFP4 weights use the four_over_six (4/6) per-block scale selection. +# Weight-only; activations BF16. MSE selects per-block M=6 (mult 1.0) vs +# M=4 (mult 1.5); the choice is baked into the amax by MSE. +# +# HF -> MCore name mapping for Nemotron-3.5-Lightning-30B-A3B (Mamba/Attn/MoE hybrid): +# - mixer.experts..{up,down}_proj -> mlp.experts.local_experts..linear_fc{1,2} +# - mixer.shared_experts.{up,down}_proj -> mlp.shared_experts.linear_fc{1,2} +# - mixer.in_proj / out_proj -> mixer.in_proj / out_proj (same name; W4A16) +# - lm_head -> output_layer (W4A16) +imports: + nvfp4_four_over_six: configs/numerics/nvfp4_four_over_six + fp8_default: configs/numerics/fp8 +metadata: + recipe_type: ptq + description: > + Lightning 3.5 W4A16 PTQ, NVFP4 four_over_six (4/6) weight scales. + Routed MoE experts, shared experts, and lm_head; FP8 Mamba in_proj/out_proj. + W4A16 use weight-only NVFP4 4/6 (static, MSE-selected). KV cache FP8. + Attention BF16. +quantize: + # 4/6: MSE selects per-block between M=6 (keep amax) and M=4 (amax x 6/4). + algorithm: + method: mse + fp8_scale_sweep: false + start_multiplier: 1.0 # M=6 (keep amax) + stop_multiplier: 1.5 # M=4 (amax x 6/4) + step_size: 0.5 # candidates [1.0, 1.5] + quant_cfg: + - quantizer_name: '*' + enable: false + # W4A16 NVFP4 4/6 weight-only (block 16, static, four_over_six: true). HF + MCore names. + - quantizer_name: '*mixer.experts*weight_quantizer' + enable: true + cfg: {$import: nvfp4_four_over_six} + - quantizer_name: '*mixer.shared_experts*weight_quantizer' + enable: true + cfg: {$import: nvfp4_four_over_six} + - quantizer_name: '*mlp.experts*weight_quantizer' + enable: true + cfg: {$import: nvfp4_four_over_six} + - quantizer_name: '*mlp.shared_experts*weight_quantizer' + enable: true + cfg: {$import: nvfp4_four_over_six} + # FP8 Mamba projections. + - quantizer_name: '*mixer.in_proj*weight_quantizer' + enable: true + cfg: {$import: fp8_default} + - quantizer_name: '*mixer.in_proj*input_quantizer' + enable: true + cfg: {$import: fp8_default} + - quantizer_name: '*mixer.out_proj*weight_quantizer' + enable: true + cfg: {$import: fp8_default} + - quantizer_name: '*mixer.out_proj*input_quantizer' + enable: true + cfg: {$import: fp8_default} + # KV cache -> FP8. + - quantizer_name: '*[kv]_bmm_quantizer' + enable: true + cfg: + num_bits: e4m3 + # lm_head (output_layer) W4A16 4/6 weight-only. + - quantizer_name: 'output_layer.weight_quantizer' + enable: true + cfg: {$import: nvfp4_four_over_six} + - quantizer_name: 'lm_head.weight_quantizer' + enable: true + cfg: {$import: nvfp4_four_over_six} + # Keep the entire MTP subtree in BF16. This rule must remain last so it + # overrides the broad expert, Mamba, KV-cache, and output-layer selectors. + - quantizer_name: 'mtp.*' + enable: false diff --git a/modelopt_recipes/huggingface/phi4mm/ptq/README.md b/modelopt_recipes/huggingface/phi4mm/ptq/README.md deleted file mode 100644 index bedaf1fcb6b..00000000000 --- a/modelopt_recipes/huggingface/phi4mm/ptq/README.md +++ /dev/null @@ -1,13 +0,0 @@ -# Phi-4-Multimodal PTQ recipes - -Phi-4-Multimodal is a multimodal model. Quantization should be applied only to -the language model; the speech, audio, image, and vision branches are kept in -full precision to avoid accuracy regressions on those modalities. - -| File | What's model-specific | -|------|-----------------------| -| `disabled_quantizers.yaml` | Reusable unit (`QuantizerCfgListConfig`). Merges the standard `default_disabled_quantizers` exclusions with Phi-4-MM ones (`*speech*`, `*audio*`, `*image*`, `*vision*`). Imported by recipes below as the single `disabled_quantizers` slot so they don't pull in two disabled-quantizer sets. | -| `nvfp4-kv_fp8_cast.yaml` | NVFP4 W4A4 model quantization + FP8 KV-cache cast (constant amax, no KV calibration). Identical numerics to the general `nvfp4` preset / `kv_fp8_cast` unit; what makes it model-specific is that it imports `disabled_quantizers.yaml` from this folder to skip the non-language branches. | - -Additional `-kv_fp8_cast.yaml` recipes can be generated for other formats -if needed; only `nvfp4-kv_fp8_cast.yaml` is shipped by default. diff --git a/modelopt_recipes/huggingface/phi4mm/ptq/nvfp4-kv_fp8_cast.yaml b/modelopt_recipes/huggingface/phi4mm/ptq/nvfp4-kv_fp8_cast.yaml deleted file mode 100644 index dfb1be1778d..00000000000 --- a/modelopt_recipes/huggingface/phi4mm/ptq/nvfp4-kv_fp8_cast.yaml +++ /dev/null @@ -1,36 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# Phi-4-Multimodal-specific PTQ recipe for the `nvfp4` quantization format. -# Equivalent to the general `nvfp4` preset with quantization disabled -# on non-language branches. - -imports: - base_disable_all: configs/ptq/units/base_disable_all - w4a4_nvfp4_nvfp4: configs/ptq/units/w4a4_nvfp4_nvfp4 - disabled_quantizers: huggingface/phi4mm/ptq/disabled_quantizers - kv_fp8_cast: configs/ptq/units/kv_fp8_cast - -metadata: - recipe_type: ptq - description: 'Phi-4-Multimodal PTQ recipe (nvfp4): same numerics as the general nvfp4 preset, applied to the language model only (speech, audio, image, - and vision branches are skipped).' -quantize: - algorithm: max - quant_cfg: - - $import: base_disable_all - - $import: w4a4_nvfp4_nvfp4 - - $import: kv_fp8_cast - - $import: disabled_quantizers diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index 22d6fbd5ff5..847117acc0c 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -232,7 +232,7 @@ that baseline. The deviations come in four kinds: |------|-------------------------------------|----------| | **Architecture-aware `quant_cfg`** | Per-sub-module format choices a single wildcard scheme can't express | `minimax_m3_vl`, `qwen3_5`, `qwen3_5_moe`, `vit`, `nemotron_llama` | | **Algorithm override** | Same numerics & scope, but the *calibration algorithm* is tweaked because the default breaks or regresses | `gemma`, `gemma4`, `mpt` | -| **Extra exclusions** | Adds disabled-quantizer patterns so non-language branches stay full precision | `nemotron_vl`, `phi4mm`, `diffusion_gemma` | +| **Extra exclusions** | Adds disabled-quantizer patterns so non-language branches stay full precision | `nemotron_vl`, `diffusion_gemma` | | **Checkpoint mirror** | A mixed-precision map reproducing one published checkpoint exactly | `models/nvidia/Nemotron-3-*`, `models/nvidia/Mistral-Medium-3.5-128B-NVFP4` | The numerics and standard exclusions are still inherited from `configs/` @@ -312,7 +312,7 @@ These quantize the **same layers** as the general recipes; only the *Why special:* identical scope/numerics to a general scheme, but a general recipe's default algorithm would overflow or regress here. -### Extra exclusions — `nemotron_vl`, `phi4mm`, `diffusion_gemma` +### Extra exclusions — `nemotron_vl`, `diffusion_gemma` Each of these is **numerically identical** to a general recipe. What makes them special is a model-local `disabled_quantizers.yaml` unit that *extends* the @@ -322,8 +322,6 @@ standard exclusions so a model-specific branch stays in full precision: `nvfp4_default-kv_fp8_cast` numerics, adding `*vision*`, `*image*`, `*radio*`, `*visual*`, `*encoder*`, `*model_encoder*` so only the language decoder is quantized. -- **`phi4mm`** (Phi-4-Multimodal) — general `nvfp4_default-kv_fp8_cast` - numerics, adding `*speech*`, `*audio*`, `*image*`, `*vision*`. - **`diffusion_gemma`** (block-diffusion encoder-decoder text LLM on a Gemma4 MoE backbone) — general `nvfp4_experts_only-kv_fp8_cast` numerics, adding `*self_conditioning*`: the self-conditioning network is text-only and never @@ -344,7 +342,7 @@ checkpoint's** quant config verbatim: `nvidia/Mistral-Medium-3.5-128B-NVFP4`: decoder MLP layers 4–86 use NVFP4 W4A4, edge MLP layers 0–3 and 87 use FP8 W8A8, and all attention projections and the KV cache use FP8. It uses max calibration. -- **`Nemotron-3-Super-120B-A12B/ptq/nvfp4-mse`** mirrors +- **`Nemotron-3-Super-120B-A12B-BF16/ptq/nvfp4-mse`** mirrors `nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4` exactly — a hybrid **Mamba-MoE** with a hand-mapped, **per-component** precision scheme: - MoE routed experts → NVFP4 W4A4, `group_size 16`, **static** weight scales @@ -355,13 +353,17 @@ checkpoint's** quant config verbatim: `nvfp4-mse.yaml` uses MSE calibration with an FP8-scale sweep (matches the release); `nvfp4-max-calib.yaml` is the identical layer map under plain `max` calibration, kept for comparison. -- **`Nemotron-3-Ultra-550B-A55B/ptq/nvfp4-4o6`** follows the same Super-style +- **`Nemotron-3-Ultra-550B-A55B-BF16/ptq/nvfp4-4o6`** follows the same Super-style component map (routed experts NVFP4 W4A4 block-16; shared experts + Mamba `in/out_proj` + KV cache FP8; everything else BF16), but the routed-expert weights use **Four-over-Six (4/6)** NVFP4: an MSE search picks each weight's amax multiplier from `[1.0, 1.5]` (M=6 vs. M=4). Activations stay dynamic NVFP4 (not MSE-calibrated). -- **`Nemotron-3-Nano-4B/ptq/nvfp4_w4a16`** mirrors the GGUF **Q4_K_M** bit +- **`Nemotron-3.5-Lightning-30B-A3B-BF16/ptq/w4a16_nvfp4_4o6`** applies + Four-over-Six NVFP4 W4A16 to routed experts, shared experts, and the language + model head; Mamba `in/out_proj` weights and inputs plus the KV cache use FP8, + while attention remains BF16. +- **`Nemotron-3-Nano-4B-BF16/ptq/nvfp4_w4a16`** mirrors the GGUF **Q4_K_M** bit allocation of the Nemotron-H hybrid, mapped onto NVFP4/FP8 **per layer**: Q4_K/Q5_0 linears → NVFP4 W4A4 (attention q/k/v/o kept uniform so export can fuse them), the Q6_K MLP `down_proj` layers → FP8 W8A8, embeddings → NVFP4 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/examples/megatron_bridge/test_prune_minitron.py b/tests/examples/megatron_bridge/test_prune_minitron.py index 32464814fef..78b30637468 100644 --- a/tests/examples/megatron_bridge/test_prune_minitron.py +++ b/tests/examples/megatron_bridge/test_prune_minitron.py @@ -19,6 +19,7 @@ from _test_utils.examples.megatron_bridge import qwen35_moe_bridge_supported from _test_utils.examples.run_command import extend_cmd_parts, run_example_command from _test_utils.torch.transformers_models import ( + create_tiny_deepseek_v3_dir, create_tiny_gemma3vl_dir, create_tiny_nemotron_h_dir, create_tiny_qwen3_5_moe_vl_dir, @@ -28,20 +29,18 @@ @pytest.mark.parametrize( - ("create_teacher", "megatron_format"), + ("create_teacher", "expected_pruned_config"), [ - # Dense Qwen3 LM, exported back to HF (reloadable to verify the pruned param count). + # Dense Qwen3 LM. pytest.param( lambda tmp_path, num_gpus: create_tiny_qwen3_dir( tmp_path, with_tokenizer=True, return_model=True, num_hidden_layers=num_gpus ), - False, + {}, id="qwen3", ), - # NemotronH (nemotron-3-nano): Mamba + attention + MoE hybrid. Saved in Megatron checkpoint - # format because HF export of a pruned NemotronH requires transformers<5. - # MTP heads are enabled so the run covers dropping them during calibration and the - # hybrid pattern MCore builds for them. + # NemotronH (nemotron-3-nano): Mamba + attention + MoE hybrid. + # MTP heads are enabled so the run covers dropping them during calibration. pytest.param( lambda tmp_path, num_gpus: create_tiny_nemotron_h_dir( tmp_path, @@ -50,26 +49,39 @@ num_nextn_predict_layers=1, mtp_hybrid_override_pattern="*E", ), - True, + {"n_shared_experts": 1}, id="nemotron_h", ), + # DeepSeek-V3: MLA (Q-LoRA) + MoE sizing the shared expert as + # n_shared_experts * moe_intermediate_size, so it covers candidate_filter end-to-end: + # without the filter the search picks a shared size that is not a multiple of the routed + # one and the reload below fails on the resulting shape mismatch. + pytest.param( + # n_group=1 so num_moe_experts stays divisible by it after expert pruning. + lambda tmp_path, num_gpus: create_tiny_deepseek_v3_dir( + tmp_path, + with_tokenizer=True, + return_model=True, + num_hidden_layers=num_gpus, + n_group=1, + topk_group=1, + ), + {"n_shared_experts": 1}, + id="deepseek_v3", + ), ], ) -def test_prune_minitron(tmp_path, num_gpus, create_teacher, megatron_format): +def test_prune_minitron(tmp_path, num_gpus, create_teacher, expected_pruned_config): teacher_hf_path, teacher_model = create_teacher(tmp_path, num_gpus) teacher_params = sum(p.numel() for p in teacher_model.parameters()) prune_target_params = int(teacher_params * 0.8) pruned_path = tmp_path / "pruned" - output_kwarg = ( - {"output_megatron_path": pruned_path} - if megatron_format - else {"output_hf_path": pruned_path} - ) # TODO: Dont enable grouped GEMM for MoE models until nemo:26.08 container prune_command_parts = extend_cmd_parts( ["torchrun", f"--nproc_per_node={num_gpus}", "prune_minitron.py", "--no_moe_grouped_gemm"], hf_model_name_or_path=teacher_hf_path, + output_hf_path=pruned_path, pp_size=num_gpus, calib_dataset_name="cnn_dailymail", calib_num_samples=8, @@ -80,17 +92,14 @@ def test_prune_minitron(tmp_path, num_gpus, create_teacher, megatron_format): ss_channel_divisor=4, hparams_to_skip="num_attention_heads", top_k=1, - **output_kwarg, ) run_example_command(prune_command_parts, example_path="megatron_bridge") - if megatron_format: - # HF reload of a pruned NemotronH needs transformers<5; just verify the Megatron checkpoint. - assert (pruned_path / "latest_checkpointed_iteration.txt").exists() - else: - assert (pruned_path / "config.json").exists() - pruned_model = AutoModelForCausalLM.from_pretrained(pruned_path) - assert sum(p.numel() for p in pruned_model.parameters()) <= prune_target_params + assert (pruned_path / "config.json").exists() + pruned_model = AutoModelForCausalLM.from_pretrained(pruned_path) + assert sum(p.numel() for p in pruned_model.parameters()) <= prune_target_params + for field, expected in expected_pruned_config.items(): + assert getattr(pruned_model.config, field) == expected @pytest.mark.parametrize( 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/gpu_megatron/torch/prune/plugins/test_mcore_mamba_minitron_pruning.py b/tests/gpu_megatron/torch/prune/plugins/test_mcore_mamba_minitron_pruning.py index 69d8c7c31ec..79f33b806f9 100644 --- a/tests/gpu_megatron/torch/prune/plugins/test_mcore_mamba_minitron_pruning.py +++ b/tests/gpu_megatron/torch/prune/plugins/test_mcore_mamba_minitron_pruning.py @@ -446,13 +446,32 @@ def _test_mcore_mamba_hybrid_pruning_nas_memory_mb(rank, size, ckpt_dir): ) memory_threshold = baseline_memory_mb * 0.7 + def _shared_is_multiple_of_routed(ss_cfg): + return ss_cfg["moe_shared_expert_intermediate_size"] % ss_cfg["moe_ffn_hidden_size"] == 0 + constraints = {"memory_mb": memory_threshold} config = { **_base_nas_config(ckpt_dir), "seq_length": sequence_length, "batch_size": 1, + # moe_shared_expert_intermediate_size is skipped, so the filter reads it from the model + # config: only routed sizes that divide it survive. + "candidate_filter": _shared_is_multiple_of_routed, } - model, searcher_state = prune_minitron(model, constraints, config, _NAS_CHANNEL_DIVISOR) + stdout_capture = io.StringIO() + with contextlib.redirect_stdout(stdout_capture): + model, searcher_state = prune_minitron(model, constraints, config, _NAS_CHANNEL_DIVISOR) + + shared_size = _NAS_MODEL_KWARGS["moe_shared_expert_intermediate_size"] + (candidates,) = searcher_state["all_candidates_per_constraint"].values() + assert all(shared_size % c.ss_config["moe_ffn_hidden_size"] == 0 for c in candidates) + if rank == 0: + # Half of the 512-combo search space: moe_ffn_hidden_size is [12, 16] and only 16 divides the + # (skipped, so unpruned) moe_shared_expert_intermediate_size of 16. + output = stdout_capture.getvalue() + match = re.search(r"Rejected (\d+) candidates", output) + assert match, f"candidate_filter rejected nothing:\n{output}" + assert int(match.group(1)) == 256, output pruned_params, _ = mcore_param_count( model.config, @@ -465,17 +484,18 @@ def _test_mcore_mamba_hybrid_pruning_nas_memory_mb(rank, size, ckpt_dir): sorted_layers = _get_sorted_layers(searcher_state) # fmt: off if sorted_layers == [1, 4, 3, 2]: + # All moe_ffn_hidden_size 16: the 12s do not divide the unpruned shared size of 16. expected_top_k = [ - [{"num_layers": 4, "hidden_size": 12, "mamba_num_heads": 8, "mamba_head_dim": 16, "num_moe_experts": 6, "moe_ffn_hidden_size": 12, "ffn_hidden_size": 24}, {"memory_mb": 0.0226287841796875}, 114], # noqa: E501 [{"num_layers": 4, "hidden_size": 12, "mamba_num_heads": 8, "mamba_head_dim": 12, "num_moe_experts": 8, "moe_ffn_hidden_size": 16, "ffn_hidden_size": 32}, {"memory_mb": 0.022613525390625}, 124], # noqa: E501 [{"num_layers": 4, "hidden_size": 12, "mamba_num_heads": 6, "mamba_head_dim": 16, "num_moe_experts": 8, "moe_ffn_hidden_size": 16, "ffn_hidden_size": 32}, {"memory_mb": 0.022556304931640625}, 126], # noqa: E501 - [{"num_layers": 4, "hidden_size": 16, "mamba_num_heads": 6, "mamba_head_dim": 12, "num_moe_experts": 7, "moe_ffn_hidden_size": 12, "ffn_hidden_size": 24}, {"memory_mb": 0.022541046142578125}, 113], # noqa: E501 - [{"num_layers": 3, "hidden_size": 16, "mamba_num_heads": 8, "mamba_head_dim": 16, "num_moe_experts": 5, "moe_ffn_hidden_size": 12, "ffn_hidden_size": 20}, {"memory_mb": 0.0225067138671875}, 112], # noqa: E501 [{"num_layers": 3, "hidden_size": 16, "mamba_num_heads": 8, "mamba_head_dim": 16, "num_moe_experts": 5, "moe_ffn_hidden_size": 16, "ffn_hidden_size": 20}, {"memory_mb": 0.0225067138671875}, 116], # noqa: E501 - [{"num_layers": 3, "hidden_size": 16, "mamba_num_heads": 8, "mamba_head_dim": 16, "num_moe_experts": 6, "moe_ffn_hidden_size": 12, "ffn_hidden_size": 20}, {"memory_mb": 0.0225067138671875}, 113], # noqa: E501 [{"num_layers": 3, "hidden_size": 16, "mamba_num_heads": 8, "mamba_head_dim": 16, "num_moe_experts": 6, "moe_ffn_hidden_size": 16, "ffn_hidden_size": 20}, {"memory_mb": 0.0225067138671875}, 117], # noqa: E501 - [{"num_layers": 3, "hidden_size": 16, "mamba_num_heads": 8, "mamba_head_dim": 16, "num_moe_experts": 7, "moe_ffn_hidden_size": 12, "ffn_hidden_size": 20}, {"memory_mb": 0.0225067138671875}, 114], # noqa: E501 [{"num_layers": 3, "hidden_size": 16, "mamba_num_heads": 8, "mamba_head_dim": 16, "num_moe_experts": 7, "moe_ffn_hidden_size": 16, "ffn_hidden_size": 20}, {"memory_mb": 0.0225067138671875}, 118], # noqa: E501 + [{"num_layers": 3, "hidden_size": 16, "mamba_num_heads": 8, "mamba_head_dim": 16, "num_moe_experts": 8, "moe_ffn_hidden_size": 16, "ffn_hidden_size": 20}, {"memory_mb": 0.0225067138671875}, 119], # noqa: E501 + [{"num_layers": 4, "hidden_size": 16, "mamba_num_heads": 6, "mamba_head_dim": 12, "num_moe_experts": 5, "moe_ffn_hidden_size": 16, "ffn_hidden_size": 28}, {"memory_mb": 0.022480010986328125}, 119], # noqa: E501 + [{"num_layers": 4, "hidden_size": 12, "mamba_num_heads": 8, "mamba_head_dim": 12, "num_moe_experts": 8, "moe_ffn_hidden_size": 16, "ffn_hidden_size": 28}, {"memory_mb": 0.022430419921875}, 120], # noqa: E501 + [{"num_layers": 4, "hidden_size": 12, "mamba_num_heads": 6, "mamba_head_dim": 16, "num_moe_experts": 8, "moe_ffn_hidden_size": 16, "ffn_hidden_size": 28}, {"memory_mb": 0.022373199462890625}, 122], # noqa: E501 + [{"num_layers": 4, "hidden_size": 12, "mamba_num_heads": 8, "mamba_head_dim": 12, "num_moe_experts": 8, "moe_ffn_hidden_size": 16, "ffn_hidden_size": 24}, {"memory_mb": 0.022247314453125}, 116], # noqa: E501 ] else: raise RuntimeError(f"FIXME: Non deterministic test, assertions may fail: {sorted_layers=}") diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py index 36f80787931..504e9cce965 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py @@ -14,9 +14,13 @@ # limitations under the License. import copy +import sys +import types from contextlib import nullcontext from functools import partial from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch import pytest import torch @@ -49,6 +53,7 @@ get_tensor_model_parallel_group, ) from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear +from megatron.core.transformer import MegatronModule, TransformerConfig from megatron.core.transformer.moe.experts import SequentialMLP, TEGroupedMLP from megatron.core.transformer.moe.router import TopKRouter @@ -58,8 +63,11 @@ from modelopt.torch.quantization.algorithms import QuantRecipe, _AutoQuantizeBaseSearcher from modelopt.torch.quantization.nn import QuantModuleRegistry from modelopt.torch.quantization.plugins.megatron import ( + _output_layer_untied, _QuantTEMCoreRowParallelLinear, + _resolve_output_layer_untied, get_mcore_layerwise_calibration_layers, + megatron_replace_quant_module_hook, ) from modelopt.torch.quantization.utils import is_quantized_linear from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector @@ -1301,3 +1309,113 @@ def test_homogeneous_sharded_state_dict_te_spec(dist_workers, tmp_path): {"transformer_impl": "transformer_engine"}, ), ) + + +def test_resolve_output_layer_untied(): + """The tiedness signal is read off the model, not from Megatron-LM global args.""" + + class _Flagged(torch.nn.Module): + def __init__(self, shared): + super().__init__() + self.share_embeddings_and_output_weights = shared + + # No signal anywhere -> unknown. + assert _resolve_output_layer_untied(torch.nn.Module()) is None + + # The root's own flag wins over any subtree. + root = _Flagged(False) + root.inner = _Flagged(True) + assert _resolve_output_layer_untied(root) is True + + # Otherwise fall back to a subtree scan. + root = torch.nn.Module() + root.language_model = _Flagged(True) + assert _resolve_output_layer_untied(root) is False + + # Subtrees that do not own the language model's output_layer are skipped: the vision tower + # and a distillation teacher, either of which may be tied differently from the student. + root = torch.nn.Module() + root.vision_model = _Flagged(True) + root._teacher_model = _Flagged(True) + root.language_model = _Flagged(False) + assert _resolve_output_layer_untied(root) is True + + +@pytest.mark.parametrize("mlm_untied", [True, False]) +def test_output_layer_untied_falls_back_to_megatron_lm_args(mlm_untied): + """With no model-derived flag, the answer comes from Megatron-LM's args.""" + + class _Config: + pass + + fake_training = types.ModuleType("megatron.training") + fake_training.get_args = lambda: SimpleNamespace(untie_embeddings_and_output_weights=mlm_untied) + + config = _Config() + with patch.dict(sys.modules, {"megatron.training": fake_training}): + assert _output_layer_untied(config) is mlm_untied + + # The model-derived flag takes precedence over the args fallback. + config.modelopt_output_layer_untied = not mlm_untied + with patch.dict(sys.modules, {"megatron.training": fake_training}): + assert _output_layer_untied(config) is (not mlm_untied) + + +def test_output_layer_untied_warns_once_when_args_unavailable(): + """Without either signal the layer is treated as tied, and the warning is not repeated.""" + + class _Config: + pass + + broken = types.ModuleType("megatron.training") # no get_args attribute + + config = _Config() + with ( + patch.dict(sys.modules, {"megatron.training": broken}), + patch("modelopt.torch.quantization.plugins.megatron.warn_rank_0") as warn, + ): + assert _output_layer_untied(config) is False + assert _output_layer_untied(config) is False + assert warn.call_count == 1 + + +def test_output_layer_untied_warns_when_args_uninitialized(): + """Megatron-LM importable but not initialized: treated as tied, warned once.""" + + class _Config: + pass + + def _uninitialized(): + raise AssertionError("args is not initialized.") + + fake_training = types.ModuleType("megatron.training") + fake_training.get_args = _uninitialized + + config = _Config() + with ( + patch.dict(sys.modules, {"megatron.training": fake_training}), + patch("modelopt.torch.quantization.plugins.megatron.warn_rank_0") as warn, + ): + assert _output_layer_untied(config) is False + assert _output_layer_untied(config) is False + assert warn.call_count == 1 + + +def test_output_layer_untied_not_stamped_onto_teacher_config(): + """A distillation teacher keeps its own tiedness; the student's answer must not leak in.""" + + def _config(): + return TransformerConfig(num_layers=1, hidden_size=8, num_attention_heads=1) + + class _Tiny(MegatronModule): + def __init__(self, config, shared): + super().__init__(config) + self.share_embeddings_and_output_weights = shared + + student = _Tiny(_config(), shared=False) # untied + student._teacher_model = _Tiny(_config(), shared=True) # tied -- must not be overwritten + + megatron_replace_quant_module_hook(student) + + assert student.config.modelopt_output_layer_untied is True + assert not hasattr(student._teacher_model.config, "modelopt_output_layer_untied") diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py index 05b3c6ace8d..578922db077 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py @@ -28,7 +28,7 @@ import vllm from vllm.v1.attention.backend import CommonAttentionMetadata from vllm.v1.attention.backends import flashinfer as flashinfer_backend -from vllm.v1.attention.backends.flash_attn import FlashAttentionImpl +from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend, FlashAttentionImpl from vllm.v1.attention.backends.flashinfer import ( FlashInferBackend, FlashInferImpl, @@ -541,6 +541,76 @@ def _make_flash_attention_impl(*, sparse=False, quantized=False): return impl +def _flash_attention_kv_cache(num_blocks, page_size, num_kv_heads, head_size): + layout = vllm_plugin._flash_attention_kv_cache_layout() + if layout == "packed": + shape = [num_blocks, num_kv_heads, page_size, 2 * head_size] + else: + shape = [num_blocks, page_size, num_kv_heads, head_size] + shape.insert(0 if layout == "kv-first" else 1, 2) + return torch.zeros(shape, dtype=torch.float16) + + +@pytest.mark.parametrize( + ("layout", "backend_shape"), + [ + ("kv-first", (2, 3, 16, 1, 16)), + ("blocks-first", (3, 2, 16, 1, 16)), + ("packed", (3, 1, 16, 32)), + ], +) +def test_flash_attention_forward_follows_backend_kv_cache_layout( + monkeypatch, layout, backend_shape +): + impl = _make_flash_attention_impl(sparse=True) + if layout == "packed": + shape = [3, impl.num_kv_heads, 16, 2 * impl.head_size] + else: + shape = [3, 16, impl.num_kv_heads, impl.head_size] + shape.insert(0 if layout == "kv-first" else 1, 2) + monkeypatch.setattr( + FlashAttentionBackend, "get_kv_cache_shape", staticmethod(lambda *_args: backend_shape) + ) + vllm_plugin._flash_attention_kv_cache_layout.cache_clear() + kv_cache = torch.zeros(shape, dtype=torch.float16) + query = torch.zeros(4, impl.num_heads, impl.head_size, dtype=torch.float16) + metadata = _flash_attention_metadata(query.shape[0], 16) + captured = {} + + def fake_attention(query, **kwargs): + captured.update(kwargs) + return torch.zeros_like(query) + + monkeypatch.setattr(vllm_plugin, "triton_attention", fake_attention) + + try: + impl.forward( + layer=None, + query=query, + key=query, + value=query, + kv_cache=kv_cache, + attn_metadata=metadata, + output=torch.empty_like(query), + ) + finally: + vllm_plugin._flash_attention_kv_cache_layout.cache_clear() + + if layout == "packed": + expected_key_cache, expected_value_cache = kv_cache.transpose(1, 2).split( + impl.head_size, dim=-1 + ) + else: + expected_key_cache, expected_value_cache = kv_cache.unbind(0 if layout == "kv-first" else 1) + assert captured["k_cache"].shape == expected_key_cache.shape + assert captured["v_cache"].shape == expected_value_cache.shape + assert captured["k_cache"].stride() == expected_key_cache.stride() + assert captured["v_cache"].stride() == expected_value_cache.stride() + assert captured["k_cache"].data_ptr() == expected_key_cache.data_ptr() + assert captured["v_cache"].data_ptr() == expected_value_cache.data_ptr() + assert captured["page_size"] == 16 + + def _flash_attention_mixed_metadata(decode_len=1, prefill_len=17): query_lens = (decode_len, prefill_len) seq_lens = (16, 34) @@ -568,7 +638,7 @@ def test_flash_attention_mixed_batch_splits_decode_and_prefill(monkeypatch, quan prefill_tokens = 17 impl = _make_flash_attention_impl(sparse=True, quantized=quantized) query = torch.zeros(1 + prefill_tokens, 2, 64, dtype=torch.float16) - kv_cache = torch.zeros(2, 4, 16, 2, 64, dtype=torch.float16) + kv_cache = _flash_attention_kv_cache(4, 16, 2, 64) metadata = _flash_attention_mixed_metadata(decode_len=1, prefill_len=prefill_tokens) layer = SimpleNamespace( _query_quant_in_kernel=quantized, @@ -761,7 +831,7 @@ def test_forward_delegates_cascade_metadata_to_vllm(monkeypatch): """Cascade/prefix-cache metadata should use vLLM's native implementation.""" impl = _clone_sparse_impl(_make_old_impl()) q = torch.zeros(1, impl.num_heads, impl.head_size, dtype=torch.float16) - kv_cache = torch.zeros(2, 1, 16, impl.num_kv_heads, impl.head_size, dtype=torch.float16) + kv_cache = _flash_attention_kv_cache(1, 16, impl.num_kv_heads, impl.head_size) output = torch.empty_like(q) attn_metadata = type("AttnMetadata", (), {"use_cascade": True})() called = {} @@ -838,9 +908,7 @@ def test_forward_delegates_launches_without_effective_sparse_work( impl = _clone_sparse_impl(_make_old_impl()) impl.sparse_kw = sparse_kw q = torch.zeros(max_query_len, impl.num_heads, impl.head_size, dtype=torch.float16) - kv_cache = torch.zeros( - 2, 1, max_seq_len, impl.num_kv_heads, impl.head_size, dtype=torch.float16 - ) + kv_cache = _flash_attention_kv_cache(1, max_seq_len, impl.num_kv_heads, impl.head_size) output = torch.empty_like(q) attn_metadata = _flash_attention_metadata(max_query_len, max_seq_len) called = {} @@ -903,7 +971,7 @@ def test_forward_resolves_calibrated_skip_softmax_threshold(monkeypatch): "target_sparse_ratio": {"prefill": 0.4, "decode": 0.6}, } q = torch.zeros(max_query_len, impl.num_heads, impl.head_size, dtype=torch.float16) - kv_cache = torch.zeros(2, 1, seq_len, impl.num_kv_heads, impl.head_size, dtype=torch.float16) + kv_cache = _flash_attention_kv_cache(1, seq_len, impl.num_kv_heads, impl.head_size) attn_metadata = _flash_attention_metadata(max_query_len, seq_len) captured = {} @@ -980,7 +1048,7 @@ def quantize_q(query): } q = torch.full((4, impl.num_heads, impl.head_size), 2.0, dtype=torch.float16) q[2:] = 10_000 - kv_cache = torch.zeros(2, 4, 16, impl.num_kv_heads, impl.head_size, dtype=torch.float16) + kv_cache = _flash_attention_kv_cache(4, 16, impl.num_kv_heads, impl.head_size) metadata = SimpleNamespace( num_actual_tokens=q.shape[0], max_query_len=1, @@ -1023,8 +1091,15 @@ def fake_decode(query, key_cache, value_cache, block_table, seq_lens, **kwargs): "v_qdq_scale": 1.0, } key_cache, value_cache, block_table, seq_lens, decode_kw = calls["decode"] - assert key_cache.data_ptr() == kv_cache[0].data_ptr() - assert value_cache.data_ptr() == kv_cache[1].data_ptr() + layout = vllm_plugin._flash_attention_kv_cache_layout() + if layout == "packed": + expected_key_cache, expected_value_cache = kv_cache.transpose(1, 2).split( + impl.head_size, dim=-1 + ) + else: + expected_key_cache, expected_value_cache = kv_cache.unbind(0 if layout == "kv-first" else 1) + assert key_cache.data_ptr() == expected_key_cache.data_ptr() + assert value_cache.data_ptr() == expected_value_cache.data_ptr() assert block_table is metadata.block_table assert seq_lens is metadata.seq_lens assert calls["query"].shape[0] == metadata.seq_lens.shape[0] @@ -1048,7 +1123,7 @@ def test_quantized_skip_softmax_decode_stays_on_shared_kernel(monkeypatch): } impl.sparse_kw = {"skip_softmax_threshold": 0.001} q = torch.zeros(1, impl.num_heads, impl.head_size, dtype=torch.float16) - kv_cache = torch.zeros(2, 1, 16, impl.num_kv_heads, impl.head_size, dtype=torch.float16) + kv_cache = _flash_attention_kv_cache(1, 16, impl.num_kv_heads, impl.head_size) metadata = _flash_attention_metadata(1, 16) captured = {} @@ -1140,7 +1215,7 @@ def test_forward_allows_chunked_prefill_metadata(monkeypatch): q_len = 4 kv_len = 10 q = torch.zeros(q_len, impl.num_heads, impl.head_size, dtype=torch.float16) - kv_cache = torch.zeros(2, 1, 16, impl.num_kv_heads, impl.head_size, dtype=torch.float16) + kv_cache = _flash_attention_kv_cache(1, 16, impl.num_kv_heads, impl.head_size) attn_metadata = _flash_attention_metadata(q_len, kv_len) captured = {} diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_plugin.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_plugin.py index fa11b144354..d08a4072b4d 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_plugin.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_plugin.py @@ -20,7 +20,7 @@ * ``query_start_loc`` -> ``b_start_loc`` / ``b_seq_len`` * ``seq_lens`` -> ``b_seq_len_k`` -* ``kv_cache.unbind(0)`` -> key_cache / value_cache (axis order) +* backend-declared K/V axis -> key_cache / value_cache * ``k_cache.shape[1]`` -> ``page_size`` Asserted against a contiguous reference call to the underlying Triton kernel. @@ -33,7 +33,7 @@ from vllm.v1.attention.backends.flash_attn import FlashAttentionImpl from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE -from modelopt.torch.sparsity.attention_sparsity.plugins.vllm import ModelOptSparseAttentionImpl +from modelopt.torch.sparsity.attention_sparsity.plugins import vllm as vllm_plugin if TRITON_KERNEL_AVAILABLE: from modelopt.torch.kernels.common.attention import attention as triton_attention @@ -46,11 +46,19 @@ } +def _make_backend_paged_cache(k_cache, v_cache): + layout = vllm_plugin._flash_attention_kv_cache_layout() + if layout == "kv-first": + return torch.stack([k_cache, v_cache], dim=0) + if layout == "blocks-first": + return torch.stack([k_cache, v_cache], dim=1) + return torch.cat([k_cache, v_cache], dim=-1).transpose(1, 2) + + def _make_paged_cache(k, v, b_start_loc, b_seq_len, num_kv_heads, head_dim, page_size): - """Scatter contiguous K/V into a paged KV cache stacked as [2, ...]. + """Scatter contiguous K/V into the installed vLLM paged-cache layout. - Returns a single ``kv_cache`` tensor (matching vLLM's layout that - ``ModelOptSparseAttentionImpl`` consumes via ``kv_cache.unbind(0)``). + Returns a single ``kv_cache`` tensor with the backend-declared K/V axis. """ batch = b_seq_len.shape[0] device, dtype = k.device, k.dtype @@ -76,14 +84,13 @@ def _make_paged_cache(k, v, b_start_loc, b_seq_len, num_kv_heads, head_dim, page v_cache[g, :n] = v[start + ts : start + te] g += 1 - # Stack on a new leading axis so kv_cache.unbind(0) recovers (k_cache, v_cache). - kv_cache = torch.stack([k_cache, v_cache], dim=0) + kv_cache = _make_backend_paged_cache(k_cache, v_cache) return kv_cache, block_table def _make_impl(num_heads, head_dim, num_kv_heads): """Construct ModelOptSparseAttentionImpl with minimal valid kwargs.""" - return ModelOptSparseAttentionImpl( + return vllm_plugin.ModelOptSparseAttentionImpl( num_heads=num_heads, head_size=head_dim, scale=1.0 / (head_dim**0.5), @@ -132,7 +139,7 @@ def test_prefill_matches_contiguous(self): **_ACTIVE_PREFILL_SPARSE_KW, ) - # Build paged kv_cache shaped [2, num_blocks, page_size, num_kv_heads, head_dim]. + # Build the paged cache using the installed backend's K/V axis. kv_cache, block_table = _make_paged_cache( k, v, b_start_loc, b_seq_len, num_kv_heads, head_dim, page_size ) @@ -174,7 +181,8 @@ def test_chunked_prefill_is_forwarded_to_kernel(self): block_table=torch.zeros(1, 1, device="cuda", dtype=torch.int32), ) q = torch.zeros(4, 2, 64, device="cuda", dtype=torch.float16) - kv_cache = torch.zeros(2, 1, 16, 2, 64, device="cuda", dtype=torch.float16) + k_cache = torch.zeros(1, 16, 2, 64, device="cuda", dtype=torch.float16) + kv_cache = _make_backend_paged_cache(k_cache, torch.zeros_like(k_cache)) out = impl.forward( layer=None, query=q, @@ -346,8 +354,7 @@ def test_page_size_inferred_from_k_cache(self): kv_cache, block_table = _make_paged_cache( k, v, b_start_loc, b_seq_len, num_kv_heads, head_dim, page_size ) - # Sanity: kv_cache axis 1 is page_size. - assert kv_cache.shape == (2, seq_len // page_size, page_size, num_kv_heads, head_dim) + assert kv_cache.shape[2] == page_size attn_metadata = SimpleNamespace( num_actual_tokens=seq_len, diff --git a/tests/unit/recipe/test_loader.py b/tests/unit/recipe/test_loader.py index e15b897a224..f88240b4dcf 100644 --- a/tests/unit/recipe/test_loader.py +++ b/tests/unit/recipe/test_loader.py @@ -213,6 +213,41 @@ def test_nvfp4_mlp_only_novit_recipe_disables_vision_quantizers(): assert {"*visual*", "*vision_tower*"} <= disabled_quantizers +@pytest.mark.parametrize( + "recipe_path", + [ + "general/ptq/nvfp4_mlp_only-kv_fp8", + "general/ptq/nvfp4_mlp_only-novit-kv_fp8", + "general/ptq/nvfp4_mlp_only-kv_fp8_cast", + "general/ptq/nvfp4_mlp_only_mse-kv_fp8_cast", + "general/ptq/nvfp4_omlp_only-kv_fp8", + "general/ptq/nvfp4_omlp_only-kv_fp8_cast", + ], +) +def test_nvfp4_mlp_only_recipes_match_nemotron_h_dense_mlp(recipe_path): + recipe = load_recipe(recipe_path) + enabled_patterns = [ + entry["quantizer_name"] + for entry in recipe.quantize.model_dump()["quant_cfg"] + if entry["enable"] + ] + + for quantizer_name in ( + "backbone.layers.0.mixer.up_proj.weight_quantizer", + "backbone.layers.0.mixer.up_proj.input_quantizer", + "backbone.layers.0.mixer.down_proj.weight_quantizer", + "backbone.layers.0.mixer.down_proj.input_quantizer", + ): + assert any(fnmatch(quantizer_name, pattern) for pattern in enabled_patterns) + + for quantizer_name in ( + "backbone.layers.0.mixer.in_proj.weight_quantizer", + "backbone.layers.0.mixer.out_proj.input_quantizer", + "backbone.layers.0.mixer.shared_experts.up_proj.weight_quantizer", + ): + assert not any(fnmatch(quantizer_name, pattern) for pattern in enabled_patterns) + + @pytest.mark.parametrize( "recipe_path", [ diff --git a/tests/unit/torch/export/test_hf_spec_rope_export.py b/tests/unit/torch/export/test_hf_spec_rope_export.py index 720082bc617..fbeb218793e 100644 --- a/tests/unit/torch/export/test_hf_spec_rope_export.py +++ b/tests/unit/torch/export/test_hf_spec_rope_export.py @@ -139,3 +139,16 @@ def test_dflash_rope_theta_inherits_base(): """rope_theta is inherited from the target/base config (draft drafts for the base).""" config = _make_dflash_exporter(base_rope_theta=5000000.0)._export_config() assert config["rope_theta"] == 5000000.0 + + +def test_dflash_rope_theta_inherits_base_rope_parameters(): + """Transformers 5 stores the target RoPE base in rope_parameters.""" + exporter = _make_dflash_exporter(base_rope_theta=None) + exporter.model.config.rope_parameters = { + "rope_type": "default", + "rope_theta": 5000000.0, + } + + config = exporter._export_config() + + assert config["rope_theta"] == 5000000.0 diff --git a/tests/unit/torch/export/test_quant_aware_conversion.py b/tests/unit/torch/export/test_quant_aware_conversion.py index 98bf6c250c2..7cbe311370c 100644 --- a/tests/unit/torch/export/test_quant_aware_conversion.py +++ b/tests/unit/torch/export/test_quant_aware_conversion.py @@ -316,6 +316,151 @@ def test_nested_text_prefix_reverse_still_applies_to_text_model(): assert mapper("model.layers.0") == "model.language_model.layers.0" +def test_scoped_submodel_prefix_change_does_not_capture_siblings(): + """A vision sub-model's ``PrefixChange`` must not prefix the whole VLM state dict. + + NVBug 6525511: ``LlavaForConditionalGeneration`` on transformers>=5.12 collects the + vision tower's own "add ``vision_model.``" prefix change. transformers scopes it to + ``model.vision_tower`` via ``scope_prefix`` and only matches keys under that prefix; + applying the raw pattern instead prefixes *every* key, so the export writes + ``vision_model.language_model.*`` / ``vision_model.lm_head.*`` and vLLM fails with + "There is no module or parameter named 'vision_model'". + """ + pytest.importorskip("transformers.core_model_loading") + # Local import: optional dependency, guarded by the importorskip above. + from transformers.core_model_loading import PrefixChange + + model = torch.nn.Module() + model.model = torch.nn.Module() + model.model.vision_tower = torch.nn.Module() + model.model.vision_tower.encoder = torch.nn.Linear(2, 2, bias=False) + model.model.language_model = torch.nn.Module() + model.model.language_model.layers = torch.nn.ModuleList([torch.nn.Linear(2, 2, bias=False)]) + model.lm_head = torch.nn.Linear(2, 2, bias=False) + + prefix_change = PrefixChange(prefix_to_remove="vision_model") + prefix_change.scope_prefix = "model.vision_tower" + prefix_change.base_model_prefix = "model" + model._weight_conversions = [prefix_change] + + state_dict = { + "model.vision_tower.encoder.weight": torch.randn(2, 2), + "model.language_model.layers.0.weight": torch.randn(2, 2), + "lm_head.weight": torch.randn(2, 2), + } + reverted = revert_weight_conversion_quant_aware(model, state_dict) + + # Only the vision tower's own subtree gains the ``vision_model.`` segment. + assert set(reverted) == { + "model.vision_tower.vision_model.encoder.weight", + "model.language_model.layers.0.weight", + "lm_head.weight", + } + # Regression guard: nothing may be moved under a bogus top-level ``vision_model``. + assert not any(k.startswith("vision_model.") for k in reverted) + + +def test_scoped_rule_maps_config_module_names_consistently(): + """``build_reverse_name_mapper`` must apply the same scoping as the weight rename. + + Otherwise ``exclude_modules`` (which lists the BF16 vision tower) lands in a + different namespace than the weights and a deployment loader silently treats an + excluded layer as quantized. + """ + pytest.importorskip("transformers.core_model_loading") + # Local import: optional dependency, guarded by the importorskip above. + from transformers.core_model_loading import PrefixChange + + model = torch.nn.Module() + model.model = torch.nn.Module() + model.model.vision_tower = torch.nn.Module() + model.model.vision_tower.encoder = torch.nn.Linear(2, 2, bias=False) + model.model.language_model = torch.nn.Module() + model.model.language_model.layers = torch.nn.ModuleList([torch.nn.Linear(2, 2, bias=False)]) + + prefix_change = PrefixChange(prefix_to_remove="vision_model") + prefix_change.scope_prefix = "model.vision_tower" + prefix_change.base_model_prefix = "model" + model._weight_conversions = [prefix_change] + + mapper = build_reverse_name_mapper(model) + assert mapper is not None + assert mapper("model.vision_tower.encoder") == "model.vision_tower.vision_model.encoder" + # Sibling namespaces are untouched. + assert mapper("model.language_model.layers.0") == "model.language_model.layers.0" + # A trailing-wildcard exclude pattern tracks the same rename its weights got, so the + # excluded (BF16) vision tower still matches the exported tensor names. + assert mapper("model.vision_tower*") == "model.vision_tower.vision_model*" + + +def test_root_scoped_rule_still_faces_shadowing_guard(): + """A ``scope_prefix == ""`` rule has whole-key-space reach and must not bypass #2032. + + ``_scope_prefixes`` keeps an empty candidate for the root scope, which + ``_sub_scoped`` matches against every key -- so such a rule is as broad as an + unscoped one. Skipping the shadowing heuristic merely because ``scope_prefixes`` is a + non-empty *tuple* would reintroduce NVBug 6525534: the nested text model's + ``^model.language_model.`` reverse would be kept and rewrite ``model.visual.*``. + """ + pytest.importorskip("transformers.core_model_loading") + # Local import: optional dependency, guarded by the importorskip above. + from transformers.core_model_loading import WeightRenaming + + model = torch.nn.Module() + model.model = torch.nn.Module() + model.model.visual = torch.nn.Module() + model.model.visual.patch_embed = torch.nn.Linear(2, 2, bias=False) + model.model.language_model = torch.nn.Module() + model.model.language_model.layers = torch.nn.ModuleList([torch.nn.Linear(2, 2, bias=False)]) + + renaming = WeightRenaming( + source_patterns=r"^model.language_model.", + target_patterns=r"^model.(?!language_model.)", + ) + # Root scope: reaches every key, exactly like an unscoped rule. + renaming.scope_prefix = "" + renaming.base_model_prefix = "" + model._weight_conversions = [renaming] + + state_dict = { + "model.visual.patch_embed.weight": torch.randn(2, 2), + "model.language_model.layers.0.weight": torch.randn(2, 2), + } + reverted = revert_weight_conversion_quant_aware(model, state_dict) + + # The sibling vision namespace must be untouched (the #2032 guarantee). + assert set(reverted) == set(state_dict) + assert not any("language_model.visual" in k for k in reverted) + + +def test_scoped_weight_converter_is_refused(): + """A scoped ``WeightConverter`` must fall back rather than emit unscoped rules. + + Converter-derived rules (expert leaf renames, dense splits) match by module suffix, + so they cannot be confined to a sub-model subtree the way an anchored rename can. + No current architecture scopes a converter -- transformers only scopes + ``WeightRenaming``/``PrefixChange`` -- so if one ever appears, refusing the whole + conversion keeps the in-memory names (a warning) instead of silently rewriting an + identically-named module in a sibling namespace. + """ + pytest.importorskip("transformers.core_model_loading") + # Local import: optional dependency, guarded by the importorskip above. + from transformers.core_model_loading import Chunk, WeightConverter + + conv = WeightConverter( + source_patterns="mlp.gate_up_proj", + target_patterns=["mlp.gate_proj", "mlp.up_proj"], + operations=[Chunk(dim=0)], + ) + conv.scope_prefix = "model.language_model" + conv.base_model_prefix = "model" + model = types.SimpleNamespace(_weight_conversions=[conv]) + + sd = _nvfp4_linear("model.language_model.layers.0.mlp.gate_up_proj", 8, 16) + with pytest.raises(QuantConversionUnsupportedError, match="scoped WeightConverter"): + revert_weight_conversion_quant_aware(model, sd) + + def test_split_collision_raises(): """A split whose target key already exists must fail instead of overwriting.""" sd = _nvfp4_linear("m.gate_up_proj", 8, 16) 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) diff --git a/tests/unit/torch/opt/plugins/test_transformers_save_load.py b/tests/unit/torch/opt/plugins/test_transformers_save_load.py index fced5734e4f..e8b3ed22151 100644 --- a/tests/unit/torch/opt/plugins/test_transformers_save_load.py +++ b/tests/unit/torch/opt/plugins/test_transformers_save_load.py @@ -18,13 +18,20 @@ import pytest import torch +import torch.nn as nn from _test_utils.torch.opt.utils import apply_mode_with_sampling from _test_utils.torch.transformers_models import ( create_tiny_llama_dir, tf_modelopt_state_and_output_tester, ) +from safetensors.torch import load_file from transformers import AutoConfig, AutoModelForCausalLM, LlamaForCausalLM +from modelopt.torch.opt.plugins.transformers import ( + _TRANSFORMERS_GE_5_0, + _legacy_tied_weights_keys_as_dict, +) + @pytest.mark.parametrize("model_cls", [LlamaForCausalLM, AutoModelForCausalLM]) def test_causal_lm_save_restore(tmp_path, model_cls): @@ -40,6 +47,61 @@ def test_causal_lm_save_restore(tmp_path, model_cls): tf_modelopt_state_and_output_tester(model_ref, model_test) +@pytest.mark.parametrize("tie_word_embeddings", [False, True]) +def test_save_pretrained_with_legacy_tied_weights_keys(tmp_path, tie_word_embeddings): + """A model declaring 4.x list-style `_tied_weights_keys` must still save (nvbug 6518665). + + transformers>=5 expects a `{target: source}` dict there and calls `.keys()` on it while + saving, which crashes for `trust_remote_code` modeling code that has not migrated yet. + + Both tying configurations are covered because normalizing the list to `{key: key}` feeds + those keys to the save-time dedup as patterns: an untied weight must survive it, and a + genuinely tied one must still be deduped down to its canonical name. + """ + tiny_llama_dir = create_tiny_llama_dir( + tmp_path, hidden_size=128, dtype=torch.float32, tie_word_embeddings=tie_word_embeddings + ) + model = AutoModelForCausalLM.from_pretrained(tiny_llama_dir) + model = apply_mode_with_sampling(model, ["quantize"]) + + model._tied_weights_keys = ["lm_head.weight"] + model.model._tied_weights_keys = ["embed_tokens.weight"] + + save_dir = tiny_llama_dir / "legacy_tied_keys_model" + model.save_pretrained(save_dir) + + # The declarations the model owned before the save are restored verbatim. + assert model._tied_weights_keys == ["lm_head.weight"] + assert model.model._tied_weights_keys == ["embed_tokens.weight"] + + # No weight is silently dropped: `lm_head.weight` is written out unless it really does + # share storage with the embedding, in which case transformers re-ties it on load. + saved_keys = set(load_file(save_dir / "model.safetensors")) + assert "model.embed_tokens.weight" in saved_keys + assert ("lm_head.weight" in saved_keys) is not tie_word_embeddings + + model_test = AutoModelForCausalLM.from_pretrained(save_dir) + tf_modelopt_state_and_output_tester(model, model_test) + + +@pytest.mark.skipif(not _TRANSFORMERS_GE_5_0, reason="list-style keys are native to transformers 4") +def test_legacy_tied_weights_keys_as_dict_restores_class_attribute(): + """The shim must not leave an instance attribute shadowing the class declaration.""" + + class _LegacyChild(nn.Module): + # How remote-code models declare it: on the class, not the instance. + _tied_weights_keys = ["lm_head.weight"] + + parent = nn.Module() + parent.child = _LegacyChild() + + with _legacy_tied_weights_keys_as_dict(parent): + assert parent.child._tied_weights_keys == {"lm_head.weight": "lm_head.weight"} + + assert _LegacyChild._tied_weights_keys == ["lm_head.weight"] + assert "_tied_weights_keys" not in parent.child.__dict__ + + def test_causal_lm_from_config(tmp_path): """Test loading a model using from_config after applying optimizations""" tiny_llama_dir = create_tiny_llama_dir(tmp_path, hidden_size=128, dtype=torch.float32) diff --git a/tests/unit/torch/speculative/plugins/test_fakebase.py b/tests/unit/torch/speculative/plugins/test_fakebase.py index 2880bf4ef1c..cf6dfe1a6bc 100644 --- a/tests/unit/torch/speculative/plugins/test_fakebase.py +++ b/tests/unit/torch/speculative/plugins/test_fakebase.py @@ -134,3 +134,31 @@ def _fake_from_pretrained(*args, **kwargs): model = load_vlm_or_llm("fake-model", use_offline_training=True, use_fake_base=False) assert captured_kwargs.get("num_hidden_layers") == 0 assert model.config.num_orig_hidden_layers == 4 + + +def test_load_vlm_or_llm_uses_transformers5_vlm_auto_class(monkeypatch): + """Transformers 5 loads VLMs through AutoModelForImageTextToText.""" + cfg = transformers.PretrainedConfig() + cfg.model_type = "qwen3_vl" + cfg.text_config = object() + monkeypatch.setattr(transformers.AutoConfig, "from_pretrained", lambda *a, **kw: cfg) + + captured = {} + + class _FakeVLM: + @staticmethod + def from_pretrained(*args, **kwargs): + captured["args"] = args + captured["kwargs"] = kwargs + return object() + + # ``transformers`` exposes auto classes lazily, so deleting this attribute + # lets its module-level ``__getattr__`` recreate the legacy class. An + # explicit ``None`` models its absence and reliably exercises the v5 + # fallback. + monkeypatch.setattr(transformers, "AutoModelForVision2Seq", None, raising=False) + monkeypatch.setattr(transformers, "AutoModelForImageTextToText", _FakeVLM) + + assert load_vlm_or_llm("qwen3-vl", dtype="auto") is not None + assert captured["args"] == ("qwen3-vl",) + assert captured["kwargs"]["torch_dtype"] == "auto" diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash.py b/tests/unit/torch/speculative/plugins/test_hf_dflash.py index ef2eec2ad07..bd243421d2c 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dflash.py +++ b/tests/unit/torch/speculative/plugins/test_hf_dflash.py @@ -35,6 +35,7 @@ import modelopt.torch.opt as mto import modelopt.torch.speculative as mtsp +import modelopt.torch.speculative.plugins.hf_dflash as hf_dflash from modelopt.torch.speculative.config import DFLASH_DEFAULT_CFG from modelopt.torch.speculative.plugins.hf_dflash import ( DFlashAttention, @@ -119,6 +120,222 @@ def test_convert_sets_mask_token_id(self): assert model.mask_token_id == 0 +def test_qwen3_vl_transformers_530_position_ids_expand_video_grid(monkeypatch): + """Only mRoPE receives a per-frame video grid on Transformers 5.3.0.""" + original_grid = torch.tensor([[3, 4, 5], [2, 6, 7]]) + expected_position_ids = torch.ones(3, 1, 12, dtype=torch.long) + get_rope_index = MagicMock(return_value=(expected_position_ids, torch.zeros(1, 1))) + compute_position_ids = MagicMock() + fake_model = SimpleNamespace( + config=SimpleNamespace(model_type="qwen3_vl"), + model=SimpleNamespace( + get_rope_index=get_rope_index, + compute_3d_position_ids=compute_position_ids, + ), + ) + monkeypatch.setattr(hf_dflash.transformers, "__version__", "5.3.0") + + position_ids = HFDFlashModel._qwen3_vl_position_ids( + fake_model, + input_ids=torch.ones(1, 12, dtype=torch.long), + attention_mask=torch.ones(1, 12, dtype=torch.long), + position_ids=None, + past_key_values=None, + inputs_embeds=None, + model_kwargs={ + "video_grid_thw": original_grid, + "mm_token_type_ids": torch.tensor([[2, 0, 2, 0, 2, 0, 2, 0, 2, 0, 0, 0]]), + }, + ) + + assert position_ids is expected_position_ids + assert not compute_position_ids.called + assert torch.equal(original_grid, torch.tensor([[3, 4, 5], [2, 6, 7]])) + assert torch.equal( + get_rope_index.call_args.kwargs["video_grid_thw"], + torch.tensor([[1, 4, 5], [1, 4, 5], [1, 4, 5], [1, 6, 7], [1, 6, 7]]), + ) + + +def test_qwen3_vl_moe_transformers_530_position_ids_expand_video_grid(monkeypatch): + """Qwen3-VL family variants use the same 5.3.0 mRoPE workaround.""" + expected_position_ids = torch.ones(3, 1, 4, dtype=torch.long) + get_rope_index = MagicMock(return_value=(expected_position_ids, torch.zeros(1, 1))) + fake_model = SimpleNamespace( + config=SimpleNamespace(model_type="qwen3_vl_moe"), + model=SimpleNamespace(get_rope_index=get_rope_index), + ) + monkeypatch.setattr(hf_dflash.transformers, "__version__", "5.3.0") + + position_ids = HFDFlashModel._qwen3_vl_position_ids( + fake_model, + input_ids=torch.ones(1, 4, dtype=torch.long), + attention_mask=torch.ones(1, 4, dtype=torch.long), + position_ids=None, + past_key_values=None, + inputs_embeds=None, + model_kwargs={ + "video_grid_thw": torch.tensor([[2, 4, 4]]), + "mm_token_type_ids": torch.tensor([[2, 0, 2, 0]]), + }, + ) + + assert position_ids is expected_position_ids + assert torch.equal( + get_rope_index.call_args.kwargs["video_grid_thw"], + torch.tensor([[1, 4, 4], [1, 4, 4]]), + ) + + +def test_qwen3_vl_transformers_53_patch_release_raises(monkeypatch): + """Avoid double expansion when a 5.3 patch backports the upstream fix.""" + fake_model = SimpleNamespace( + config=SimpleNamespace(model_type="qwen3_vl"), + model=SimpleNamespace(get_rope_index=MagicMock()), + ) + monkeypatch.setattr(hf_dflash.transformers, "__version__", "5.3.1") + + with pytest.raises(RuntimeError, match=r"5\.3\.0 or >=5\.4\.0"): + HFDFlashModel._qwen3_vl_position_ids( + fake_model, + input_ids=torch.ones(1, 4, dtype=torch.long), + attention_mask=torch.ones(1, 4, dtype=torch.long), + position_ids=None, + past_key_values=None, + inputs_embeds=None, + model_kwargs={ + "video_grid_thw": torch.tensor([[1, 4, 4]]), + "mm_token_type_ids": torch.tensor([[2, 0, 0, 0]]), + }, + ) + + +def test_qwen3_vl_transformers_54_uses_native_position_ids(monkeypatch): + """Transformers 5.4+ performs the grid expansion inside get_rope_index.""" + get_rope_index = MagicMock() + fake_model = SimpleNamespace( + config=SimpleNamespace(model_type="qwen3_vl"), + model=SimpleNamespace(get_rope_index=get_rope_index), + ) + monkeypatch.setattr(hf_dflash.transformers, "__version__", "5.4.0") + + position_ids = HFDFlashModel._qwen3_vl_position_ids( + fake_model, + input_ids=torch.ones(1, 4, dtype=torch.long), + attention_mask=torch.ones(1, 4, dtype=torch.long), + position_ids=None, + past_key_values=None, + inputs_embeds=None, + model_kwargs={ + "video_grid_thw": torch.tensor([[1, 4, 4]]), + "mm_token_type_ids": torch.tensor([[2, 0, 0, 0]]), + }, + ) + + assert position_ids is None + assert not get_rope_index.called + + +def test_qwen3_vl_transformers_530_rejects_bad_video_frame_groups(monkeypatch): + """Fail before mRoPE construction when processor and video-grid contracts differ.""" + fake_model = SimpleNamespace( + config=SimpleNamespace(model_type="qwen3_vl"), + model=SimpleNamespace(get_rope_index=MagicMock()), + ) + monkeypatch.setattr(hf_dflash.transformers, "__version__", "5.3.0") + + with pytest.raises(ValueError, match="video frame groups"): + HFDFlashModel._qwen3_vl_position_ids( + fake_model, + input_ids=torch.ones(1, 4, dtype=torch.long), + attention_mask=torch.ones(1, 4, dtype=torch.long), + position_ids=None, + past_key_values=None, + inputs_embeds=None, + model_kwargs={ + "video_grid_thw": torch.tensor([[2, 4, 4]]), + "mm_token_type_ids": torch.tensor([[2, 2, 0, 0]]), + }, + ) + + +def test_multimodal_forward_kwargs_exclude_non_model_inputs(): + """Do not forward Trainer or collator-only fields to Hugging Face models.""" + pixel_values = torch.ones(1) + mm_token_type_ids = torch.zeros(1, 4, dtype=torch.long) + + forwarded = hf_dflash._multimodal_forward_kwargs( + { + "pixel_values": pixel_values, + "mm_token_type_ids": mm_token_type_ids, + "assistant_masks": torch.ones(1, 4), + "loss_mask": torch.ones(1, 4), + "num_items_in_batch": 4, + "unexpected_dataset_column": "drop me", + } + ) + + assert set(forwarded) == {"pixel_values", "mm_token_type_ids"} + assert forwarded["pixel_values"] is pixel_values + assert forwarded["mm_token_type_ids"] is mm_token_type_ids + + +def test_eval_does_not_precompute_qwen3_vl_position_ids(monkeypatch): + """Evaluation delegates mRoPE construction to the base model and its cache.""" + model = get_tiny_llama(num_hidden_layers=4) + mtsp.convert(model, [("dflash", _get_dflash_config())]) + precompute_position_ids = MagicMock() + monkeypatch.setattr(model, "_qwen3_vl_position_ids", precompute_position_ids) + + model.eval() + model(input_ids=torch.tensor([[1, 2, 3, 4]])) + + precompute_position_ids.assert_not_called() + + +def test_qwen3_vl_transformers_53_position_ids_require_mm_token_types(monkeypatch): + """Never silently fall back to one-dimensional positions for a visual batch.""" + fake_model = SimpleNamespace( + config=SimpleNamespace(model_type="qwen3_vl"), + model=SimpleNamespace(get_rope_index=MagicMock()), + ) + monkeypatch.setattr(hf_dflash.transformers, "__version__", "5.3.0") + + with pytest.raises(ValueError, match="mm_token_type_ids"): + HFDFlashModel._qwen3_vl_position_ids( + fake_model, + input_ids=torch.ones(1, 12, dtype=torch.long), + attention_mask=torch.ones(1, 12, dtype=torch.long), + position_ids=None, + past_key_values=None, + inputs_embeds=None, + model_kwargs={"image_grid_thw": torch.tensor([[1, 4, 4]])}, + ) + + +def test_qwen3_vl_transformers_53_position_ids_reject_bad_mm_token_shape(monkeypatch): + """Keep processor-produced modality ids aligned with the padded text sequence.""" + fake_model = SimpleNamespace( + config=SimpleNamespace(model_type="qwen3_vl"), + model=SimpleNamespace(get_rope_index=MagicMock()), + ) + monkeypatch.setattr(hf_dflash.transformers, "__version__", "5.3.0") + + with pytest.raises(ValueError, match="same shape as input_ids"): + HFDFlashModel._qwen3_vl_position_ids( + fake_model, + input_ids=torch.ones(1, 12, dtype=torch.long), + attention_mask=torch.ones(1, 12, dtype=torch.long), + position_ids=None, + past_key_values=None, + inputs_embeds=None, + model_kwargs={ + "image_grid_thw": torch.tensor([[1, 4, 4]]), + "mm_token_type_ids": torch.zeros(1, 11, dtype=torch.long), + }, + ) + + class TestDPaceWeights: """Test the D-PACE position-weighting objective (arXiv:2605.18810).""" diff --git a/tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py b/tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py index 2deefadd9e3..5abaa124d68 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py +++ b/tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py @@ -30,7 +30,7 @@ import pytest import torch -from _test_utils.torch.transformers_models import get_tiny_llama +from _test_utils.torch.transformers_models import get_tiny_llama, get_tiny_tokenizer import modelopt.torch.speculative as mtsp from modelopt.torch.speculative.eagle.default_config import default_eagle_config @@ -38,6 +38,7 @@ EagleOfflineDataCollator, OfflineSupervisedDataset, ) +from modelopt.torch.utils.plugins import transformers_dataset _mock_scripts = types.ModuleType("scripts") _mock_ar = types.ModuleType("scripts.ar_validate") @@ -57,6 +58,58 @@ make_speculative_data_module = _eagle_utils.make_speculative_data_module +# --------------------------------------------------------------------------- +# online VLM data-module wiring +# --------------------------------------------------------------------------- + + +def test_vlm_data_module_passes_dflash_label_mode(monkeypatch): + """VLM batches must use unshifted labels for DFlash and preserve OSL settings.""" + data_args = argparse.Namespace( + mode="online", + data_path="unused.jsonl", + vlm_processor="dummy-vlm-processor", + vlm_img_dir="/images", + chat_template=None, + ) + collator = MagicMock() + monkeypatch.setattr(_eagle_utils, "ShardedDataset", MagicMock()) + monkeypatch.setattr(_eagle_utils, "VisionLanguageDataCollator", collator) + + module = make_speculative_data_module( + MagicMock(), data_args, train_len=16, answer_only_loss=True, shift_labels=False + ) + + collator.assert_called_once_with( + processor="dummy-vlm-processor", + train_len=16, + local_image_path="/images", + return_labels=True, + answer_only_loss=True, + shift_labels=False, + chat_template=None, + ) + assert module["data_collator"] is collator.return_value + + +def test_vlm_data_collator_accepts_unshifted_labels(monkeypatch): + """The real VLM collator must support DFlash's unshifted labels.""" + processor = types.SimpleNamespace(tokenizer=get_tiny_tokenizer()) + monkeypatch.setattr( + transformers_dataset.transformers.AutoProcessor, + "from_pretrained", + lambda *_args, **_kwargs: processor, + ) + + collator = transformers_dataset.VisionLanguageDataCollator( + processor="dummy-vlm-processor", + chat_template="{{ messages }}", + shift_labels=False, + ) + + assert collator.shift_labels is False + + # --------------------------------------------------------------------------- # sample_size truncation tests # --------------------------------------------------------------------------- diff --git a/tools/launcher/common/megatron_lm/quantize/quantize.sh b/tools/launcher/common/megatron_lm/quantize/quantize.sh index 083ef7399f0..c6b72094c28 100755 --- a/tools/launcher/common/megatron_lm/quantize/quantize.sh +++ b/tools/launcher/common/megatron_lm/quantize/quantize.sh @@ -34,7 +34,7 @@ if [[ -z ${HF_MODEL_CKPT} ]]; then fi # Persist PTQ ckpt + HF export under /cicd ($SLURM_JOB_DIR/cicd) so later # experiments can re-use them. -export MLM_MODEL_SAVE="/cicd/megatron-lm/${MLM_MODEL_CFG}" +export MLM_MODEL_SAVE="${MLM_MODEL_SAVE:-/cicd/megatron-lm/${MLM_MODEL_CFG}}" # If QUANT_CFG is a recipe path, collapse to a flat tag (strip dirs + .yaml/.yml). _QUANT_CFG_TAG="$(basename "${QUANT_CFG}")" _QUANT_CFG_TAG="${_QUANT_CFG_TAG%.yaml}" diff --git a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16/megatron_lm_ptq.yaml b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16/megatron_lm_ptq.yaml index 05f6d986f43..6ec49dd7a0a 100644 --- a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16/megatron_lm_ptq.yaml +++ b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16/megatron_lm_ptq.yaml @@ -29,7 +29,7 @@ pipeline: - --calib-size 32 environment: - MLM_MODEL_CFG: nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16 - - QUANT_CFG: huggingface/models/nvidia/Nemotron-3-Super-120B-A12B/ptq/nvfp4-mse + - QUANT_CFG: huggingface/models/nvidia/Nemotron-3-Super-120B-A12B-BF16/ptq/nvfp4-mse - HF_MODEL_CKPT: /hf-local/nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16 # MMLU + Export run as separate tasks; quantize.sh does quantize only. - RUN_MMLU: "false" @@ -52,7 +52,7 @@ pipeline: script: common/megatron_lm/export/export.sh environment: - MLM_MODEL_CFG: nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16 - - QUANT_CFG: huggingface/models/nvidia/Nemotron-3-Super-120B-A12B/ptq/nvfp4-mse + - QUANT_CFG: huggingface/models/nvidia/Nemotron-3-Super-120B-A12B-BF16/ptq/nvfp4-mse - HF_MODEL_CKPT: /hf-local/nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16 - TP: "1" - PP: "4" diff --git a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16/megatron_lm_ptq.yaml b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16/megatron_lm_ptq.yaml index d297a1c4b50..dfe363da634 100644 --- a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16/megatron_lm_ptq.yaml +++ b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16/megatron_lm_ptq.yaml @@ -30,7 +30,7 @@ pipeline: - --calib-size 32 environment: - MLM_MODEL_CFG: nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16 - - QUANT_CFG: huggingface/models/nvidia/Nemotron-3-Ultra-550B-A55B/ptq/nvfp4-4o6 + - QUANT_CFG: huggingface/models/nvidia/Nemotron-3-Ultra-550B-A55B-BF16/ptq/nvfp4-4o6 - HF_MODEL_CKPT: /hf-local/nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16 # MMLU + Export run as separate tasks; quantize.sh does quantize only. - RUN_MMLU: "false" @@ -53,7 +53,7 @@ pipeline: script: common/megatron_lm/export/export.sh environment: - MLM_MODEL_CFG: nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16 - - QUANT_CFG: huggingface/models/nvidia/Nemotron-3-Ultra-550B-A55B/ptq/nvfp4-4o6 + - QUANT_CFG: huggingface/models/nvidia/Nemotron-3-Ultra-550B-A55B-BF16/ptq/nvfp4-4o6 - HF_MODEL_CKPT: /hf-local/nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16 - TP: "1" - PP: "12" diff --git a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/megatron_lm_qad.yaml b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/megatron_lm_qad.yaml new file mode 100644 index 00000000000..b9201c27784 --- /dev/null +++ b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/megatron_lm_qad.yaml @@ -0,0 +1,151 @@ +# NVIDIA Nemotron 3.5 Lightning 30B-A3B NVFP4 quantization-aware distillation (QAD). +# +# The pipeline converts the HuggingFace BF16 model to an MCore teacher, +# quantizes a separate MCore student, distills the student for 400 iterations, +# and exports the resulting checkpoint. The training task uses an explicit +# Nemotron-Post-Training-Dataset-v2 chat shard so Hugging Face Datasets does not +# prepare the repository's other large splits. +# +# PTQ topology: 1 B200 node x 4 GPUs, TP=1, PP=1, CP=1, EP=4, ETP=1. +# QAD topology: 2 B200 nodes x 4 GPUs, TP=2, PP=1, CP=1, EP=4, ETP=1. +# With micro-batch-size=1 and global-batch-size=16, train-samples=6400 produces +# 400 iterations. To use another sequence length, change both --seq-length and +# --max-position-embeddings. Our final QAD run used a sequence length of 524,288. +# +# Requirements: +# - The BF16 model is mounted at /hf-local/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16. +# - HF_TOKEN can access the gated nvidia/Nemotron-Post-Training-Dataset-v2 dataset. +# +# Usage from tools/launcher: +# source .env-slurm +# uv run launch.py --yaml examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/megatron_lm_qad.yaml --yes + +job_name: Nemotron-3.5-Lightning-30B-A3B_QAD_32k_400iter +pipeline: + allow_to_fail: false + skip: false + note: "NVFP4 TEGroupedMLP QAD at 32K for 400 iterations on one explicit Nemotron post-training chat shard" + + # Import the BF16 Hugging Face checkpoint as the MCore teacher checkpoint. + task_0: + script: common/megatron_bridge/import/import.sh + environment: + - HF_MODEL_ID: /hf-local/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 + - OUTPUT_DIR: /cicd/megatron-lm-bf16/nvidia + - TORCH_DTYPE: bfloat16 + slurm_config: + _factory_: "slurm_factory" + container: nvcr.io/nvidia/nemo:26.06.00 + modelopt_install_path: /opt/venv/lib/python3.12/site-packages/modelopt + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 4 + + # Quantize the teacher checkpoint into the NVFP4 student checkpoint. + task_1: + script: common/megatron_lm/quantize/quantize.sh + args: + - --seq-length 32768 --max-position-embeddings 32768 + - --calib-size 768 + - --skip-generate + - --export-default-te-spec + environment: + - MLM_MODEL_CFG: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 + - QUANT_CFG: huggingface/models/nvidia/Nemotron-3.5-Lightning-30B-A3B-BF16/ptq/w4a16_nvfp4_4o6 + - MLM_MODEL_CKPT: /cicd/megatron-lm-bf16/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16-MCore + - MLM_MODEL_SAVE: /cicd/megatron-lm/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16-W4A16 + - HF_MODEL_CKPT: /hf-local/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 + - RUN_MMLU: "false" + - RUN_EXPORT: "false" + - DP: "1" + - CP: "1" + - TP: "1" + - PP: "1" + - EP: "4" + - ETP: "1" + slurm_config: + _factory_: "slurm_factory" + container: nvcr.io/nvidia/nemo:26.06.00 + modelopt_install_path: /opt/venv/lib/python3.12/site-packages/modelopt + nodes: 1 + ntasks_per_node: 4 + gpus_per_node: 4 + + # Distill the quantized student from the BF16 teacher on chat-template data. + task_2: + script: common/megatron_lm/train/sft.sh + args: + # Data + - --seq-length 32768 --max-position-embeddings 32768 + - --micro-batch-size 1 --global-batch-size 16 + - --train-samples 6400 + - --lr-decay-samples 6400 + - --lr-warmup-samples 0 + - --split 99,1,0 + - --finetune-data-split chat + - --finetune-data-files data/chat-00000-of-00012.parquet + # QAD + - --modelopt-enabled + - --export-default-te-spec + - --export-kd-teacher-load /cicd/megatron-lm-bf16/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16-MCore + - --attention-dropout 0.0 --hidden-dropout 0.0 + - --no-check-for-nan-in-loss-and-grad + - --recompute-granularity selective + - --recompute-modules layernorm moe + - --sequence-parallel + - --ckpt-fully-parallel-load --ckpt-fully-parallel-save + # Optimizer + - --lr 5.0e-6 + - --lr-decay-style constant + - --clip-grad 1.0 --weight-decay 0.0 + - --adam-beta1 0.9 --adam-beta2 0.95 + - --init-method-std 0.010 + - --use-distributed-optimizer + # Evaluation and checkpoints + - --eval-iters 2 --eval-interval 25 + - --save-interval 50 --log-interval 10 + - --dist-ckpt-strictness log_all + environment: + - MLM_MODEL_CFG: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 + - MLM_MODEL_CKPT: /cicd/megatron-lm/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16-W4A16 + - MLM_MODEL_SAVE: /cicd/megatron-lm-qad/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 + - HF_MODEL_CKPT: /hf-local/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 + - MLM_TRAIN_SCRIPT: finetune + - DATASET: nvidia/Nemotron-Post-Training-Dataset-v2 + - DP: "1" + - CP: "1" + - TP: "2" + - PP: "1" + - EP: "4" + - ETP: "1" + slurm_config: + _factory_: "slurm_factory" + container: nvcr.io/nvidia/nemo:26.06.00 + modelopt_install_path: /opt/venv/lib/python3.12/site-packages/modelopt + nodes: 2 + ntasks_per_node: 4 + gpus_per_node: 4 + + task_3: + script: common/megatron_lm/export/export.sh + args: + - --export-default-te-spec + - --dist-ckpt-strictness log_all + environment: + - MLM_MODEL_CFG: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 + - MLM_MODEL_CKPT: /cicd/megatron-lm-qad/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 + - EXPORT_DIR: /cicd/export/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16_NVFP4_QAD + - HF_MODEL_CKPT: /hf-local/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 + - DP: "1" + - CP: "1" + - TP: "1" + - PP: "4" + - EP: "1" + - ETP: "1" + slurm_config: + _factory_: "slurm_factory" + container: nvcr.io/nvidia/nemo:26.06.00 + modelopt_install_path: /opt/venv/lib/python3.12/site-packages/modelopt + nodes: 1 + ntasks_per_node: 4 + gpus_per_node: 4