diff --git a/.agents/skills/ptq/references/slurm-setup-ptq.md b/.agents/skills/ptq/references/slurm-setup-ptq.md index 635e5e1a4a8..c642c4aacf2 100644 --- a/.agents/skills/ptq/references/slurm-setup-ptq.md +++ b/.agents/skills/ptq/references/slurm-setup-ptq.md @@ -36,13 +36,17 @@ pip install -U transformers For unlisted models that need unreleased transformers (e.g., from git), see `references/unsupported-models.md` Step A. -**Prefer `PYTHONPATH`** to use the synced ModelOpt source instead of installing inside the container — this avoids risking dependency conflicts (e.g., `pip install -U nvidia-modelopt[hf]` can upgrade PyTorch and break other packages): +**Prefer `pip install -e ".[hf]" --no-build-isolation`** (run from the Model-Optimizer repo root) to make the synced ModelOpt source importable in the container — this matches how `examples/hf_ptq/slurm/multinode_fsdp2_ptq.slurm` sets up the job, and unlike `PYTHONPATH` it surfaces packaging/build issues instead of masking them. Avoid `pip install -U nvidia-modelopt[hf]` from PyPI, which can upgrade PyTorch and break other packages. ```bash -export PYTHONPATH=/path/to/Model-Optimizer:$PYTHONPATH +pip install -e ".[hf]" --no-build-isolation ``` -If `PYTHONPATH` doesn't work due to missing compiled extensions, fall back to `pip install -e ".[hf]" --no-build-isolation` (run from the Model-Optimizer repo root). +If you specifically need to leave the container's installed packages untouched (e.g. to sidestep a dependency conflict), fall back to `PYTHONPATH` — but note it skips the editable install, so a missing compiled extension only surfaces at import time: + +```bash +export PYTHONPATH=/path/to/Model-Optimizer:$PYTHONPATH +``` **Watch for pip dependency conflicts** — NGC containers set `PIP_CONSTRAINT` to pin versions, causing `ResolutionImpossible` errors. Unset it first so pip can resolve freely: @@ -63,23 +67,11 @@ pip install -U transformers --no-deps Estimate GPU count from model size and available GPU memory. `hf_ptq.py` uses `device_map="auto"` so it fills GPUs automatically — request only as many as needed. -For multi-node PTQ (200B+ params), use `examples/hf_ptq/multinode_ptq.py` with FSDP2 and accelerate: - -```bash -accelerate launch \ - --config_file examples/hf_ptq/fsdp2.yaml \ - --num_machines $NUM_NODES \ - --num_processes $((NUM_NODES * GPUS_PER_NODE)) \ - --main_process_ip $MASTER_ADDR \ - --main_process_port $MASTER_PORT \ - --machine_rank $SLURM_PROCID \ - examples/hf_ptq/multinode_ptq.py \ - --pyt_ckpt_path \ - --qformat \ - --export_path -``` +For multi-node PTQ (200B+ params), use `hf_ptq.py --use_fsdp2`. For the launch commands (`sbatch` +and manual `torchrun`) and the `--recipe` format, see the *Multi-Node Post-Training Quantization with +FSDP2* section of `examples/hf_ptq/README.md`. -The `num_machines`, `num_processes`, `main_process_ip`, and `machine_rank` are overridden on the command line — no need to edit `fsdp2.yaml`. Only update `fsdp_transformer_layer_cls_to_wrap` in the YAML if the model uses a non-default decoder layer class. +Sizing guidance specific to this path: when the per-rank decoder shard approaches GPU capacity (200B+ at low rank count), either add more nodes (more ranks → smaller shard per rank) or add `--cpu_offload`. Layer detection is automatic; no YAML config needed. Use the multi-node template from `skills/common/slurm-setup.md` section 4 as the job script wrapper. diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 2ea1b2f3db6..232f0684237 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -1,48 +1,30 @@ Changelog ========= -Experimental +0.47 (2026-xx-xx) ^^^^^^^^^^^^^^^^^ -- Add pruning examples for Qwen3.5-9B and Nemotron3-Nano using the `new experimental puzzletron branch `_, this branch uses `AutoModel `_ for better parallelization and efficiency. +**New Features** -0.46 (2026-08-xx) -^^^^^^^^^^^^^^^^^ +*Misc* -**Backward Breaking Changes** +- Add ``onnxsim`` as an alternative ONNX simplification backend for ``modelopt.onnx.quantization.quantize(..., simplify=True)`` (and the ``--simplify`` CLI flag). The new ``simplify_backend`` argument / ``--simplify_backend`` flag selects between ``"onnxslim"`` (default, unchanged) and ``"onnxsim"``; both produce an equivalent simplified model. ``onnxsim>=0.7.0`` now ships wheels for Python 3.12+ and aarch64. -- 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. +**Backward Breaking Changes** **Deprecations** -- ``examples/hf_ptq`` AutoQuantize is now driven by an **AutoQuantize recipe** (``--recipe``). The ``--auto_quantize_bits``, ``--auto_quantize_method``, ``--auto_quantize_score_size``, ``--auto_quantize_cost_model``, and ``--auto_quantize_active_moe_expert_ratio`` flags are **deprecated** but still work: they are converted into an ``AutoQuantizeConfig`` on the fly (emitting a ``DeprecationWarning``) and will be removed in a future release. Prefer a recipe under ``modelopt_recipes/general/auto_quantize/``. See ``examples/hf_ptq/README.md``. +**Bug Fixes** -- Renamed ``examples/llm_ptq`` to ``examples/hf_ptq`` to reflect that it covers Hugging Face LLM **and** VLM PTQ. A relative symlink ``examples/llm_ptq`` -> ``hf_ptq`` keeps existing paths and commands working; it will be removed in a future release. Please update references to the new ``examples/hf_ptq`` path. -- Consolidated ``examples/vlm_ptq`` into ``examples/hf_ptq``. Vision-language model PTQ now shares the ``hf_ptq.py`` entry point and ``scripts/huggingface_example.sh``; pass ``--vlm`` to run the TensorRT-LLM multimodal quickstart smoke test. The ``examples/vlm_ptq/scripts/huggingface_example.sh`` entry point is deprecated: it now prints a warning and forwards to the ``hf_ptq`` script with ``--vlm``, and will be removed in a future release. See `examples/hf_ptq/README.md `__. -- Dropped VILA / NVILA vision-language model support in ``examples/hf_ptq``. VILA's modeling code requires ``transformers<=4.50.0``, which conflicts with ModelOpt's minimum supported ``transformers`` version. The VILA-specific bootstrap (repo clone, ``requirements-vila.txt``) and loading paths in ``example_utils.py`` have been removed. -- Bump minimum nemo container requirement to ``nemo:26.04`` (recommended ``nemo:26.06``) for Megatron-Bridge / Megatron-LM optimization features. +0.46 (2026-08-17) +^^^^^^^^^^^^^^^^^ **New Features** -- Add support for retaining all Megatron-Bridge distillation checkpoints via ``distill.py --checkpoint_keep_last -1`` and exporting all or selected iterations with ``export_distilled_megatron_to_hf.py --export_iterations``. -- Add the ``prepare_megatron_data_blend`` utility to prepare weighted Megatron data blends from YAML configs, including optional token-budgeted subsets for distillation workflows. See the `Megatron data preparation guide `_. +*Quantization* + - Add Learned Scale Quantization (LSQ) and Dual-LSQ support for quantization-aware distillation, including learnable ``amax`` parameters, tied-scale and pre-scale options, focused NVFP4 recipes, and scale-only training. -- Add the **D-PACE** loss objective for DFlash speculative-decoding training (`arXiv:2605.18810 `_) and make it the default (``dflash_loss_objective: dpace``). It replaces the static exponential position decay with dynamic, confidence-derived per-position weights that adapt to whichever block positions currently limit acceptance. Smoothing is controlled by ``dflash_dpace_alpha`` (default 0.5); set ``dflash_loss_objective: decay`` to restore the previous static schedule. Training-only and detached from the gradient (no architecture or inference change). -- Add the ``day0-release`` agent skill (``.agents/skills/day0-release/``), a deterministic end-to-end driver that chains the PTQ → evaluation → comparison skills (the evaluation stage deploys the checkpoint itself) with an enforced gate after each stage and returns a publish decision (ACCEPT / REGRESSION / ANOMALOUS / INFEASIBLE). Ships three GPU-free, unit-tested gate scripts (``gate_ptq.py``, ``gate_run.py``, ``gate_compare.py``) that validate checkpoint coverage, evaluation-run completeness, and baseline-vs-candidate accuracy threshold. v1 reports and stops on regression; the recipe-search loop is deferred. -- Add **streaming** speculative-decoding training (EAGLE3 / DFlash): the draft trains on base-model hidden states produced on the fly by a co-located ``vllm serve`` (no disk dump), moved trainer-side over NIXL RDMA, scaling to multi-node (dedicated serve replicas + DDP trainers). New launcher examples for NVFP4 Kimi-K2.5 / K2.6 on GB200/aarch64 under ``tools/launcher/examples/moonshotai/``. - Add a fused Triton fast path for ``local_hessian`` NVFP4 weight-scale search (the Hessian-weighted FP8-E4M3 scale sweep). For each NVFP4 block it minimizes ``dwᵀ H dw`` over the 126 candidate scales using the per-cin-block local Hessian on tensor cores, replacing the per-weight Python reference sweep — roughly **34x** faster on a single 8192x4096 weight and bit-exact with the reference for fp32/fp16 weights. Used automatically during ``local_hessian`` calibration for both dense and fused-MoE expert weights; falls back to the reference sweep on CPU, when Triton is unavailable, or via ``MODELOPT_NVFP4_TRITON_SWEEP=0``. -- Add **context-parallel (CP)** and **data-parallel (DP)** support to the shared Megatron-Core inference/calibration utilities. Under CP, ``get_megatron_calibration_forward_loop`` and ``megatron_mmlu`` partition each sequence across CP ranks (zigzag load-balanced), ``megatron_prefill`` accepts a CP-partitioned ``position_ids`` and lets the CP-aware causal attention build the mask, and MMLU gathers per-rank logits back to the full sequence for last-token scoring. Under DP, calibration shards the dataset across data-parallel ranks (``DistributedSampler``; amax is max-reduced across the DP group inside ``mtq``) and ``megatron_mmlu`` shards whole batches across DP ranks and all-reduces the per-subject counts. DP is implicit (``world_size / (tp * pp * cp)``); ``examples/megatron_bridge/quantize.py`` gains a ``--cp_size`` flag. -- Add Minitron pruning support for Megatron-Core models with the following new attention and MoE variants. For these, only ``hidden_size`` is pruned (alongside the usual ``ffn_hidden_size`` / ``num_layers`` / MoE dimensions); the variant-internal dimensions noted below are not pruned: -- Add support for ONNX Q/DQ node placement for DLA via the new flag ``--target_dla``. - - - **GatedDeltaNet** (linear attention) and **gated attention** (``attention_output_gate``), such as Qwen3.5 (hybrid GatedDeltaNet + gated-attention) language models, including MoE variants — attention / linear-attention heads are not pruned. - - **Multi-Latent Attention (MLA)**, such as DeepSeek — MLA latent ranks are not pruned. - - **Latent MoE**, such as Nemotron-3-Super — ``hidden_size`` pruning resizes the latent projections while the MoE latent dim itself is not pruned. -- Add Minitron pruning support for the language model part of vision-language models (e.g. Qwen3.5-VL, Gemma3-VL) via ``examples/megatron_bridge/prune_minitron.py``. The language model is pruned while the vision tower is left intact and the full VLM is saved back; ``hidden_size`` is not pruned if it is shared with the vision projector. Pruning importance is estimated from image-text calibration (the full VLM forward over vision-conditioned activations) by default, or from a text dataset for text-only ablations. -- Add PTQ support for the language model part of vision-language models (e.g. Qwen3.5-VL, Gemma3-VL) via ``examples/megatron_bridge/quantize.py``. Only the language model is quantized (vision tower + projector left in full precision) and the full VLM is saved as a Megatron checkpoint. The calibration modality is inferred from ``--calib_dataset_name``: an image-text dataset drives the full VLM forward (vision-conditioned activations), while a text dataset runs text-only calibration of the language model. Image-text calibration shards across data-parallel ranks (context parallelism is supported only for text-only calibration). HuggingFace unified export of a quantized VLM is not yet supported. -- Add Megatron-Bridge distillation and QAD support for the language model part of vision-language models (e.g. Qwen3.5-VL, Gemma3-VL) via ``examples/megatron_bridge/distill.py``. - Add NVFP4 Four-Over-Six (4/6) weight quantization (``mtq.NVFP4_FOUR_OVER_SIX_CFG``): MSE weight calibration picks, per block, between an M=6 and an M=4 dynamic range (the choice is folded into the FP8 per-block scales), with the ``four_over_six: true`` flag normalizing those scales by 256 (vs 448) for M=4 headroom. Supported via ``mtq.quantize`` and HF / Megatron export only -- **not** ``mtq.compress``, which does not preserve the per-block M=4/M=6 choice - Add dLLM (tied-weight PTQ and HF-checkpoint export) support for diffusion-based encoder-decoder LLMs (e.g. DiffusionGemma) whose encoder/decoder stacks share parameters via HF ``_tied_weights_keys``. @@ -51,7 +33,6 @@ Experimental - 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. - ``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 **Domino** speculative-decoding training: the parallel DFlash draft backbone plus a lightweight GRU causal correction head, selected via ``dflash_architecture_config.projector_type=domino``. Trained with a base/final dual loss whose ``dflash_lambda_base_start``/``dflash_lambda_base_decay_ratio`` curriculum decays the base-loss weight 1→0. Exports in the z-lab drafter format; recipe at ``modelopt_recipes/general/speculative_decoding/domino.yaml``. Training only — the inference path is not wired up yet. - 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/``. - Add module-specific AutoQuantize search spaces through ``mtq.auto_quantize(..., module_search_spaces=...)`` and recipe-level ``auto_quantize.module_search_spaces``. Glob-matched runtime decision groups can override global candidate formats and control whether BF16/no-quant is solver-selectable with ``allow_no_quant``. Recipes can instead reuse a normal PTQ ``quantize`` config as the fixed baseline and explicitly list only genuinely searched modules; fixed and searched groups remain in one calibration, scoring, effective-bits, checkpoint, and export flow. Rules cannot partially split runtime-fused groups, fixed groups are isolated from unrelated calibration algorithms, infeasible resolved budgets fail before calibration, and checkpoint replay validates the fixed baseline, resolved groups, candidate choices, scoring boundaries, and cost weights before reusing calibration or sensitivity state. @@ -60,6 +41,48 @@ Experimental - Add ``MaxCalibConfig.skip_forward_without_activation_calib`` (opt-in, default ``False``): when enabled, max calibration skips the ``forward_loop`` if no enabled quantizer needs data-driven activation statistics — e.g. an experts-only recipe whose activation quantizers all use ``constant_amax`` / ``use_constant_amax``, or dynamic / MX (MXFP4/MXFP8) quantization (and none has a static ``bias_calibrator``). Weight calibration still runs on the weight tensors directly, so the quantized weights are unchanged; only the wasted forward is avoided. It is opt-in because the ``forward_loop`` can carry caller-side effects (notably materializing sharded parameters under DeepSpeed ZeRO-3). The advanced algorithms that always need activations (MSE, local Hessian, SmoothQuant, SVDQuant, GPTQ) call ``max_calibrate`` directly and are unaffected. The ``nvfp4_experts_only_input_scale1-kv_fp8_cast`` recipe enables it. - Add ``examples/minimax_m3/hf_ptq_mixed_mxfp8_nvfp4.py`` for streaming MiniMax-M3 export and a model-specific ``hf_ptq.py`` recipe that produces an MXFP8 language-model base with MSE-calibrated NVFP4 routed experts directly from BF16. The NVFP4 expert ``input_scale`` is fixed to 1.0. +*Speculative Decoding* + +- Add the **D-PACE** loss objective for DFlash speculative-decoding training (`arXiv:2605.18810 `_) and make it the default (``dflash_loss_objective: dpace``). It replaces the static exponential position decay with dynamic, confidence-derived per-position weights that adapt to whichever block positions currently limit acceptance. Smoothing is controlled by ``dflash_dpace_alpha`` (default 0.5); set ``dflash_loss_objective: decay`` to restore the previous static schedule. Training-only and detached from the gradient (no architecture or inference change). +- Add **streaming** speculative-decoding training (EAGLE3 / DFlash): the draft trains on base-model hidden states produced on the fly by a co-located ``vllm serve`` (no disk dump), moved trainer-side over NIXL RDMA, scaling to multi-node (dedicated serve replicas + DDP trainers). New launcher examples for NVFP4 Kimi-K2.5 / K2.6 on GB200/aarch64 under ``tools/launcher/examples/moonshotai/``. +- Add **Domino** speculative-decoding training: the parallel DFlash draft backbone plus a lightweight GRU causal correction head, selected via ``dflash_architecture_config.projector_type=domino``. Trained with a base/final dual loss whose ``dflash_lambda_base_start``/``dflash_lambda_base_decay_ratio`` curriculum decays the base-loss weight 1→0. Exports in the z-lab drafter format; recipe at ``modelopt_recipes/general/speculative_decoding/domino.yaml``. Training only — the inference path is not wired up yet. + +*Megatron Framework (M-LM / M-Bridge)* + +- Add Minitron pruning support for Megatron-Core models with the following new attention and MoE variants. For these, only ``hidden_size`` is pruned (alongside the usual ``ffn_hidden_size`` / ``num_layers`` / MoE dimensions); the variant-internal dimensions noted below are not pruned: + + - **GatedDeltaNet** (linear attention) and **gated attention** (``attention_output_gate``), such as Qwen3.5 (hybrid GatedDeltaNet + gated-attention) language models, including MoE variants — attention / linear-attention heads are not pruned. + - **Multi-Latent Attention (MLA)**, such as DeepSeek — MLA latent ranks are not pruned. + - **Latent MoE**, such as Nemotron-3-Super — ``hidden_size`` pruning resizes the latent projections while the MoE latent dim itself is not pruned. +- Optimize Minitron pruning support for MoE models using the fused **grouped GEMM** experts (``TEGroupedMLP``) in addition to the existing ``SequentialMLP`` path. ``examples/megatron_bridge/prune_minitron.py`` now uses grouped GEMM by default (pass ``--no_moe_grouped_gemm`` to fall back to ``TESequentialMLP``). +- Add Minitron pruning support for the language model part of vision-language models (e.g. Qwen3.5-VL, Gemma3-VL) via ``examples/megatron_bridge/prune_minitron.py``. The language model is pruned while the vision tower is left intact and the full VLM is saved back; ``hidden_size`` is not pruned if it is shared with the vision projector. Pruning importance is estimated from image-text calibration (the full VLM forward over vision-conditioned activations) by default, or from a text dataset for text-only ablations. +- Add PTQ support for the language model part of vision-language models (e.g. Qwen3.5-VL, Gemma3-VL) via ``examples/megatron_bridge/quantize.py``. Only the language model is quantized (vision tower + projector left in full precision) and the full VLM is saved as a Megatron checkpoint. The calibration modality is inferred from ``--calib_dataset_name``: an image-text dataset drives the full VLM forward (vision-conditioned activations), while a text dataset runs text-only calibration of the language model. Image-text calibration shards across data-parallel ranks (context parallelism is supported only for text-only calibration). HuggingFace unified export of a quantized VLM is not yet supported. +- Add Megatron-Bridge distillation and Quantization-Aware Distillation (QAD) support for the language model part of vision-language models (e.g. Qwen3.5-VL, Gemma3-VL) via ``examples/megatron_bridge/distill.py``. +- Add **context-parallel (CP)** and **data-parallel (DP)** support to the shared Megatron-Core inference/calibration utilities. Under CP, ``get_megatron_calibration_forward_loop`` and ``megatron_mmlu`` partition each sequence across CP ranks (zigzag load-balanced), ``megatron_prefill`` accepts a CP-partitioned ``position_ids`` and lets the CP-aware causal attention build the mask, and MMLU gathers per-rank logits back to the full sequence for last-token scoring. Under DP, calibration shards the dataset across data-parallel ranks (``DistributedSampler``; amax is max-reduced across the DP group inside ``mtq``) and ``megatron_mmlu`` shards whole batches across DP ranks and all-reduces the per-subject counts. DP is implicit (``world_size / (tp * pp * cp)``); ``examples/megatron_bridge/quantize.py`` gains a ``--cp_size`` flag. +- Add support for retaining all Megatron-Bridge distillation checkpoints via ``distill.py --checkpoint_keep_last -1`` and exporting all or selected iterations with ``export_distilled_megatron_to_hf.py --export_iterations``. +- Add the ``prepare_megatron_data_blend`` utility to prepare weighted Megatron data blends from YAML configs, including optional token-budgeted subsets for distillation workflows. See the `Megatron data preparation guide `_. + +*Misc* + +- Add the ``day0-release`` agent skill (``.agents/skills/day0-release/``), a deterministic end-to-end driver that chains the PTQ → evaluation → comparison skills (the evaluation stage deploys the checkpoint itself) with an enforced gate after each stage and returns a publish decision (ACCEPT / REGRESSION / ANOMALOUS / INFEASIBLE). Ships three GPU-free, unit-tested gate scripts (``gate_ptq.py``, ``gate_run.py``, ``gate_compare.py``) that validate checkpoint coverage, evaluation-run completeness, and baseline-vs-candidate accuracy threshold. v1 reports and stops on regression; the recipe-search loop is deferred. +- Add support for ONNX Q/DQ node placement for DLA via the new flag ``--target_dla``. +- (Experimental) Add pruning examples for Qwen3.5-9B and Nemotron3-Nano using the `new experimental puzzletron branch `_, this branch uses `AutoModel `_ for better parallelization and efficiency. + +**Backward Breaking Changes** + +- 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. + +**Deprecations** + +- ``examples/hf_ptq`` AutoQuantize is now driven by an **AutoQuantize recipe** (``--recipe``). The ``--auto_quantize_bits``, ``--auto_quantize_method``, ``--auto_quantize_score_size``, ``--auto_quantize_cost_model``, and ``--auto_quantize_active_moe_expert_ratio`` flags are **deprecated** but still work: they are converted into an ``AutoQuantizeConfig`` on the fly (emitting a ``DeprecationWarning``) and will be removed in a future release. Prefer a recipe under ``modelopt_recipes/general/auto_quantize/``. See ``examples/hf_ptq/README.md``. +- Renamed ``examples/llm_ptq`` to ``examples/hf_ptq`` to reflect that it covers Hugging Face LLM **and** VLM PTQ. A relative symlink ``examples/llm_ptq`` to ``hf_ptq`` keeps existing paths and commands working; it will be removed in a future release. Please update references to the new ``examples/hf_ptq`` path. +- Consolidated ``examples/vlm_ptq`` into ``examples/hf_ptq``. Vision-language model PTQ now shares the ``hf_ptq.py`` entry point and ``scripts/huggingface_example.sh``; pass ``--vlm`` to run the TensorRT-LLM multimodal quickstart smoke test. The ``examples/vlm_ptq/scripts/huggingface_example.sh`` entry point is deprecated: it now prints a warning and forwards to the ``hf_ptq`` script with ``--vlm``, and will be removed in a future release. See `examples/hf_ptq/README.md `__. +- Dropped VILA / NVILA vision-language model support in ``examples/hf_ptq``. VILA's modeling code requires ``transformers<=4.50.0``, which conflicts with ModelOpt's minimum supported ``transformers`` version. The VILA-specific bootstrap (repo clone, ``requirements-vila.txt``) and loading paths in ``example_utils.py`` have been removed. +- Bump minimum nemo container requirement to ``nemo:26.04`` (recommended ``nemo:26.06``) for Megatron-Bridge / Megatron-LM optimization features. +- Python 3.10 support will be dropped in the next release as it is reaching EOL. + **Bug Fixes** - 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. diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 950f3640aaf..a3967565ae0 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -471,33 +471,37 @@ mtq.calibrate(model, algorithm="max", forward_loop=calibrate_loop) ## Multi-Node Post-Training Quantization with FSDP2 -ModelOpt enables quantization of LLMs across multiple GPU nodes using various quantization formats. It leverages HuggingFace's Accelerate library and FSDP2 for distributed model sharding and calibration. +ModelOpt enables quantization of LLMs across multiple GPU nodes using FSDP2 for distributed model sharding and calibration, exposed via the `--use_fsdp2` flag on the standard `hf_ptq.py` entry point. ### Usage -For distributed execution across multiple nodes, use the `accelerate` library. A template configuration file (`fsdp2.yaml`) is provided and can be customized for user specific requirements. +#### Slurm (recommended) -On each node run the following command: +Slurm orchestrates launching the job on every node for you, so this is the easiest way to run a multi-node PTQ. A ready-to-run example that quantizes Nemotron-3-Super to NVFP4 is provided in [`slurm/multinode_fsdp2_ptq.slurm`](./slurm/multinode_fsdp2_ptq.slurm). Edit the `CONFIG` block (container image, model path, export path, recipe) and submit: ```bash -accelerate launch --config_file fsdp2.yaml \ - --num_machines= \ - --machine_rank= \ - --main_process_ip= \ - --main_process_port= \ - --fsdp_transformer_layer_cls_to_wrap= - multinode_ptq.py \ +sbatch --nodes=2 slurm/multinode_fsdp2_ptq.slurm +``` + +#### Manual (run on each node) + +Without Slurm, start `torchrun` on every node yourself: + +```bash +torchrun \ + --nnodes= --node_rank= \ + --master_addr= --master_port= \ + --nproc_per_node= \ + hf_ptq.py \ --pyt_ckpt_path \ - --qformat \ - --kv_cache_qformat \ + --recipe general/ptq/nvfp4_default-kv_fp8_cast \ --batch_size \ --calib_size \ - --dataset \ --export_path \ - --trust_remote_code + --use_fsdp2 ``` -The exported checkpoint can be deployed using TensorRT-LLM/ vLLM/ SGLang. For more details refer to the [deployment section](#deployment) of this document. +See [Recipe-based Quantization](#recipe-based-quantization) for the recipe format and built-in recipe names. The exported checkpoint can be deployed using TensorRT-LLM/ vLLM/ SGLang. For more details refer to the [deployment section](#deployment) of this document. > *Performance Note: FSDP2 is designed for training workloads and may result in longer calibration and export times. For faster calibration, maximize the batch size based on available GPU memory and choose the right number of GPUs to avoid unnecessary communication.* diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index 83a54849110..d74ffb34efb 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -23,6 +23,8 @@ import shutil import warnings from collections.abc import Callable, Iterable +from dataclasses import dataclass +from datetime import timedelta from pathlib import Path from typing import Any @@ -48,11 +50,68 @@ except ImportError: snapshot_download = None +from modelopt.torch.utils import distributed as dist_utils + logger = logging.getLogger(__name__) SPECULATIVE_MODEL_LIST = ["Eagle", "Medusa"] +@dataclass +class DistributedState: + """Example-local distributed state for model loading, dataloader sharding, and rank-0 output.""" + + rank: int + world_size: int + device: torch.device | str + is_main: bool + + +def setup_distributed_args(args): + """Initialize and attach ``args.dist_state`` (single-process if FSDP2 off).""" + if getattr(args, "use_fsdp2", False): + # Raise the collective timeout above NCCL's 30-min default: rank 0's checkpoint write can + # exceed it, and PyTorch 2.8 has no per-call barrier() timeout (must be set at PG creation). + dist_utils.setup(timeout=timedelta(hours=2)) + rank = dist_utils.rank() + args.dist_state = DistributedState( + rank=rank, + world_size=dist_utils.size(), + device=torch.device(f"cuda:{dist_utils.local_rank()}"), + is_main=rank == 0, + ) + else: + args.dist_state = DistributedState(rank=0, world_size=1, device=args.device, is_main=True) + + +def cleanup_distributed(args): + """Destroy the process group if ``--use_fsdp2`` set it up.""" + if getattr(args, "use_fsdp2", False): + dist_utils.cleanup() + + +def validate_fsdp2_supported(args, config): + """Raise ``NotImplementedError`` for model/CLI combos the FSDP2 path doesn't support yet.""" + issues = [] + if "vila" in args.pyt_ckpt_path.lower(): + issues.append("VILA (custom builder + non-standard layer layout)") + if is_nemotron_vl(config) or _is_multimodal_config(config): + issues.append("multimodal / VL models (decoder layers not auto-detectable)") + if getattr(config, "quantization_config", None) is not None: + issues.append("pack-quantized / compressed-tensors checkpoints") + if getattr(args, "specdec_offline_dataset", None) is not None: + issues.append("speculative decoding (--specdec_offline_dataset)") + if getattr(args, "low_memory_mode", False): + issues.append("--low_memory_mode (redundant with FSDP2)") + + if issues: + raise NotImplementedError( + "--use_fsdp2 does not support:\n - " + + "\n - ".join(issues) + + "\nRemove --use_fsdp2 or use a standard causal-LM checkpoint." + ) + + def run_nemotron_vl_preview( full_model, tokenizer, @@ -372,6 +431,19 @@ def _apply_to_model_state_dict( return out_state_dict +def mtp_layer_prefixes_from_checkpoint(model_path: str) -> list[str]: + """MTP exclude-prefixes from a checkpoint's safetensors index (``[]`` if none); reads no tensors. + + Local-index-only, matching :func:`load_mtp_weights`, so detection and re-attach stay in sync. + """ + index_file = Path(model_path) / "model.safetensors.index.json" + if not index_file.exists(): + return [] + weight_map = json.load(open(index_file))["weight_map"] + mtp_keys = [k for k, v in weight_map.items() if "mtp" in k or "mtp" in v] + return list(_keys_to_prefixes(mtp_keys)) + + def load_mtp_weights( model: torch.nn.Module, model_path: str ) -> tuple[list[str], dict[str, torch.Tensor]]: diff --git a/examples/hf_ptq/fsdp2.yaml b/examples/hf_ptq/fsdp2.yaml deleted file mode 100644 index 09977835561..00000000000 --- a/examples/hf_ptq/fsdp2.yaml +++ /dev/null @@ -1,30 +0,0 @@ -# ============================================================================= -# FSDP Configuration for running LLM PTQ on multinode setup. This file is consumed by examples/hf_ptq/multinode_ptq.py -# ============================================================================= - -compute_environment: LOCAL_MACHINE -debug: false -distributed_type: FSDP -downcast_bf16: 'no' -enable_cpu_affinity: false -fsdp_config: - fsdp_activation_checkpointing: false - fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP - fsdp_cpu_ram_efficient_loading: true - fsdp_offload_params: false - fsdp_reshard_after_forward: true - fsdp_state_dict_type: FULL_STATE_DICT - fsdp_transformer_layer_cls_to_wrap: LlamaDecoderLayer - fsdp_use_orig_params: true - fsdp_version: 2 -machine_rank: 0 -main_training_function: main -mixed_precision: 'no' -num_machines: 2 -num_processes: 16 -rdzv_backend: c10d -same_network: true -tpu_env: [] -tpu_use_cluster: false -tpu_use_sudo: false -use_cpu: false diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 57a3dd6e264..6a4bd476984 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -15,6 +15,7 @@ import argparse import copy +import os import random import time import warnings @@ -29,6 +30,7 @@ from example_utils import ( _resolve_model_path, build_quant_cfg, + cleanup_distributed, copy_custom_model_files, create_vlm_calibration_loop, get_model, @@ -37,9 +39,12 @@ is_enc_dec, is_nemotron_vl, load_mtp_weights, + mtp_layer_prefixes_from_checkpoint, needs_checkpoint_path_update, resolve_checkpoint_dir, run_nemotron_vl_preview, + setup_distributed_args, + validate_fsdp2_supported, ) from torch.utils.data import DataLoader from transformers import ( @@ -75,6 +80,7 @@ EagleOfflineDataCollator, OfflineSupervisedDataset, ) +from modelopt.torch.utils import print_rank_0 from modelopt.torch.utils.dataset_utils import ( create_forward_loop, get_dataset_dataloader, @@ -82,6 +88,7 @@ get_supported_datasets, ) from modelopt.torch.utils.memory_monitor import launch_memory_monitor +from modelopt.torch.utils.plugins.model_load_utils import parallel_load_and_prepare_fsdp2 from modelopt.torch.utils.speech_dataset_utils import get_speech_dataset_dataloader from modelopt.torch.utils.vlm_dataset_utils import get_vlm_dataset_dataloader @@ -251,6 +258,12 @@ def make_calib_dataloader( max_sample_length=args.calib_seq, device=device, include_labels=include_labels, + distributed=args.use_fsdp2, + sampler_kwargs=( + {"num_replicas": args.dist_state.world_size, "rank": args.dist_state.rank} + if args.use_fsdp2 + else None + ), ) return calib_dataloader, first_text_speech_dataset @@ -438,6 +451,13 @@ def auto_quantize( "Auto Quantization is not supported for pipeline parallel size > 1" ) + if args.use_fsdp2: + warnings.warn( + "AutoQuantize with --use_fsdp2 has not been validated end-to-end yet " + "(distributed calibration, sensitivity scoring, and recipe/checkpoint " + "synchronization across ranks); use at your own risk." + ) + inputs = _mtq_inputs_from_auto_quantize_config( aq_config, args, fixed_quantize_config=fixed_quantize_config ) @@ -530,10 +550,30 @@ def _recipe_is_auto_quantize(recipe: str | None) -> bool: def load_model(args: argparse.Namespace): # If low memory mode is enabled, we compress the model while loading the HF checkpoint. calibration_only = False - if args.specdec_offline_dataset is not None or not args.low_memory_mode: + if args.use_fsdp2: + hf_config = AutoConfig.from_pretrained( + args.pyt_ckpt_path, trust_remote_code=args.trust_remote_code + ) + validate_fsdp2_supported(args, hf_config) + full_model = parallel_load_and_prepare_fsdp2( + args.pyt_ckpt_path, + args.dist_state.device, + args.dist_state.rank, + args.dist_state.world_size, + trust_remote_code=args.trust_remote_code, + cpu_offload=args.cpu_offload, + attn_implementation=args.attn_implementation, + hf_config=hf_config, + ) + # The FSDP2 loader drops MTP weights (re-attached BF16 at export); flag their prefixes now + # so the pre-quant exclusion below skips any MTP module from_config did build. + mtp_prefixes = mtp_layer_prefixes_from_checkpoint(args.pyt_ckpt_path) + if mtp_prefixes: + full_model._mtp_layer_prefixes = mtp_prefixes + elif args.specdec_offline_dataset is not None or not args.low_memory_mode: full_model = get_model( args.pyt_ckpt_path, - args.device, + args.dist_state.device, gpu_mem_percentage=args.gpu_max_mem_percentage, trust_remote_code=args.trust_remote_code, use_seq_device_map=args.use_seq_device_map, @@ -565,9 +605,12 @@ def load_model(args: argparse.Namespace): model_type = get_model_type(full_model) - device = full_model.device - if hasattr(full_model, "model"): - device = full_model.model.device + if args.use_fsdp2: + device = args.dist_state.device + else: + device = full_model.device + if hasattr(full_model, "model"): + device = full_model.model.device processor = None tokenizer = None language_model = full_model @@ -861,7 +904,6 @@ def export_quantized( mtp_layer_prefixes, mtp_state_dict = load_mtp_weights( full_model, args.pyt_ckpt_path ) - if mtp_layer_prefixes: full_model._mtp_layer_prefixes = mtp_layer_prefixes @@ -882,16 +924,18 @@ def export_quantized( tokenizer.padding_side = default_padding_side if default_pad_token is not None: tokenizer.pad_token = default_pad_token - tokenizer.save_pretrained(export_path) + if args.dist_state.is_main: + tokenizer.save_pretrained(export_path) # Copy custom model files (Python files and JSON configs) if trust_remote_code is used. # This must run AFTER tokenizer.save_pretrained() so original tokenizer files # from the source checkpoint take precedence over regenerated ones (which may # differ in format due to newer transformers versions). - copy_custom_model_files(args.pyt_ckpt_path, export_path, args.trust_remote_code) + if args.dist_state.is_main: + copy_custom_model_files(args.pyt_ckpt_path, export_path, args.trust_remote_code) end_time = time.time() - print( + print_rank_0( f"Quantized model exported to: {export_path}. Total time used {end_time - start_time}s" ) @@ -992,7 +1036,7 @@ def post_quantize( ) return - if args.verbose: + if args.verbose and args.dist_state.is_main: try: mtq.print_quant_summary(full_model, args.export_path) save_expert_token_count_table(full_model, args.export_path) @@ -1453,6 +1497,20 @@ def parse_args() -> argparse.Namespace: default=False, action="store_true", ) + parser.add_argument( + "--use_fsdp2", + action="store_true", + help=( + "Run calibration under PyTorch FSDP2 (requires torchrun); takes precedence over " + "--use_seq_device_map. v1: standard causal-LM only (no VILA / pack-quantized / " + "speculative / auto-quantize / sparsity / VLM / MTP)." + ), + ) + parser.add_argument( + "--cpu_offload", + action="store_true", + help="With --use_fsdp2, keep decoder shards on CPU between forwards (frees GPU memory, adds PCIe traffic).", + ) parser.add_argument( "--verbose", help="Print verbose output (e.g. quantization summary). Disable by --no-verbose.", @@ -1583,6 +1641,19 @@ def parse_args() -> argparse.Namespace: "--low_memory_mode does not support --recipe or AutoQuantize (--auto_quantize_bits); " "the low-memory loader initializes quantizers from --qformat/--kv_cache_qformat." ) + if args.use_fsdp2 and args.use_seq_device_map: + warnings.warn("--use_seq_device_map is ignored when --use_fsdp2 is set.") + args.use_seq_device_map = False + if args.use_fsdp2 and os.environ.get("RANK") is None: + parser.error("--use_fsdp2 requires launching with torchrun") + if args.cpu_offload and not args.use_fsdp2: + parser.error("--cpu_offload requires --use_fsdp2") + if args.use_fsdp2 and args.sparsity_fmt != "dense": + parser.error(f"--use_fsdp2 does not support --sparsity_fmt {args.sparsity_fmt}.") + if args.use_fsdp2 and args.vllm_fakequant_export: + parser.error("--use_fsdp2 does not support --vllm_fakequant_export.") + if args.use_fsdp2 and args.cast_mxfp4_to_nvfp4: + parser.error("--use_fsdp2 does not support --cast_mxfp4_to_nvfp4.") return args @@ -1594,31 +1665,16 @@ def main(args: argparse.Namespace): random.seed(RAND_SEED) np.random.seed(RAND_SEED) - # launch a memory monitor to read the currently used GPU memory. - launch_memory_monitor() + setup_distributed_args(args) - # Force eager execution for all model types. - torch.compiler.set_stance("force_eager") + try: + # launch a memory monitor to read the currently used GPU memory. + launch_memory_monitor() - ( - full_model, - language_model, - model_type, - calibration_only, - processor, - tokenizer, - default_padding_side, - default_pad_token, - device, - ) = load_model(args) + # Force eager execution for all model types. + torch.compiler.set_stance("force_eager") - if args.sparsity_fmt != "dense": - # Sparse - sparsity_main(args, full_model, tokenizer, device) - else: - # Quantize - quantize_main( - args, + ( full_model, language_model, model_type, @@ -1628,7 +1684,27 @@ def main(args: argparse.Namespace): default_padding_side, default_pad_token, device, - ) + ) = load_model(args) + + if args.sparsity_fmt != "dense": + # Sparse + sparsity_main(args, full_model, tokenizer, device) + else: + # Quantize + quantize_main( + args, + full_model, + language_model, + model_type, + calibration_only, + processor, + tokenizer, + default_padding_side, + default_pad_token, + device, + ) + finally: + cleanup_distributed(args) if __name__ == "__main__": diff --git a/examples/hf_ptq/multinode_ptq.py b/examples/hf_ptq/multinode_ptq.py deleted file mode 100644 index 12e6c04e535..00000000000 --- a/examples/hf_ptq/multinode_ptq.py +++ /dev/null @@ -1,369 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 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. - -"""Multi-node PTQ (Post-Training Quantization) with FSDP2 support.""" - -import argparse -import json -import os -import random -import time -import warnings -from pathlib import Path - -import numpy as np -import torch -import torch.nn as nn -from accelerate import Accelerator -from example_utils import build_quant_cfg, get_tokenizer -from tqdm import tqdm -from transformers import AutoModelForCausalLM, PreTrainedTokenizer, PreTrainedTokenizerFast - -import modelopt.torch.opt as mto -import modelopt.torch.quantization as mtq -from modelopt.recipe.presets import KV_CACHE_NONE, KV_QUANT_CFG_CHOICES, QUANT_CFG_CHOICES -from modelopt.torch.export import get_model_type -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.quantization.config import need_calibration -from modelopt.torch.quantization.utils import patch_fsdp_mp_dtypes -from modelopt.torch.utils.dataset_utils import get_dataset_dataloader, get_supported_datasets - -# Constants -RAND_SEED = 1234 - - -# Enable HuggingFace checkpointing -mto.enable_huggingface_checkpointing() - - -def parse_args(): - """Parse command line arguments.""" - parser = argparse.ArgumentParser(description="Multi-node post-training quantization with FSDP2") - - parser.add_argument( - "--pyt_ckpt_path", - required=True, - help="Path to PyTorch checkpoint", - ) - parser.add_argument( - "--qformat", - default="fp8", - choices=list(QUANT_CFG_CHOICES), - help="Quantization format", - ) - parser.add_argument( - "--kv_cache_qformat", - default="fp8", - choices=[KV_CACHE_NONE, *KV_QUANT_CFG_CHOICES], - help="KV cache quantization format", - ) - parser.add_argument( - "--batch_size", - type=int, - default=1, - help="Batch size for calibration", - ) - parser.add_argument( - "--calib_size", - type=str, - default="512", - help="Comma-separated list of calibration sizes per dataset", - ) - parser.add_argument( - "--dataset", - help=( - f"name of a dataset, or a comma separated list of datasets. " - f"dataset choices are {get_supported_datasets()}" - ), - type=str, - default=None, - ) - parser.add_argument( - "--export_path", - default="exported_model", - help="Directory to export the quantized model", - ) - parser.add_argument( - "--trust_remote_code", - action="store_true", - help="Trust remote code for HuggingFace models", - ) - parser.add_argument("--awq_block_size", default=0, type=int) - - args = parser.parse_args() - - # Parse comma-separated lists - args.dataset = args.dataset.split(",") if args.dataset else None - args.calib_size = [int(x) for x in args.calib_size.split(",")] - - return args - - -def load_and_prepare_model( - model_path: str, - calib_dataloader: torch.utils.data.DataLoader, - accelerator: Accelerator, - trust_remote_code: bool = False, -) -> tuple[nn.Module, str, list[str], torch.utils.data.DataLoader]: - """Load model and prepare it for FSDP2 distributed execution. - - Args: - model_path: Path to the HuggingFace model - calibration_dataloader: Calibration dataloader to be sharded for calibration - accelerator: Accelerate's Accelerator instance - trust_remote_code: Whether to trust remote code - - Returns: - Tuple of (prepared_model, model_type, original_architectures, calibration_dataloader) - """ - model = AutoModelForCausalLM.from_pretrained( - model_path, dtype="auto", trust_remote_code=trust_remote_code - ) - model.eval() - model_type = get_model_type(model) - # Need the original architectures for export - # FSDP prefix is added to the architectures for FSDP2 wrapped models - original_architectures = model.config.architectures - - # FSDP2 requires an optimizer to be prepared together with the model - dummy_optimizer = torch.optim.SGD(model.parameters(), lr=0.0) - model, _, calibration_dataloader = accelerator.prepare(model, dummy_optimizer, calib_dataloader) - - return model, model_type, original_architectures, calibration_dataloader - - -def create_calibration_dataloader( - tokenizer: PreTrainedTokenizer | PreTrainedTokenizerFast, - dataset_names: list[str], - calib_sizes: list[int], - batch_size: int, -) -> torch.utils.data.DataLoader: - """Create calibration dataloader from dataset. - - Args: - tokenizer: HuggingFace tokenizer - dataset_names: List of dataset names (defaults to cnn_dailymail) - calib_sizes: Number of samples for each dataset - batch_size: Batch size for calibration - - Returns: - DataLoader for calibration - """ - - return get_dataset_dataloader( - dataset_name=dataset_names, - tokenizer=tokenizer, - batch_size=batch_size, - num_samples=calib_sizes, - device=None, # Keep data on CPU, calibration loop handles device transfer - include_labels=False, - ) - - -def create_fsdp2_calibration_loop( - model: nn.Module, - dataloader: torch.utils.data.DataLoader, - accelerator: Accelerator, -): - """Create calibration loop compatible with FSDP2. - - For FSDP2, we need to use the outer FSDP-wrapped model instead of - the parameter passed by mtq.quantize to properly handle DTensor. - - Args: - model: FSDP2-wrapped model - dataloader: Calibration dataloader - accelerator: Accelerator instance for device management - - Returns: - Calibration function compatible with mtq.quantize - """ - - def calibrate(unwrapped_model): - """Calibration loop that uses the FSDP-wrapped model.""" - for batch in tqdm(dataloader, desc="Calibrating"): - if isinstance(batch, dict): - batch = { - k: v.to(accelerator.device) if isinstance(v, torch.Tensor) else v - for k, v in batch.items() - } - # Use outer model (FSDP-wrapped), not the parameter - # Important: We should forward pass using the unwrapped model - # mtq.quantize will unwrap the model & pass to the forward_loop - model(**batch) - - return calibrate - - -def export_model( - model: nn.Module, - accelerator: Accelerator, - export_path: str | Path, - architectures: list[str], -): - """Export quantized model to HuggingFace format. - - Args: - model: Quantized model - accelerator: Accelerator instance for state dict gathering - export_path: Directory to export model to - """ - export_dir = Path(export_path) - export_dir.mkdir(parents=True, exist_ok=True) - - post_state_dict, hf_quant_config = _export_transformers_checkpoint( - model, torch.bfloat16, accelerator=accelerator - ) - - if accelerator.is_main_process: - # Save hf_quant_config.json for backward compatibility - with open(f"{export_dir}/hf_quant_config.json", "w") as file: - json.dump(hf_quant_config, file, indent=4) - - hf_quant_config = convert_hf_quant_config_format(hf_quant_config) - - # Save model - model.save_pretrained(export_dir, state_dict=post_state_dict, save_modelopt_state=False) - - original_config = f"{export_dir}/config.json" - config_data = {} - - with open(original_config) as file: - config_data = json.load(file) - - config_data["quantization_config"] = hf_quant_config - # Update config architectures to use original architectures that does not have FSDP prefix - config_data["architectures"] = architectures - - with open(original_config, "w") as file: - json.dump(config_data, file, indent=4) - - -def main(args): - """Main quantization workflow.""" - # Validate GPU availability - if not torch.cuda.is_available(): - raise OSError("GPU is required for quantization.") - - # Validate quantization format - if args.qformat not in QUANT_CFG_CHOICES: - raise ValueError( - f"Quantization format {args.qformat} not supported. Choose from: {list(QUANT_CFG_CHOICES)}" - ) - - # Set random seeds - random.seed(RAND_SEED) - np.random.seed(RAND_SEED) - torch.manual_seed(RAND_SEED) - - # Initialize accelerator - accelerator = Accelerator() - - print(f"Rank: {os.environ.get('RANK', 'Not set')}") - print(f"World Size: {os.environ.get('WORLD_SIZE', 'Not set')}") - print(f"Local Rank: {os.environ.get('LOCAL_RANK', 'Not set')}") - - # Load tokenizer - tokenizer = get_tokenizer(args.pyt_ckpt_path, trust_remote_code=args.trust_remote_code) - default_padding_side = tokenizer.padding_side - tokenizer.padding_side = "left" # Left padding for better calibration - - # Set default dataset if not provided - if args.dataset is None: - args.dataset = ["cnn_dailymail", "nemotron-post-training-dataset-v2"] - warnings.warn( - "No dataset specified. Defaulting to cnn_dailymail and nemotron-post-training-dataset-v2." - ) - # Adjust calib_size to match dataset length by extending or truncating as needed - args.calib_size = (args.calib_size + [args.calib_size[-1]] * len(args.dataset))[ - : len(args.dataset) - ] - - # Create calibration dataloader with max batch size - calib_dataloader = create_calibration_dataloader( - tokenizer=tokenizer, - dataset_names=args.dataset, - calib_sizes=args.calib_size, - batch_size=args.batch_size, - ) - - # Load and prepare model - model, model_type, original_architectures, calib_dataloader = load_and_prepare_model( - model_path=args.pyt_ckpt_path, - calib_dataloader=calib_dataloader, - accelerator=accelerator, - trust_remote_code=args.trust_remote_code, - ) - - quant_cfg = QUANT_CFG_CHOICES[args.qformat] - - quant_cfg = build_quant_cfg( - quant_cfg, - args.awq_block_size, - ) - - enable_quant_kv_cache = args.kv_cache_qformat != KV_CACHE_NONE - print(f"{'Enable' if enable_quant_kv_cache else 'Disable'} KV cache quantization") - - # Check if any bmm_quantizer is in the quant_cfg. If so, we need to enable the bmm_quantizer. - if enable_quant_kv_cache: - quant_cfg = mtq.update_quant_cfg_with_kv_cache_quant( - quant_cfg, - KV_QUANT_CFG_CHOICES[args.kv_cache_qformat]["quant_cfg"], - ) - - # Quantize the model - if accelerator.is_main_process: - print("Starting quantization...") - - start_time = time.time() - - if need_calibration(quant_cfg): - calibrate_fn = create_fsdp2_calibration_loop(model, calib_dataloader, accelerator) - else: - calibrate_fn = None - warnings.warn("Dynamic quantization. Calibration skipped.") - - with torch.no_grad(): - model = mtq.quantize(model, quant_cfg, forward_loop=calibrate_fn) - - elapsed = time.time() - start_time - - if accelerator.is_main_process: - print(f"Quantization completed in {elapsed:.2f}s") - mtq.print_quant_summary(model) - - start_time = time.time() - export_model(model, accelerator, args.export_path, original_architectures) - elapsed = time.time() - start_time - - if accelerator.is_main_process: - # Restore default padding and export the tokenizer as well. - if tokenizer is not None: - tokenizer.padding_side = default_padding_side - tokenizer.save_pretrained(args.export_path) - # Export the model - print(f"Export completed in {elapsed:.2f}s") - print(f"Model exported to {args.export_path}") - - print("Unpatching FSDP2 MP dtypes") - - -if __name__ == "__main__": - args = parse_args() - # This context manager can be removed once the update to FSDP2 function is reflected in torch - with patch_fsdp_mp_dtypes(): - main(args) diff --git a/examples/hf_ptq/slurm/multinode_fsdp2_ptq.slurm b/examples/hf_ptq/slurm/multinode_fsdp2_ptq.slurm new file mode 100644 index 00000000000..96ed4e2dbe0 --- /dev/null +++ b/examples/hf_ptq/slurm/multinode_fsdp2_ptq.slurm @@ -0,0 +1,90 @@ +#!/bin/bash + +# SPDX-FileCopyrightText: Copyright (c) 2025 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. +# +# Multi-node post-training quantization with FSDP2, ready to run under Slurm. +# +# Slurm allocates the nodes and launches one task per node; each task starts a `torchrun` that +# spawns one process per GPU. `hf_ptq.py --use_fsdp2` then shards the model across all ranks +# (world_size = num_nodes * gpus_per_node) for distributed loading, calibration, and export. +# +# The defaults below quantize nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16 to NVFP4 with an FP8 +# KV cache. Edit the CONFIG block for your cluster/model, then submit with e.g.: +# +# sbatch --nodes=2 multinode_fsdp2_ptq.slurm +# +# For a single-node run, `--nodes=1` works unchanged. + +#SBATCH --job-name=fsdp2-ptq +#SBATCH --account={account} +#SBATCH --partition={partition} +#SBATCH --nodes=2 +#SBATCH --ntasks-per-node=1 # one torchrun launcher per node; it fans out to the GPUs +#SBATCH --gpus-per-node=8 +#SBATCH --exclusive +#SBATCH --time=04:00:00 +#SBATCH --output=%x_%j.log + +set -euo pipefail + +# --------------------------------------------------------------------------- +# CONFIG — edit these for your cluster and model. They are exported so `srun` +# propagates them into the container task below. +# --------------------------------------------------------------------------- +export CONTAINER_IMAGE={container_image} # e.g. nvcr.io#nvidia/pytorch:25.10-py3 +export MODELOPT_PATH={path_to_modelopt_repo} # host clone of TensorRT-Model-Optimizer +export MODEL_PATH=nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16 # HF repo id (auto-downloaded) or a local dir +export EXPORT_PATH={path_to_export_dir} # where the quantized checkpoint is written +export HF_HOME={path_to_hf_cache} # HF cache; persist so a repo id isn't re-downloaded each run +# export HF_TOKEN={hf_token} # required for gated repos (or `huggingface-cli login` on host) +export RECIPE=general/ptq/nvfp4_default-kv_fp8_cast # built-in recipe name or /path/to/recipe.yaml +export CALIB_SIZE=512 +export BATCH_SIZE=4 + +# Rendezvous: node 0 is the master; all ranks meet at MASTER_ADDR:MASTER_PORT. +export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -1) +export MASTER_PORT=29531 + +# Mount the repo, export dir, and HF cache; add the model dir only when MODEL_PATH is a local path +# (a repo id is downloaded into HF_HOME instead). Run from the hf_ptq example dir. +CONTAINER_MOUNTS="${MODELOPT_PATH}:/modelopt,${EXPORT_PATH}:${EXPORT_PATH},${HF_HOME}:${HF_HOME}" +if [ -d "${MODEL_PATH}" ]; then + CONTAINER_MOUNTS="${CONTAINER_MOUNTS},${MODEL_PATH}:${MODEL_PATH}" +fi +srun --container-image="${CONTAINER_IMAGE}" \ + --container-mounts="${CONTAINER_MOUNTS}" \ + --container-workdir=/modelopt/examples/hf_ptq \ + bash -c ' + set -euo pipefail + export PYTHONUNBUFFERED=1 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True + # ModelOpt + its HF deps (transformers/accelerate/datasets/...) from the mounted source, then example extras. + pip install -q -e "/modelopt[hf]" --no-build-isolation + pip install -q -r requirements.txt + + torchrun \ + --nnodes="${SLURM_NNODES}" \ + --node_rank="${SLURM_NODEID}" \ + --nproc_per_node="$(nvidia-smi -L | wc -l)" \ + --rdzv_backend=c10d \ + --rdzv_endpoint="${MASTER_ADDR}:${MASTER_PORT}" \ + hf_ptq.py \ + --pyt_ckpt_path "${MODEL_PATH}" \ + --recipe "${RECIPE}" \ + --calib_size "${CALIB_SIZE}" \ + --batch_size "${BATCH_SIZE}" \ + --export_path "${EXPORT_PATH}" \ + --use_fsdp2 + ' diff --git a/examples/megatron_bridge/prune_minitron.py b/examples/megatron_bridge/prune_minitron.py index 49cc8beee87..9f938495d4c 100644 --- a/examples/megatron_bridge/prune_minitron.py +++ b/examples/megatron_bridge/prune_minitron.py @@ -107,6 +107,14 @@ def get_args() -> argparse.Namespace: parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) parser.add_argument("--hf_model_name_or_path", type=str, required=True) parser.add_argument("--trust_remote_code", action="store_true") + parser.add_argument( + "--no_moe_grouped_gemm", + action="store_true", + help=( + "Use SequentialMLP for MoE experts instead of the (default) efficient fused " + "TEGroupedMLP (grouped GEMM). Only affects MoE models." + ), + ) target_group = parser.add_mutually_exclusive_group(required=True) target_group.add_argument( @@ -382,7 +390,7 @@ def main(args: argparse.Namespace): "mtp_num_layers": 0, # MTP is not supported during calibration }, init_model_parallel=True, - moe_grouped_gemm=False, + moe_grouped_gemm=not args.no_moe_grouped_gemm, ) # TODO: Support pruning with MTP heads enabled (e.g. Qwen3.5 mtp_num_hidden_layers=1). @@ -503,18 +511,9 @@ def main(args: argparse.Namespace): ) match = re.fullmatch(r"mmlu_(\d+)pct_bs(\d+)", args.prune_score_func) - legacy_match = re.fullmatch(r"mmlu_(\d+)pct", args.prune_score_func) if match: mmlu_frac = float(match.group(1)) / 100.0 batch_size = int(match.group(2)) - elif legacy_match: - warn_rank_0( - f"Score function '{args.prune_score_func}' uses the deprecated format " - "'mmlu_pct'. Use 'mmlu_pct_bs' to specify the evaluation batch size. " - "Falling back to batch_size=1." - ) - mmlu_frac = float(legacy_match.group(1)) / 100.0 - batch_size = 1 else: raise ValueError( f"Invalid score function: {args.prune_score_func}. " diff --git a/examples/speculative_decoding/requirements.txt b/examples/speculative_decoding/requirements.txt index 0c679b55024..c13df8ca017 100644 --- a/examples/speculative_decoding/requirements.txt +++ b/examples/speculative_decoding/requirements.txt @@ -1,3 +1,3 @@ accelerate>=1.12.0 peft==0.18.1 -transformers>=5.0,<5.4 +transformers>=5.0,<5.13 diff --git a/modelopt/onnx/quantization/__main__.py b/modelopt/onnx/quantization/__main__.py index edf05df30e3..a42eb82953e 100644 --- a/modelopt/onnx/quantization/__main__.py +++ b/modelopt/onnx/quantization/__main__.py @@ -327,6 +327,16 @@ def get_parser() -> argparse.ArgumentParser: action="store_true", help="If True, the given ONNX model will be simplified before quantization is performed.", ) + argparser.add_argument( + "--simplify_backend", + type=str, + default="onnxslim", + choices=["onnxslim", "onnxsim"], + help=( + "ONNX simplification package to use when --simplify is set. " + "Both produce an equivalent simplified model (default: onnxslim)." + ), + ) argparser.add_argument( "--calibrate_per_node", action="store_true", @@ -552,6 +562,7 @@ def main(): use_zero_point=args.use_zero_point, passes=args.passes, simplify=args.simplify, + simplify_backend=args.simplify_backend, calibrate_per_node=args.calibrate_per_node, direct_io_types=args.direct_io_types, opset=args.opset, diff --git a/modelopt/onnx/quantization/quantize.py b/modelopt/onnx/quantization/quantize.py index d8b2471f127..339cf21555e 100755 --- a/modelopt/onnx/quantization/quantize.py +++ b/modelopt/onnx/quantization/quantize.py @@ -122,6 +122,7 @@ def _preprocess_onnx( trt_plugins_precision: list[str] | None, override_shapes: str, simplify: bool = False, + simplify_backend: str = "onnxslim", quantize_mode: str = "int8", opset: int | None = None, ) -> tuple[str, onnx.ModelProto, list[str], bool, bool, bool, dict, dict]: @@ -218,10 +219,32 @@ def _preprocess_onnx( # Simplify model if requested if simplify: - logger.info("Attempting to simplify model") + logger.info(f"Attempting to simplify model with '{simplify_backend}'") + + # Resolve the backend before attempting simplification so that a missing + # optional dependency or an unknown backend name fails loudly instead of + # being silently swallowed by the graceful fallback below. + if simplify_backend == "onnxsim": + try: + import onnxsim + except ModuleNotFoundError as e: + logger.warning( + "onnxsim is not installed. Please install it with 'pip install onnxsim'." + ) + raise e + elif simplify_backend != "onnxslim": + raise ValueError( + f"Unsupported simplify_backend '{simplify_backend}'. " + "Choose one of: 'onnxslim', 'onnxsim'." + ) + try: - model_simp = onnxslim.slim(onnx_model, skip_fusion_patterns=["FusionGemm"]) - if model_simp: + if simplify_backend == "onnxslim": + model_simp = onnxslim.slim(onnx_model, skip_fusion_patterns=["FusionGemm"]) + check = model_simp is not None + else: + model_simp, check = onnxsim.simplify(onnx_model) + if check: onnx_model = model_simp onnx_path = os.path.join(output_dir, f"{model_name}_simp.onnx") save_onnx(onnx_model, onnx_path, use_external_data_format) @@ -375,6 +398,7 @@ def quantize( use_zero_point: bool = False, passes: list[str] = ["concat_elimination"], simplify: bool = False, + simplify_backend: str = "onnxslim", calibrate_per_node: bool = False, input_shapes_profile: Sequence[dict[str, str]] | None = None, model_id: str | None = None, @@ -477,6 +501,9 @@ def quantize( List of optimization passes name, if set, appropriate pre/post-processing passes will be invoked. simplify: Simplify the given model before quantization. + simplify_backend: + ONNX simplification package to use when ``simplify`` is set. One of ``"onnxslim"`` + (default) or ``"onnxsim"``; both produce an equivalent simplified model. calibrate_per_node: Calibrate the model node by node instead of calibrating the entire model. This allows calibration with a lower system memory with the cost of longer calibration time. @@ -627,6 +654,7 @@ def quantize( trt_plugins_precision, override_shapes, # type: ignore[arg-type] simplify, + simplify_backend, quantize_mode, opset, ) diff --git a/modelopt/torch/export/registry.py b/modelopt/torch/export/registry.py index 8e2cda63df9..260cb32eea3 100644 --- a/modelopt/torch/export/registry.py +++ b/modelopt/torch/export/registry.py @@ -32,6 +32,8 @@ import torch import torch.nn as nn +from modelopt.torch.utils.distributed import is_fsdp2_model + __all__ = [ "ExportContext", "ExportHandler", @@ -54,8 +56,17 @@ class ExportContext: model: nn.Module dtype: torch.dtype is_modelopt_qlora: bool = False - tied_cache: dict[int, nn.Module] = field(default_factory=dict) - moe_tied_cache: dict[tuple[int, int], nn.Module] = field(default_factory=dict) + tied_cache: dict[int, nn.Module] | None = field(default_factory=dict) + moe_tied_cache: dict[tuple[int, int], nn.Module] | None = field(default_factory=dict) + + def __post_init__(self) -> None: + # FSDP2 may recycle data_ptr() values as modules are resharded, so pointer-keyed dedup can + # falsely alias distinct weights. Disable it for FSDP2; consequently, legitimately tied + # packed weights and scale buffers are not re-aliased and may be stored as duplicates. + # TODO: replace this with stable, name-based tied-group deduplication. + if is_fsdp2_model(self.model): + self.tied_cache = None + self.moe_tied_cache = None ExportHandler = Callable[[str, nn.Module, ExportContext], None] diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index cee64c22c05..f51eea17b1a 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -51,6 +51,7 @@ except ImportError: HAS_DIFFUSERS = False +from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict from torch.distributed.fsdp import FSDPModule from modelopt.torch.quantization import set_quantizer_by_cfg_context @@ -58,6 +59,7 @@ from modelopt.torch.quantization.qtensor import MXFP8QTensor, NVFP4QTensor from modelopt.torch.quantization.utils import fsdp2_aware_weight_update, quantizer_attr_names from modelopt.torch.utils.dataset_utils import _disable_use_cache +from modelopt.torch.utils.distributed import is_fsdp2_model try: from modelopt.torch.sparsity.attention_sparsity.conversion import export_sparse_attention_config @@ -840,7 +842,6 @@ def _export_transformers_checkpoint( Args: model: the full torch model to export. The actual quantized model may be a submodule. dtype: the weights data type to export the unquantized layers or the default model data type if None. - accelerator: the accelerator instance in case of distributed export setup. Returns: post_state_dict: Dict containing quantized weights @@ -854,8 +855,6 @@ def _export_transformers_checkpoint( f"({dtype}), which may lead to numerical errors." ) - accelerator = kwargs.get("accelerator") - # Handle input quantizers of experts that are not calibrated. Each MoE block is # dispatched by its experts container to the matching preparation handler. prepare_ctx = ExportContext(model=model, dtype=dtype, is_modelopt_qlora=is_modelopt_qlora) @@ -925,10 +924,14 @@ def _export_transformers_checkpoint( _reconstruct_fused_moe_linear(model) - if accelerator is not None: - # Gather state_dict from all ranks - quantized_state_dict = accelerator.get_state_dict(model) + if is_fsdp2_model(model): + # FSDP2: gather the full (unsharded) state_dict to CPU on rank 0. + quantized_state_dict = get_model_state_dict( + model, + options=StateDictOptions(full_state_dict=True, cpu_offload=True), + ) else: + # Non-FSDP2: assumes a replicated model (rank 0 has the full state dict). quantized_state_dict = model.state_dict() # We define kv cache scale as amax / 448 for both FP8 and NVFP4 KV cache quantization. @@ -1397,6 +1400,10 @@ def export_hf_checkpoint( This function automatically detects whether the model is from transformers or diffusers and applies the appropriate export logic. + Under ``torch.distributed`` (e.g. FSDP2), all ranks participate in the + collective state-dict gather inside ``_export_transformers_checkpoint``; + only rank 0 writes files. A final barrier syncs the other ranks. + Args: model: The full torch model to export. The actual quantized model may be a submodule. Supports both transformers models (e.g., LlamaForCausalLM) and diffusers @@ -1430,6 +1437,11 @@ def export_hf_checkpoint( ) return + is_distributed = ( + torch.distributed.is_available() + and torch.distributed.is_initialized() + and is_fsdp2_model(model) + ) try: post_state_dict, hf_quant_config = _export_transformers_checkpoint(model, dtype, **kwargs) @@ -1461,6 +1473,10 @@ def export_hf_checkpoint( "names may not match the original HF hub checkpoint." ) + # Under torch.distributed only rank 0 writes; others sync at the finally barrier. + if is_distributed and torch.distributed.get_rank() != 0: + return + # Only treat the export as quantized when at least one quant_algo field is set. # get_quant_config always returns a dict (even for sparsity-only or unmodified models), # so emitting hf_quant_config.json unconditionally produces a file with @@ -1489,6 +1505,7 @@ def export_hf_checkpoint( _sanitize_generation_config_for_save(model) + # TODO: parallelize the disk write across ranks (avoid single-process speed + rank-0 OOM). try: model.save_pretrained( export_dir, @@ -1525,3 +1542,6 @@ def export_hf_checkpoint( " can be saved with torch.save for further inspection." ) raise e + finally: + if is_distributed: + torch.distributed.barrier() diff --git a/modelopt/torch/nas/plugins/megatron.py b/modelopt/torch/nas/plugins/megatron.py index 860e6c35d70..b48485591c8 100644 --- a/modelopt/torch/nas/plugins/megatron.py +++ b/modelopt/torch/nas/plugins/megatron.py @@ -19,15 +19,18 @@ import types from abc import ABC from collections.abc import Callable, Sequence +from functools import partial import torch import torch.nn as nn import transformer_engine as te from megatron.core.extensions.transformer_engine import ( + TEColumnParallelGroupedLinear, TEColumnParallelLinear, TEDotProductAttention, TELayerNormColumnParallelLinear, TELinear, + TERowParallelGroupedLinear, TERowParallelLinear, ) from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding @@ -44,7 +47,7 @@ from megatron.core.transformer.identity_op import IdentityOp from megatron.core.transformer.mlp import MLP from megatron.core.transformer.moe import moe_utils -from megatron.core.transformer.moe.experts import SequentialMLP +from megatron.core.transformer.moe.experts import SequentialMLP, TEGroupedMLP from megatron.core.transformer.moe.moe_layer import MoELayer from megatron.core.transformer.moe.router import TopKRouter from megatron.core.transformer.moe.shared_experts import SharedExpertMLP @@ -759,6 +762,11 @@ def _setup(self, *, hidden_size: TracedHp): for expert in self.local_experts: DMRegistry.convert(expert, hidden_size=hidden_size, hp_name="moe_ffn_hidden_size") + def modify(self, ffn_hidden_size_divisor: int = 1, **kwargs) -> None: + """Modify each expert's moe_ffn_hidden_size hparam choices based on search space config.""" + for expert in self.local_experts: + expert.modify(ffn_hidden_size_divisor=ffn_hidden_size_divisor) + def export(self) -> torch.nn.Module: """Export the dynamic module to a standard SequentialMLP.""" for expert in self.local_experts: @@ -767,6 +775,130 @@ def export(self) -> torch.nn.Module: return super().export() +@DMRegistry.register( + { + TEColumnParallelGroupedLinear: ( + "megatron.core.extensions.transformer_engine.TEColumnParallelGroupedLinear" + ), + TERowParallelGroupedLinear: ( + "megatron.core.extensions.transformer_engine.TERowParallelGroupedLinear" + ), + } +) +class _DynamicTEGroupedLinear(DynamicModule): + """A TEGroupedLinear (column/row parallel) with dynamic hyperparams for grouped-GEMM MoE. + + TEGroupedMLP fuses all local experts into two grouped linears, each storing the per-expert + weights as separate ``weight0..weight{num_gemms-1}`` params (shape ``[out, in]``, optional + ``bias{i}``). ``moe_ffn_hidden_size`` / ``hidden_size`` slice each expert weight by + ``output_size`` (rows) / ``input_size`` (cols) like a normal linear; ``num_local_experts`` + reorders/drops experts by remapping position ``j`` to the ``j``-th most important expert and + exposing ``num_gemms = num_local_experts.active`` so TE only reads the kept experts. + """ + + def _setup(self, *, input_size: TracedHp, output_size: TracedHp, num_local_experts: TracedHp): + assert not self.single_grouped_weight, ( + "moe_single_grouped_weight=True is not supported for grouped-GEMM pruning yet." + ) + # input_size/output_size/num_local_experts are all shared with the sibling grouped linear + # (and num_local_experts additionally with the router) via _DynamicTEGroupedMLP. + self._register_hparam("input_size", input_size) + self._register_hparam("output_size", output_size) + self._register_hparam("num_local_experts", num_local_experts) + + self._register_dynamic_attribute("num_gemms", lambda mod, val: num_local_experts.active) + self._register_dynamic_attribute("in_features", lambda mod, val: input_size.active) + self._register_dynamic_attribute("out_features", lambda mod, val: output_size.active) + for j in range(self.num_gemms): + self._register_dynamic_attribute(f"weight{j}", partial(self._get_expert_param, pos=j)) + if self.use_bias: + self._register_dynamic_attribute(f"bias{j}", partial(self._get_expert_param, pos=j)) + + @staticmethod + def _get_expert_param(mod: "_DynamicTEGroupedLinear", val: torch.Tensor, *, pos: int): + """Dynamic getter for weight{pos}/bias{pos}: map position -> ranked expert, then slice.""" + hp = mod.get_hparam("num_local_experts") + max_experts = hp.max + assert isinstance(max_experts, int) + order = hp._slice_order.tolist() if hp._slice_order is not None else range(max_experts) + e = order[pos] + is_weight = val.dim() == 2 + raw = mod._parameters[f"{'weight' if is_weight else 'bias'}{e}"] + slices = [mod.get_hparam("output_size").active_slice] + if is_weight: + slices.append(mod.get_hparam("input_size").active_slice) + return get_sliced_tensor_by_slices(raw, slices) + + def export(self) -> torch.nn.Module: + """Export to a standard TEGroupedLinear with the kept experts sliced + reordered in place.""" + # Read all sliced/reordered params (via the dynamic getters) before mutating any, then drop + # the per-expert weight/bias attrs so the base export only folds num_gemms/in/out_features. + active = self.get_hparam("num_local_experts").active + assert isinstance(active, int) + weights = [getattr(self, f"weight{j}").detach().clone() for j in range(active)] + biases = [ + getattr(self, f"bias{j}").detach().clone() for j in range(active) if self.use_bias + ] + for name in [n for n in list(self._parameters) if n.startswith(("weight", "bias"))]: + delattr(self, name) + + super().export() # num_gemms -> active, in/out_features -> sliced sizes, class un-patched + + for j, weight in enumerate(weights): + self.register_parameter(f"weight{j}", torch.nn.Parameter(weight)) + for j, bias in enumerate(biases): + self.register_parameter(f"bias{j}", torch.nn.Parameter(bias)) + return self + + +@DMRegistry.register({TEGroupedMLP: "megatron.core.transformer.moe.experts.TEGroupedMLP"}) +class _DynamicTEGroupedMLP(DynamicModule): + """A TEGroupedMLP (grouped-GEMM MoE experts) with dynamic hyperparams. + + Mirrors ``_DynamicSequentialMLP`` but the experts are two fused ``TEGroupedLinear`` layers rather + than an ``nn.ModuleList`` of per-expert MLPs. Since Minitron prunes homogeneously, all experts + share a single ``moe_ffn_hidden_size`` hparam (unlike the SequentialMLP path which registers one per expert). + """ + + def _setup(self, *, hidden_size: TracedHp): + """Setup the TEGroupedMLP dynamic module with global hidden_size hparam.""" + num_local_experts = TracedHp(list(range(1, self.num_local_experts + 1))) + self._register_hparam("num_local_experts", num_local_experts) + + moe_ffn_hidden_size = TracedHp(list(range(1, self.config.moe_ffn_hidden_size + 1))) + self._register_hparam("moe_ffn_hidden_size", moe_ffn_hidden_size) + + linear_fc1_output_size = ( + build_concat_hp([moe_ffn_hidden_size] * 2) + if self.config.gated_linear_unit + else moe_ffn_hidden_size + ) + DMRegistry.convert( # _DynamicTEGroupedLinear + self.linear_fc1, + input_size=hidden_size, + output_size=linear_fc1_output_size, + num_local_experts=num_local_experts, + ) + DMRegistry.convert( # _DynamicTEGroupedLinear + self.linear_fc2, + input_size=moe_ffn_hidden_size, + output_size=hidden_size, + num_local_experts=num_local_experts, + ) + + def modify(self, ffn_hidden_size_divisor: int = 1, **kwargs) -> None: + """Modify the shared moe_ffn_hidden_size hparam choices based on search space config.""" + hp = self.get_hparam("moe_ffn_hidden_size") + choices = {int(make_divisible(c, ffn_hidden_size_divisor)) for c in hp.choices} # type: ignore[arg-type] + hp.choices = list(set(hp.choices) & choices | {hp.original}) + + def export(self) -> torch.nn.Module: + """Export the dynamic module to a standard TEGroupedMLP.""" + self.linear_fc1.export() + self.linear_fc2.export() + return super().export() + + @DMRegistry.register({MoELayer: "megatron.core.transformer.moe.moe_layer.MoELayer"}) class _DynamicMoELayer(DynamicModule): """A MoELayer with dynamic hyperparams.""" @@ -837,8 +969,7 @@ def modify( expert_hp.choices = list(set(expert_hp.choices) & choices | {expert_hp.original}) # Modify expert FFN hparam choices - for expert in self.experts.local_experts: - expert.modify(ffn_hidden_size_divisor=ffn_hidden_size_divisor) + self.experts.modify(ffn_hidden_size_divisor=ffn_hidden_size_divisor) if self.use_shared_expert: self.shared_experts.modify(ffn_hidden_size_divisor) diff --git a/modelopt/torch/opt/utils.py b/modelopt/torch/opt/utils.py index 96c0dd555f3..787b96ceddd 100644 --- a/modelopt/torch/opt/utils.py +++ b/modelopt/torch/opt/utils.py @@ -113,6 +113,7 @@ def _lazy_init_retain_mesh_info(self): fsdp_state._lazy_init = types.MethodType(_lazy_init_retain_mesh_info, fsdp_state) fsdp_states.append(fsdp_state) + yield for fsdp_state in fsdp_states: if fsdp_state._fsdp_param_group and hasattr(fsdp_state, "_post_forward_mesh_info_after"): diff --git a/modelopt/torch/prune/plugins/mcore_minitron.py b/modelopt/torch/prune/plugins/mcore_minitron.py index ac6fe138b2d..876432fee01 100644 --- a/modelopt/torch/prune/plugins/mcore_minitron.py +++ b/modelopt/torch/prune/plugins/mcore_minitron.py @@ -67,6 +67,7 @@ _DynamicMoELayer, _DynamicSelfAttention, _DynamicSequentialMLP, + _DynamicTEGroupedMLP, _DynamicTransformerLayer, ) from modelopt.torch.nas.plugins.megatron_model_stats import ( @@ -974,6 +975,8 @@ def __init__(self, model: DynamicModule): _register_mlp_importance(module, self) elif isinstance(module, _DynamicSequentialMLP): _register_sequential_mlp_importance(module, self) + elif isinstance(module, _DynamicTEGroupedMLP): + _register_grouped_mlp_importance(module, self) elif isinstance(module, _DynamicMambaMixer): _register_mamba_mixer_importance(module, self) @@ -1402,6 +1405,46 @@ def _estimate_expert_importance(mod): ) +def _register_grouped_mlp_importance( + module: _DynamicTEGroupedMLP, registry: ImportanceEstimatorRegistry +) -> None: + """Register importance estimators for TEGroupedMLP (grouped-GEMM MoE experts) modules. + + Mirrors the SequentialMLP path: ``num_local_experts`` reuses the expert-L2 hook (TEGroupedMLP + shares SequentialMLP's forward signature), and ``moe_ffn_hidden_size`` is a single shared score + from the fused ``linear_fc2`` input activations (all experts' tokens), since experts prune + homogeneously. + """ + # Expert importance for num_local_experts; also creates module._activations (the dict saved and + # restored by the per-rank score checkpoint). We stash the ffn score in it so re-pruning from a + # checkpoint recovers it without re-running the forward loop. + _register_sequential_mlp_importance(module, registry) + module._activations["ffn_activations"] = None + + def _grouped_fc2_forward_hook(mod, module_inner, input, output): + """Collect ffn-channel activations from the fused linear_fc2 input (all experts' tokens).""" + # input[0] is the permuted intermediate [total_tokens, moe_ffn_hidden_size] (no batch dim) + acts = gather_from_tensor_model_parallel_region(input[0]).detach()[:, None, :] + acts = acts.to(torch.float32).abs().mean(dim=0).pow(2).sum(dim=0) # [moe_ffn_hidden_size] + prev = mod._activations["ffn_activations"] + mod._activations["ffn_activations"] = acts if prev is None else prev + acts + + def _estimate_grouped_ffn_importance(mod): + """Return the activation magnitude-based importance (L2 norm) of moe_ffn_hidden_size.""" + acts = mod._activations["ffn_activations"] + assert acts is not None, "No activations collected for importance estimation." + return acts.pow(0.5) + + registry.register_hook( + module.linear_fc2, + partial(_grouped_fc2_forward_hook, module), + hook_type="forward", + ) + registry.register_importance( + module, "moe_ffn_hidden_size", lambda: _estimate_grouped_ffn_importance(module) + ) + + def _register_mamba_mixer_importance( module: _DynamicMambaMixer, registry: ImportanceEstimatorRegistry ) -> None: diff --git a/modelopt/torch/quantization/utils/core_utils.py b/modelopt/torch/quantization/utils/core_utils.py index 56eb373a7f6..1bdf23da64a 100644 --- a/modelopt/torch/quantization/utils/core_utils.py +++ b/modelopt/torch/quantization/utils/core_utils.py @@ -427,11 +427,14 @@ def _get_fsdp2_mesh(module: nn.Module): return None fsdp_state = _get_module_state(module) - if ( - fsdp_state._fsdp_param_group - and fsdp_state._fsdp_param_group.post_forward_mesh_info is not None - ): - return fsdp_state._fsdp_param_group.post_forward_mesh_info.mesh + pg = fsdp_state._fsdp_param_group + if pg is None: + return None + # A root FSDP module has reshard_after_forward=False by default, so its + # post_forward_mesh_info is None; fall back to the sharding mesh (mesh_info), + # which is the same FSDP shard mesh (post_forward_mesh_info is only the reshard target). + mesh_info = pg.post_forward_mesh_info or pg.mesh_info + return mesh_info.mesh if mesh_info is not None else None def _get_module_name(module: nn.Module, root_model: nn.Module, name_to_module: dict | None = None): @@ -514,8 +517,6 @@ def fsdp2_weight_access_and_writeback_context( If TP is implemented with DTensor, the weight will be a local tensor of the TP DTensor under this context. """ - assert isinstance(root_model, torch.distributed.fsdp.FSDPModule), "We only support FSDP2" - assert not hasattr(module, "_hf_hook"), "We dont support FSDP2 with HF accelerate hooks" fsdp_module = _get_enclosing_fsdp_module(module, root_model) assert fsdp_module is not None, "Module is not wrapped by FSDP" @@ -538,28 +539,48 @@ def fsdp2_weight_access_and_writeback_context( assert ( fsdp_device_mesh.mesh_dim_names == original_device_mesh.mesh_dim_names[:fsdp_dim] ), "FSDP2 mesh should be a slice of DTensor's device mesh." - collected = param.redistribute( + unsharded_dtensor = param.redistribute( placements=[Replicate()] * fsdp_dim + list(original_placements[fsdp_dim:]), device_mesh=original_device_mesh, ) - originals[name] = (param, collected, original_placements, original_device_mesh) - _set_parameter(module, name, nn.Parameter(collected.to_local())) - - yield - - # Write back and restore original DTensor parameters. - for name, ( - original_param, - collected, - original_placements, - original_device_mesh, - ) in originals.items(): - original_param.to_local().data.copy_( - collected.redistribute( - placements=original_placements, device_mesh=original_device_mesh - ).to_local() + unsharded_tensor = unsharded_dtensor.to_local() + # cpu_offload: gathered shard is on CPU; mirror to GPU for forward. + needs_gpu_copy = unsharded_tensor.device.type == "cpu" and torch.cuda.is_available() + gpu_tensor = ( + unsharded_tensor.to(torch.cuda.current_device()) if needs_gpu_copy else unsharded_tensor ) - _set_parameter(module, name, original_param) + cpu_writeback_tensor = unsharded_tensor if needs_gpu_copy else None + originals[name] = ( + param, + unsharded_dtensor, + original_placements, + original_device_mesh, + cpu_writeback_tensor, + gpu_tensor, + ) + _set_parameter(module, name, nn.Parameter(gpu_tensor)) + + try: + yield + finally: + # Write back and restore original DTensor parameters. Runs on both success + # and exception so the module never lingers with the temporary local params. + for name, ( + original_param, + unsharded_dtensor, + original_placements, + original_device_mesh, + cpu_writeback_tensor, + gpu_tensor, + ) in originals.items(): + if cpu_writeback_tensor is not None: + cpu_writeback_tensor.data.copy_(gpu_tensor.data.to(cpu_writeback_tensor.device)) + original_param.to_local().data.copy_( + unsharded_dtensor.redistribute( + placements=original_placements, device_mesh=original_device_mesh + ).to_local() + ) + _set_parameter(module, name, original_param) @contextmanager diff --git a/modelopt/torch/speculative/plugins/hf_streaming_dataset.py b/modelopt/torch/speculative/plugins/hf_streaming_dataset.py index aebb726e27f..0f9713948fd 100644 --- a/modelopt/torch/speculative/plugins/hf_streaming_dataset.py +++ b/modelopt/torch/speculative/plugins/hf_streaming_dataset.py @@ -41,6 +41,7 @@ import base64 import os +import re import time from typing import TypedDict @@ -105,6 +106,33 @@ def _tokenize_with_loss_mask( recovery = None if answer_only_loss and not getattr(tokenizer, "is_fast", False): recovery = get_loss_mask_recovery(tokenizer) + if answer_only_loss and recovery is None: + # Fail loudly on a template without {% generation %} tags: transformers only + # warns and returns an ALL-ZERO assistant mask, which trains every sample at + # zero loss with no other symptom. (The regex matches the tag itself, not the + # unrelated `add_generation_prompt` identifier most templates contain.) + template = getattr(tokenizer, "chat_template", None) or "" + if not re.search(r"\{%-?\s*generation\b", template): + raise RuntimeError( + "answer_only_loss=True needs assistant masks, but the tokenizer's chat " + "template has no {% generation %} tags, so apply_chat_template would " + "return an all-zero mask and training would silently run at zero loss. " + "Pass a tagged template copy via data.chat_template (see " + "tools/launcher/examples/MiniMaxAI/MiniMax-M3/m3_chat_template_generation.jinja " + "for a worked example), register a loss-mask recovery " + "(modelopt.torch.utils.loss_mask), or set answer_only_loss=false." + ) + if not getattr(tokenizer, "is_fast", False): + # A tagged template is not enough on a slow tokenizer: assistant-mask + # alignment needs the fast tokenizer's char_to_token, so apply_chat_template + # would fail downstream with an unrelated-looking error. + raise RuntimeError( + "answer_only_loss=True needs assistant masks, but the tokenizer is not a " + "fast tokenizer, so apply_chat_template cannot align {% generation %} tags " + "to tokens (char_to_token is unavailable). Use a fast tokenizer, register " + "a loss-mask recovery (modelopt.torch.utils.loss_mask), or set " + "answer_only_loss=false." + ) out = tokenizer.apply_chat_template( conversations, tokenize=True, diff --git a/modelopt/torch/speculative/plugins/modeling_fakebase.py b/modelopt/torch/speculative/plugins/modeling_fakebase.py index f896abab696..2b5fe989c03 100644 --- a/modelopt/torch/speculative/plugins/modeling_fakebase.py +++ b/modelopt/torch/speculative/plugins/modeling_fakebase.py @@ -204,7 +204,9 @@ def from_source(cls, source: str, trust_remote_code: bool = False) -> "FakeBaseM intermediate_size=getattr(base_cfg, "intermediate_size", None), rms_norm_eps=getattr(base_cfg, "rms_norm_eps", 1e-6), rope_theta=getattr(base_cfg, "rope_theta", None), - final_norm_type=_select_final_norm_type(getattr(base_cfg, "model_type", None)), + final_norm_type=_select_final_norm_type( + getattr(base_cfg, "model_type", None), base_cfg + ), ) model = cls(config) # Load lm_head, embed_tokens, and (for known models) the final norm into the model. diff --git a/modelopt/torch/speculative/plugins/modeling_final_norm.py b/modelopt/torch/speculative/plugins/modeling_final_norm.py index 85852bc2bef..718d591b662 100644 --- a/modelopt/torch/speculative/plugins/modeling_final_norm.py +++ b/modelopt/torch/speculative/plugins/modeling_final_norm.py @@ -22,6 +22,8 @@ only when we know which one the model uses. """ +from collections.abc import Callable + import torch from transformers.models.llama.modeling_llama import LlamaRMSNorm @@ -41,11 +43,37 @@ def __init__(self, hidden_size, eps=1e-6, dtype=torch.bfloat16): self.to(dtype) +class _FinalGemmaRMSNorm(torch.nn.Module): + """Gemma-style RMSNorm: fp32 normalize, scale by ``(1 + weight)``, multiply-then-cast. + + Mirrors MiniMax M3's ``MiniMaxM3VLRMSNorm`` (``use_gemma_norm=True`` configs): the stored + ``weight`` is a zero-centered delta, the effective scale is ``1 + weight``, and the multiply + happens in fp32 BEFORE casting back (``(x_hat * (1 + w)).type_as(x)``), unlike LlamaRMSNorm's + cast-then-multiply. Reusing ``_FinalRMSNorm`` here would silently drop the ``+1`` and corrupt + the reconstructed logits. ``weight`` is loaded from the base checkpoint. + """ + + def __init__(self, hidden_size, eps=1e-6, dtype=torch.bfloat16): + super().__init__() + self.eps = eps + self.weight = torch.nn.Parameter(torch.zeros(hidden_size, dtype=dtype)) + + def forward(self, x): + output = x.float() + output = output * torch.rsqrt(output.pow(2).mean(-1, keepdim=True) + self.eps) + output = output * (1.0 + self.weight.float()) + return output.type_as(x) + + def extra_repr(self): + return f"{tuple(self.weight.shape)}, eps={self.eps}" + + # Registry of self-implemented final-norm variants. We deliberately reimplement these # (rather than importing the base model's actual module) to keep FakeBaseModel lightweight. # Only a small, explicit set is supported; add a class here when a new type is needed. -_FINAL_NORM_CLASSES = { +_FINAL_NORM_CLASSES: dict[str, Callable[..., torch.nn.Module]] = { "rmsnorm": _FinalRMSNorm, + "gemma_rmsnorm": _FinalGemmaRMSNorm, } # Base ``model_type`` → final-norm type. ONLY listed models get a norm — applying the wrong or @@ -62,17 +90,28 @@ def __init__(self, hidden_size, eps=1e-6, dtype=torch.bfloat16): "deepseek_v3": "rmsnorm", "kimi_k2": "rmsnorm", # Kimi-K2 / K2-Thinking (DeepSeek-V3 arch) report model_type "kimi_k2" "kimi_k25": "rmsnorm", # Kimi-K2.5 / K2.6 / K2.7 all report model_type "kimi_k25" + # M3's final norm is always gemma-style; map it here too so a config that lost its + # use_gemma_norm flag still gets the correct flavor instead of silently dropping the +1. + "minimax_m3_vl_text": "gemma_rmsnorm", # gpt_oss intentionally DISABLED: GptOssRMSNorm uses an fp32 weight + multiply-then-cast, # unlike _FinalRMSNorm's bf16 weight, so reusing it would silently bias reconstructed logits. # Re-enable once a gpt_oss-style class (fp32 weight, multiply-then-cast) is in _FINAL_NORM_CLASSES. } -def _select_final_norm_type(model_type: str | None) -> str | None: +def _select_final_norm_type(model_type: str | None, base_cfg=None) -> str | None: """Return the final-norm type for a base ``model_type``, or ``None`` if unknown. ``None`` means we don't know the model's final norm, so FakeBaseModel builds no norm. + + ``base_cfg`` (the resolved text config, optional) takes precedence over the model_type + table when it carries an explicit ``use_gemma_norm=True`` flag: only MiniMax M2.x/M3 set + it, and their ``model_type`` is unreliable — MiniMax's VL remote code coerces a + model_type-less ``text_config`` to Mixtral (whose table entry is plain ``rmsnorm``), so + keying off model_type alone would silently apply the wrong norm flavor. """ + if base_cfg is not None and getattr(base_cfg, "use_gemma_norm", False): + return "gemma_rmsnorm" return _FINAL_NORM_TYPE_BY_MODEL_TYPE.get(model_type or "") diff --git a/modelopt/torch/utils/dataset_utils.py b/modelopt/torch/utils/dataset_utils.py index a3d44064683..a0beb92afbc 100644 --- a/modelopt/torch/utils/dataset_utils.py +++ b/modelopt/torch/utils/dataset_utils.py @@ -1013,7 +1013,9 @@ def _get_free_gpu_mem(): free_mem_before, max_allocated_before = _get_free_gpu_mem() is_enc_dec = model_type_is_enc_dec(model) - infer_method = model.generate if is_enc_dec else model.forward + # Call the module (not .forward) so nn.Module.__call__ runs pre/post-forward hooks — this is how + # FSDP2 unshards/reshards a sharded root. generate() also calls the module internally. + infer_method = model.generate if is_enc_dec else model if sample_input_single_batch is None: sample_input_single_batch = ( @@ -1163,7 +1165,9 @@ def _forward_loop( """ with _disable_use_cache(model), torch.no_grad(): is_enc_dec = model_type_is_enc_dec(model) - infer_method = model.generate if is_enc_dec else model.forward + # Call the module (not .forward) so nn.Module.__call__ runs pre/post-forward hooks — this is + # how FSDP2 unshards/reshards a sharded root. generate() also calls the module internally. + infer_method = model.generate if is_enc_dec else model max_working_batch_size = None # Initialize max working batch size as None for _, data in enumerate(tqdm(dataloader)): diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 7922b688052..12287865b7f 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -27,6 +27,7 @@ import torch import torch.distributed +from torch.distributed.fsdp import CPUOffloadPolicy, FSDPModule, fully_shard from torch.distributed.tensor import DTensor __all__ = [ @@ -34,7 +35,9 @@ "ParallelState", "backend", "barrier", + "fsdp2_wrap", "is_available", + "is_fsdp2_model", "is_initialized", "is_master", "rank", @@ -216,6 +219,82 @@ def cleanup(): torch.distributed.destroy_process_group() +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()) + + +def fsdp2_wrap(model, shard_root=True, mp_policy=None, cpu_offload: bool = False): + """Auto-detect a HF causal-LM's decoder layers and FSDP2 ``fully_shard`` each one. + + By default (``shard_root=True``) the root module is wrapped too, so embed/lm_head/norm are + sharded instead of replicated per rank; pass ``shard_root=False`` to leave the root replicated + (only decoder layers sharded). Returns the detected decoder layers so callers can reuse the + detection result. + """ + # Lazy import: layerwise_calib imports this module at top level (circular). + from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector + + decoder_layers = LayerActivationCollector.get_decoder_layers(model) + if decoder_layers is None: + raise RuntimeError( + "Could not auto-detect decoder layers; FSDP2 wrap requires a standard HF causal-LM layout." + ) + + fsdp_kwargs: dict[str, Any] = {"reshard_after_forward": True} + if mp_policy is not None: + fsdp_kwargs["mp_policy"] = mp_policy + if cpu_offload: + fsdp_kwargs["offload_policy"] = CPUOffloadPolicy() + + # Snapshot/restore config.architectures: some HF builders mutate it during fully_shard. + config = getattr(model, "config", None) + architectures = list(getattr(config, "architectures", []) or []) + for layer in decoder_layers: + fully_shard(layer, **fsdp_kwargs) + if shard_root: + fully_shard(model, **fsdp_kwargs) + if config is not None and architectures: + config.architectures = architectures + + return decoder_layers + + +def broadcast_state_dict( + state_dict_or_none: dict | None, + src: int, + device: torch.device, + pg=None, +) -> dict: + """Broadcast a dict of CPU tensors from rank ``src`` to all ranks. + + Two phases: (1) broadcast metadata (key list + shape/dtype) via + ``broadcast_object_list``, (2) broadcast each tensor via ``dist.broadcast``. + Source rank passes the populated dict; non-source ranks pass ``None``. + Returns a dict of tensors on ``device`` on every rank. + """ + is_src = torch.distributed.get_rank() == src + meta: list[Any] = ( + [{name: (tuple(t.shape), t.dtype) for name, t in state_dict_or_none.items()}] + if is_src and state_dict_or_none is not None + else [None] + ) + torch.distributed.broadcast_object_list(meta, src=src, group=pg) + meta_dict = meta[0] + assert meta_dict is not None, f"src rank {src} passed no state dict to broadcast" + + src_state_dict = state_dict_or_none or {} + out: dict[str, torch.Tensor] = {} + for name, (shape, dtype) in meta_dict.items(): + if is_src: + t = src_state_dict[name].to(device) + else: + t = torch.empty(shape, dtype=dtype, device=device) + torch.distributed.broadcast(t, src=src, group=pg) + out[name] = t + return out + + class DistributedProcessGroup: """A convenient wrapper around torch.distributed.ProcessGroup objects.""" diff --git a/modelopt/torch/utils/plugins/mbridge.py b/modelopt/torch/utils/plugins/mbridge.py index b6432ca5f51..d0e712e15a5 100644 --- a/modelopt/torch/utils/plugins/mbridge.py +++ b/modelopt/torch/utils/plugins/mbridge.py @@ -95,10 +95,9 @@ def load_mbridge_model_from_hf( assert hasattr(provider, key), f"{type(provider)} does not have attribute {key}" setattr(provider, key, value) - # Pruning does not support grouped GEMM yet, so disable it for MoE models. Set the flag on the - # provider (the bridge's native, possibly custom/hybrid spec reads it at build time) rather than - # replacing the whole layer spec -- overwriting it would drop custom layers (e.g. Qwen3.5's - # GatedDeltaNet + gated-attention or Gemma3's custom spec). + # Set moe_grouped_gemm on the provider (the bridge's native, possibly custom/hybrid spec reads + # it at build time) rather than replacing the whole layer spec -- overwriting it would drop + # custom layers (e.g. Qwen3.5's GatedDeltaNet + gated-attention or Gemma3's custom spec). if HAS_HYBRID and isinstance(provider, (HybridModelProvider)): provider.hybrid_stack_spec = get_te_hybrid_stack_spec(moe_grouped_gemm=moe_grouped_gemm) provider.moe_grouped_gemm = moe_grouped_gemm diff --git a/modelopt/torch/utils/plugins/megatron_generate.py b/modelopt/torch/utils/plugins/megatron_generate.py index d868f1ba65c..1fcf90c1d73 100644 --- a/modelopt/torch/utils/plugins/megatron_generate.py +++ b/modelopt/torch/utils/plugins/megatron_generate.py @@ -88,6 +88,25 @@ def cp_gather_logits(local_logits: torch.Tensor, cp_group, global_seq_len: int) return torch.cat(chunks, dim=1) +def _assert_mamba_within_int32_indexing(model: MegatronModule, batch_size: int, seq_length: int): + """Guard Mamba2 calibration against int32 activation-indexing overflow. + + The SSD Triton kernels index activations with int32, so ``batch * seq * mamba in_proj`` must stay + < 2**31 or they hit a cryptic CUDA 'illegal memory access'. Fail fast with the max safe batch. + """ + d = getattr(model, "_modelopt_max_mamba_in_proj", None) + if d is None: + d = max( + (m.in_proj.weight.shape[0] for m in model.modules() if hasattr(m, "in_proj")), default=0 + ) + model._modelopt_max_mamba_in_proj = d + if d and batch_size * seq_length * d >= 2**31: + raise ValueError( + f"{batch_size=} x {seq_length=} x mamba in_proj={d} overflows int32 in the Mamba2 kernels. " + f"Reduce calibration batch size to <= {2**31 // (seq_length * d)}." + ) + + def get_current_memory_info(): """Get current memory usage.""" remaining_mem, total_mem = torch.cuda.mem_get_info() @@ -157,6 +176,9 @@ def megatron_prefill( padded_seq_len = tokens.shape[-1] + # Fail fast with a clear message if the Mamba activation would overflow int32 kernel indexing. + _assert_mamba_within_int32_indexing(model, batch_size, padded_seq_len) + cp_size = mpu.get_context_parallel_world_size() # Under CP a local triu mask would be wrong for the per-rank zigzag chunks; pass None and let diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py new file mode 100644 index 00000000000..cd66567fa9a --- /dev/null +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -0,0 +1,425 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""HuggingFace-coupled FSDP2 model loading helpers.""" + +import json +import logging +import os +import re +from collections.abc import Callable +from itertools import chain +from typing import Any + +import torch +import torch.nn as nn +from accelerate import init_empty_weights +from huggingface_hub import snapshot_download +from safetensors import safe_open +from torch.distributed.checkpoint.state_dict import StateDictOptions, set_model_state_dict +from torch.distributed.tensor import DTensor +from transformers import AutoConfig, AutoModelForCausalLM + +try: + from transformers.conversion_mapping import get_model_conversion_mapping + from transformers.core_model_loading import WeightConverter, dot_natural_key, rename_source_key +except ImportError: # transformers<5 has no weight-conversion engine + get_model_conversion_mapping = rename_source_key = WeightConverter = dot_natural_key = None + +from modelopt.torch.utils.distributed import ( + barrier, + broadcast_state_dict, + fsdp2_wrap, + is_initialized, +) + +logger = logging.getLogger(__name__) + + +def read_safetensors_subset( + ckpt_path: str, + weight_map: dict, + select: Callable[[str], bool], +) -> dict: + """Read tensors whose name satisfies ``select`` from safetensors files. + + Groups param names by file to avoid re-opening. Returns CPU tensors. + Uses ``safe_open`` so only the requested tensors' bytes are read. + + ``get_tensor`` returns a zero-copy view into the mmap'd file; the bytes are + not actually read from disk until first touched. We ``clone()`` here to force + the read eagerly, while this function runs (each rank reading its own layers + in parallel). Without it the read is deferred to the later per-source + broadcast (``.to(device)``), which is serialized across ranks and silently + destroys the read parallelism this loader exists to provide. + """ + by_file: dict[str, list[str]] = {} + for name, file in weight_map.items(): + if select(name): + by_file.setdefault(file, []).append(name) + + state: dict[str, torch.Tensor] = {} + for file, names in by_file.items(): + with safe_open(os.path.join(ckpt_path, file), framework="pt", device="cpu") as f: + for name in names: + state[name] = f.get_tensor(name).clone() + return state + + +def weight_map_for(ckpt_path: str) -> dict[str, str]: + """Return the ``param_name → safetensors_file`` map for a local checkpoint directory.""" + index_path = os.path.join(ckpt_path, "model.safetensors.index.json") + single_file = os.path.join(ckpt_path, "model.safetensors") + if os.path.exists(index_path): + with open(index_path) as f: + return json.load(f)["weight_map"] + if os.path.exists(single_file): + with safe_open(single_file, framework="pt", device="cpu") as f: + return dict.fromkeys(f.keys(), "model.safetensors") + raise RuntimeError( + f"No safetensors checkpoint at {ckpt_path} " + "(expected model.safetensors or model.safetensors.index.json)." + ) + + +def _resolve_checkpoint_dir(ckpt_path: str, rank: int) -> str: + """Local dir for ``ckpt_path``; resolves an HF Hub ID (rank 0 downloads, others wait).""" + if os.path.isdir(ckpt_path): + return ckpt_path + if rank == 0: + snapshot_download(ckpt_path) + if is_initialized(): + barrier() + return snapshot_download(ckpt_path) + + +def _materialize_meta_model(model: nn.Module, device: torch.device) -> None: + """Replace meta params/buffers with empty real ones on ``device``; move real buffers there. + + Goes through ``model._apply`` so FSDP2's override refreshes its internal + ``_sharded_param_data`` pointers via ``reset_sharded_param``. + """ + model._apply(lambda t: torch.empty_like(t, device=device) if t.is_meta else t.to(device)) + + +def _promote_non_dtensor_to_gpu(model: nn.Module, device: torch.device) -> None: + """Move all non-DTensor params + buffers in ``model`` to ``device`` in-place. + + Used after CPU-offload loading: decoder DTensor shards stay on CPU (FSDP2 + streams them to GPU per layer), while root-level plain params and buffers + need to live on GPU so forwards work. + """ + for module in model.modules(): + for name, param in list(module._parameters.items()): + if param is None or isinstance(param, DTensor): + continue + module._parameters[name] = nn.Parameter( + param.data.to(device), requires_grad=param.requires_grad + ) + for name, buf in list(module._buffers.items()): + if buf is None or isinstance(buf, DTensor): + continue + module._buffers[name] = buf.to(device) + + +def _conversion_plan(model: nn.Module) -> dict | None: + """Transformers' own conversion mapping for ``model``, or ``None`` if nothing needs converting. + + ``legacy_renames`` (``_checkpoint_conversion_mapping``) covers transformers<5; on 5+ the + ``renamings``/``converters`` from HF's engine drive renaming + MoE weight fusion directly. + """ + legacy_renames = dict(getattr(model, "_checkpoint_conversion_mapping", None) or {}) + renamings, converters = [], [] + for entry in get_model_conversion_mapping(model) if get_model_conversion_mapping else []: + (converters if isinstance(entry, WeightConverter) else renamings).append(entry) + if not (legacy_renames or renamings or converters): + return None + return { + "legacy_renames": legacy_renames, + "renamings": renamings, + "converters": converters, + "prefix": model.base_model_prefix, + "meta_state_dict": model.state_dict(), + } + + +def _resolve_target(plan: dict, key: str) -> tuple[str, str | None]: + """Resolve a checkpoint key to ``(target param name, matched converter source pattern)``. + + No tensors are read. ``source_pattern`` is ``None`` for a plain (non-fused) key. + """ + for old, new in plan["legacy_renames"].items(): + key = re.sub(old, new, key) + if rename_source_key is None: # transformers<5: legacy renames only, no converters + return key, None + return rename_source_key( + key, plan["renamings"], plan["converters"], plan["prefix"], plan["meta_state_dict"] + ) + + +def _convert_keys(plan: dict, state: dict) -> dict: + """Rename 1:1 keys and fuse per-expert keys by driving transformers' own conversion ops.""" + if rename_source_key is None: # transformers<5: legacy renames only, no fusion + return {_resolve_target(plan, k)[0]: v for k, v in state.items()} + result: dict = {} + collected: dict = {} # target -> (converter, {source_pattern: [(sort_key, tensor)]}) + for key in sorted(state, key=dot_natural_key): + renamed, source_pattern = _resolve_target(plan, key) + if source_pattern is None: # plain rename, no fusion + result[renamed] = state[key] + continue + conv = next(c for c in plan["converters"] if source_pattern in c.source_patterns) + collected.setdefault(renamed, (conv, {}))[1].setdefault(source_pattern, []).append( + (dot_natural_key(key), state[key]) + ) + for target, (conv, by_src) in collected.items(): + tensors = { + sp: [t for _, t in sorted(lst, key=lambda kt: kt[0])] for sp, lst in by_src.items() + } + for op in conv.operations: # transformers runs the fusion math (any op, no whitelist) + tensors = op.convert( + tensors, source_patterns=conv.source_patterns, target_patterns=conv.target_patterns + ) + if len(tensors) != 1: + raise NotImplementedError( + f"Only many-to-one conversions supported; got {list(tensors)}" + ) + result[target] = next(iter(tensors.values())) + return result + + +def build_meta_causal_lm( + ckpt_path: str, + trust_remote_code: bool, + attn_implementation: str | None, + hf_config=None, +): + """Build a meta-init causal LM (no real storage allocated).""" + if hf_config is None: + config_kwargs: dict[str, Any] = {"trust_remote_code": trust_remote_code} + if attn_implementation is not None: + config_kwargs["attn_implementation"] = attn_implementation + hf_config = AutoConfig.from_pretrained(ckpt_path, **config_kwargs) + elif attn_implementation is not None: + # Honor the override even when the caller passed in a pre-fetched config. + hf_config._attn_implementation = attn_implementation + dtype = getattr(hf_config, "torch_dtype", None) or torch.bfloat16 + with init_empty_weights(include_buffers=False): + model = AutoModelForCausalLM.from_config( + hf_config, torch_dtype=dtype, trust_remote_code=trust_remote_code + ) + model.eval() + return model + + +def _layers_for_rank(n_layers: int, world_size: int, r: int) -> list[int]: + return [i for i in range(n_layers) if i % world_size == r] + + +def _read_and_convert( + resolved_path: str, weight_map: dict, keyset: set[str], plan: dict | None +) -> dict: + raw = read_safetensors_subset(resolved_path, weight_map, lambda k: k in keyset) + return _convert_keys(plan, raw) if plan else raw + + +# One decoder layer's converted weights (param-name suffix -> tensor); the outer dict is keyed +# by decoder-layer index. +LayerStateDict = dict[str, torch.Tensor] +OwnedLayerStateDicts = dict[int, LayerStateDict] + + +def _read_owned_layers( + resolved_path: str, + weight_map: dict, + layer_sources: dict, + owned_layer_indices: list[int], + plan: dict | None, +) -> OwnedLayerStateDicts: + """Read + convert this rank's owned decoder layers from disk (ranks read in parallel).""" + return { + layer_idx: _read_and_convert(resolved_path, weight_map, set(layer_sources[layer_idx]), plan) + for layer_idx in owned_layer_indices + } + + +def _broadcast_load_group( + layer_indices: list[int], + source_rank: int, + current_rank: int, + owned_layer_state_dicts: OwnedLayerStateDicts, + decoder_layers: list[nn.Module], + layer_prefixes: list[str], + device: torch.device, + cpu_offload: bool, +) -> None: + """Broadcast ``layer_indices`` from ``source_rank`` to all ranks and reshard into FSDP2 shards. + + The owner assembles the group's full tensors; every rank receives them, reshards its local slice, + then frees the full copy (capping the transient GPU peak). The owner drops its read copy after. + """ + group_state_dict: dict | None = None + if current_rank == source_rank: + group_state_dict = {} + for layer_idx in layer_indices: + group_state_dict.update(owned_layer_state_dicts[layer_idx]) + broadcasted_state_dict = broadcast_state_dict(group_state_dict, src=source_rank, device=device) + for layer_idx in layer_indices: + prefix = layer_prefixes[layer_idx] + layer_state_dict = { + k[len(prefix) :]: v for k, v in broadcasted_state_dict.items() if k.startswith(prefix) + } + if cpu_offload: + layer_state_dict = {k: v.cpu() for k, v in layer_state_dict.items()} + set_model_state_dict( + decoder_layers[layer_idx], + layer_state_dict, + options=StateDictOptions(full_state_dict=True, broadcast_from_rank0=False), + ) + del layer_state_dict + del broadcasted_state_dict + if current_rank == source_rank: + for layer_idx in layer_indices: + del owned_layer_state_dicts[layer_idx] + + +def _group_sources_by_layer( + weight_map: dict, plan: dict | None, model_param_names: set[str], layer_prefixes: list[str] +) -> tuple[dict[int, list[str]], list[str], int]: + """Bucket checkpoint keys by the decoder layer their converted target lives in. + + Returns ``(layer_sources, non_layer_sources, skipped)``: ``layer_sources[i]`` holds the keys + targeting decoder layer ``i``, ``non_layer_sources`` holds root (embed/lm_head/norm) keys, and + ``skipped`` counts keys whose target isn't in the model (aux weights, e.g. an MTP head). + """ + layer_sources: dict[int, list[str]] = {i: [] for i in range(len(layer_prefixes))} + non_layer_sources: list[str] = [] + skipped = 0 + for ckpt_key in weight_map: + target = _resolve_target(plan, ckpt_key)[0] if plan else ckpt_key + if target not in model_param_names: + skipped += 1 + continue + for i, prefix in enumerate(layer_prefixes): + if target.startswith(prefix): + layer_sources[i].append(ckpt_key) + break + else: + non_layer_sources.append(ckpt_key) + return layer_sources, non_layer_sources, skipped + + +def parallel_load_and_prepare_fsdp2( + ckpt_path: str, + device: torch.device, + rank: int, + world_size: int, + trust_remote_code: bool = False, + mp_policy=None, + cpu_offload: bool = False, + attn_implementation: str | None = None, + hf_config=None, + broadcast_chunk_size: int | None = 8, +) -> nn.Module: + """Load and FSDP2-shard a HuggingFace causal LM via parallel safetensors reads. + + Round-robin assigns decoder layers to ranks; each rank reads only its owned + layers' weights from disk in parallel, then broadcasts to the others. Non-decoder + weights (embed, lm_head, norm) are read on rank 0 and broadcast. + + Requires an initialized ``torch.distributed`` process group (FSDP2's ``fully_shard`` + and the per-layer broadcasts both need it). A 1-rank PG (e.g. ``torchrun + --nproc_per_node=1``) is allowed; bare single-process is not. + + Pass ``hf_config`` if the caller has already fetched it (skips a redundant fetch). + + ``broadcast_chunk_size`` sets how many of a source's owned layers are broadcast per collective: + a smaller value lowers the peak transient GPU memory at the cost of more collectives (default 8; + pass ``None`` to broadcast all of a source's layers at once). + """ + resolved_path = _resolve_checkpoint_dir(ckpt_path, rank) + weight_map = weight_map_for(resolved_path) + + model = build_meta_causal_lm(resolved_path, trust_remote_code, attn_implementation, hf_config) + + # fsdp2_wrap shards each decoder layer + the root (embed/lm_head/norm sharded, not replicated). + decoder_layers = fsdp2_wrap(model, mp_policy=mp_policy, cpu_offload=cpu_offload) + module_to_name = {m: n for n, m in model.named_modules()} + layer_prefixes = [module_to_name[layer] + "." for layer in decoder_layers] + + # transformers>=5 fuses/renames checkpoint keys so they no longer match param names 1:1 + # (None => the pre-5.x identity path). + plan = _conversion_plan(model) + + # Valid targets; keys converting to anything else are aux weights (e.g. an MTP head) we skip. + model_param_names = {n for n, _ in chain(model.named_parameters(), model.named_buffers())} + + # Bucket each checkpoint key by its target's decoder layer (root params go to non_layer_sources). + layer_sources, non_layer_sources, skipped = _group_sources_by_layer( + weight_map, plan, model_param_names, layer_prefixes + ) + if skipped: + logger.debug( + "skipping %d checkpoint keys not present in the model (e.g. MTP head)", skipped + ) + + _materialize_meta_model(model, torch.device("cpu") if cpu_offload else device) + + owned_layer_indices = _layers_for_rank(len(decoder_layers), world_size, rank) + owned_layer_state_dicts = _read_owned_layers( + resolved_path, weight_map, layer_sources, owned_layer_indices, plan + ) + + # Smaller broadcast_chunk_size lowers the transient GPU peak (more, smaller collectives). + for source_rank in range(world_size): + source_layer_indices = _layers_for_rank(len(decoder_layers), world_size, source_rank) + if not source_layer_indices: + continue + chunk = broadcast_chunk_size or len(source_layer_indices) + for start in range(0, len(source_layer_indices), chunk): + _broadcast_load_group( + source_layer_indices[start : start + chunk], + source_rank, + rank, + owned_layer_state_dicts, + decoder_layers, + layer_prefixes, + device, + cpu_offload, + ) + + # Non-decoder params: rank 0 reads + broadcasts; resharded into the root below. + # TODO: layerwise support. + non_layer = None + if rank == 0: + non_layer = _read_and_convert(resolved_path, weight_map, set(non_layer_sources), plan) + non_layer = broadcast_state_dict(non_layer, src=0, device=device) + if cpu_offload: + non_layer = {k: v.cpu() for k, v in non_layer.items()} + # shard_root=True makes the root params sharded DTensors, so reshard the full tensors via + # set_model_state_dict. strict=False: decoder keys are absent here (loaded above). + set_model_state_dict( + model, + non_layer, + options=StateDictOptions(full_state_dict=True, broadcast_from_rank0=False, strict=False), + ) + + if cpu_offload: + # Loaded on CPU for set_model_state_dict; FSDP2 streams decoder shards per forward, but + # the unwrapped root must live on GPU, so promote it. + _promote_non_dtensor_to_gpu(model, device) + if hasattr(model, "tie_weights"): + model.tie_weights() + return model diff --git a/pyproject.toml b/pyproject.toml index ccec4e1b82b..3abfc21fb22 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -69,6 +69,7 @@ onnx = [ "onnxruntime~=1.24.2; python_version > '3.10' and (platform_machine == 'aarch64' or platform_system == 'Darwin')", "onnxruntime-gpu~=1.24.2; python_version > '3.10' and platform_machine != 'aarch64' and platform_system != 'Darwin' and platform_system != 'Windows'", "onnxscript", + "onnxsim>=0.7.0", "onnxslim>=0.1.76", "polygraphy>=0.49.22", ] diff --git a/tests/examples/megatron_bridge/test_prune_minitron.py b/tests/examples/megatron_bridge/test_prune_minitron.py index f21f25d5209..f468e8abcba 100644 --- a/tests/examples/megatron_bridge/test_prune_minitron.py +++ b/tests/examples/megatron_bridge/test_prune_minitron.py @@ -60,8 +60,9 @@ def test_prune_minitron(tmp_path, num_gpus, create_teacher, megatron_format): 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"], + ["torchrun", f"--nproc_per_node={num_gpus}", "prune_minitron.py", "--no_moe_grouped_gemm"], hf_model_name_or_path=teacher_hf_path, pp_size=num_gpus, calib_dataset_name="cnn_dailymail", @@ -126,8 +127,9 @@ def test_prune_minitron_vlm(tmp_path, num_gpus, create_teacher): prune_target_params = int(language_model_params * 0.7) pruned_model_path = tmp_path / "pruned" + # 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"], + ["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_model_path, pp_size=num_gpus, diff --git a/tests/gpu/torch/quantization/test_fsdp2.py b/tests/gpu/torch/quantization/test_fsdp2.py index 1ad88d087d1..f9fec0d2a4c 100644 --- a/tests/gpu/torch/quantization/test_fsdp2.py +++ b/tests/gpu/torch/quantization/test_fsdp2.py @@ -331,3 +331,136 @@ def _test_persistent_materialization(rank, size): def test_persistent_materialization(dist_workers): dist_workers.run(_test_persistent_materialization) + + +def _test_writeback_root_unwrapped(rank, size): + """Writeback works when only the decoder layers are FSDP2-wrapped and the root is unsharded. + + The root is only the search boundary: ``enable_weight_access_and_writeback(layer, model)`` + walks ``layer[0]``'s ancestors to the sharded decoder layer and gathers/writes back its + DTensor, so the root needs no FSDP state. Covers the ``shard_root=False`` / nested-FSDP case + (``fsdp2_wrap`` now defaults to ``shard_root=True``, wrapping the root too). Regression guard + for the old ``isinstance(root_model, FSDPModule)`` assert that wrongly required a wrapped root. + """ + from modelopt.torch.quantization.utils import enable_weight_access_and_writeback + + dim = 32 + torch.manual_seed(1) + # Root is a plain container; model[0] stands in for a decoder layer. + model = nn.Sequential(nn.Sequential(nn.Linear(dim, dim), nn.Linear(dim, dim))).cuda(rank) + synchronize_state_dict(model) + + # Wrap ONLY the "decoder layer" -- intentionally NO ``fully_shard(model)`` on the root, + # mirroring fsdp2_wrap. ``root_model`` (model) is therefore not an FSDPModule. + fully_shard(model[0]) + layer = model[0] + inputs = torch.randn(2, dim).cuda(rank) + + # Warmup forward to trigger FSDP2's lazy_init (mirrors layerwise calibration). + model(inputs) + + # This is the exact call save()/full_restore() make. Before the fix it tripped the + # ``assert isinstance(root_model, FSDPModule)`` because the root is unwrapped — that's + # the regression we guard. The DTensor-shape checks are not portable across torch + # versions when the root is not FSDP-wrapped, so we just verify the writeback path + # runs and mutations persist. + with enable_weight_access_and_writeback(layer[0], model): + ref_weight = layer[0].weight.clone() + layer[0].weight.data.add_(1.0) # mutate -> exercises the writeback path + + # Modification was written back into the shards. + with enable_weight_access_and_writeback(layer[0], model): + assert torch.allclose(layer[0].weight, ref_weight + 1.0) + + +def test_writeback_root_unwrapped(dist_workers): + dist_workers.run(_test_writeback_root_unwrapped) + + +def _test_writeback_cpu_offload(rank, size): + """Writeback round-trip when the FSDP2 shard is CPU-resident (``CPUOffloadPolicy``). + + Regression guard for the CPU↔GPU mirror added to + ``fsdp2_weight_access_and_writeback_context``: the gathered shard is on CPU, + so the helper mirrors it to GPU for in-context mutation and must copy + modifications back to the CPU shard on exit. + """ + from torch.distributed.fsdp import CPUOffloadPolicy + + from modelopt.torch.quantization.utils import enable_weight_access_and_writeback + + dim = 32 + torch.manual_seed(1) + model = nn.Sequential(nn.Sequential(nn.Linear(dim, dim), nn.Linear(dim, dim))).cuda(rank) + synchronize_state_dict(model) + + # Wrap the "decoder layer" with cpu_offload; root stays unwrapped. + fully_shard(model[0], offload_policy=CPUOffloadPolicy()) + layer = model[0] + + # Warmup forward triggers FSDP2's lazy_init. + model(torch.randn(2, dim).cuda(rank)) + + # Regression guard for the CPU→GPU mirror in fsdp2_weight_access_and_writeback_context: + # if the helper handed back a CPU tensor under cpu_offload, calibration ops would crash + # on the in-context mutation below (GPU activations vs CPU weight). The fact that this + # block runs and the mutation persists is the evidence the mirror trip worked. + with enable_weight_access_and_writeback(layer[0], model): + ref_weight = layer[0].weight.clone() + layer[0].weight.data.add_(1.0) + + # Mutation written back to the CPU shard. + with enable_weight_access_and_writeback(layer[0], model): + assert torch.allclose(layer[0].weight, ref_weight + 1.0) + + +def test_writeback_cpu_offload(dist_workers): + dist_workers.run(_test_writeback_cpu_offload) + + +class _EmbedRootModel(nn.Module): + """Root owns embed/norm params plus a decoder block. Mirrors the sharded-root layout + where ``model(**batch)`` must fire the root's FSDP2 hook to unshard embed for the forward.""" + + def __init__(self, vocab=16, dim=32): + super().__init__() + self.embed = nn.Embedding(vocab, dim) + self.block = _DecoderBlock(dim) + self.norm = nn.LayerNorm(dim) + + def forward(self, input_ids=None, **kwargs): + return self.norm(self.block(self.embed(input_ids))) + + +def _test_sharded_root_calibration(rank, size): + """Calibration through the standard forward loop works with a *sharded* FSDP2 root. + + Regression guard for removing ``materialize_fsdp2_root``: ``_forward_loop`` now calls + ``model(**batch)`` (not ``model.forward``), so the root's FSDP2 pre/post-forward hooks + unshard embed/norm for the forward and reshard them after — no manual materialization. + With the old ``model.forward`` bypass this hit ``aten.embedding: mixed Tensor and DTensor``. + """ + from modelopt.torch.utils.dataset_utils import _forward_loop + + dim = 32 + torch.manual_seed(1) + model = _EmbedRootModel(dim=dim).cuda(rank) + synchronize_state_dict(model) + + # Shard the decoder block AND the root -> the root's own params (embed/norm) are sharded DTensors. + fully_shard(model.block) + model = fully_shard(model) + assert isinstance(model.embed.weight, DTensor) + + batches = [{"input_ids": torch.randint(0, 16, (2, 8), device=rank)} for _ in range(2)] + mtq.quantize(model, mtq.INT8_DEFAULT_CFG, lambda m: _forward_loop(m, batches)) + + # Root params are resharded after calibration (needed for export / get_model_state_dict), + # and the model still runs. + assert isinstance(model.embed.weight, DTensor) + assert isinstance(model.norm.weight, DTensor) + model(input_ids=batches[0]["input_ids"]) + + +def test_sharded_root_calibration(dist_workers): + dist_workers.run(_test_sharded_root_calibration) diff --git a/tests/gpu/torch/utils/test_model_load_utils.py b/tests/gpu/torch/utils/test_model_load_utils.py new file mode 100644 index 00000000000..8371f40d2d7 --- /dev/null +++ b/tests/gpu/torch/utils/test_model_load_utils.py @@ -0,0 +1,129 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""GPU/distributed tests for the FSDP2 load path and its helpers.""" + +import json +import os +import tempfile +from functools import partial + +import pytest +import torch +import torch.distributed as dist +from torch.distributed.tensor import DTensor + + +def _test_broadcast_state_dict_roundtrip(rank, size): + """Round-trip from every rank as source (matches the per-layer rotation in the loader).""" + from modelopt.torch.utils.distributed import broadcast_state_dict + + device = torch.device(f"cuda:{rank}") + # Distinct payload per source rank so a wrong-src result would fail content checks. + for source in range(size): + src_dict = { + "w": torch.full((2, 4), float(source)), + "b": torch.tensor([float(source), float(source) + 1.0]), + } + out = broadcast_state_dict(src_dict if rank == source else None, src=source, device=device) + assert set(out.keys()) == {"w", "b"} + assert out["w"].device == device + assert torch.equal(out["w"].cpu(), src_dict["w"]) + assert torch.equal(out["b"].cpu(), src_dict["b"]) + + +def test_broadcast_state_dict_roundtrip(dist_workers): + dist_workers.run(_test_broadcast_state_dict_roundtrip) + + +def _build_tiny_llama_checkpoint(path: str) -> None: + """Write a tiny LlamaForCausalLM checkpoint (config + safetensors) to ``path``.""" + from transformers import LlamaConfig, LlamaForCausalLM + + config = LlamaConfig( + vocab_size=64, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + max_position_embeddings=32, + torch_dtype="bfloat16", + ) + model = LlamaForCausalLM(config).to(torch.bfloat16) + model.save_pretrained(path) + + +def _test_parallel_load_and_export(rank, size, cpu_offload): + """Load a tiny Llama via the FSDP2 loader, forward, then export — config.architectures preserved. + + Parametrized over ``cpu_offload`` to cover both shard placements: + - off: decoder DTensor shards on GPU, plain root on GPU. + - on: decoder DTensor shards on CPU (streamed per layer), root promoted to GPU + via ``_promote_non_dtensor_to_gpu``. + """ + from modelopt.torch.export.unified_export_hf import export_hf_checkpoint + from modelopt.torch.utils.plugins.model_load_utils import parallel_load_and_prepare_fsdp2 + + suffix = "offload" if cpu_offload else "noffload" + ckpt_dir = os.path.join(tempfile.gettempdir(), f"_test_parallel_load_{suffix}_{os.getpid()}") + if rank == 0: + os.makedirs(ckpt_dir, exist_ok=True) + _build_tiny_llama_checkpoint(ckpt_dir) + dist.barrier() + + device = torch.device(f"cuda:{rank}") + model = parallel_load_and_prepare_fsdp2( + ckpt_dir, + device, + rank, + size, + cpu_offload=cpu_offload, + ) + + # Decoder layers AND root params (embed/lm_head) are sharded DTensors under shard_root=True. + decoder_params = list(model.model.layers[0].parameters()) + assert any(isinstance(p, DTensor) for p in decoder_params) + assert isinstance(model.model.embed_tokens.weight, DTensor) + if not cpu_offload: + # Non-offload: the root's local shard lives on GPU. + assert model.model.embed_tokens.weight.to_local().device.type == "cuda" + if cpu_offload: + # Under cpu_offload the decoder shards live on CPU between forwards. + decoder_dtensors = [p for p in decoder_params if isinstance(p, DTensor)] + assert all(p.to_local().device.type == "cpu" for p in decoder_dtensors) + + # Forward exercises FSDP2 hooks + (under cpu_offload) the per-layer CPU↔GPU stream. + input_ids = torch.randint(0, 64, (1, 8), device=device) + out = model(input_ids=input_ids).logits + assert out.shape == (1, 8, 64) + + # Export and verify the saved config.json retains the original architectures. + export_dir = os.path.join( + tempfile.gettempdir(), f"_test_parallel_export_{suffix}_{os.getpid()}" + ) + if rank == 0: + os.makedirs(export_dir, exist_ok=True) + dist.barrier() + export_hf_checkpoint(model, export_dir=export_dir, dtype=torch.bfloat16) + + if rank == 0: + with open(os.path.join(export_dir, "config.json")) as f: + cfg = json.load(f) + assert cfg["architectures"] == ["LlamaForCausalLM"] + + +@pytest.mark.parametrize("cpu_offload", [False, True]) +def test_parallel_load_and_export(dist_workers, cpu_offload): + dist_workers.run(partial(_test_parallel_load_and_export, cpu_offload=cpu_offload)) diff --git a/tests/gpu_megatron/torch/nas/plugins/test_megatron_gpt_dynamic_modules.py b/tests/gpu_megatron/torch/nas/plugins/test_megatron_gpt_dynamic_modules.py index 158b6cafacd..5df4c2fa79e 100644 --- a/tests/gpu_megatron/torch/nas/plugins/test_megatron_gpt_dynamic_modules.py +++ b/tests/gpu_megatron/torch/nas/plugins/test_megatron_gpt_dynamic_modules.py @@ -36,6 +36,8 @@ _DynamicMoELayer, _DynamicSelfAttention, _DynamicSequentialMLP, + _DynamicTEGroupedLinear, + _DynamicTEGroupedMLP, _DynamicTELayerNormColumnParallelLinear, _DynamicTEProjRowParallelLinear, _DynamicTEQKVLayerNormColumnParallelLinear, @@ -231,7 +233,7 @@ def test_gpt_self_attention_head_sorting(distributed_setup_size_1): destroy_model_parallel() -def _test_gpt_moe_search_space(rank, size): +def _test_gpt_moe_search_space(moe_grouped_gemm, rank, size): channel_divisor = 4 num_layers = min(size * 2, 8) @@ -258,6 +260,7 @@ def _test_gpt_moe_search_space(rank, size): activation_func="squared_relu", transformer_impl="transformer_engine", num_moe_experts=num_moe_experts, + moe_grouped_gemm=moe_grouped_gemm, moe_ffn_hidden_size=moe_ffn_hidden_size, moe_shared_expert_intermediate_size=moe_shared_expert_intermediate_size, ).cuda() @@ -280,11 +283,16 @@ def _test_gpt_moe_search_space(rank, size): moe = model.decoder.layers[0].mlp assert isinstance(moe, _DynamicMoELayer) assert isinstance(moe.router, _DynamicTopKRouter) - assert isinstance(moe.experts, _DynamicSequentialMLP) - assert isinstance(moe.experts.local_experts, DynamicModuleList) - for expert in moe.experts.local_experts: - assert isinstance(expert, _DynamicMLP) assert isinstance(moe.shared_experts, _DynamicMLP) + if moe_grouped_gemm: + assert isinstance(moe.experts, _DynamicTEGroupedMLP) + assert isinstance(moe.experts.linear_fc1, _DynamicTEGroupedLinear) + assert isinstance(moe.experts.linear_fc2, _DynamicTEGroupedLinear) + else: + assert isinstance(moe.experts, _DynamicSequentialMLP) + assert isinstance(moe.experts.local_experts, DynamicModuleList) + for expert in moe.experts.local_experts: + assert isinstance(expert, _DynamicMLP) # NOTE: `search_space_size` does not reduce across TP/PP groups ss_size_per_pp = search_space_size(model) @@ -293,15 +301,12 @@ def _test_gpt_moe_search_space(rank, size): moe_shared_ffn_choices = moe_shared_expert_intermediate_size // channel_divisor hidden_size_choices = hidden_size // channel_divisor num_layers_per_pp = num_layers // size - # SequentialMLP has per-expert moe_ffn_hidden_size hparams + # SequentialMLP has one moe_ffn_hidden_size hparam per expert (moe_ffn_choices**num_moe_experts); + # TEGroupedMLP shares a single one (moe_ffn_choices). + moe_ffn_ss = moe_ffn_choices if moe_grouped_gemm else moe_ffn_choices**num_moe_experts assert ( ss_size_per_pp - == ( - num_heads_choices - * num_moe_experts - * moe_ffn_choices**num_moe_experts - * moe_shared_ffn_choices - ) + == (num_heads_choices * num_moe_experts * moe_ffn_ss * moe_shared_ffn_choices) ** num_layers_per_pp * num_layers * hidden_size_choices @@ -319,5 +324,6 @@ def _test_gpt_moe_search_space(rank, size): assert not any(named_dynamic_modules(model)) -def test_gpt_moe_search_space(dist_workers): - dist_workers.run(_test_gpt_moe_search_space) +@pytest.mark.parametrize("moe_grouped_gemm", [False, True]) +def test_gpt_moe_search_space(dist_workers, moe_grouped_gemm): + dist_workers.run(partial(_test_gpt_moe_search_space, moe_grouped_gemm)) diff --git a/tests/gpu_megatron/torch/prune/plugins/test_mcore_gpt_minitron_pruning.py b/tests/gpu_megatron/torch/prune/plugins/test_mcore_gpt_minitron_pruning.py index 3b3425f6717..348f6799ca8 100644 --- a/tests/gpu_megatron/torch/prune/plugins/test_mcore_gpt_minitron_pruning.py +++ b/tests/gpu_megatron/torch/prune/plugins/test_mcore_gpt_minitron_pruning.py @@ -356,7 +356,7 @@ def test_mcore_gpt_pruning( ) -def _test_mcore_gpt_moe_parameter_sorting(rank, size): +def _test_mcore_gpt_moe_parameter_sorting(moe_grouped_gemm, rank, size): set_seed(SEED) # Use relatively bigger model here for more accurate test for sorting channel_divisor = 64 @@ -385,6 +385,7 @@ def _test_mcore_gpt_moe_parameter_sorting(rank, size): activation_func="squared_relu", transformer_impl="transformer_engine", num_moe_experts=num_moe_experts, + moe_grouped_gemm=moe_grouped_gemm, moe_ffn_hidden_size=moe_ffn_hidden_size, moe_shared_expert_intermediate_size=moe_shared_expert_intermediate_size, bf16=False, @@ -408,9 +409,11 @@ def _test_mcore_gpt_moe_parameter_sorting(rank, size): sortable_per_pp = [ n for n, hp in dynamic_space.named_hparams(configurable=True) if hp.importance is not None ] - # (num_moe_experts + 3) hps per layer + 1 for hidden_size (num_layers is not sorted!) - # Per layer: num_attention_heads, num_moe_experts, moe_ffn (per expert), moe_shared_ffn - assert len(sortable_per_pp) == (num_moe_experts + 3) * num_layers // size + 1 + # (moe_ffn_count + 3) hps per layer + 1 for hidden_size (num_layers is not sorted!) + # Per layer: num_attention_heads, num_moe_experts, moe_ffn, moe_shared_ffn. + # SequentialMLP registers one moe_ffn hparam per expert; TEGroupedMLP shares a single one. + moe_ffn_count = 1 if moe_grouped_gemm else num_moe_experts + assert len(sortable_per_pp) == (moe_ffn_count + 3) * num_layers // size + 1 # sanity check if the model functionality is preserved after sorting export_searchspace(model, mtn.get_subnet_config(model)) @@ -418,11 +421,12 @@ def _test_mcore_gpt_moe_parameter_sorting(rank, size): compare_outputs(y1, y2, rtol=1e-5, atol=1e-3) -def test_mcore_gpt_moe_parameter_sorting(dist_workers): - dist_workers.run(_test_mcore_gpt_moe_parameter_sorting) +@pytest.mark.parametrize("moe_grouped_gemm", [False, True]) +def test_mcore_gpt_moe_parameter_sorting(dist_workers, moe_grouped_gemm): + dist_workers.run(partial(_test_mcore_gpt_moe_parameter_sorting, moe_grouped_gemm)) -def _test_mcore_gpt_pruning_moe(ckpt_dir, rank, size): +def _test_mcore_gpt_pruning_moe(ckpt_dir, moe_grouped_gemm, rank, size): channel_divisor = 4 num_layers = size @@ -446,6 +450,7 @@ def _get_model(initialize_megatron=True): activation_func="squared_relu", transformer_impl="transformer_engine", num_moe_experts=num_moe_experts, + moe_grouped_gemm=moe_grouped_gemm, moe_ffn_hidden_size=moe_ffn_hidden_size, moe_shared_expert_intermediate_size=moe_shared_expert_intermediate_size, ).cuda() @@ -483,10 +488,24 @@ def _get_model(initialize_megatron=True): assert moe.router.expert_bias.shape == (pruned_num_moe_experts,) assert moe.router.weight.shape == (pruned_num_moe_experts, pruned_hidden_size) assert moe.experts.num_local_experts == pruned_num_moe_experts - assert len(moe.experts.local_experts) == pruned_num_moe_experts - for expert in moe.experts.local_experts: - assert expert.linear_fc1.weight.shape == (pruned_moe_ffn, pruned_hidden_size) - assert expert.linear_fc2.weight.shape == (pruned_hidden_size, pruned_moe_ffn) + if moe_grouped_gemm: + # TEGroupedMLP fuses experts into two grouped linears with per-expert weight{i} params + assert moe.experts.linear_fc1.num_gemms == pruned_num_moe_experts + assert moe.experts.linear_fc2.num_gemms == pruned_num_moe_experts + for i in range(pruned_num_moe_experts): + assert getattr(moe.experts.linear_fc1, f"weight{i}").shape == ( + pruned_moe_ffn, + pruned_hidden_size, + ) + assert getattr(moe.experts.linear_fc2, f"weight{i}").shape == ( + pruned_hidden_size, + pruned_moe_ffn, + ) + else: + assert len(moe.experts.local_experts) == pruned_num_moe_experts + for expert in moe.experts.local_experts: + assert expert.linear_fc1.weight.shape == (pruned_moe_ffn, pruned_hidden_size) + assert expert.linear_fc2.weight.shape == (pruned_hidden_size, pruned_moe_ffn) assert moe.shared_experts.linear_fc1.weight.shape == ( pruned_moe_shared_ffn, pruned_hidden_size, @@ -519,8 +538,11 @@ def _get_model(initialize_megatron=True): ) -def test_mcore_gpt_pruning_moe(dist_workers, tmp_path): - dist_workers.run(partial(_test_mcore_gpt_pruning_moe, tmp_path / "minitron_scores")) +@pytest.mark.parametrize("moe_grouped_gemm", [False, True]) +def test_mcore_gpt_pruning_moe(dist_workers, tmp_path, moe_grouped_gemm): + dist_workers.run( + partial(_test_mcore_gpt_pruning_moe, tmp_path / "minitron_scores", moe_grouped_gemm) + ) def _build_and_prune_variant(size, export_config, *, num_attention_heads=4, **model_kwargs): 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 a080ab66bd8..69d8c7c31ec 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 @@ -258,13 +258,14 @@ def test_mcore_mamba_hybrid_pruning(dist_workers, tmp_path): } -def _make_nas_hybrid_model(size): +def _make_nas_hybrid_model(size, moe_grouped_gemm=False): return get_mcore_mamba_hybrid_model( tensor_model_parallel_size=1, pipeline_model_parallel_size=size, initialize_megatron=True, transformer_impl="transformer_engine", bf16=False, + moe_grouped_gemm=moe_grouped_gemm, **_NAS_MODEL_KWARGS, ).cuda() @@ -335,7 +336,8 @@ def _assert_top_k_candidates(searcher_state, constraint_key, expected_top_k, k=1 def _test_mcore_mamba_hybrid_pruning_nas_params(rank, size, ckpt_dir): set_seed(SEED) - model = _make_nas_hybrid_model(size) + # Covers grouped-GEMM (TEGroupedMLP) MoE pruning; the memory_mb test below covers SequentialMLP. + model = _make_nas_hybrid_model(size, moe_grouped_gemm=True) baseline_params, baseline_active = mcore_param_count( model.config, @@ -430,7 +432,8 @@ def _test_mcore_mamba_hybrid_pruning_nas_memory_mb(rank, size, ckpt_dir): set_seed(SEED) dtype_bytes = 2 sequence_length = 128 - model = _make_nas_hybrid_model(size) + # Covers SequentialMLP MoE pruning; the params test above covers grouped-GEMM (TEGroupedMLP). + model = _make_nas_hybrid_model(size, moe_grouped_gemm=False) _, _, _, baseline_memory_mb = mcore_memory_footprint_mb( model.config, diff --git a/tests/unit/onnx/quantization/test_quantize_api.py b/tests/unit/onnx/quantization/test_quantize_api.py index f350d5d89f4..e7c595d23a8 100644 --- a/tests/unit/onnx/quantization/test_quantize_api.py +++ b/tests/unit/onnx/quantization/test_quantize_api.py @@ -89,6 +89,7 @@ def fake_preprocess( trt_plugins_precision, override_shapes, simplify, + simplify_backend, quantize_mode, opset, ): diff --git a/tests/unit/torch/speculative/plugins/test_hf_streaming_dataset.py b/tests/unit/torch/speculative/plugins/test_hf_streaming_dataset.py index 0cb581f0296..efd77267b3c 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_streaming_dataset.py +++ b/tests/unit/torch/speculative/plugins/test_hf_streaming_dataset.py @@ -484,3 +484,62 @@ def handler(request: httpx.Request) -> httpx.Response: ) ds[0] assert seen_ports == {advertised_port} + + +# --------------------------------------------------------------------------- +# answer_only_loss template guard +# --------------------------------------------------------------------------- + + +def _fast_tokenizer_with_template(template: str, seq: int = 8) -> MagicMock: + """Fast-tokenizer mock with a given chat template; returns ids + assistant_masks.""" + tok = MagicMock() + tok.is_fast = True + tok.chat_template = template + tok.apply_chat_template.return_value = { + "input_ids": torch.arange(seq, dtype=torch.long).unsqueeze(0), + "assistant_masks": torch.ones(1, seq, dtype=torch.long), + } + return tok + + +_CONV = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}] + + +def test_answer_only_loss_rejects_template_without_generation_tags(): + """A fast tokenizer whose template lacks {% generation %} tags fails loudly. + + Without the guard transformers only warns and returns an ALL-ZERO assistant + mask -- every sample then trains at zero loss with no other symptom. + """ + tok = _fast_tokenizer_with_template("{{ bos }}{% if add_generation_prompt %}x{% endif %}") + with pytest.raises(RuntimeError, match="generation"): + hf_streaming_dataset._tokenize_with_loss_mask(tok, _CONV, answer_only_loss=True) + + +def test_answer_only_loss_accepts_tagged_template(): + """Templates carrying {% generation %} (either whitespace-control form) pass.""" + for tag in ("{% generation %}", "{%- generation -%}"): + tok = _fast_tokenizer_with_template("{{ bos }}" + tag + "{{ c }}") + ids, mask = hf_streaming_dataset._tokenize_with_loss_mask(tok, _CONV, answer_only_loss=True) + assert mask.sum() == ids.shape[-1] + + +def test_full_loss_skips_template_guard(): + """answer_only_loss=False never consults the template (mask is all ones).""" + tok = _fast_tokenizer_with_template("{{ bos }}") + ids, mask = hf_streaming_dataset._tokenize_with_loss_mask(tok, _CONV, answer_only_loss=False) + assert mask.sum() == ids.shape[-1] + + +def test_answer_only_loss_rejects_slow_tokenizer_without_recovery(): + """A slow tokenizer with no registered recovery fails loudly even on a tagged template. + + Assistant-mask alignment needs the fast tokenizer's char_to_token; without the + guard apply_chat_template fails downstream with an unrelated-looking error. + """ + tok = _fast_tokenizer_with_template("{{ bos }}{% generation %}{{ c }}") + tok.is_fast = False + tok.convert_tokens_to_ids.return_value = None # defeat recovery detect()s + with pytest.raises(RuntimeError, match="fast tokenizer"): + hf_streaming_dataset._tokenize_with_loss_mask(tok, _CONV, answer_only_loss=True) diff --git a/tests/unit/torch/utils/test_dataset_utils.py b/tests/unit/torch/utils/test_dataset_utils.py index 49fceecb828..59433644d7e 100644 --- a/tests/unit/torch/utils/test_dataset_utils.py +++ b/tests/unit/torch/utils/test_dataset_utils.py @@ -313,7 +313,7 @@ def test_get_max_batch_size_oom_retry_shrinks_input(): seen_batch_sizes: list[int] = [] - def fake_forward(x): + def fake_call(x): seen_batch_sizes.append(x.shape[0]) # First call is the single-batch probe — succeeds. # Second call is the target-batch attempt — OOMs. @@ -322,7 +322,9 @@ def fake_forward(x): raise torch.cuda.OutOfMemoryError model = Mock(spec=torch.nn.Module) - model.forward = fake_forward + # get_max_batch_size calls the module (model(...)) so FSDP2 hooks fire, not model.forward, + # so route the mock's __call__ via side_effect. + model.side_effect = fake_call model.__class__.__name__ = "DummyModel" # not enc/dec free_before = 1000 @@ -344,7 +346,7 @@ def fake_forward(x): sample_input_single_batch=sample_input, ) - # Forward calls: probe(1), retry-at-target(10), retry-after-halve(5) + # Model calls: probe(1), retry-at-target(10), retry-after-halve(5) assert seen_batch_sizes == [1, 10, 5] # Final batch is 5 -> regulated to 4 (5 // 4 * 4 = 4). assert result == 4 diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py new file mode 100644 index 00000000000..323fb92a568 --- /dev/null +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -0,0 +1,156 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""Pure-function tests for ``modelopt.torch.utils.plugins.model_load_utils``.""" + +import json + +import pytest +import torch +from packaging.version import Version +from safetensors.torch import save_file + +pytest.importorskip("accelerate") + +from modelopt.torch.utils.plugins.model_load_utils import ( + _conversion_plan, + _convert_keys, + _resolve_target, + read_safetensors_subset, + weight_map_for, +) + + +def test_weight_map_for_sharded(tmp_path): + save_file({"a.weight": torch.zeros(2)}, str(tmp_path / "shard1.safetensors")) + save_file({"b.weight": torch.zeros(2)}, str(tmp_path / "shard2.safetensors")) + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps( + {"weight_map": {"a.weight": "shard1.safetensors", "b.weight": "shard2.safetensors"}} + ) + ) + + assert weight_map_for(str(tmp_path)) == { + "a.weight": "shard1.safetensors", + "b.weight": "shard2.safetensors", + } + + +def test_weight_map_for_single_file(tmp_path): + save_file( + {"a.weight": torch.zeros(2), "b.weight": torch.zeros(2)}, + str(tmp_path / "model.safetensors"), + ) + + assert weight_map_for(str(tmp_path)) == { + "a.weight": "model.safetensors", + "b.weight": "model.safetensors", + } + + +def test_weight_map_for_missing(tmp_path): + with pytest.raises(RuntimeError, match="No safetensors checkpoint"): + weight_map_for(str(tmp_path)) + + +def test_read_safetensors_subset(tmp_path): + save_file( + {"a.weight": torch.tensor([1.0, 2.0]), "a.bias": torch.tensor([3.0])}, + str(tmp_path / "shard1.safetensors"), + ) + save_file({"b.weight": torch.tensor([4.0])}, str(tmp_path / "shard2.safetensors")) + weight_map = { + "a.weight": "shard1.safetensors", + "a.bias": "shard1.safetensors", + "b.weight": "shard2.safetensors", + } + + result = read_safetensors_subset(str(tmp_path), weight_map, lambda n: n.startswith("a.")) + + assert set(result.keys()) == {"a.weight", "a.bias"} + assert torch.equal(result["a.weight"], torch.tensor([1.0, 2.0])) + assert torch.equal(result["a.bias"], torch.tensor([3.0])) + + +def _build_tiny_qwen3_moe(): + """A tiny meta-init Qwen3-MoE (fused ``gate_up_proj`` + ``down_proj`` experts) for converter tests.""" + transformers = pytest.importorskip("transformers") + if Version(transformers.__version__) < Version("5.0"): + pytest.skip("multi-source fused-MoE conversion needs transformers>=5") + from accelerate import init_empty_weights + from transformers import AutoConfig, AutoModelForCausalLM + + cfg = AutoConfig.for_model( + "qwen3_moe", + hidden_size=8, + intermediate_size=16, + moe_intermediate_size=6, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + head_dim=4, + num_experts=4, + num_experts_per_tok=2, + vocab_size=32, + max_position_embeddings=16, + ) + try: + with init_empty_weights(): + model = AutoModelForCausalLM.from_config(cfg) + except Exception as e: # modeling class unavailable in this transformers build + pytest.skip(f"qwen3_moe modeling unavailable: {e}") + return model, cfg + + +def test_checkpoint_key_converter_multisource_expert_fusion(): + """A multi-source fused MoE (Qwen3: gate_proj+up_proj -> gate_up_proj) converts correctly. + + Exercises the multi-source path (gate/up concatenated after expert stacking) AND the + single-source path (down_proj is a plain expert stack) in one model. + """ + model, cfg = _build_tiny_qwen3_moe() + plan = _conversion_plan(model) + assert plan is not None + + names = dict(model.named_parameters()) + gname = next(n for n in names if n.endswith("mlp.experts.gate_up_proj")) + prefix = gname[: -len("mlp.experts.gate_up_proj")] + n_exp, inter, hidden = cfg.num_experts, cfg.moe_intermediate_size, cfg.hidden_size + + gate = [torch.randn(inter, hidden) for _ in range(n_exp)] + up = [torch.randn(inter, hidden) for _ in range(n_exp)] + down = [torch.randn(hidden, inter) for _ in range(n_exp)] + state = {} + for e in range(n_exp): + state[f"{prefix}mlp.experts.{e}.gate_proj.weight"] = gate[e] + state[f"{prefix}mlp.experts.{e}.up_proj.weight"] = up[e] + state[f"{prefix}mlp.experts.{e}.down_proj.weight"] = down[e] + + out = _convert_keys(plan, state) + + # Multi-source: experts stacked (dim 0) then gate|up concatenated (dim 1), gate first. + gate_up = out[f"{prefix}mlp.experts.gate_up_proj"] + assert gate_up.shape == (n_exp, 2 * inter, hidden) == tuple(names[gname].shape) + assert torch.equal(gate_up, torch.cat([torch.stack(gate), torch.stack(up)], dim=1)) + + # Single-source (regression): down_proj is a plain expert stack. + down_proj = out[f"{prefix}mlp.experts.down_proj"] + assert down_proj.shape == (n_exp, hidden, inter) + assert torch.equal(down_proj, torch.stack(down)) + + # Name-only mapping routes every source key to the fused target. + for e in range(n_exp): + assert _resolve_target(plan, f"{prefix}mlp.experts.{e}.gate_proj.weight")[0] == gname + assert _resolve_target(plan, f"{prefix}mlp.experts.{e}.up_proj.weight")[0] == gname diff --git a/tools/launcher/common/eagle3/train_eagle_streaming.sh b/tools/launcher/common/eagle3/train_eagle_streaming.sh index 2f1e165062c..72b7388c3b1 100755 --- a/tools/launcher/common/eagle3/train_eagle_streaming.sh +++ b/tools/launcher/common/eagle3/train_eagle_streaming.sh @@ -47,6 +47,9 @@ # SERVE_CPU_OFFLOAD_GB GB/GPU offloaded to host RAM (fits big models on too-few GPUs; slower) # SERVE_MAX_MODEL_LEN cap context length (trims KV/activation) # SERVE_MAX_NUM_SEQS cap concurrent sequences (trims KV/activation) +# SERVE_BLOCK_SIZE KV-cache block size (e.g. 128 for MiniMax-M3 MSA sparse attention). +# Needs its own knob: multi-token SERVE_EXTRA_ARGS values are mangled +# by nemo_run's unquoted env export, so "--block-size 128" cannot ride it. # SERVE_HOST single-node: bind/connect host. default 127.0.0.1 # SERVE_GPU single-node: CUDA_VISIBLE_DEVICES for vllm. default "0" # SERVE_TP tensor-parallel size. default 1 single-node / all serve-node GPUs @@ -153,6 +156,7 @@ launch_vllm() { [ -n "${SERVE_CPU_OFFLOAD_GB:-}" ] && opt_args+=(--cpu-offload-gb "$SERVE_CPU_OFFLOAD_GB") [ -n "${SERVE_MAX_MODEL_LEN:-}" ] && opt_args+=(--max-model-len "$SERVE_MAX_MODEL_LEN") [ -n "${SERVE_MAX_NUM_SEQS:-}" ] && opt_args+=(--max-num-seqs "$SERVE_MAX_NUM_SEQS") + [ -n "${SERVE_BLOCK_SIZE:-}" ] && opt_args+=(--block-size "$SERVE_BLOCK_SIZE") # --no-enable-chunked-prefill / --no-enable-prefix-caching: connector captures hidden states during prefill; both skip recomputing cached/partial prefixes, yielding short/empty hidden_states. Required. # --no-enable-flashinfer-autotune: on NVFP4 MoE the autotuner re-tunes on the first serving step and stalls a worker past vLLM's execute-model timeout, killing EngineCore. # Hidden states move serve -> trainer over NIXL RDMA (no disk round-trip): one diff --git a/tools/launcher/examples/MiniMaxAI/MiniMax-M3/hf_streaming_dspark_multi_node.yaml b/tools/launcher/examples/MiniMaxAI/MiniMax-M3/hf_streaming_dspark_multi_node.yaml new file mode 100644 index 00000000000..d3e6290b13a --- /dev/null +++ b/tools/launcher/examples/MiniMaxAI/MiniMax-M3/hf_streaming_dspark_multi_node.yaml @@ -0,0 +1,134 @@ +# DSpark streaming speculative-decoding training for MiniMax-M3 (multi-node). +# DSpark = the DFlash backbone + a lightweight Markov head + a confidence head, +# generating a causal block semi-autoregressively; see dspark.yaml for the head +# and loss config. Runs the shared streaming pipeline +# (common/eagle3/train_eagle_streaming.sh) with the M3-specific base, draft dims, +# mask token and chat template; trained from scratch. A starting point for +# reproduction — tune node counts, batch, steps and serve limits for your cluster. +# +# MiniMax-M3 specifics this yaml encodes (each was a silent failure mode): +# * SERVE_BLOCK_SIZE=128: M3's MSA sparse attention (sparse_block_size=128) +# requires KV block 128 or vLLM dies at engine init ("No common block size"). +# * data.chat_template: M3 ships a FAST tokenizer whose template has no +# {% generation %} tags -> assistant_masks come back ALL-ZERO and +# answer_only_loss training silently runs at zero loss. The tagged template +# copy next to this yaml wraps the assistant turn (think prefix + content + +# tool calls + eos) in {% generation %} tags. +# * The draft does NOT inherit base GQA/FFN dims (set explicitly below), and +# M3's base rope_theta (5e6) is pinned onto the draft. +# * EAGLE_CAPTURE_IDS = draft default target_layer_ids+1 (6 aux) + final (60). +# * M3 is a VLM wrapper (text_config nested): the base's gemma-style final +# norm is selected via the config's use_gemma_norm flag (see +# modeling_final_norm.py) — its text_config coerces to a mixtral model_type, +# so the model_type table alone would pick the wrong norm. +# +# Run ON the cluster login node (paramiko can't reach it through the login proxy): +# export SLURM_HOST=localhost SLURM_ACCOUNT= \ +# SLURM_PARTITION= \ +# SLURM_HF_LOCAL= \ +# SLURM_JOB_DIR= \ +# NEMORUN_HOME=$PWD +# uv run launch.py --yaml examples/MiniMaxAI/MiniMax-M3/hf_streaming_dspark_multi_node.yaml \ +# identity=$HOME/.ssh/id_ecdsa detach=True --yes +# +# The export lands in /scratchspace/export. + +job_name: MiniMax-M3_DSpark_streaming_multi_node +pipeline: + allow_to_fail: false + skip: false + note: + + global_vars: + hf_model: /hf-local/MiniMaxAI/MiniMax-M3 + + # Build /scratchspace/data/train.jsonl. Point data.data_path at the full + # Spec-Decoding-Dataset-v2 corpus to reproduce; eagle_utils also accepts a + # directory of *.jsonl shards directly. + task_0: + script: common/eagle3/make_dataset.sh + args: + - -f modules/Model-Optimizer/examples/dataset/example_data_config.yaml + - --full-conversations + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 8 + container: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc10 + + task_1: + script: common/eagle3/train_eagle_streaming.sh + args: + - --config modules/Model-Optimizer/modelopt_recipes/general/speculative_decoding/dspark.yaml + - model.model_name_or_path=<> + - model.use_fake_base_for_offline=true + - model.trust_remote_code=true + - data.mode=streaming + - data.data_path=/scratchspace/data/train.jsonl + # M3's own template has no {% generation %} tags; without this tagged copy + # answer_only_loss trains on an all-zero mask (see header). + - data.chat_template=examples/MiniMaxAI/MiniMax-M3/m3_chat_template_generation.jinja + - training.output_dir=/scratchspace/dspark + - training.training_seq_len=4096 + - training.disable_tqdm=true + - training.ar_validate_steps=500000 + - training.num_train_epochs=1 + - training.per_device_train_batch_size=4 + - training.gradient_accumulation_steps=1 + - training.save_steps=1000 + - training.logging_steps=20 + - training.learning_rate=1.0e-4 + - training.warmup_steps=2000 + - training.answer_only_loss=true + # The vLLM serve container has no tensorboard -> trainer init crash. + - training.report_to=none + # The DSpark draft does NOT inherit the base GQA/FFN dims, so set them + # explicitly to match the M3 backbone (else a silently wrong-shape draft). + # intermediate_size matches M3's dense FFN (dense_intermediate_size). + - dflash.dflash_architecture_config.num_hidden_layers=6 + - dflash.dflash_architecture_config.num_key_value_heads=8 + - dflash.dflash_architecture_config.intermediate_size=12288 + # Pin the base's rope_theta onto the draft (M3 uses 5e6, not the Qwen3 + # default 1e6; a mismatch trains rope into the weights and caps AL). + - dflash.dflash_architecture_config.rope_theta=5000000 + # Semi-AR generation block (dspark.yaml ships 16; Kimi/M3 runs use 8). + - dflash.dflash_block_size=8 + # M3 has no mask token; vocab 200064, added tokens end at 200060 -> 200063 free. + - dflash.dflash_mask_token_id=200063 + environment: + - HF_MODEL_CKPT: <> + # 6 aux capture ids = the draft's default target_layer_ids+1, plus the true + # final hidden (60). Requires the aux-capture fix vllm#46788 (in-tree in + # recent nightlies); mismatched ids silently skew train vs inference. + - EAGLE_CAPTURE_IDS: "[2,13,24,36,47,58,60]" + - SERVE_NODES: "4" + - SERVE_TP: "8" + - STREAMING_NUM_WORKERS: "4" + # M3's custom-modeling base needs trust_remote_code at export and serve. + - EXPORT_EXTRA_ARGS: "--trust_remote_code" + - SERVE_EXTRA_ARGS: "--trust-remote-code" + # REQUIRED for M3: MSA sparse attention needs KV block 128 (dedicated knob — + # multi-token SERVE_EXTRA_ARGS values are mangled by nemo_run's unquoted + # env export, so "--block-size 128" cannot ride it). + - SERVE_BLOCK_SIZE: "128" + - SERVE_MAX_MODEL_LEN: "4160" + - SERVE_MAX_NUM_SEQS: "32" + - SERVE_GPU_MEM_UTIL: "0.9" + - SERVE_READY_TIMEOUT: "3600" + - VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: "1200" + - VLLM_ENGINE_ITERATION_TIMEOUT_S: "1200" + # RDMA transport is UCX (InfiniBand) by default. On AWS EFA, uncomment — + # and note UCX SEGFAULTS at agent init on EFA nodes (it detects the EFA + # devices), so LIBFABRIC is required there even for single-node runs: + # - NIXL_BACKENDS: "LIBFABRIC" + # - FI_PROVIDER: "efa" + # - NCCL_IB_DISABLE: "1" + slurm_config: + _factory_: "slurm_factory" + nodes: 6 + ntasks_per_node: 1 + gpus_per_node: 8 + # vLLM x86_64 build with native MiniMax-M3 support (vllm/models/minimax_m3) + # and the aux-capture fix (vllm#46788). + container: diff --git a/tools/launcher/examples/MiniMaxAI/MiniMax-M3/m3_chat_template_generation.jinja b/tools/launcher/examples/MiniMaxAI/MiniMax-M3/m3_chat_template_generation.jinja new file mode 100644 index 00000000000..95488f5483c --- /dev/null +++ b/tools/launcher/examples/MiniMaxAI/MiniMax-M3/m3_chat_template_generation.jinja @@ -0,0 +1,255 @@ +{# MiniMax-M3 chat template with {% generation %} tags for answer_only_loss training. + Adapted from https://huggingface.co/MiniMaxAI/MiniMax-M3/blob/main/chat_template.jinja + with {% generation %} / {% endgeneration %} wrapping assistant content (think prefix, + content, and eos), so apply_chat_template can return the assistant token mask. +-#} +{# ---------- special token variables ---------- #} +{%- set ns_token = ']<]minimax[>[' -%} +{%- set bod_token = ']~!b[' -%} +{%- set bos_token = ']~b]' -%} +{%- set eos_token = '[e~[' -%} +{%- set toolcall_begin_token = ns_token ~ '' -%} +{%- set toolcall_end_token = ns_token ~ '' -%} +{%- set think_begin_token = '' -%} +{%- set think_end_token = '' -%} +{%- set image_token = ']<]image[>[' -%} +{%- set video_token = ']<]video[>[' -%} +{#- Thinking mode: "enabled" / "disabled" / "adaptive" / not defined -#} +{#- Recursive XML renderer for tool_call arguments ======================== -#} +{#- None values are intentionally skipped in mapping iteration so that + `null` (which would round-trip to the literal string "null") + never appears in the rendered tool_call. The convention is: omit the + field entirely. The top-level `_args` loop applies the same rule. + The `val is none` branch below is a safety net only — upstream cleaning + (drop_none_in_tool_arguments) should ensure no None ever reaches here. -#} +{%- macro to_xml(val, ns) -%} +{%- if val is mapping -%} +{%- for k, v in val.items() if v is not none -%} +{{ ns }}<{{ k }}>{{ to_xml(v, ns) }}{{ ns }} +{%- endfor -%} +{%- elif val is iterable and val is not string -%} +{%- for item in val -%} +{{ ns }}{{ to_xml(item, ns) }}{{ ns }} +{%- endfor -%} +{%- elif val is none -%} +{#- Should be unreachable when upstream cleaning is applied. -#} +{%- elif val is boolean -%} +{{ val | tojson }} +{%- else -%} +{{ val }} +{%- endif -%} +{%- endmacro -%} +{#- Tool Rendering Functions ============================================== -#} +{%- macro render_tool_namespace(namespace_name, tool_list) -%} +{%- for tool in tool_list -%} +{{ tool.function | tojson(ensure_ascii=False) }} +{% endfor -%} +{%- endmacro -%} +{%- macro visible_text(content) -%} + {%- if content is string -%} + {{ content }} + {%- elif content is iterable and content is not mapping -%} + {%- for item in content -%} + {%- if item is mapping and item.type == 'text' -%} + {{- item.text }} + {%- elif item is mapping and item.type == 'image' -%} + {{- image_token }} + {%- elif item is mapping and item.type == 'video' -%} + {{- video_token}} + {%- elif item is string -%} + {{- item }} + {%- endif -%} + {%- endfor -%} + {%- elif content is none -%} + {{- '' }} + {%- else -%} + {{- content }} + {%- endif -%} +{%- endmacro -%} +{#- System Message Construction ============================================ -#} +{%- macro build_system_message(system_message) -%} + {%- if system_message and system_message.content -%} + {{- visible_text(system_message.content) }} + {%- else -%} + {{- 'Your model version is MiniMax-M3, developed by MiniMax. Knowledge cutoff: January 2026. Founded in early 2022, MiniMax is a global AI foundation model company committed to advancing the frontiers of AI towards AGI.' }} + {%- endif -%} + + {#- Thinking mode instructions -#} + {{- '\n\n\n' }} + {{- 'You have a thinking capability that allows you to reason step by step before responding. When thinking is enabled, wrap your reasoning in ' ~ think_begin_token ~ think_end_token ~ ' tags before your response. When thinking is disabled, begin your response directly after the ' ~ think_end_token ~ ' prefix. When thinking is adaptive, decide on your own whether to think for the current turn.\n' }} + {%- if thinking_mode is defined -%} + {%- if thinking_mode == "enabled" -%} + {{- 'Current thinking mode: enabled. You MUST think step by step before every response, including after receiving function/tool results.\n' }} + {%- elif thinking_mode == "disabled" -%} + {{- 'Current thinking mode: disabled. Do not output any thinking process.\n' }} + {%- elif thinking_mode == "adaptive" -%} + {{- 'Current thinking mode: adaptive. You are encouraged to think for complex decision-making, multi-step reasoning, or when analyzing function/tool results.\n' }} + {%- endif -%} + {%- else -%} + {{- 'Current thinking mode: adaptive. You are encouraged to think for complex decision-making, multi-step reasoning, or when analyzing function/tool results.\n' }} + {%- endif -%} + {{- '' }} +{%- endmacro -%} +{%- macro build_developer_message(developer_message) -%} + {%- if developer_message and developer_message.content -%} + {{- visible_text(developer_message.content) }} + {%- else -%} + {%- if model_identity is not defined -%} + {%- set model_identity = "You are a helpful assistant." -%} + {%- endif -%} + {{- model_identity }} + {%- endif -%} +{%- endmacro -%} +{#- Main Template Logic ================================================= -#} +{#- Role mapping: root -> system sp (high priority), system/developer -> developer sp (low priority) -#} +{%- set system_message = none -%} +{%- set developer_message = none -%} +{%- set conversation_messages = messages -%} +{%- if messages and messages[0].role == "root" -%} + {%- set system_message = messages[0] -%} + {%- set conversation_messages = messages[1:] -%} + {%- if conversation_messages and conversation_messages[0].role in ["system", "developer"] -%} + {%- set developer_message = conversation_messages[0] -%} + {%- set conversation_messages = conversation_messages[1:] -%} + {%- endif -%} +{%- elif messages and messages[0].role in ["system", "developer"] -%} + {%- set developer_message = messages[0] -%} + {%- set conversation_messages = messages[1:] -%} +{%- endif -%} +{#- Render system sp (higher priority, root role only) -#} +{{- bod_token ~ bos_token ~ 'system' ~ '\n' }} +{{- build_system_message(system_message) }} +{{- eos_token ~ '\n' }} + +{#- Render developer sp (lower priority: system/developer role + tools) -#} +{{- bos_token ~ 'developer' ~ '\n' }} +{{- build_developer_message(developer_message) }} +{%- if tools -%} + {{- '\n\n' ~ '# Tools' ~ '\n' ~ 'You may call one or more tools to assist with the user query.\nHere are the tools available in JSONSchema format:' ~ '\n' }} + {{- '\n' ~ '' ~ '\n' }} + {{- render_tool_namespace("functions", tools) }} + {{- '' ~ '\n\n' }} + {{- 'To call tools, wrap all invocations in a single ' ~ toolcall_begin_token ~ toolcall_end_token ~ ' block. Parameter values containing nested objects or arrays are recursively expanded into XML elements. Example:\n' }} + {{- '\n' ~ toolcall_begin_token ~ '\n' }} + {{- ns_token + '' }} + {{- ns_token + 'value-1' + ns_token + '' }} + {{- ns_token + '' }} + {{- ns_token + '' }} + {{- ns_token + 'val-a' + ns_token + '' }} + {{- ns_token + 'val-b' + ns_token + '' }} + {{- ns_token + '' }} + {{- ns_token + '' }} + {{- ns_token + '\n' }} + {{- ns_token + '' }} + {{- ns_token + 'value-1' + ns_token + '' }} + {{- ns_token + '\n' }} + {{- toolcall_end_token }} +{%- endif -%} +{{- eos_token ~ '\n' }} + +{#- Render messages -#} +{%- set last_tool_call = namespace(name=none) -%} +{%- for message in conversation_messages -%} + {%- if message.role == 'assistant' -%} + {{- bos_token ~ 'ai' ~ '\n' }} + {%- generation -%} + + {%- set reasoning_content = '' %} + {%- set content = visible_text(message.content) %} + {%- if message.reasoning_content is string %} + {%- set reasoning_content = message.reasoning_content %} + {%- else %} + {%- if think_end_token in content %} + {%- set reasoning_content = content.split(think_end_token)[0].strip('\n').split(think_begin_token)[-1].strip('\n') %} + {%- set content = content.split(think_end_token)[-1].strip('\n') %} + {%- endif %} + {%- endif %} + + {%- if reasoning_content -%} + {#- Render thinking for every assistant turn (all-turn visible) -#} + {{- think_begin_token ~ reasoning_content ~ think_end_token }} + {%- else -%} + {#- No thinking rendered → prefix with think_end_token -#} + {{- think_end_token }} + {%- endif -%} + + {%- if content -%} + {{- content }} + {%- endif -%} + {%- if message.tool_calls -%} + {{- toolcall_begin_token ~ '\n' }} + + {%- for tool_call in message.tool_calls -%} + {%- if tool_call.function -%} + {%- set tool_call = tool_call.function -%} + {%- endif -%} +{{- ns_token + '' }} +{%- set _args = tool_call.arguments -%} +{%- for k, v in _args.items() if v is not none %} +{{- ns_token + '<' + k + '>' -}} +{{- to_xml(v, ns_token) -}} +{{- ns_token + '' }} +{%- endfor -%} +{{- ns_token + '' ~ '\n' }} + {%- endfor -%} + + {{- toolcall_end_token }} + {%- if message.tool_calls[-1].function -%} + {%- set last_tool_call.name = message.tool_calls[-1].function.name -%} + {%- else -%} + {%- set last_tool_call.name = message.tool_calls[-1].name -%} + {%- endif -%} + {%- else -%} + {%- set last_tool_call.name = none -%} + {%- endif -%} + {{- eos_token }} + {%- endgeneration -%} + {{- '\n' }} + + {%- elif message.role == 'tool' -%} + {%- if last_tool_call.name is none -%} + {{- raise_exception("Message has tool role, but there was no previous assistant message with a tool call!") }} + {%- endif -%} + {%- if loop.first or (conversation_messages[loop.index0 - 1].role != 'tool') -%} + {{- bos_token ~ 'tool' }} + {%- endif -%} + {{- '\n' }} + {%- if message.content is string -%} + {{- message.content }} + {%- else -%} + {%- for tr in message.content -%} + {%- if tr is mapping and tr.type is defined and tr.type == 'image' -%} + {{- image_token }} + {%- elif tr is mapping and tr.type is defined and tr.type == 'video' -%} + {{- video_token }} + {%- else -%} + {{- tr.output if tr.output is defined else (tr.text if tr.type == 'text' and tr.text is defined else tr) }} + {%- endif -%} + {%- endfor -%} + {%- endif -%} + {{- '' }} + {%- if loop.last or (conversation_messages[loop.index0 + 1].role != 'tool') -%} + {{- eos_token ~ '\n' -}} + {%- endif -%} + + {%- elif message.role == 'user' -%} + {{- bos_token ~ 'user' ~ '\n' }} + {{- visible_text(message.content) }} + {{- eos_token ~ '\n' }} + {%- endif -%} +{%- endfor -%} + +{#- Generation prompt -#} +{%- if add_generation_prompt -%} +{{- bos_token ~ 'ai' ~ '\n' }} +{%- if thinking_mode is defined and thinking_mode == "disabled" -%} + {{- think_end_token }} +{%- elif thinking_mode is defined and thinking_mode == "adaptive" -%} + {#- adaptive: no prefix, let model decide -#} +{%- elif thinking_mode is defined and thinking_mode == "enabled" -%} + {#- enabled or not defined: default to think -#} + {{- think_begin_token }} +{%- else -%} + {#- adaptive: no prefix, let model decide -#} +{%- endif -%} +{%- endif -%} diff --git a/uv.lock b/uv.lock index ebe5c12d2fe..4f044327f3a 100644 --- a/uv.lock +++ b/uv.lock @@ -2273,6 +2273,7 @@ all = [ { name = "onnxruntime-gpu", version = "1.22.0", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and platform_machine != 'aarch64' and sys_platform != 'darwin') or sys_platform == 'win32'" }, { name = "onnxruntime-gpu", version = "1.24.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and platform_machine != 'aarch64' and sys_platform != 'darwin' and sys_platform != 'win32'" }, { name = "onnxscript" }, + { name = "onnxsim" }, { name = "onnxslim" }, { name = "pandas", version = "2.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "pandas", version = "3.0.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, @@ -2312,6 +2313,7 @@ dev = [ { name = "onnxruntime-gpu", version = "1.22.0", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and platform_machine != 'aarch64' and sys_platform != 'darwin') or sys_platform == 'win32'" }, { name = "onnxruntime-gpu", version = "1.24.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and platform_machine != 'aarch64' and sys_platform != 'darwin' and sys_platform != 'win32'" }, { name = "onnxscript" }, + { name = "onnxsim" }, { name = "onnxslim" }, { name = "pandas", version = "2.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "pandas", version = "3.0.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, @@ -2396,6 +2398,7 @@ onnx = [ { name = "onnxruntime-gpu", version = "1.22.0", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and platform_machine != 'aarch64' and sys_platform != 'darwin') or sys_platform == 'win32'" }, { name = "onnxruntime-gpu", version = "1.24.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and platform_machine != 'aarch64' and sys_platform != 'darwin' and sys_platform != 'win32'" }, { name = "onnxscript" }, + { name = "onnxsim" }, { name = "onnxslim" }, { name = "polygraphy" }, ] @@ -2445,6 +2448,7 @@ requires-dist = [ { name = "onnxruntime-gpu", marker = "python_full_version < '3.11' and platform_machine != 'aarch64' and sys_platform != 'darwin' and sys_platform != 'win32' and extra == 'onnx'", specifier = "~=1.22.0" }, { name = "onnxruntime-gpu", marker = "sys_platform == 'win32' and extra == 'onnx'", specifier = "==1.22.0" }, { name = "onnxscript", marker = "extra == 'onnx'" }, + { name = "onnxsim", marker = "extra == 'onnx'", specifier = ">=0.7.0" }, { name = "onnxslim", marker = "extra == 'onnx'", specifier = ">=0.1.76" }, { name = "packaging" }, { name = "pandas", marker = "extra == 'puzzletron'" }, @@ -2758,6 +2762,31 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/dd/bd/a0c8e737b6afda10e42a597787d53d5b66e00268df6f59184701eeae37d9/onnxscript-0.7.1-py3-none-any.whl", hash = "sha256:544763b7fdef49940cdd9412ff5135cbae96d59ac6bc1921457f21280f40f4b7", size = 721970, upload-time = "2026-06-29T23:33:23.298Z" }, ] +[[package]] +name = "onnxsim" +version = "0.7.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "onnx" }, + { name = "rich" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/29/d4/0def0ec1f143963137419fc2cdfc8521724216861abd8ca2cc42f5d257c2/onnxsim-0.7.0.tar.gz", hash = "sha256:9ee396257785ce07c5a468ba6cc0a8aba4eee0c6ed73f945deb7dc14b5c892c0", size = 2976263, upload-time = "2026-07-26T08:14:26.164Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/01/6d/4eb3b9284bbcf43ec01e85a483241a85562d31d6065a5dd79ef24d01c9ee/onnxsim-0.7.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:e5c123acbee25b4d6b7aa8f273afb46d9adda1e2fa32f6d70cc9ea68345d63fb", size = 2590443, upload-time = "2026-07-26T08:14:03.104Z" }, + { url = "https://files.pythonhosted.org/packages/b6/31/898cbc058b77947aedc060a3fc631f5d15d4e877229a52d60960b27b2021/onnxsim-0.7.0-cp310-cp310-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a03b84c3e527a9516bf261194beb7f0a545183c53baabd347a2cb815327ecd4f", size = 2385717, upload-time = "2026-07-26T08:14:05.181Z" }, + { url = "https://files.pythonhosted.org/packages/b6/b4/dcd6f799f26f62b4242f3f0a2fa3a77b105349af1fc85aa053ce17f7675c/onnxsim-0.7.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:26c12ae184d308c75c01644bdb1c53da52cb8e4738e0a2e325f1238fcee63ad9", size = 2658534, upload-time = "2026-07-26T08:14:06.812Z" }, + { url = "https://files.pythonhosted.org/packages/c7/8a/125afeed244cbc37427568d56748c0dd72c2a43f9d1805dc170da8fc9777/onnxsim-0.7.0-cp310-cp310-win_amd64.whl", hash = "sha256:1803fbe03709896a90d0fce9203e1b0ae1e2b6c911126f45851658b39b776fea", size = 3085963, upload-time = "2026-07-26T08:14:09.212Z" }, + { url = "https://files.pythonhosted.org/packages/10/50/40fb9043d5161a287e7df88b28c7b813cf893f9b1fa35317c0b9f67e73f7/onnxsim-0.7.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:5902ccc6cc3f78249be3136a9e0e4c7c4d4b405cdb9a0390edcd814683b4fe50", size = 2589879, upload-time = "2026-07-26T08:14:10.828Z" }, + { url = "https://files.pythonhosted.org/packages/98/5e/4251b34ce2305b10f818050317b83948ff39a68482bd0e756d2f77560b92/onnxsim-0.7.0-cp311-cp311-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f1409944a76b93b2a20c43b9cd7405a36b7482e44a84317d9a6ee9808f7d5561", size = 2385366, upload-time = "2026-07-26T08:14:12.453Z" }, + { url = "https://files.pythonhosted.org/packages/f4/f9/de22837a77f0dabfb5075cecbc7c1f5a2d2f2a7f741d9fe33097882edb3b/onnxsim-0.7.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ad35ab2066e44b5b839b553e59d135268faf3f7331f2c5e575d336128eaba171", size = 2658256, upload-time = "2026-07-26T08:14:14.094Z" }, + { url = "https://files.pythonhosted.org/packages/0d/06/5d516c4bdd77391e778206243e6d5538b6edda9084e364b7ab8d24143bbc/onnxsim-0.7.0-cp311-cp311-win_amd64.whl", hash = "sha256:c807b308fd7249ee7a31949efb42781a969dc705078d7342d94e3a6f3d482a86", size = 3085705, upload-time = "2026-07-26T08:14:15.89Z" }, + { url = "https://files.pythonhosted.org/packages/f0/65/583263cca89c18abd68192a40a0be25506e67f1bc8d53df7edb1fa3f3a6a/onnxsim-0.7.0-cp312-abi3-macosx_11_0_arm64.whl", hash = "sha256:ba0fc54e1f8f0bac1643ff8a3c00516428913ae799ab7717b771cf23baa46680", size = 2589625, upload-time = "2026-07-26T08:14:18.053Z" }, + { url = "https://files.pythonhosted.org/packages/8a/bc/743235ea06f3ef7aa55c515ed4e28316fc851d696d1acc686171529fb48f/onnxsim-0.7.0-cp312-abi3-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a10fc08b853ce60d3b5b8076b26e2bda1322c1e01c808b8db908e73f9a047a9d", size = 2381363, upload-time = "2026-07-26T08:14:19.486Z" }, + { url = "https://files.pythonhosted.org/packages/89/88/a771a9e5cfcd0210fdabdb242bb56857bd7beee59ba96f1178c1b702c17d/onnxsim-0.7.0-cp312-abi3-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0aa412583982475f7a5b45882d46353447b15694f7870737e43c59ed2300f472", size = 2654139, upload-time = "2026-07-26T08:14:21.148Z" }, + { url = "https://files.pythonhosted.org/packages/92/03/8c7913bc77c6f29fd122e92381fac85da5918a939efd68f396474072500b/onnxsim-0.7.0-cp312-abi3-win_amd64.whl", hash = "sha256:c4cfe7b2df2feeef8daf255537378d9028123d4ef627842ef19db9a2f6bd68e0", size = 3084638, upload-time = "2026-07-26T08:14:22.815Z" }, + { url = "https://files.pythonhosted.org/packages/94/d7/27430466aaff6e97d4ed4d1d335281363167f33ff5bc5726c6a3906e896d/onnxsim-0.7.0-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:477f9c8d9894e144d698a1216d7e94eee6043be6ed4669440063daa6c1616a68", size = 2658112, upload-time = "2026-07-26T08:14:24.386Z" }, +] + [[package]] name = "onnxslim" version = "0.1.94"