diff --git a/.github/actions/ubuntu-setup/action.yml b/.github/actions/ubuntu-setup/action.yml index 58fdc12b494..6a9c9fd81f5 100644 --- a/.github/actions/ubuntu-setup/action.yml +++ b/.github/actions/ubuntu-setup/action.yml @@ -1,5 +1,5 @@ name: Ubuntu runner setup -description: Standard setup for ubuntu-latest runners (setup Python, cleanup disk space) +description: Standard setup for ubuntu-24.04 runners (setup Python, cleanup disk space) inputs: python-version: diff --git a/.github/workflows/_pr_gate.yml b/.github/workflows/_pr_gate.yml index e3605202f7c..3720a3cd824 100644 --- a/.github/workflows/_pr_gate.yml +++ b/.github/workflows/_pr_gate.yml @@ -23,7 +23,7 @@ on: jobs: check-file-changes: - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 outputs: any_changed: ${{ steps.changed-tests.outputs.any_modified || steps.non-pr.outputs.any_changed }} steps: @@ -53,7 +53,7 @@ jobs: if: >- startsWith(github.ref, 'refs/heads/pull-request/') && needs.check-file-changes.outputs.any_changed == 'true' - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 permissions: checks: read steps: diff --git a/.github/workflows/bump_uv_lock.yml b/.github/workflows/bump_uv_lock.yml index ae624ddb211..01b8e50cc84 100644 --- a/.github/workflows/bump_uv_lock.yml +++ b/.github/workflows/bump_uv_lock.yml @@ -12,7 +12,7 @@ permissions: jobs: bump-uv-lock: - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 timeout-minutes: 15 steps: - uses: actions/checkout@v6 diff --git a/.github/workflows/check_model_hub_orgs.yml b/.github/workflows/check_model_hub_orgs.yml index fe4df1c5385..39d190b57a8 100644 --- a/.github/workflows/check_model_hub_orgs.yml +++ b/.github/workflows/check_model_hub_orgs.yml @@ -17,7 +17,7 @@ concurrency: jobs: check-model-hub-orgs: - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 timeout-minutes: 5 steps: - uses: actions/checkout@v6 diff --git a/.github/workflows/claude.yml b/.github/workflows/claude.yml index 46d0996515f..850383e66e3 100644 --- a/.github/workflows/claude.yml +++ b/.github/workflows/claude.yml @@ -47,7 +47,7 @@ jobs: contains(fromJson('["OWNER", "MEMBER", "COLLABORATOR"]'), github.event.issue.author_association) ) ) - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 timeout-minutes: 10 permissions: contents: write diff --git a/.github/workflows/claude_review.yml b/.github/workflows/claude_review.yml index 46f27063de8..b7abc076638 100644 --- a/.github/workflows/claude_review.yml +++ b/.github/workflows/claude_review.yml @@ -22,7 +22,7 @@ jobs: github.event.issue.pull_request && contains(github.event.comment.body, '/claude review') && contains(fromJson('["OWNER", "MEMBER", "COLLABORATOR"]'), github.event.comment.author_association) - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 timeout-minutes: 30 permissions: contents: read diff --git a/.github/workflows/close_inactive_issues_pr.yml b/.github/workflows/close_inactive_issues_pr.yml index 5c63c8316b9..53656b5d44a 100644 --- a/.github/workflows/close_inactive_issues_pr.yml +++ b/.github/workflows/close_inactive_issues_pr.yml @@ -7,7 +7,7 @@ on: jobs: stale: - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 permissions: issues: write pull-requests: write diff --git a/.github/workflows/code_quality.yml b/.github/workflows/code_quality.yml index 4743fef960b..80996cdc772 100644 --- a/.github/workflows/code_quality.yml +++ b/.github/workflows/code_quality.yml @@ -16,7 +16,7 @@ concurrency: jobs: code-quality: - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 timeout-minutes: 30 steps: - uses: actions/checkout@v6 diff --git a/.github/workflows/example_tests.yml b/.github/workflows/example_tests.yml index 8cc432a22ba..33cd33d3308 100644 --- a/.github/workflows/example_tests.yml +++ b/.github/workflows/example_tests.yml @@ -21,7 +21,7 @@ jobs: # One changed-files pass decides which lanes run. Lane granularity, not per-example: a job's # `if` cannot read `matrix`, so gating an individual example would need a job per example. pr-gate: - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 permissions: checks: read contents: read @@ -237,7 +237,7 @@ jobs: # Run even if example tests are skipped if: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && always() }} needs: [pr-gate, torch, vllm, trtllm, megatron, diffusers, onnx] - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 steps: - name: Required example tests did not succeed if: | diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml index a0389da7b62..91577ca15b4 100644 --- a/.github/workflows/gpu_tests.yml +++ b/.github/workflows/gpu_tests.yml @@ -104,7 +104,7 @@ jobs: # Run even if gpu-tests is skipped if: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && always() }} needs: [pr-gate, gpu-tests] - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 steps: - name: Required GPU tests did not succeed if: ${{ needs.pr-gate.result != 'success' || (needs.pr-gate.outputs.any_changed == 'true' && needs.gpu-tests.result != 'success') }} diff --git a/.github/workflows/pages.yml b/.github/workflows/pages.yml index 2132cdd302e..7dd889f88ad 100644 --- a/.github/workflows/pages.yml +++ b/.github/workflows/pages.yml @@ -24,7 +24,7 @@ permissions: jobs: build-docs: if: github.event.action != 'closed' - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 timeout-minutes: 10 steps: - uses: actions/checkout@v6 @@ -40,7 +40,7 @@ jobs: changes: if: github.event_name == 'pull_request' && github.event.action != 'closed' - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 outputs: docs: ${{ steps.filter.outputs.docs }} steps: @@ -62,7 +62,7 @@ jobs: github.event.pull_request.head.repo.full_name == github.repository && (github.event.action == 'closed' || needs.changes.outputs.docs == 'true') needs: [build-docs, changes] - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 timeout-minutes: 10 # Per-PR concurrency without cancel-in-progress so 'closed' cleanup always runs concurrency: @@ -83,7 +83,7 @@ jobs: deploy-gh-pages: if: github.event_name == 'push' || github.event_name == 'schedule' || github.event_name == 'workflow_dispatch' needs: build-docs - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 steps: - uses: actions/checkout@v6 - name: Download docs artifact diff --git a/.github/workflows/regression_tests.yml b/.github/workflows/regression_tests.yml index 001c4e6af4a..a28ac16611c 100644 --- a/.github/workflows/regression_tests.yml +++ b/.github/workflows/regression_tests.yml @@ -74,7 +74,7 @@ jobs: # Run even if regression-tests is skipped if: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && always() }} needs: [pr-gate, regression-tests] - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 steps: - name: Required regression tests did not succeed if: | diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 002a08850aa..dbdb208f2ba 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -9,7 +9,7 @@ permissions: jobs: build-upload: - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 steps: - uses: actions/checkout@v6 - uses: ./.github/actions/ubuntu-setup diff --git a/.github/workflows/unit_tests.yml b/.github/workflows/unit_tests.yml index 6fbca185d9f..090c7065e12 100644 --- a/.github/workflows/unit_tests.yml +++ b/.github/workflows/unit_tests.yml @@ -28,7 +28,7 @@ concurrency: jobs: check-file-changes: - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 outputs: any_changed: ${{ steps.changed.outputs.any_modified || steps.non-pr.outputs.any_changed }} steps: @@ -58,7 +58,7 @@ jobs: .agents/skills/** plugins/modelopt/skills/** linux: - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 timeout-minutes: 20 steps: - uses: actions/checkout@v6 @@ -120,7 +120,7 @@ jobs: multi-version: if: needs.check-file-changes.outputs.any_changed == 'true' needs: [linux, check-file-changes] - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 timeout-minutes: 20 strategy: fail-fast: false @@ -150,7 +150,7 @@ jobs: partial-install: if: needs.check-file-changes.outputs.any_changed == 'true' needs: [linux, check-file-changes] - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 timeout-minutes: 15 strategy: fail-fast: false @@ -164,7 +164,7 @@ jobs: launcher: if: needs.check-file-changes.outputs.any_changed == 'true' needs: [linux, check-file-changes] - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 timeout-minutes: 15 steps: - uses: actions/checkout@v6 @@ -181,7 +181,7 @@ jobs: mcp: if: needs.check-file-changes.outputs.any_changed == 'true' needs: [linux, check-file-changes] - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 timeout-minutes: 15 steps: - uses: actions/checkout@v6 @@ -204,7 +204,7 @@ jobs: skills: if: needs.check-file-changes.outputs.any_changed == 'true' needs: [linux, check-file-changes] - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 timeout-minutes: 10 steps: - uses: actions/checkout@v6 @@ -222,7 +222,7 @@ jobs: # Run even if some jobs are skipped if: ${{ github.event_name == 'pull_request' && always() }} needs: [check-file-changes, linux, windows, multi-version, partial-install, launcher, mcp, skills] - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 steps: - name: Required unit tests did not succeed # windows is deliberately absent from this condition. It still runs and is still visible, diff --git a/CHANGELOG.rst b/CHANGELOG.rst index ef059b41228..4058290db06 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -111,6 +111,7 @@ Changelog - ``get_te_hybrid_stack_spec`` was removed from ``modelopt.torch.nas.plugins.megatron``; it had no use outside tests. Use ``modelopt.torch.utils.plugins.megatron_layer_specs.te_hybrid_stack_spec_sequential_mlp`` for the SequentialMLP layout, or ``megatron.core.models.hybrid.hybrid_layer_specs.hybrid_stack_spec`` for grouped GEMM. - Unified HuggingFace export now fails with ``NotImplementedError`` when it meets an MoE block whose expert projection names it does not know, instead of assuming Mixtral's ``w1``/``w2``/``w3``. If you hit this, register a ``ModelSpec`` for the model under ``modelopt/torch/models/``. Every MoE architecture ModelOpt exported correctly before this change is registered, so no supported model regresses. - ``--recipe`` (and ``modelopt.recipe.load_recipe``) now resolve a recipe path **filesystem-first**: a recipe of the same relative path in the current working directory takes precedence over the shipped built-in of that name, matching how recipe ``$import`` paths already resolve. Previously the built-in won. +- Diffusion HF exports now omit attention quantizer buffers by default, except for FLUX.1, which retains enabled Q/K/V ``_amax`` keys for runtime compatibility; other consumers of those keys can opt in with ``export_attention_quantizers=True`` (CLI: ``--extra-param export_attention_quantizers=true``). Disabled attention quantizers remain excluded, and diffusion components, including LTX-2, no longer emit KV-cache quantization metadata. **Deprecations** @@ -132,12 +133,18 @@ Changelog **Bug Fixes** +- Fix LTX-2 distillation full-mode checkpoints to use the safetensors format expected by their file extension and loaders. +- Fix corrupted ONNX external tensor data when exporting models on Windows by upgrading ONNX. Regenerate affected models from clean weights. +- Harden the default locations of the TensorRT deployment runtime's build artifacts and the ONNX AutoTune timing cache: they no longer sit at fixed paths under the shared system temp directory, where another local user could pre-create them to plant symlinks or swap a built engine. Both now live in a private temporary directory that is created once per process; pass ``timing_cache_path`` (runtime) or ``--timing_cache`` (AutoTune) to keep a timing cache across processes. +- Fix ``import modelopt.torch`` failing with ``ModuleNotFoundError`` for ``requests`` or ``huggingface_hub`` when ``nvidia-modelopt`` is installed without the ``[hf]`` extra. - Harden the distributed object broadcast/allgather helpers and the Megatron expert-parallel export gather so a compromised peer rank can no longer execute code through a crafted pickle payload: payloads load with ``weights_only=True`` and the export gathers raw bytes instead of pickled objects. +- Harden ``examples/vllm_serve/Dockerfile``: ``/workspace`` is no longer world-writable, so no other unprivileged user in the container can replace the CUDA extensions or the ModelOpt sources that the vLLM server loads. The directory is now owned by the ``vllm`` user; rebuild the image, and to run it as a different user point ``TORCH_EXTENSIONS_DIR`` at a directory that user can write, since ModelOpt otherwise cannot load its CUDA extensions and falls back to slower eager fake quantization. - Fix LSQ clamping small, representable FP8 block scales too aggressively, which could zero quantized values and suppress scale gradients. - Fix ONNX AutoTune benchmark initialization with TensorRT 11 while preserving strongly typed networks on TensorRT 10. - Keep vision-block Linears in high precision in the ``qwen3_vl`` and ``qwen3_5`` FP8 vision recipes so exported checkpoints avoid a vLLM vision-attention shape failure. The recipes still quantize primary and deepstack merger Linears; re-quantize affected checkpoints before deployment. - Fix unified HF export writing a weight name the source checkpoint does not use when a model declares a backwards-compatibility *load* rename as a reversible ``WeightRenaming``. Reversing such a rule rewrote the canonical name to a legacy spelling, and a deployment runtime keyed on the canonical name then failed to load the quantized checkpoint. Each exported name is now resolved against the names the source checkpoint actually carries, for both the exported tensors and the quantization config's module references. - Fix exports using the built-in Transformers NemotronH implementation retaining incompatible legacy remote-code config and model mappings. Re-export affected checkpoints before deployment. +- Fix ``examples/hf_ptq`` failing with ``KeyError: 'mlp'`` on dense Nemotron-H checkpoints (Nano-9B/12B-v2, H-8B/47B/56B) with Transformers 5.5, whose built-in ``NemotronHForCausalLM`` supports only the MoE variant. ``get_model`` now uses the checkpoint's bundled code when ``--trust_remote_code`` is set and the built-in config can't be built. - Fix ``export_hf_checkpoint`` dropping Diffusers pipeline settings, such as Wan 2.2 ``boundary_ratio``, from ``model_index.json`` when ``components`` is given. Re-export affected checkpoints to a new directory, since component exports do not overwrite an existing ``model_index.json``. - Fix INT8 real quantization (``mtq.compress``) of FP16 weights zeroing or clipping channels with a small amax: their scales underflowed in FP16. Scales of newly compressed INT8 weights are now stored in FP32. - Fix Hugging Face exports classifying registered Qwen VLM parameters as unplaced source weights, which overwrote quantized tensors with original weights and incorrectly excluded their modules from quantization. Re-export affected checkpoints. @@ -186,6 +193,19 @@ Changelog - Fix GPTQ with Megatron data parallelism: each data-parallel replica of a module now builds its Hessian from the whole calibration set instead of only its own shard, so all replicas reach the same quantized weights. - Fix ONNX quantization with the default ``concat_elimination`` pass crashing in calibration with ``KeyError: '_ReduceMin'`` when a ``Concat`` input is a constant, such as the class token of a ViT. - Fix PTQ of native Transformers FP8 experts to apply source weight scales during calibration and export. Re-quantize affected checkpoints from the original FP8 source. +- The default PTQ recipes now leave multi-token-prediction (MTP) layers unquantized in Megatron vision-language models too (``language_model.mtp.*``), as they already did at the model root. +- Fix ``examples/llm_sparsity/weight_sparsity/finetune.py`` saving a sharded state dict under FSDP2, the default FSDP version in Transformers 5, which made ``export_hf_ckpt.py`` fail to restore it. Re-run finetuning to produce a loadable ``finetuned_modelopt_state.pth``. +- Fix GPTQ with ``fused: true`` failing with a CUDA illegal memory access, or a Triton pointer error, on a model that ``device_map="auto"`` spreads over several GPUs once calibration reaches a layer outside the current GPU. +- Diffusers ``export_hf_checkpoint`` now honors ``save_modelopt_state=False`` (the default) when HuggingFace checkpointing is enabled. Pass ``save_modelopt_state=True`` to retain the ModelOpt state file. +- Fix ONNX INT4 quantization ending the Python process when the INT4 packing extension fails to build, for example without a C++ compiler. It now falls back to the Python implementation. +- Fix Megatron Top-K/Top-P distillation gradients being multiplied by the tensor-parallel size when tensor parallelism is enabled; reported loss values are unchanged. +- Fix ``fsdp2_aware_weight_update`` masking an error raised during its unshard (e.g. a CUDA OOM) with a spurious ``UnboundLocalError``. +- Nested ``only_student_forward()`` and ``only_teacher_forward()`` contexts now preserve the enclosing execution mode, including when an inner context is disabled or raises an exception. +- Fix an illegal memory access in static NVFP4 fake quantization for tensors with more than 2^31 elements, such as stacked MoE expert weights. The kernel is also about 14x faster. +- Preserve accumulated histogram counts when computing an entropy calibration threshold with NumPy histograms. +- Fix KV-cache quantization through ``register_attention_for_kv_quant``, the fallback for Hugging Face attention classes that do not use the attention interface, attaching the ``q``, ``k`` and ``v`` bmm quantizers to the wrong tensors for most attention layouts. Recalibrate and re-export models quantized through this fallback. Attention classes that compute their scores with ``torch.baddbmm`` (e.g. ``BloomAttention``, ``GPTNeoXJapaneseAttention``) are no longer patched, since their operands were always wired wrong (or raised ``IndexError``), and now get the existing warning to register a custom quantized attention class. +- Fix the diffusers skip-softmax HF export to emit the schema TensorRT-LLM VisualGen loads (``threshold_scale_factor.coefficients`` and a scalar ``target_sparsity``). Re-export checkpoints that failed to load. +- EAGLE training with ``eagle_decoder_type="kimik2"`` now downloads the Kimi-K2 modeling code from a pinned ``moonshotai/Kimi-K2-Thinking`` revision and verifies its SHA-256 hashes before importing it. Offline runs need that revision in the Hugging Face cache. 0.47.0 (2026-09-23) ^^^^^^^^^^^^^^^^^^^ diff --git a/README.md b/README.md index dc838483e88..532bd1669e2 100644 --- a/README.md +++ b/README.md @@ -139,7 +139,7 @@ more fine-grained control on installed dependencies or for alternative docker im |------------|----------------| | LLM / VLM Quantization | [View Support Matrix](./examples/hf_ptq/README.md#support-matrix) | | Diffusers Quantization | [View Support Matrix](./examples/diffusers/README.md#support-matrix) | -| ONNX Quantization | [View Support Matrix](./examples/torch_onnx/README.md#onnx-export-supported-llm-models) | +| ONNX Quantization | [View Support Matrix](./examples/torch_onnx/README.md#supported-models) | | Windows Quantization | [View Support Matrix](./examples/windows/README.md#support-matrix) | | Quantization Aware Training | [View Support Matrix](./examples/llm_qat/README.md#support-matrix) | | Pruning | [View Support Matrix](./examples/pruning/README.md#support-matrix) | diff --git a/docs/source/getting_started/windows/_installation_windows_arm64.rst b/docs/source/getting_started/windows/_installation_windows_arm64.rst index 2bafcf0be4e..e50f039a303 100644 --- a/docs/source/getting_started/windows/_installation_windows_arm64.rst +++ b/docs/source/getting_started/windows/_installation_windows_arm64.rst @@ -101,7 +101,7 @@ explicit dependency installation avoids selecting the Windows x64 ``onnxruntime- & $PythonExe -m pip install ` cppimport cupy-cuda13x lief ml_dtypes ` - "onnx~=1.21.0" "onnx-graphsurgeon>=0.6.1" ` + "onnx~=1.23.2" "onnx-graphsurgeon>=0.6.1" ` "onnxconverter-common~=1.16.0" onnxscript ` "onnxslim>=0.1.76" "polygraphy>=0.49.22" ` "onnxruntime>=1.24.2" onnxruntime-ep-nv-tensorrt-rtx-cu13 diff --git a/docs/source/guides/7_nas.rst b/docs/source/guides/7_nas.rst index 14566bfcf95..c86dc3cc198 100644 --- a/docs/source/guides/7_nas.rst +++ b/docs/source/guides/7_nas.rst @@ -326,7 +326,7 @@ The conversion process ---------------------- ModelOpt will automatically generate a search space for you from your custom PyTorch model. -This is a one time process process performed during pruning and NAS. Once a model is converted, you +This is a one time process performed during pruning and NAS. Once a model is converted, you can save and restore it for downstream tasks like training, inference, and fine-tuning. To help you better understand how the search space is derived from your model, we go through the diff --git a/docs/source/guides/_pytorch_quantization.rst b/docs/source/guides/_pytorch_quantization.rst index f8f12b068ba..1e7a2b05c09 100644 --- a/docs/source/guides/_pytorch_quantization.rst +++ b/docs/source/guides/_pytorch_quantization.rst @@ -142,7 +142,7 @@ quantization and weight & activation quantization). ``effective_bits`` constrain You may specify a ``effective_bits`` constraint such as 4.8 for mixed precision quantization using :attr:`NVFP4_DEFAULT_CFG` & :attr:`FP8_DEFAULT_CFG`. ``AutoQuantize`` will automatically quantize highly sensitive layers in :attr:`FP8_DEFAULT_CFG` while keeping less sensitive layers in :attr:`NVFP4_DEFAULT_CFG` (and even skip quantization for any extremely sensitive layers) so that -the the final mixed precision quantized model has an effective quantized bits of 4.8. +the final mixed precision quantized model has an effective quantized bits of 4.8. This model would give a better accuracy than the model quantized with vanilla :attr:`NVFP4_DEFAULT_CFG` since the more aggressive :attr:`NVFP4_DEFAULT_CFG` quantization was not applied for the highly sensitive layers. diff --git a/examples/deepseek/README.md b/examples/deepseek/README.md index 33848ed06aa..ccab5a09ef0 100644 --- a/examples/deepseek/README.md +++ b/examples/deepseek/README.md @@ -26,7 +26,7 @@ export HF_FP4_PATH={path_to_save_the_final_FP4_checkpoint} ### DeepSeek V3 R1 V3.1 ```bash -# download the FP8 checkpoint from Hugginface. This is an example of DeepSeek-R1 +# download the FP8 checkpoint from Hugging Face. This is an example of DeepSeek-R1 huggingface-cli download deepseek-ai/DeepSeek-R1 --local-dir $HF_FP8_CKPT # clone DeepSeek-V3 (base model of R1) Github repository for FP8 inference, @@ -36,7 +36,7 @@ git clone https://github.com/deepseek-ai/DeepSeek-V3.git && cd DeepSeek-V3 && gi ### [Experimental] DeepSeek V3.2 ```bash -# download the FP8 checkpoint from Hugginface. +# download the FP8 checkpoint from Hugging Face. huggingface-cli download deepseek-ai/DeepSeek-V3.2-Exp --local-dir $HF_FP8_CKPT # clone DeepSeek-V3.2 Github repository for FP8 inference, diff --git a/examples/diffusers/README.md b/examples/diffusers/README.md index 8d6e61e9644..b67d9630b8d 100644 --- a/examples/diffusers/README.md +++ b/examples/diffusers/README.md @@ -16,7 +16,7 @@ Cache Diffusion is a technique that reuses cached outputs from previous diffusio | :------------: | :------------: | :------------: | :------------: | | Pre-Requisites | Required & optional packages to use this technique | \[[Link](#pre-requisites)\] | | | Getting Started | Learn how to optimize your models using quantization/cache diffusion to reduce precision and improve inference efficiency | \[[Link](#getting-started)\] | \[[docs](https://nvidia.github.io/Model-Optimizer/guides/1_quantization.html)\] | -| Support Matrix | View the support matrix to see quantization/cahce diffusion compatibility and feature availability across different models | \[[Link](#support-matrix)\] | \[[docs](https://nvidia.github.io/Model-Optimizer/guides/1_quantization.html)\] | +| Support Matrix | View the support matrix to see quantization/cache diffusion compatibility and feature availability across different models | \[[Link](#support-matrix)\] | \[[docs](https://nvidia.github.io/Model-Optimizer/guides/1_quantization.html)\] | | Sparse Attention (Skip-Softmax) | Skip-softmax sparse attention for diffusion models | \[[Link](#sparse-attention-skip-softmax)\] | | | Cache Diffusion | Caching technique to accelerate inference without compromising quality | \[[Link](#cache-diffusion)\] | | | Post Training Quantization (PTQ) | Example scripts on how to run PTQ on diffusion models | \[[Link](#post-training-quantization-ptq)\] | \[[docs](https://nvidia.github.io/Model-Optimizer/guides/1_quantization.html)\] | @@ -97,7 +97,7 @@ We support calibration for INT8, FP8 and FP4 precision and for both weights and We also provide instructions on deploying and running E2E diffusion pipelines with Model Optimizer quantized INT8 and FP8 Backbone to generate images and measure latency on target GPUs. Note, Jetson devices are not supported at this time due to the incompatibility of the software. > [!NOTE] -> Model calibration requires relatively more GPU computing power then deployment. It does not need to be on the same GPUs as the deployment target GPUs. Deprecated ONNX export and TensorRT engine instructions remain available in [`quantization/ONNX-TRT-Deployment.md`](./quantization/ONNX-TRT-Deployment.md). +> Model calibration requires relatively more GPU computing power than deployment. It does not need to be on the same GPUs as the deployment target GPUs. Deprecated ONNX export and TensorRT engine instructions remain available in [`quantization/ONNX-TRT-Deployment.md`](./quantization/ONNX-TRT-Deployment.md). ### Quantize scripts @@ -124,6 +124,17 @@ python quantize.py \ --hf-ckpt-dir ./hf_ckpt ``` +HF exports omit attention quantizer state by default, including when calibration uses +`--quantize-mha`, except for FLUX.1. Qwen-Image's TRT-LLM loader requires this clean +checkpoint. FLUX.1 retains enabled Q/K/V `_amax` keys by default for compatibility with +the static FP8 attention path in [TensorRT-LLM PR #18396](https://github.com/NVIDIA/TensorRT-LLM/pull/18396). +Use `--extra-param export_attention_quantizers=true` to retain enabled Q/K/V quantizers +explicitly, or `export_attention_quantizers=false` to omit them, including for FLUX.1. +Disabled quantizers are always omitted. The equivalent Python option is +`export_hf_checkpoint(pipe, export_dir="./hf_ckpt", export_attention_quantizers=True)`. +Diffusion attention quantization, including native LTX-2 components, is not exported +as KV-cache quantization. Transformers components retain their KV-cache metadata. + #### Wan 2.2 VAE NVFP4 (Conv3D Implicit GEMM) The Wan 2.2 VAE (`AutoencoderKLWan`, shared between the 5B and 14B pipelines) is built from 3D convolutions. When quantizing the VAE with NVFP4, the `Conv3d` layers are automatically dispatched through a custom BF16 WMMA implicit-GEMM kernel with fused FP4 activation quantization. Requires SM80+ (Ampere or newer). See [`modelopt/torch/kernels/quantization/conv/README.md`](../../modelopt/torch/kernels/quantization/conv/README.md) for kernel details. @@ -154,6 +165,14 @@ python quantize.py \ > (including quantized or distilled checkpoints) remain subject to the LTX Community License > Agreement and are **not** covered by Apache 2.0. +> [!NOTE] +> This example targets the LTX-2 packages at commit +> [`28c3c73`](https://github.com/Lightricks/LTX-2/commit/28c3c73fe557666c3de176e1e50a5220152ccfca) +> (`QuantizationPolicy.fp8_cast()` and `checkpoint_path`/`gemma_root` pipeline arguments), e.g. +> `pip install "ltx-core @ git+https://github.com/Lightricks/LTX-2.git@28c3c73#subdirectory=packages/ltx-core" "ltx-pipelines @ git+https://github.com/Lightricks/LTX-2.git@28c3c73#subdirectory=packages/ltx-pipelines"`. +> Newer LTX-2 releases removed `fp8_cast()` and changed the pipeline constructor, so this example does +> not run on them yet. + This example produces three outputs: a PyTorch checkpoint (`--quantized-torch-ckpt-save-path`), a Hugging Face checkpoint (`--hf-ckpt-dir`), and a ComfyUI-compatible merged safetensor (`--extra-param merged_base_safetensor_path`). ```sh @@ -361,7 +380,7 @@ By following these steps, your PEFT LoRA model should be efficiently quantized u ## Sparse Attention (Skip-Softmax) -Skip-softmax sparse attention skips KV tiles whose attention scores are negligible during the softmax computation, reducing FLOPs without retraining. An exponential model (`scale_factor = a * exp(b * target_sparsity)`) is calibrated once, then the target sparsity can be adjusted at runtime without recalibration. Calibrated coefficients can be exported as a Hugging Face checkpoint (embedded in each component's `config.json` under `sparse_attention_config`) consumed directly by TRT-LLM's `SkipSoftmaxAttentionConfig.resolve_for_target_sparsity` — no extra conversion needed downstream. +Skip-softmax sparse attention skips KV tiles whose attention scores are negligible during the softmax computation, reducing FLOPs without retraining. An exponential model (`scale_factor = a * exp(b * target_sparsity)`) is calibrated once, then the target sparsity can be adjusted at runtime without recalibration. Calibrated coefficients can be exported as a Hugging Face checkpoint (embedded in each component's `config.json` under `sparse_attention_config`) in the schema TRT-LLM VisualGen loads (`threshold_scale_factor.coefficients` plus a scalar `target_sparsity`) — no extra conversion needed downstream. ### Getting Started @@ -478,7 +497,7 @@ compile( engine_path=Path("./engine"), ) -cachify.prepare(pipe, num_inference_steps, SDXL_DEFAULT_CONFIG) +cachify.prepare(pipe, SDXL_DEFAULT_CONFIG) ``` Afterward, use it as a standard cache diffusion pipeline to generate the image. @@ -497,7 +516,7 @@ SDXL_DEFAULT_CONFIG = [ } ] -cachify.prepare(pipe, num_inference_steps, SDXL_DEFAULT_CONFIG) +cachify.prepare(pipe, SDXL_DEFAULT_CONFIG) ``` Two parameters are essential: `wildcard_or_filter_func` and `select_cache_step_func`. diff --git a/examples/diffusers/distillation/distillation_trainer.py b/examples/diffusers/distillation/distillation_trainer.py index 38908f5c8f3..c1b4cc60e9a 100644 --- a/examples/diffusers/distillation/distillation_trainer.py +++ b/examples/diffusers/distillation/distillation_trainer.py @@ -1138,7 +1138,7 @@ def _save_checkpoint(self) -> Path | None: k: v.to(save_dtype) if isinstance(v, Tensor) else v for k, v in full_state_dict.items() } - self._accelerator.save(full_state_dict, saved_weights_path) + self._accelerator.save(full_state_dict, saved_weights_path, safe_serialization=True) rel_path = saved_weights_path.relative_to(self._config.output_dir) logger.info(f"Model weights for step {self._global_step} saved in {rel_path}") @@ -1504,6 +1504,7 @@ def train( # Calculate how many raw steps to skip and how many to run total_raw_steps = cfg.optimization.steps * cfg.optimization.gradient_accumulation_steps skip_raw_steps = resume_step * cfg.optimization.gradient_accumulation_steps + saved_this_step = False with progress: # Initial validation before training starts (skip if resuming) @@ -1616,6 +1617,7 @@ def train( ) if not saved_this_step: self._save_training_state() + saved_this_step = True # Break out of the training loop; post-loop code # will collect stats and return. break @@ -1700,8 +1702,9 @@ def train( global_batch_size=cfg.optimization.batch_size * self._accelerator.num_processes, ) - # Save final training state (for potential resume) - self._save_training_state() + # Interval and time-limit saves may already have checkpointed the final step. + if not saved_this_step: + self._save_training_state() # Save inference-ready model weights (standalone safetensors file) saved_path = self._save_checkpoint() diff --git a/examples/diffusers/quantization/calib/plugin_calib.py b/examples/diffusers/quantization/calib/plugin_calib.py index b8d82734f21..7490ad15ad6 100644 --- a/examples/diffusers/quantization/calib/plugin_calib.py +++ b/examples/diffusers/quantization/calib/plugin_calib.py @@ -26,6 +26,7 @@ def __init__(self, num_bits=8, axis=None, unsigned=False, track_amax=False, **kw self.total_step = kwargs["total_step"] self.collect_method = kwargs["collect_method"] self.data = {} + self._counts = {} self.i = 0 def collect(self, x): @@ -49,6 +50,7 @@ def collect(self, x): self.data[_cur_step] = torch.max(self.data[_cur_step], local_amax) else: self.data[_cur_step] += local_amax + self._counts[_cur_step] = self._counts.get(_cur_step, 0) + 1 if self._track_amax: raise NotImplementedError self.i += 1 @@ -57,7 +59,7 @@ def compute_amax(self): """Return the absolute max of all tensors collected.""" up_lim = int(self.total_step * self.percentile) if self.collect_method == "min-mean": - amaxs_values = [self.data[i] / self.total_step for i in range(up_lim)] + amaxs_values = [self.data[i] / self._counts[i] for i in range(up_lim)] else: amaxs_values = [self.data[i] for i in range(up_lim)] if self.collect_method == "mean-max": diff --git a/examples/diffusers/quantization/pipeline_manager.py b/examples/diffusers/quantization/pipeline_manager.py index af89ed568ff..c03af7b3ce9 100644 --- a/examples/diffusers/quantization/pipeline_manager.py +++ b/examples/diffusers/quantization/pipeline_manager.py @@ -217,13 +217,24 @@ def _create_ltx2_pipeline(self) -> Any: distilled_lora_strength = params.pop("distilled_lora_strength", 0.8) spatial_upsampler_path = params.pop("spatial_upsampler_path", None) gemma_root = params.pop("gemma_root", None) - fp8_quantization = params.pop("fp8_quantization", None) or params.pop( - "fp8transformer", False - ) + # Pop both aliases: `or` would short-circuit and leave `fp8transformer` in the kwargs. + fp8_quantization = params.pop("fp8_quantization", None) + fp8_quantization = params.pop("fp8transformer", False) or fp8_quantization params.pop("merged_base_safetensor_path", None) params.pop("enable_swizzle_layout", None) params.pop("padding_strategy", None) params.pop("enable_layerwise_quant_metadata", None) + # Calibration-only knobs read in calibration.py; the pipeline constructor rejects them. + for key in ( + "height", + "width", + "num_frames", + "frame_rate", + "seed", + "images", + "tiling_config", + ): + params.pop(key, None) if not checkpoint_path: raise ValueError("Missing required extra_param: checkpoint_path.") diff --git a/examples/diffusers/quantization/quantize.py b/examples/diffusers/quantization/quantize.py index 502801157fe..73ce6808df3 100644 --- a/examples/diffusers/quantization/quantize.py +++ b/examples/diffusers/quantization/quantize.py @@ -389,7 +389,11 @@ def export_hf_ckpt(self, pipe: Any, model_config: ModelConfig | None = None) -> self.logger.info(f"Merging base safetensors from {merged_path} for LTX2 export") kwargs["merged_base_safetensor_path"] = merged_path if model_config: - for key in ("enable_swizzle_layout", "enable_layerwise_quant_metadata"): + for key in ( + "enable_swizzle_layout", + "enable_layerwise_quant_metadata", + "export_attention_quantizers", + ): val = model_config.extra_params.get(key) if val is not None: normalized = str(val).strip().lower() diff --git a/examples/diffusers/sparsity/README.md b/examples/diffusers/sparsity/README.md index e102f1b6d5d..bc1188d2ea0 100644 --- a/examples/diffusers/sparsity/README.md +++ b/examples/diffusers/sparsity/README.md @@ -65,10 +65,9 @@ python wan22_skip_softmax.py \ `--export-dir` writes a Hugging Face checkpoint with the calibrated `threshold_scale_factor` block embedded in each component's `config.json` -(under the `sparse_attention_config` key). TRT-LLM's -`SkipSoftmaxAttentionConfig.resolve_for_target_sparsity` reads the -`(a, b)` directly via `coeffs['a'] * math.exp(coeffs['b'] * sparsity)` — -no extra conversion needed downstream. +(under the `sparse_attention_config` key), in the schema TRT-LLM VisualGen +loads: `(a, b)` under `threshold_scale_factor.coefficients` and a scalar +`target_sparsity`. No extra conversion is needed downstream. ## Threshold Modes diff --git a/examples/gpt-oss/README.md b/examples/gpt-oss/README.md index be3dea6b558..6779d7e29ba 100644 --- a/examples/gpt-oss/README.md +++ b/examples/gpt-oss/README.md @@ -115,7 +115,7 @@ Note: Model Optimizer currently exports quantized checkpoints in formats other t
Deployment on TensorRT-LLM -To setup TensorRT-LLM, follow the official guide: [Deploying GPT-OSS on TensorRT-LLM](https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/blogs/tech_blog/blog9_Deploying_GPT_OSS_on_TRTLLM.md) +To setup TensorRT-LLM, follow the official guide: [Deploying GPT-OSS on TensorRT-LLM](https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/blogs/tech_blog/blog09_Deploying_GPT_OSS_on_TRTLLM.md) Once installed, launch an OpenAI-compatible endpoint using: ```bash diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 7e3c7d8fc44..a0877b52e39 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -363,7 +363,7 @@ weight-only and weight & activation quantization) — the effective number of bi You may specify an `effective_bits` target such as 5.4 for mixed precision quantization using `NVFP4_DEFAULT_CFG` & `FP8_DEFAULT_CFG`. `AutoQuantize` will automatically quantize highly sensitive layers in `FP8_DEFAULT_CFG` while keeping less sensitive layers in `NVFP4_DEFAULT_CFG` (and even skip quantization for any extremely sensitive layers) so that -the the final mixed precision quantized model has an effective quantized bits of 5.4. This model would give a better accuracy than the model quantized with vanilla `NVFP4_DEFAULT_CFG` configuration since the more aggressive `NVFP4_DEFAULT_CFG` quantization was not applied for the highly sensitive layers. +the final mixed precision quantized model has an effective quantized bits of 5.4. This model would give a better accuracy than the model quantized with vanilla `NVFP4_DEFAULT_CFG` configuration since the more aggressive `NVFP4_DEFAULT_CFG` quantization was not applied for the highly sensitive layers. Here is an example usage for `AutoQuantize` algorithm (Please see [auto_quantize](https://nvidia.github.io/Model-Optimizer/reference/generated/modelopt.torch.quantization.model_quant.html#modelopt.torch.quantization.model_quant.auto_quantize) API for more details): @@ -732,7 +732,7 @@ with torch.inference_mode(): ### Build the TensorRT-LLM engines -After the TensorRT-LLM checkpoint export, you can use the `trtllm-build` build command to build the engines from the exported checkpoints. Please check the [TensorRT-LLM Build API](https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/architecture/workflow.md#build-apis) documentation for reference. +After the TensorRT-LLM checkpoint export, you can use the `trtllm-build` build command to build the engines from the exported checkpoints. Please check the [TensorRT-LLM Build API](https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/legacy/architecture/workflow.md#build-apis) documentation for reference. ## Pre-Quantized Checkpoints diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index 52241079d39..dc774b12da1 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -47,7 +47,11 @@ from modelopt.torch.export import has_spec_opt from modelopt.torch.export.model_utils import is_multimodal_model from modelopt.torch.models import hf_model_type -from modelopt.torch.models.hf import checkpoint_has_mtp, prepare_model_for_loading +from modelopt.torch.models.hf import ( + checkpoint_has_mtp, + prefers_bundled_code, + prepare_model_for_loading, +) from modelopt.torch.utils import distributed as dist_utils from modelopt.torch.utils.mlflow import Tool, resolved_recipe_texts, tracked_run from modelopt.torch.utils.plugins.model_load_utils import record_unplaced_source_keys @@ -722,7 +726,15 @@ def has_pack_quantized_config(config): # DeepSeek ships bundled modeling code, but the built-in class is what the # disk-offload and streaming-export paths are validated against. - use_bundled_code = trust_remote_code and "Deepseek" in architecture + # Also use it when the model plugin says the built-in class can't represent the checkpoint + # (dense Nemotron-H on Transformers 5.5, whose built-in class is MoE-only). + use_bundled_code = trust_remote_code and ( + "Deepseek" in architecture + or ( + hasattr(transformers, architecture) + and prefers_bundled_code(hf_config.model_type, hf_config, ckpt_path, config_kwargs) + ) + ) if not hasattr(transformers, architecture) or use_bundled_code: if not hasattr(transformers, architecture): diff --git a/examples/llm_distill/README.md b/examples/llm_distill/README.md index 834ec15ada9..08f800a62bb 100644 --- a/examples/llm_distill/README.md +++ b/examples/llm_distill/README.md @@ -11,7 +11,7 @@ This section focuses on demonstrating how to apply Model Optimizer to perform kn | **Section** | **Description** | **Link** | **Docs** | | :------------: | :------------: | :------------: | :------------: | | Pre-Requisites | Required & optional packages to use this technique | \[[Link](#pre-requisites)\] | | -| Getting Started | Learn how to optimize your models using distillation to produce more intellegant smaller models | \[[Link](#getting-started)\] | \[[docs](https://nvidia.github.io/Model-Optimizer/guides/4_distillation.html)\] | +| Getting Started | Learn how to optimize your models using distillation to produce more intelligent smaller models | \[[Link](#getting-started)\] | \[[docs](https://nvidia.github.io/Model-Optimizer/guides/4_distillation.html)\] | | Support Matrix | View the support matrix to see compatibility and feature availability across different models | \[[Link](#support-matrix)\] | | | Distillation with Megatron-Bridge | Learn how to distill your models with Megatron-Bridge Framework | \[[Link](#knowledge-distillation-kd-in-nvidia-megatron-bridge-framework)\] | \[[docs](https://nvidia.github.io/Model-Optimizer/guides/4_distillation.html)\] | | Distillation with Megatron-LM | Learn how to distill your models with Megatron-LM Framework | \[[Link](#knowledge-distillation-kd-in-nvidia-megatron-lm-framework)\] | | diff --git a/examples/llm_qat/README.md b/examples/llm_qat/README.md index cd88ae30498..0925d4dd1d9 100644 --- a/examples/llm_qat/README.md +++ b/examples/llm_qat/README.md @@ -39,7 +39,7 @@ The Qwen3-8B example below requires a minimum of **2 x 80GB GPUs**. ## Run End-to-End QAT/QAD Example All arguments can be set via YAML, CLI, or both (CLI overrides YAML). See -[ARGUMENTS.md](ARGUMENTS.md), `--help`, and [Configuration](#advanced-configuration). +[ARGUMENTS.md](ARGUMENTS.md), `--help`, and [Advanced Topics](#advanced-topics). ### Quantization Recipes diff --git a/examples/llm_qat/notebooks/QAT_QAD_Walkthrough.ipynb b/examples/llm_qat/notebooks/QAT_QAD_Walkthrough.ipynb index 6e65d2dcc9b..a9e486b39f6 100644 --- a/examples/llm_qat/notebooks/QAT_QAD_Walkthrough.ipynb +++ b/examples/llm_qat/notebooks/QAT_QAD_Walkthrough.ipynb @@ -23,7 +23,7 @@ "id": "c3f7f931-ac38-494e-aea8-ca2cd6d05794", "metadata": {}, "source": [ - "## Installing Prerequisites and Dependancies" + "## Installing Prerequisites and Dependencies" ] }, { @@ -31,14 +31,14 @@ "id": "d7d4f25f-e569-42cf-8022-bb7cc6f9ea6e", "metadata": {}, "source": [ - "If you haven't already, install the required dependencies for this notebook. Key dependancies include:\n", + "If you haven't already, install the required dependencies for this notebook. Key dependencies include:\n", "\n", "- nvidia-modelopt\n", "- torch\n", "- transformers\n", "- jupyterlab\n", "\n", - "This repo contains a `examples/llm_qat/notebooks/requirements.txt` file that can be used to install all required dependancies." + "This repo contains a `examples/llm_qat/notebooks/requirements.txt` file that can be used to install all required dependencies." ] }, { @@ -398,7 +398,7 @@ "id": "e471ef6c-1346-4e5e-8782-5e9f2bc38f8a", "metadata": {}, "source": [ - "Once you have quantized the model you can now start the post-training process for QAT. The training process will calculate validation loss at 50 step intervals and save the model. These can be controled by adjusting the `eval_steps` and `output_dir` above along with other `training_args`." + "Once you have quantized the model you can now start the post-training process for QAT. The training process will calculate validation loss at 50 step intervals and save the model. These can be controlled by adjusting the `eval_steps` and `output_dir` above along with other `training_args`." ] }, { @@ -779,7 +779,7 @@ "id": "eca645fb-d8d2-4c98-9cb2-afbae7d30d6c", "metadata": {}, "source": [ - "## Stop the TensorRT-LLM Docker contrainer" + "## Stop the TensorRT-LLM Docker container" ] }, { diff --git a/examples/llm_qat/requirements.txt b/examples/llm_qat/requirements.txt index 8b2ca580c07..1b9061b5596 100644 --- a/examples/llm_qat/requirements.txt +++ b/examples/llm_qat/requirements.txt @@ -1,5 +1,4 @@ -# Accelerate 1.14 regresses FSDP2 handling of shared parameters. -accelerate>=1.0.0,<1.14 +accelerate>=1.15.0 flash-attn>=2.6.0 liger-kernel>=0.5.0; platform_system != 'Darwin' and platform_system != 'Windows' # peft 0.21 stops writing adapter_model.safetensors where export.py expects it. diff --git a/examples/llm_sparsity/weight_sparsity/README.md b/examples/llm_sparsity/weight_sparsity/README.md index 42bac030ce9..2330595a270 100644 --- a/examples/llm_sparsity/weight_sparsity/README.md +++ b/examples/llm_sparsity/weight_sparsity/README.md @@ -133,7 +133,7 @@ The output should be similar to the following: ROUGE scores: {'rouge1': 42.174, 'rouge2': 19.2724, 'rougeL': 28.6989, 'rougeLsum': 39.1394} ``` -Please refer to [link](../llm_eval/README.md#Evaluation-scripts-for-LLM-tasks) for more details of how to evaluate the sparsified models on other benchmarks, such as MMLU and HumanEval. +Please refer to [link](../../llm_eval/README.md#evaluation-scripts-for-llm-tasks) for more details of how to evaluate the sparsified models on other benchmarks, such as MMLU and HumanEval. ## Export and Serve a Hugging Face Checkpoint diff --git a/examples/llm_sparsity/weight_sparsity/finetune.py b/examples/llm_sparsity/weight_sparsity/finetune.py index 6eb199adc50..892e08577af 100644 --- a/examples/llm_sparsity/weight_sparsity/finetune.py +++ b/examples/llm_sparsity/weight_sparsity/finetune.py @@ -30,6 +30,7 @@ from dataclasses import dataclass, field import torch +import torch.distributed.checkpoint.state_dict as dcp_state_dict import transformers import utils from torch.utils.data import Dataset @@ -206,11 +207,13 @@ def preprocess( def get_model_state_dict(trainer: transformers.Trainer): - """Collects the state dict.""" - state_dict = trainer.model.state_dict() - cpu_state_dict = {key: value.cpu() for key, value in state_dict.items()} - del state_dict - return cpu_state_dict + """Collects the full (unsharded) state dict on rank 0's CPU; must be called on all ranks.""" + # Under FSDP2 (the transformers 5.x default), model.state_dict() holds sharded DTensors that + # cannot be loaded into an unsharded model, so gather full tensors instead. + return dcp_state_dict.get_model_state_dict( + trainer.model, + options=dcp_state_dict.StateDictOptions(full_state_dict=True, cpu_offload=True), + ) class SupervisedDataset(Dataset): diff --git a/examples/model_hub/README.md b/examples/model_hub/README.md index 2653503c56f..0f70f0dc629 100644 --- a/examples/model_hub/README.md +++ b/examples/model_hub/README.md @@ -1,6 +1,6 @@ # Deploy quantized models from Nvidia's Hugging Face model hub with TensorRT-LLM, vLLM, and SGLang -The provided bash scripts show an example to deploy and run the [quantized LLama 3.1 8B Instruct FP8 model](https://huggingface.co/nvidia/Llama-3.1-8B-Instruct-FP8) from Nvidia's Hugging Face model hub on TensorRT-LLM, vLLM and SGLang respectively. +The provided bash scripts show an example to deploy and run the [quantized Llama 3.1 8B Instruct FP8 model](https://huggingface.co/nvidia/Llama-3.1-8B-Instruct-FP8) from Nvidia's Hugging Face model hub on TensorRT-LLM, vLLM and SGLang respectively. Before running the bash scripts, please make sure you have setup the environment properly: diff --git a/examples/onnx_ptq/autotune/README.md b/examples/onnx_ptq/autotune/README.md index d8a86b8bae0..14cc9772af7 100644 --- a/examples/onnx_ptq/autotune/README.md +++ b/examples/onnx_ptq/autotune/README.md @@ -2,7 +2,7 @@ This example demonstrates automated Q/DQ (Quantize/Dequantize) node placement optimization for ONNX models using TensorRT performance measurements. -> **Warning:** This example uses the direct Autotune entry point for lower-level Q/DQ placement experiments. If you are starting from an unquantized ONNX model and care about **accuracy**, please use the ONNX PTQ workflow with `--autotune` enabled and with representative calibration data. See [../README#optimize-qdq-node-placement-with-autotune](../README#optimize-qdq-node-placement-with-autotune). +> **Warning:** This example uses the direct Autotune entry point for lower-level Q/DQ placement experiments. If you are starting from an unquantized ONNX model and care about **accuracy**, please use the ONNX PTQ workflow with `--autotune` enabled and with representative calibration data. See [../README.md#optimize-qdq-node-placement-with-autotune](../README.md#optimize-qdq-node-placement-with-autotune). ## Table of Contents @@ -265,7 +265,7 @@ All examples above use the command-line interface. For **low-level programmatic **See the API Reference documentation for low-level usage:** -- [`docs/source/reference/2_qdq_placement.rst`](../../docs/source/reference/2_qdq_placement.rst) +- [`modelopt.onnx.quantization.autotune`](https://nvidia.github.io/Model-Optimizer/reference/generated/modelopt.onnx.quantization.autotune.html) The API docs include detailed examples of: @@ -278,13 +278,13 @@ The API docs include detailed examples of: For comprehensive documentation on QDQ placement optimization, see: -- **User Guide**: [`docs/source/guides/9_qdq_placement.rst`](../../docs/source/guides/9_qdq_placement.rst) +- **User Guide**: [`9_autotune`](https://nvidia.github.io/Model-Optimizer/guides/9_autotune.html) - Detailed explanations of how the autotuner works - Advanced usage patterns and best practices - Configuration options and performance tuning - Troubleshooting common issues -- **API Reference**: [`docs/source/reference/2_qdq_placement.rst`](../../docs/source/reference/2_qdq_placement.rst) +- **API Reference**: [`modelopt.onnx.quantization.autotune`](https://nvidia.github.io/Model-Optimizer/reference/generated/modelopt.onnx.quantization.autotune.html) - Complete API documentation for all classes and functions - Low-level usage examples - State management and pattern cache details diff --git a/examples/pruning/cifar_resnet.ipynb b/examples/pruning/cifar_resnet.ipynb index c5a34cb1167..a3fd850cac8 100644 --- a/examples/pruning/cifar_resnet.ipynb +++ b/examples/pruning/cifar_resnet.ipynb @@ -450,7 +450,7 @@ "source": [ "We have now established a baseline model and accuracy that we will be comparing with using Model Optimizer.\n", "\n", - "So far, we have seen a regular PyTorch model trained without anything new. Now, lets optimize the model for our target constraints using Model Optimizer!\n", + "So far, we have seen a regular PyTorch model trained without anything new. Now, let's optimize the model for our target constraints using Model Optimizer!\n", "\n", "## FastNAS Pruning with Model Optimizer\n", "\n", @@ -482,15 +482,15 @@ "source": [ "### Prune base model and store pruned net\n", "\n", - "Using [mtp.prune](../reference/generated/modelopt.torch.prune.pruning.rst#modelopt.torch.prune.pruning.prune) you can\n", + "Using [mtp.prune](https://nvidia.github.io/Model-Optimizer/reference/generated/modelopt.torch.prune.pruning.html#modelopt.torch.prune.pruning.prune) you can\n", "\n", "* generate a search space for pruning from your base model;\n", "* prune the model;\n", "* obtain a valid pytorch model that can be used for fine-tuning.\n", "\n", - "Let's say you have the ResNet20 model as our base model to prune from and we are looking for a model with at most 60M FLOPs. We can provide search constraints for `flops` and/or `params` by an upper bound. The values can either be absolute numbers (e.g. `60e6`) or a string percentage (e.g. `\"75%\"`). In addition, we should also provide our training data loader to [mtp.prune](../reference/generated/modelopt.torch.prune.pruning.rst#modelopt.torch.prune.pruning.prune). The training data loader will be used to calibrate the normalization layers in the model. Finally, we will also specify a custom config for configuring the pruning search space to get a more fine-grained selection of pruned nets.\n", + "Let's say you have the ResNet20 model as our base model to prune from and we are looking for a model with at most 60M FLOPs. We can provide search constraints for `flops` and/or `params` by an upper bound. The values can either be absolute numbers (e.g. `60e6`) or a string percentage (e.g. `\"75%\"`). In addition, we should also provide our training data loader to [mtp.prune](https://nvidia.github.io/Model-Optimizer/reference/generated/modelopt.torch.prune.pruning.html#modelopt.torch.prune.pruning.prune). The training data loader will be used to calibrate the normalization layers in the model. Finally, we will also specify a custom config for configuring the pruning search space to get a more fine-grained selection of pruned nets.\n", "\n", - "Finally, we can store the pruned architecture and weights using [mto.save](../reference/generated/modelopt.torch.opt.conversion.rst#modelopt.torch.opt.conversion.save).\n", + "Finally, we can store the pruned architecture and weights using [mto.save](https://nvidia.github.io/Model-Optimizer/reference/generated/modelopt.torch.opt.conversion.html#modelopt.torch.opt.conversion.save).\n", "\n", ".. note:: We are optimizing a relatively smaller model here. A finer-grained search could be more effective in such a case. This is why we are specifying custom configs. In general however, it is recommended to convert models with the default config itself." ] @@ -677,7 +677,7 @@ "source": [ "As we can see, the best subnet (59.3M FLOPs) has fitted our constraint of 60M FLOPs. We can also see a drop in validation accuracy for the searched model. This is very common after pruning and fine-tuning is necessary for this model.\n", "\n", - "#### Restore the pruned subnet using [mto.restore](../reference/generated/modelopt.torch.opt.conversion.rst#modelopt.torch.opt.conversion.restore)" + "#### Restore the pruned subnet using [mto.restore](https://nvidia.github.io/Model-Optimizer/reference/generated/modelopt.torch.opt.conversion.html#modelopt.torch.opt.conversion.restore)" ] }, { diff --git a/examples/puzzletron/main.py b/examples/puzzletron/main.py index f093e5b7e68..09bae07de9d 100644 --- a/examples/puzzletron/main.py +++ b/examples/puzzletron/main.py @@ -156,7 +156,7 @@ def run_mip_only(hydra_config_path: str): mtpz.mip.run_mip_sweep(hydra_cfg) else: # mip_and_realize_models (distributed processing) - # TODO: How to make it part of mnt.search() api, similarly to run_full_puzzletron() API + # TODO: How to make it part of mtn.search() api, similarly to run_full_puzzletron() API mtpz.tools.mprint("Puzzletron Progress 7/8: running MIP and realizing models (multi-gpu)") mtpz.mip.launch_mip_and_realize_model(hydra_cfg) diff --git a/examples/specdec_bench/README.md b/examples/specdec_bench/README.md index 9a2d2c464df..f8e69d47416 100644 --- a/examples/specdec_bench/README.md +++ b/examples/specdec_bench/README.md @@ -2,7 +2,7 @@ ## Installation -This benchmark is meant to be a lightweight layer ontop of an existing vLLM/SGLang/TRTLLM installation. For example, no install +This benchmark is meant to be a lightweight layer on top of an existing vLLM/SGLang/TRTLLM installation. For example, no install is required if one is running in the following dockers: `vllm/vllm-openai:v0.11.0` (vLLM), `lmsysorg/sglang:v0.5.4.post2` (SGLang), or `nvcr.io/nvidia/tensorrt-llm/release:1.2.0` (TRT-LLM). diff --git a/examples/speculative_decoding/README.md b/examples/speculative_decoding/README.md index 23ba03f6e39..d1811659e04 100644 --- a/examples/speculative_decoding/README.md +++ b/examples/speculative_decoding/README.md @@ -92,9 +92,9 @@ The saved modelopt checkpoint is similar in architecture to HF models. It can be For large models, you can export intermediate hidden states to disk and train only the draft model. This significantly reduces GPU memory requirements, but requires several to tens of terabytes of disk storage depending on dataset size. -### Dumpping Hidden States to Disk +### Dumping Hidden States to Disk -We support two backends for generating base model hidden states. For better effciency, it is recommended to use TRT-LLM: +We support two backends for generating base model hidden states. For better efficiency, it is recommended to use TRT-LLM: ```bash python collect_hidden_states/compute_hidden_states_trtllm.py \ @@ -187,7 +187,7 @@ Please refer to [TRT-LLM Doc: Speculative Decoding](https://nvidia.github.io/Ten Please refer to [VLLM Doc: Speculative Decoding](https://docs.vllm.ai/en/latest/features/spec_decode/) for detailed usage. -Optionally, you can convert the exported checkpoint to contain target model information, which is accepted by vLLM to simplify depployment: +Optionally, you can convert the exported checkpoint to contain target model information, which is accepted by vLLM to simplify deployment: ```bash python scripts/convert_to_vllm_ckpt.py --input --verifier --output diff --git a/examples/vllm_serve/Dockerfile b/examples/vllm_serve/Dockerfile index 8bb3f8c052d..2f6af8e62ac 100644 --- a/examples/vllm_serve/Dockerfile +++ b/examples/vllm_serve/Dockerfile @@ -27,13 +27,13 @@ RUN cd Model-Optimizer && \ # Llama4 requires this RUN pip install flash-attn==2.7.4.post1 --no-build-isolation -# Pre-compile CUDA extensions into a world-accessible directory so the vllm -# user can use the cache at runtime. +# Pre-compile CUDA extensions so the vllm user can use the cache at runtime. ENV TORCH_EXTENSIONS_DIR=/workspace/torch_extensions RUN python3 -c "import modelopt.torch.quantization.extensions as ext; ext.precompile()" || true -# Allow the non-root vllm user to access the workspace -RUN chmod -R 777 /workspace +# Only the vllm user (and root) may write here, since this code is loaded into the serving process. +# The base image's vllm user has GID 0 and no vllm group exists, so chown the user only. +RUN chown -R vllm /workspace && chmod -R go-w /workspace USER vllm diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index 8eaea64700d..05ebd4f4e1b 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -23,6 +23,17 @@ docker build -f examples/vllm_serve/Dockerfile -t vllm-modelopt:v0.30.0 . The Dockerfile defaults to vLLM 0.30.0. Override `VLLM_VERSION` only when building for a different tested release. +The image runs as the `vllm` user, which owns `/workspace`. To run it as a different user, point +`TORCH_EXTENSIONS_DIR` at a directory that user can write, for example a mounted volume. ModelOpt +then builds its CUDA extensions there on first use rather than under `/workspace`; without it, +ModelOpt warns and falls back to slower eager fake quantization. + +```bash +mkdir -p ext_cache +docker run --user $(id -u) -v "$PWD/ext_cache:/ext_cache" -e TORCH_EXTENSIONS_DIR=/ext_cache \ + vllm-modelopt:v0.30.0 +``` + For a direct installation from the ModelOpt repository root, install the tested vLLM release and the ModelOpt extras used by this example: diff --git a/examples/windows/diffusers/qad_example/README.md b/examples/windows/diffusers/qad_example/README.md index af579325062..b102ad6db48 100644 --- a/examples/windows/diffusers/qad_example/README.md +++ b/examples/windows/diffusers/qad_example/README.md @@ -14,7 +14,7 @@ > (including quantized or distilled checkpoints) remain subject to the LTX Community License > Agreement and are **not** covered by Apache 2.0. -**Note:** This is a **sample script for illustrating the QAD pipeline**. It has been verified to run on a **Linux RTX 5090** system, but runs into **OOM (Out of Memory)** on that configuration. +**Note:** This is a **sample script for illustrating the QAD pipeline**. It has been tested on a **Linux RTX 5090** system and currently runs into **OOM (Out of Memory)** on that configuration. This example demonstrates **Quantization-Aware Distillation (QAD)** for [LTX-2](https://github.com/Lightricks/LTX-2) using the native LTX training loop and [NVIDIA ModelOpt](https://github.com/NVIDIA/Model-Optimizer). It combines: @@ -77,7 +77,7 @@ pip install torch accelerate safetensors pyyaml ### 1. Prepare your dataset -Run the LTX preprocessing script to extract latents and text embeddings from your videos. Use `preprocess_dataset.py` with the following arguments (matching the LTX training pipeline): +Run the LTX preprocessing script to extract latents and text embeddings from your videos. Use `process_dataset.py` with the following arguments (matching the LTX training pipeline): ```bash python scripts/process_dataset.py /path/to/dataset.json \ diff --git a/examples/windows/onnx_ptq/genai_llm/README.md b/examples/windows/onnx_ptq/genai_llm/README.md index a03d4f94fb9..be64ad1f64b 100644 --- a/examples/windows/onnx_ptq/genai_llm/README.md +++ b/examples/windows/onnx_ptq/genai_llm/README.md @@ -128,7 +128,7 @@ The quantization strategy selects which layers to quantize to INT8 vs INT4: 1. **INT8 Layers** (Higher Precision): Important layers that significantly impact model quality. Quantized per-channel -2. **INT4 Layers** (Maximum Compression): All other layers. Qunatized blockwise. +2. **INT4 Layers** (Maximum Compression): All other layers. Quantized blockwise. This strategy preserves accuracy for the most sensitive layers while maintaining aggressive compression elsewhere. @@ -244,4 +244,4 @@ Refer to the following example scripts and tutorials for deployment: 1. **Error - Invalid Position-IDs input to the ONNX model** - The ONNX models produced using ONNX GenerativeAI (GenAI) have different IO bindings for models produced using different execution-providers (EPs). For instance, model built with DML EP has position-ids input in the ONNX model but models builts using CUDA EP or NvTensorRtRtx EP don't have position-ids inputs. So, if base model requires, use `--add_position_ids` command-line argument for enabling position_ids calibration input or set "add_position_ids" variable to `True` value (hard-code) in the quantize script if required. + The ONNX models produced using ONNX GenerativeAI (GenAI) have different IO bindings for models produced using different execution-providers (EPs). For instance, model built with DML EP has position-ids input in the ONNX model but models built using CUDA EP or NvTensorRtRtx EP don't have position-ids inputs. So, if base model requires, use `--add_position_ids` command-line argument for enabling position_ids calibration input or set "add_position_ids" variable to `True` value (hard-code) in the quantize script if required. diff --git a/examples/windows/onnx_ptq/sam2/README.md b/examples/windows/onnx_ptq/sam2/README.md index 8d9e5fed3eb..77ea8032072 100644 --- a/examples/windows/onnx_ptq/sam2/README.md +++ b/examples/windows/onnx_ptq/sam2/README.md @@ -1,6 +1,6 @@ # Scripts for ONNX SAM2 model -This repository contains an example to demontrate 8-bit quantization of SAM2 ONNX model. +This repository contains an example to demonstrate 8-bit quantization of SAM2 ONNX model. ## Table of Contents diff --git a/examples/windows/onnx_ptq/whisper/README.md b/examples/windows/onnx_ptq/whisper/README.md index 82ae7822001..4d257735adb 100644 --- a/examples/windows/onnx_ptq/whisper/README.md +++ b/examples/windows/onnx_ptq/whisper/README.md @@ -1,6 +1,6 @@ # Scripts for ONNX Whisper model -This repository contains an example to demontrate 8-bit quantization of Whisper ONNX model. +This repository contains an example to demonstrate 8-bit quantization of Whisper ONNX model. ## Table of Contents diff --git a/modelopt/onnx/op_types.py b/modelopt/onnx/op_types.py index 637c0ad7a45..aa2ac4ef374 100644 --- a/modelopt/onnx/op_types.py +++ b/modelopt/onnx/op_types.py @@ -175,7 +175,7 @@ def is_non_reshape_copy_op(op_type: str): def is_irregular_mem_access_op(op_type: str): - """Returns whether the given op type is of Irreggular mem access category or not.""" + """Returns whether the given op type is of irregular mem access category or not.""" return op_type in [ "Gather", "GatherElements", diff --git a/modelopt/onnx/quantization/autotune/__main__.py b/modelopt/onnx/quantization/autotune/__main__.py index ccb826c291b..7c0357ca93e 100644 --- a/modelopt/onnx/quantization/autotune/__main__.py +++ b/modelopt/onnx/quantization/autotune/__main__.py @@ -17,10 +17,10 @@ import argparse import sys -import tempfile from pathlib import Path from modelopt.onnx.logging_config import logger +from modelopt.onnx.quantization.autotune.benchmark import default_timing_cache_file from modelopt.onnx.quantization.autotune.utils import ( DEFAULT_NUM_SCHEMES, DEFAULT_TIMING_RUNS, @@ -38,7 +38,6 @@ DEFAULT_OUTPUT_DIR = "./autotuner_output" DEFAULT_QUANT_TYPE = "int8" DEFAULT_DQ_DTYPE = "float32" -DEFAULT_TIMING_CACHE = str(Path(tempfile.gettempdir()) / "trtexec_timing.cache") def apply_mode_presets(args) -> None: @@ -94,6 +93,7 @@ def run_autotune() -> int: """ args = get_parser().parse_args() apply_mode_presets(args) + args.timing_cache = args.timing_cache or default_timing_cache_file() model_path = validate_file_path(args.onnx_path, "Model file") validate_file_path(args.qdq_baseline, "QDQ baseline model") output_dir = Path(args.output_dir) @@ -279,8 +279,8 @@ def get_parser() -> argparse.ArgumentParser: trt_group.add_argument( "--timing_cache", type=str, - default=DEFAULT_TIMING_CACHE, - help=f"TensorRT timing cache file (default: {DEFAULT_TIMING_CACHE})", + default=None, + help="TensorRT timing cache file (default: a file in a private temporary directory)", ) trt_group.add_argument( "--warmup_runs", diff --git a/modelopt/onnx/quantization/autotune/benchmark.py b/modelopt/onnx/quantization/autotune/benchmark.py index 6c62ee4e628..d16220b8094 100644 --- a/modelopt/onnx/quantization/autotune/benchmark.py +++ b/modelopt/onnx/quantization/autotune/benchmark.py @@ -27,6 +27,7 @@ """ import ctypes +import functools import importlib.util import os import re @@ -50,6 +51,16 @@ TORCH_CUDA_AVAILABLE = torch.cuda.is_available() +@functools.cache +def default_timing_cache_file() -> str: + """Return the default timing-cache path: a file in one private temporary directory per process. + + A fixed path under the shared temp dir would let another local user pre-create the file or plant + a symlink there. Pass ``timing_cache_file`` to keep a cache across runs. + """ + return str(Path(tempfile.mkdtemp(prefix="modelopt_autotune_")) / "trtexec_timing.cache") + + def _validate_shape_range(min_shape: list, opt_shape: list, max_shape: list) -> None: """Raise ValueError if shape lengths differ or if min <= opt <= max fails at any dimension.""" if len(min_shape) != len(opt_shape) or len(opt_shape) != len(max_shape): @@ -92,7 +103,8 @@ def __init__( Args: timing_cache_file: Path to timing cache file to accelerate engine builds. - If None, uses '/tmp/trtexec_timing.cache' as default. + If None, uses 'trtexec_timing.cache' in a private temporary directory + that is created once per process. warmup_runs: Number of warmup iterations before timing measurements. timing_runs: Number of iterations for latency measurement. Results are averaged across these runs. @@ -100,7 +112,7 @@ def __init__( These plugins will be loaded during engine building. If None, no custom plugins are loaded. """ - self.timing_cache_file = timing_cache_file or "/tmp/trtexec_timing.cache" # nosec B108 + self.timing_cache_file = timing_cache_file or default_timing_cache_file() self.warmup_runs = warmup_runs self.timing_runs = timing_runs self.plugin_libraries = plugin_libraries or [] @@ -328,7 +340,8 @@ def __init__( Args: timing_cache_file: Path to TensorRT timing cache file. If None, - defaults to '/tmp/trtexec_timing.cache'. + defaults to 'trtexec_timing.cache' in a private temporary directory + that is created once per process. warmup_runs: Number of warmup iterations before timing measurements. timing_runs: Number of iterations for latency measurement. plugin_libraries: List of paths to TensorRT plugin shared libraries (.so files). diff --git a/modelopt/onnx/quantization/autotune/workflows.py b/modelopt/onnx/quantization/autotune/workflows.py index a863d323527..21ee87313d9 100644 --- a/modelopt/onnx/quantization/autotune/workflows.py +++ b/modelopt/onnx/quantization/autotune/workflows.py @@ -29,7 +29,11 @@ from modelopt.onnx.logging_config import logger from modelopt.onnx.quantization.autotune.autotuner import QDQAutotuner -from modelopt.onnx.quantization.autotune.benchmark import TensorRTPyBenchmark, TrtExecBenchmark +from modelopt.onnx.quantization.autotune.benchmark import ( + TensorRTPyBenchmark, + TrtExecBenchmark, + default_timing_cache_file, +) from modelopt.onnx.quantization.autotune.common import Config, PatternCache from modelopt.onnx.quantization.qdq_utils import get_quantized_tensors @@ -97,7 +101,8 @@ def init_benchmark_instance( These plugins will be loaded by trtexec or TensorRT Python API during engine building. If None, no custom plugins are loaded. timing_cache_file: Path to TensorRT timing cache file for faster engine builds. - If None, uses default "trtexec_timing.cache" (default: None) + If None, uses "trtexec_timing.cache" in a private temporary directory + (default: None) warmup_runs: Number of warmup inference iterations before measurement. Allows GPU to reach stable performance state (default: 5) timing_runs: Number of timed inference iterations for latency measurement. @@ -126,7 +131,8 @@ def init_benchmark_instance( logger.info("TensorRT Python API benchmark initialized") logger.debug( f"Settings: warmup={warmup_runs}, timing={timing_runs}, " - f"cache={timing_cache_file or 'trtexec_timing.cache'}, plugin_libraries={plugin_libraries}" + f"cache={timing_cache_file or default_timing_cache_file()}, " + f"plugin_libraries={plugin_libraries}" ) return _benchmark_instance except Exception as e: diff --git a/modelopt/onnx/quantization/extensions.py b/modelopt/onnx/quantization/extensions.py index ba2c0e137dd..4bafe482e67 100644 --- a/modelopt/onnx/quantization/extensions.py +++ b/modelopt/onnx/quantization/extensions.py @@ -29,7 +29,7 @@ sys.path.append(path) round_and_pack_ext = cppimport.imp("modelopt_round_and_pack_ext") sys.path.remove(path) -except Exception as e: +except (Exception, SystemExit) as e: # setuptools exits on a failed build (e.g. no C++ compiler) logger.warning( f"{e}\nUnable to load `modelopt_round_and_pack_ext', falling back to python based optimized version. " ) diff --git a/modelopt/torch/_deploy/_runtime/tensorrt/constants.py b/modelopt/torch/_deploy/_runtime/tensorrt/constants.py index d9ace1645a3..eecc7073e2c 100644 --- a/modelopt/torch/_deploy/_runtime/tensorrt/constants.py +++ b/modelopt/torch/_deploy/_runtime/tensorrt/constants.py @@ -31,9 +31,6 @@ ONE_MEBI_IN_BYTES = 1 << 20 ONE_GIBI_IN_BYTES = 1 << 30 -# TensorRT conversion tool names -DEFAULT_ARTIFACT_DIR = "modelopt_build/trt_artifacts" - # Default conversion params DEFAULT_VALIDATION_THRESHOLD = 1e-4 DEFAULT_MAX_BATCH_SIZE = 1 diff --git a/modelopt/torch/_deploy/_runtime/tensorrt/engine_builder.py b/modelopt/torch/_deploy/_runtime/tensorrt/engine_builder.py index bb8bbd292b8..d00272ce158 100644 --- a/modelopt/torch/_deploy/_runtime/tensorrt/engine_builder.py +++ b/modelopt/torch/_deploy/_runtime/tensorrt/engine_builder.py @@ -13,17 +13,17 @@ # See the License for the specific language governing permissions and # limitations under the License. +import functools import logging import shutil import sys from pathlib import Path -from tempfile import TemporaryDirectory, gettempdir +from tempfile import TemporaryDirectory, mkdtemp from ..._runtime.common import read_bytes, timeit, write_bytes, write_string from ..._runtime.tensorrt.layerwise_profiling import process_layerwise_result from ...utils import OnnxBytes from .constants import ( - DEFAULT_ARTIFACT_DIR, DEFAULT_NUM_INFERENCE_PER_RUN, SHA_256_HASH_LENGTH, TRT_MODE_FLAGS, @@ -135,6 +135,27 @@ def _update_dynamic_shapes(dynamic_shapes: dict, cmd: list[str]) -> None: cmd.extend([f"--{shape}={convert_shape_to_string(dynamic_shapes[shape])}"]) +@functools.cache +def _default_artifact_dir() -> Path: + """Return this process's private temporary directory for builds without an ``output_dir``. + + The default must not be a fixed path under the shared temp dir: another local user could + pre-create it, plant symlinks in it, or swap the engine between its write and read-back. One + unpredictable directory per process keeps timing-cache reuse across calls and bounds what is + left behind to one directory, however many engines the process builds. + """ + return Path(mkdtemp(prefix="modelopt_trt_")) + + +def _artifact_dir(output_dir: Path | None) -> Path: + """Return ``output_dir`` (created if missing), or this process's private temporary directory.""" + if not output_dir: + return _default_artifact_dir() + path = Path(output_dir) + path.mkdir(parents=True, exist_ok=True) + return path + + @timeit def build_engine( onnx_bytes: OnnxBytes, @@ -183,7 +204,8 @@ def build_engine( For more details, please refer to: https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/python-api/infer/Core/BuilderConfig.html output_dir: Directory to save the engine file and trtexec artifacts. - If not provided, the artifacts will be saved in a temporary directory. + If not provided, this process's private temporary directory is used (created on first + use, not removed at exit). verbose: Set trtexec command to verbose or not. Returns: @@ -230,8 +252,7 @@ def _setup_files_and_paths( onnx_bytes.write_to_disk(str(tmp_onnx_dir)) onnx_path = tmp_onnx_dir / f"{onnx_bytes.model_name}.onnx" - final_output_dir = Path(output_dir or Path(gettempdir()) / DEFAULT_ARTIFACT_DIR) - final_output_dir.mkdir(parents=True, exist_ok=True) + final_output_dir = _artifact_dir(output_dir) engine_path = ( Path(engine_path) if engine_path @@ -294,7 +315,8 @@ def profile_engine( "maxShapes": {"input": [32,3,244,244]} } output_dir: Directory to save the engine file and trtexec artifacts. - If not provided, the artifacts will be saved in a temporary directory. + If not provided, this process's private temporary directory is used (created on first + use, not removed at exit). Returns: Layerwise profiling output as a json string. Stdout log produced by trtexec tool. @@ -322,8 +344,7 @@ def _setup_files_and_paths(tmp_dir_path: Path, engine_hash: str) -> tuple[Path, engine_path = tmp_dir_path / f"{engine_hash}-model.engine" write_bytes(engine_bytes[SHA_256_HASH_LENGTH:], engine_path) - final_output_dir = Path(output_dir or Path(gettempdir()) / DEFAULT_ARTIFACT_DIR) - final_output_dir.mkdir(parents=True, exist_ok=True) + final_output_dir = _artifact_dir(output_dir) profile_path = final_output_dir / f"{engine_hash}-profile.json" layer_info_path = final_output_dir / f"{engine_hash}-layerInfo.json" diff --git a/modelopt/torch/_deploy/_runtime/trt_client.py b/modelopt/torch/_deploy/_runtime/trt_client.py index a9c300eca5f..1135dc1a205 100644 --- a/modelopt/torch/_deploy/_runtime/trt_client.py +++ b/modelopt/torch/_deploy/_runtime/trt_client.py @@ -180,9 +180,6 @@ def initialize_input_output_tensors(self, engine): tensor_name, output_tensors[idx - len(input_tensors)].data_ptr(), ) - assert self.execution_context.all_shape_inputs_specified, ( - "Not all shape inputs are specified." - ) # Set selected profile idx self.execution_context.set_optimization_profile_async(0, self.stream.cuda_stream) diff --git a/modelopt/torch/distill/distillation_model.py b/modelopt/torch/distill/distillation_model.py index 3d4fc1cb6a3..2bd2a2d4094 100644 --- a/modelopt/torch/distill/distillation_model.py +++ b/modelopt/torch/distill/distillation_model.py @@ -166,22 +166,25 @@ def hide_loss_modules(self, enable=True): @contextmanager def only_teacher_forward(self, enable=True): """Context manager to temporarily disable forward passes on the student model.""" + # restore the previous mode so nested contexts keep working + only_teacher_fwd = self._only_teacher_fwd if enable: self._only_teacher_fwd = True try: yield finally: - self._only_teacher_fwd = False + self._only_teacher_fwd = only_teacher_fwd @contextmanager def only_student_forward(self, enable=True): """Context manager to temporarily run forward passes only on the student model.""" + only_student_fwd = self._only_student_fwd if enable: self._only_student_fwd = True try: yield finally: - self._only_student_fwd = False + self._only_student_fwd = only_student_fwd def train(self, mode: bool = True): """Override to prevent warnings of stored intermediate outputs in future forwards.""" diff --git a/modelopt/torch/distill/layerwise_distillation_model.py b/modelopt/torch/distill/layerwise_distillation_model.py index 4a98f5060c5..a8d52b8d988 100644 --- a/modelopt/torch/distill/layerwise_distillation_model.py +++ b/modelopt/torch/distill/layerwise_distillation_model.py @@ -34,7 +34,7 @@ class LayerwiseDistillationModel(DistillationModel): """Meta-model wrapper to support layerwise-enabled knowledge-distillation learning. The LayerwiseDistillationModel is a subclass of the DistillationModel that injects teacher inputs - into the corresponding student layers. This accomodates the case where the student model is the + into the corresponding student layers. This accommodates the case where the student model is the teacher with specific submodules replaced, which now need to be trained to mimic the original submodule in the teacher. """ diff --git a/modelopt/torch/distill/loss_balancers.py b/modelopt/torch/distill/loss_balancers.py index 0ceaca3e3c2..0f199d860e7 100644 --- a/modelopt/torch/distill/loss_balancers.py +++ b/modelopt/torch/distill/loss_balancers.py @@ -66,7 +66,7 @@ def forward(self, loss: dict[str, torch.Tensor]) -> torch.Tensor: ``{"student_loss": torch.tensor(...), "MSELoss_0": torch.tensor(...), "MSELoss_1": torch.tensor(...)}``. Returns: - The total loss after balancing student and kd loss loss components. + The total loss after balancing student and kd loss components. """ raise NotImplementedError @@ -116,7 +116,7 @@ def forward(self, loss: dict[str, torch.Tensor]) -> torch.Tensor: loss: The loss dict to aggregate. Returns: - The total loss after balancing student and kd loss loss components. + The total loss after balancing student and kd loss components. """ kd_loss_dict = loss.copy() diff --git a/modelopt/torch/distill/plugins/megatron.py b/modelopt/torch/distill/plugins/megatron.py index 26b3c9cc5c8..63745ed48e0 100644 --- a/modelopt/torch/distill/plugins/megatron.py +++ b/modelopt/torch/distill/plugins/megatron.py @@ -549,6 +549,11 @@ def forward(self, predictions: Tensor, targets: Tensor) -> Tensor: kl = q.exp() * (q - p) loss = torch.sum(mask * kl, dim=-1) + if tp_size > 1: + # Autograd collectives sum gradients from identical losses on every TP rank. + # Average only backward, preserving reported loss and loss-balancer weights. + loss = loss.detach() + (loss - loss.detach()) / tp_size + # No need to reduce since all ranks compute same global Top-K return self.post_forward(loss, tp_reduce=False) diff --git a/modelopt/torch/export/diffusers_utils.py b/modelopt/torch/export/diffusers_utils.py index 0309c83e1f6..de4d6966a61 100644 --- a/modelopt/torch/export/diffusers_utils.py +++ b/modelopt/torch/export/diffusers_utils.py @@ -295,6 +295,7 @@ def _wan_inputs() -> dict[str, torch.Tensor]: # Requires: hidden_states, encoder_hidden_states, timestep in_channels = getattr(cfg, "in_channels", 16) text_dim = getattr(cfg, "text_dim", 4096) + image_dim = getattr(cfg, "image_dim", None) max_seq_len = getattr(cfg, "rope_max_seq_len", 512) patch_dtype = getattr(getattr(model, "patch_embedding", None), "weight", None) @@ -308,13 +309,14 @@ def _wan_inputs() -> dict[str, torch.Tensor]: # Wan expects num_frames = 4 * n + 1; keep n small for dummy forward num_frames = 5 - text_seq_len = min(max_seq_len, 512) + # I2V attention splits image context from the final 512 text tokens. + text_seq_len = 512 if image_dim is not None else min(max_seq_len, 512) # Keep spatial dims small and divisible by patch size (default 2x2) height = 8 width = 8 - return { + dummy_inputs = { "hidden_states": torch.randn( batch_size, in_channels, num_frames, height, width, device=device, dtype=patch_dtype ), @@ -324,6 +326,21 @@ def _wan_inputs() -> dict[str, torch.Tensor]: "timestep": torch.randint(0, 1000, (batch_size,), device=device), "return_dict": False, } + if image_dim is not None: + image_embedder = model.condition_embedder.image_embedder + # The projection bias retains its compute dtype when weights are packed. + image_dtype = image_embedder.ff.net[0].proj.bias.dtype + image_batch_size = batch_size + image_seq_len = 257 + pos_embed_seq_len = getattr(cfg, "pos_embed_seq_len", None) + if pos_embed_seq_len is not None: + # FLF2V concatenates first/last-frame features inside the image embedder. + image_batch_size *= 2 + image_seq_len = pos_embed_seq_len // 2 + dummy_inputs["encoder_hidden_states_image"] = torch.randn( + image_batch_size, image_seq_len, image_dim, device=device, dtype=image_dtype + ) + return dummy_inputs def _qwen_inputs() -> dict[str, Any]: # QwenImageTransformer2DModel does NOT take the standard @@ -717,7 +734,7 @@ def get_diffusion_components( @contextmanager -def hide_quantizers_from_state_dict(model: nn.Module): +def hide_quantizers_from_state_dict(model: nn.Module, export_attention_quantizers: bool = False): """Context manager that temporarily removes quantizer modules from the model. This allows save_pretrained to save the model without quantizer buffers like _amax. @@ -725,6 +742,8 @@ def hide_quantizers_from_state_dict(model: nn.Module): Args: model: The model with quantizers to temporarily hide. + export_attention_quantizers: Keep enabled Q/K/V BMM quantizers for runtimes + that consume their legacy ``*_bmm_quantizer._amax`` keys. Yields: None - the model can be saved within the context. @@ -739,8 +758,23 @@ def hide_quantizers_from_state_dict(model: nn.Module): # first since we mutate the module tree while iterating. for name, module in list(model.named_modules()): backup = {} - for attr in ["weight_quantizer", "input_quantizer", "output_quantizer"]: + for attr in [ + "weight_quantizer", + "input_quantizer", + "output_quantizer", + "q_bmm_quantizer", + "k_bmm_quantizer", + "v_bmm_quantizer", + "softmax_quantizer", + "bmm2_output_quantizer", + ]: child = getattr(module, attr, None) + if ( + export_attention_quantizers + and attr in ("q_bmm_quantizer", "k_bmm_quantizer", "v_bmm_quantizer") + and getattr(child, "is_enabled", False) + ): + continue if isinstance(child, nn.Module): backup[attr] = child delattr(module, attr) diff --git a/modelopt/torch/export/layerwise_export.py b/modelopt/torch/export/layerwise_export.py index 92f31011ea3..2ff3abe966f 100644 --- a/modelopt/torch/export/layerwise_export.py +++ b/modelopt/torch/export/layerwise_export.py @@ -52,9 +52,11 @@ _fuse_shared_input_modules, _prepare_moe_inputs, _resolve_export_dtype, + _source_checkpoint, _write_hf_export_config, collect_shared_input_modules, save_non_weight_artifacts, + unplaced_keys_for, ) from .unified_export_hf_streaming import _assert_no_split_rules @@ -181,6 +183,7 @@ def __init__( self._dtype = dtype self._bound = False self._finalized = False + self._source_key_snapshot: tuple[str, list[str]] | None = None self._announced_on: list[nn.Module] = [] # A VLM calibrates its language model but exports the whole thing, so announce on # both: whichever of the two mtq.quantize is handed will find this exporter. @@ -272,6 +275,11 @@ def bind(self, calibrated_layers: list[nn.Module]) -> None: "match the original HF hub checkpoint." ) + # Packing removes fused source parameters; account for them while the tree is intact. + source = _source_checkpoint(model) + if source and Path(source).is_dir(): + self._source_key_snapshot = (str(source), unplaced_keys_for(model, source)) + self._bound = True def export_layer( diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index 88092f689dc..f19142cb650 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -897,6 +897,10 @@ def to_quantized_weight( # Clear GPU cache to avoid pontential GPU OOM issues for large models. clear_cuda_cache() return (weight / weights_scaling_factor.unsqueeze(-1)).to(torch.float8_e4m3fn) + # A 0-d scale would not promote a bf16 weight, so the division would round in bf16 instead of + # float32 like fake quantization. A 1-d scale promotes in-kernel, without a float32 copy. + if weights_scaling_factor.dim() == 0: + weights_scaling_factor = weights_scaling_factor.reshape(1) return (weight / weights_scaling_factor).to(torch.float8_e4m3fn) if quantization in [QUANTIZATION_INT8_SQ, QUANTIZATION_INT8_WO]: @@ -1758,6 +1762,8 @@ def _get_unquantized_moe_router_names(model: nn.Module) -> list[str]: def get_quant_config( model: nn.Module, is_modelopt_qlora: bool = False, + *, + include_kv_cache: bool = True, ) -> dict[str, Any]: """Generate quantization config for a model. @@ -1767,6 +1773,8 @@ def get_quant_config( Args: model: The PyTorch model to make config for. is_modelopt_qlora: Whether the model is a modelopt-trained QLoRA model. + include_kv_cache: Infer KV-cache metadata from attention quantizers. Disable + for diffusion models, whose attention quantizers do not represent a cache. Returns: Dictionary containing the quantization configuration @@ -1883,7 +1891,7 @@ def get_quant_config( for quantizer in (k_quantizer, v_quantizer, output_quantizer) if is_kv_boundary and quantizer is not None and quantizer.is_enabled ] - if enabled_kv_quantizers and is_language_model_module: + if include_kv_cache and enabled_kv_quantizers and is_language_model_module: module_kv_quant = get_kv_cache_dtype(module) if module_kv_quant != QUANTIZATION_NONE: kv_cache_formats.add(module_kv_quant) diff --git a/modelopt/torch/export/trtllm/layer_utils.py b/modelopt/torch/export/trtllm/layer_utils.py index 8f02ad7298f..006355f8e70 100755 --- a/modelopt/torch/export/trtllm/layer_utils.py +++ b/modelopt/torch/export/trtllm/layer_utils.py @@ -26,7 +26,7 @@ try: from transformers.activations import ACT2FN except Exception: - warn("Cannot find transformers package. Hugginface modules cannot be exported.") + warn("Cannot find transformers package. Hugging Face modules cannot be exported.") from modelopt.torch.models import is_moe from modelopt.torch.utils import distributed as dist diff --git a/modelopt/torch/export/trtllm/model_config.py b/modelopt/torch/export/trtllm/model_config.py index c3d630f9fc9..1bee58298f7 100755 --- a/modelopt/torch/export/trtllm/model_config.py +++ b/modelopt/torch/export/trtllm/model_config.py @@ -108,7 +108,7 @@ class LinearConfig: # For methods like W4A8 AWQ, we have two quantizers for weights # For W4A8, the first quantizer is for INT4 quantization and the second quantizer is for FP8 quantization - # `weight_scaling_factor_2` is the scaling factor the the second FP8 quantizer + # `weight_scaling_factor_2` is the scaling factor for the second FP8 quantizer weights_scaling_factor_2: torch.Tensor = None prequant_scaling_factor: torch.Tensor = None @@ -319,7 +319,7 @@ class RelativeAttentionTableConfig: class AttentionConfig: """The attention layer config.""" - # QKV can either be stored as splitted (for easier postprocessing) + # QKV can either be stored as split (for easier postprocessing) # or merged (for TRT LLM export) qkv: QKVConfig | LinearConfig = None dense: LinearConfig = None diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index f605ddaf636..842ce4dacce 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -64,6 +64,13 @@ except ImportError: HAS_DIFFUSERS = False +try: + from transformers import PreTrainedModel + + _TRANSFORMERS_MODEL_TYPES: tuple[type, ...] = (PreTrainedModel,) +except ImportError: + _TRANSFORMERS_MODEL_TYPES = () + from modelopt.torch.opt.conversion import ModeloptStateManager, modelopt_state from modelopt.torch.opt.plugins.huggingface import _MODELOPT_STATE_SAVE_NAME @@ -1242,6 +1249,8 @@ def _export_diffusers_checkpoint( export_dir: Path, components: list[str] | None, max_shard_size: int | str = "10GB", + export_attention_quantizers: bool | None = None, + save_modelopt_state: bool = False, **kwargs, ) -> None: """Internal: Export diffusion(-like) model/pipeline checkpoint. @@ -1259,6 +1268,11 @@ def _export_diffusers_checkpoint( max_shard_size: Maximum size of each shard file. If the model exceeds this size, it will be sharded into multiple files and a .safetensors.index.json will be created. Use smaller values like "5GB" or "2GB" to force sharding. + export_attention_quantizers: Preserve enabled Q/K/V BMM quantizer state for + runtimes that consume legacy attention ``_amax`` keys, such as static FP8 FLUX. + None preserves these keys for FLUX.1 and omits them for other diffusion + components. Transformers components always retain them for KV-cache scales. + save_modelopt_state: Whether to save ModelOpt state with HF checkpointing enabled. **kwargs: Runtime-specific post-processing options forwarded to :func:`_postprocess_safetensors`. See its docstring for details. """ @@ -1303,6 +1317,14 @@ def _export_diffusers_checkpoint( # Infer dtype if not provided component_dtype = dtype if dtype is not None else infer_dtype_from_model(component) + is_transformers_component = isinstance(component, _TRANSFORMERS_MODEL_TYPES) + # TRT-LLM's static FP8 FLUX.1 path consumes the legacy Q/K/V amax keys. + # Transformers components need these keys for their declared KV-cache scheme. + keep_attention_quantizers = is_transformers_component or ( + isinstance(component, diffusers.FluxTransformer2DModel) + if export_attention_quantizers is None + else export_attention_quantizers + ) if is_quantized: # Fuse QKV linears that share the same input (unify amax values) @@ -1314,7 +1336,13 @@ def _export_diffusers_checkpoint( # dummy forward fails for it, fail loudly rather than silently skipping # fusion (which would export un-unified amax values). is_qwen_component = "qwen" in type(component).__name__.lower() - _fuse_qkv_linears_diffusion(component, strict=is_qwen_component) + # Image-conditioned Wan also requires a successful dummy forward for fusion. + is_wan_i2v_component = ( + HAS_DIFFUSERS + and isinstance(component, getattr(diffusers, "WanTransformer3DModel", ())) + and getattr(getattr(component, "config", None), "image_dim", None) is not None + ) + _fuse_qkv_linears_diffusion(component, strict=is_qwen_component or is_wan_i2v_component) # Process quantized modules (convert weights, register scales) _process_quantized_modules(component, component_dtype, is_modelopt_qlora=False) @@ -1329,7 +1357,12 @@ def _export_diffusers_checkpoint( # promoted buffers are always removed, even if save / post-process / config # update raises (keeps the live module reusable for a repeated export). try: - quant_config = get_quant_config(component, is_modelopt_qlora=False) + quant_config = get_quant_config( + component, + is_modelopt_qlora=False, + # Native diffusion components (including LTX-2) have no KV cache. + include_kv_cache=is_transformers_component, + ) if quant_config: quantization_details = quant_config.get("quantization", {}) # Record the SVDQuant low-rank size so consumers know the LoRA shape. @@ -1345,12 +1378,14 @@ def _export_diffusers_checkpoint( # - diffusers ModelMixin.save_pretrained does NOT accept state_dict parameter # - for non-diffusers modules (e.g., LTX-2 transformer), fall back to torch.save if hasattr(component, "save_pretrained"): - with hide_quantizers_from_state_dict(component): + with hide_quantizers_from_state_dict(component, keep_attention_quantizers): component.save_pretrained( - component_export_dir, max_shard_size=max_shard_size + component_export_dir, + max_shard_size=max_shard_size, + save_modelopt_state=save_modelopt_state, ) else: - with hide_quantizers_from_state_dict(component): + with hide_quantizers_from_state_dict(component, keep_attention_quantizers): _save_component_state_dict_safetensors(component, component_export_dir) # Post-process — merge, metadata, padding, swizzle @@ -1376,13 +1411,17 @@ def _export_diffusers_checkpoint( _remove_promoted_quantizer_tensors(component) # Non-quantized component: just save as-is elif hasattr(component, "save_pretrained"): - component.save_pretrained(component_export_dir, max_shard_size=max_shard_size) + component.save_pretrained( + component_export_dir, + max_shard_size=max_shard_size, + save_modelopt_state=save_modelopt_state, + ) else: _save_component_state_dict_safetensors(component, component_export_dir) # Update config.json with sparse attention info (both quantized and non-quantized) if export_sparse_attention_config is not None: - sparse_attn_config = export_sparse_attention_config(component) + sparse_attn_config = export_sparse_attention_config(component, for_diffusion=True) if sparse_attn_config is not None: config_path = component_export_dir / "config.json" if config_path.exists(): @@ -1750,6 +1789,14 @@ def unplaced_keys_for(model: nn.Module, ckpt: "str | Path") -> list[str]: module scope. Where those are absent the recorded set stands alone -- degraded, but no worse than before this function existed. """ + # Local import breaks the cycle: layerwise_export imports this module. + from .layerwise_export import LAYERWISE_EXPORTER_ATTR + + exporter = getattr(model, LAYERWISE_EXPORTER_ATTR, None) + snapshot = getattr(exporter, "_source_key_snapshot", None) + if snapshot is not None and snapshot[0] == str(ckpt): + return list(snapshot[1]) + recorded = getattr(model, "_modelopt_unplaced_source_keys", None) or [] structural: list[str] = [] try: @@ -1897,7 +1944,12 @@ def export_hf_checkpoint( max_shard_size: Maximum size of each safetensors shard file. Defaults to "10GB". **kwargs: Runtime-specific post-processing options forwarded to :func:`_postprocess_safetensors` for diffusion model exports. - See its docstring for supported keys. + See its docstring for supported keys. Diffusion exports also accept + ``export_attention_quantizers=True`` to retain enabled Q/K/V BMM quantizer + state for runtimes using legacy ``*_bmm_quantizer._amax`` keys, or False + to omit it. When unspecified, FLUX.1 retains enabled Q/K/V quantizers for + runtime compatibility; other diffusion components omit them. + Transformers components always retain them for KV-cache scales. """ # Weights the model never loaded (MTP head, auxiliary tower, ...) are copied straight from the # source so the exported checkpoint is the complete model. Merged here, ahead of the path @@ -1950,6 +2002,7 @@ def export_hf_checkpoint( export_dir, components, max_shard_size, + save_modelopt_state=save_modelopt_state, **kwargs, ) return diff --git a/modelopt/torch/kernels/quantization/gemm/fp4_kernel.py b/modelopt/torch/kernels/quantization/gemm/fp4_kernel.py index c9499b612f6..7f084ed089a 100644 --- a/modelopt/torch/kernels/quantization/gemm/fp4_kernel.py +++ b/modelopt/torch/kernels/quantization/gemm/fp4_kernel.py @@ -190,6 +190,9 @@ def fp4_dequantize( return output +_STATIC_FP4_BLOCKS_PER_PROGRAM = 64 + + @triton.jit def static_blockwise_fp4_fake_quant_kernel( x_ptr, # [NUM_FP4_BLOCKS * BLOCK_SIZE] @@ -197,21 +200,21 @@ def static_blockwise_fp4_fake_quant_kernel( scale_ptr, # [NUM_FP4_BLOCKS] NUM_FP4_BLOCKS, BLOCK_SIZE: tl.constexpr, + BLOCKS_PER_PROGRAM: tl.constexpr, OUT_DTYPE: tl.constexpr, ): - pid = tl.program_id(axis=0) - if pid >= NUM_FP4_BLOCKS: - return - - block_offset = pid * BLOCK_SIZE - idx = block_offset + tl.arange(0, BLOCK_SIZE) + # int64 offsets: stacked MoE weights can exceed 2**31 elements. + pid = tl.program_id(axis=0).to(tl.int64) + block_idx = pid * BLOCKS_PER_PROGRAM + tl.arange(0, BLOCKS_PER_PROGRAM) + block_mask = block_idx < NUM_FP4_BLOCKS + idx = block_idx[:, None] * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)[None, :] - scale = tl.load(scale_ptr + pid).to(tl.float32) - x = tl.load(x_ptr + idx).to(tl.float32) + scale = tl.load(scale_ptr + block_idx, mask=block_mask, other=0.0).to(tl.float32) + x = tl.load(x_ptr + idx, mask=block_mask[:, None], other=0.0).to(tl.float32) - x_quant = nvfp4_scalar_quant(x, scale, BLOCK_SIZE) + x_quant = nvfp4_scalar_quant(x, scale[:, None], BLOCK_SIZE) - tl.store(y_ptr + idx, x_quant.to(OUT_DTYPE)) + tl.store(y_ptr + idx, x_quant.to(OUT_DTYPE), mask=block_mask[:, None]) def compute_fp4_scales( @@ -301,7 +304,7 @@ def static_blockwise_fp4_fake_quant( tl_out_dtype = _torch_dtype_to_tl(out_dtype) - grid = (NUM_FP4_BLOCKS,) + grid = (triton.cdiv(NUM_FP4_BLOCKS, _STATIC_FP4_BLOCKS_PER_PROGRAM),) with torch.cuda.device(x.device): static_blockwise_fp4_fake_quant_kernel[grid]( @@ -310,6 +313,7 @@ def static_blockwise_fp4_fake_quant( scale_flat, NUM_FP4_BLOCKS, BLOCK_SIZE, + BLOCKS_PER_PROGRAM=_STATIC_FP4_BLOCKS_PER_PROGRAM, OUT_DTYPE=tl_out_dtype, ) diff --git a/modelopt/torch/kernels/quantization/gemm/gptq_fused_kernel.py b/modelopt/torch/kernels/quantization/gemm/gptq_fused_kernel.py index fb13386f851..3b410e97bae 100644 --- a/modelopt/torch/kernels/quantization/gemm/gptq_fused_kernel.py +++ b/modelopt/torch/kernels/quantization/gemm/gptq_fused_kernel.py @@ -119,18 +119,21 @@ def gptq_fused_block_scalar( qw_block = torch.empty_like(w_block) err_block = torch.empty_like(w_block) - _gptq_scalar_kernel[(num_rows,)]( - w_block.contiguous(), - qw_block, - err_block, - block_amax.contiguous(), - global_scale, - h_inv_cho_blk.contiguous(), - num_rows, - block_amax.shape[1], - quant_block_size, - block_start, - BLOCK_SIZE=block_size, - ) + # Triton launches on the current device; a model sharded with device_map="auto" keeps later + # layers on other GPUs, where launching from the current one reads another GPU's memory. + with torch.cuda.device(w_block.device): + _gptq_scalar_kernel[(num_rows,)]( + w_block.contiguous(), + qw_block, + err_block, + block_amax.contiguous(), + global_scale, + h_inv_cho_blk.contiguous(), + num_rows, + block_amax.shape[1], + quant_block_size, + block_start, + BLOCK_SIZE=block_size, + ) return qw_block, err_block diff --git a/modelopt/torch/models/hf.py b/modelopt/torch/models/hf.py index f9bdd49a931..3b765ddd2ba 100644 --- a/modelopt/torch/models/hf.py +++ b/modelopt/torch/models/hf.py @@ -20,6 +20,7 @@ __all__ = [ "checkpoint_has_mtp", + "prefers_bundled_code", "prepare_model_for_calibration", "prepare_model_for_loading", ] @@ -53,6 +54,14 @@ def checkpoint_has_mtp(model_type, checkpoint_path: str) -> bool: return plugin is not None and plugin.has_mtp_weights(checkpoint_path) +def prefers_bundled_code(model_type, hf_config, checkpoint_path: str, config_kwargs: dict) -> bool: + """Use the model plugin to tell whether the checkpoint's bundled code must replace the built-in class.""" + plugin = _get_plugin(model_type) + return plugin is not None and plugin.prefers_bundled_code( + hf_config, checkpoint_path, config_kwargs + ) + + def prepare_model_for_calibration(model) -> None: """Run the optional model-specific hook that augments the default calibration forward.""" plugin = _get_plugin(getattr(getattr(model, "config", None), "model_type", None)) diff --git a/modelopt/torch/models/nemotron_h/mtp.py b/modelopt/torch/models/nemotron_h/mtp.py index b3a234ac21f..736ee4a5322 100644 --- a/modelopt/torch/models/nemotron_h/mtp.py +++ b/modelopt/torch/models/nemotron_h/mtp.py @@ -29,6 +29,7 @@ import torch from huggingface_hub import snapshot_download from safetensors import SafetensorError +from transformers import AutoConfig from transformers.dynamic_module_utils import get_class_from_dynamic_module from transformers.masking_utils import create_causal_mask @@ -42,6 +43,7 @@ __all__ = [ "has_mtp_weights", + "prefers_bundled_code", "prepare_for_calibration", "prepare_for_loading", ] @@ -67,6 +69,29 @@ def has_mtp_weights(checkpoint_path: str) -> bool: return _get_mtp_layout(_resolve_checkpoint_path(checkpoint_path)) is not None +def prefers_bundled_code(hf_config, checkpoint_path: str, config_kwargs: dict) -> bool: + """True for a remote-code config the built-in Nemotron-H class cannot represent. + + Built-in Nemotron-H implements only the MoE variant on some Transformers releases, so a + dense checkpoint's config fails to parse as a built-in one and needs its bundled code. + """ + if getattr(hf_config, "model_type", None) != "nemotron_h": + return False + if not type(hf_config).__module__.startswith("transformers_modules"): + return False + builtin_kwargs = {k: v for k, v in config_kwargs.items() if k != "trust_remote_code"} + try: + builtin_config = AutoConfig.from_pretrained(checkpoint_path, **builtin_kwargs) + except (KeyError, ValueError, TypeError, AttributeError): + # Config-parsing incompatibilities only; e.g. an OSError from the Hub must propagate. + return True + # Parsing can succeed while model construction fails on a block type the built-in lacks. + # Native Nemotron-H is optional: older Transformers can still load unrelated models. + from transformers.models.nemotron_h.modeling_nemotron_h import MIXER_TYPES + + return any(t not in MIXER_TYPES for t in getattr(builtin_config, "layers_block_type", ())) + + def _normalize_mtp_block_types(block_types, mixer_types): """Normalize checkpoint attention names to the installed Transformers registry.""" aliases = {"attention": "full_attention", "full_attention": "attention"} diff --git a/modelopt/torch/models/specs.py b/modelopt/torch/models/specs.py index aa9110223ee..0a203780b7f 100644 --- a/modelopt/torch/models/specs.py +++ b/modelopt/torch/models/specs.py @@ -351,7 +351,7 @@ def hf_model_type(model) -> str | None: def match_moe_block(module: "nn.Module", model_type: str | None = None) -> "MoESpec | None": - """Return the MoE layout layout for ``module``, resolved by model type. + """Return the MoE layout for ``module``, resolved by model type. ``model_type`` (the root ``model.config.model_type``) is a strict filter: only that model's own spec is consulted, and an unregistered model type resolves to diff --git a/modelopt/torch/nas/autonas.py b/modelopt/torch/nas/autonas.py index 00393657c97..2cdf6349eb4 100644 --- a/modelopt/torch/nas/autonas.py +++ b/modelopt/torch/nas/autonas.py @@ -476,7 +476,7 @@ class RandomSearcher(IterativeSearcher): """An iterative searcher that samples subnets randomly.""" def sample(self) -> dict[str, Any]: - """Random sample new subset during each steo.""" + """Random sample new subnet during each step.""" return self._sample() diff --git a/modelopt/torch/opt/__init__.py b/modelopt/torch/opt/__init__.py index 950be0b986f..c46bf64687c 100644 --- a/modelopt/torch/opt/__init__.py +++ b/modelopt/torch/opt/__init__.py @@ -22,7 +22,7 @@ * export a model back to a regular pytorch model after optimization; * save, restore, and manage checkpoints from which the model modifications can be resumed. -Please refer to each individual sub-module to learn more about the various concepts wihin +Please refer to each individual sub-module to learn more about the various concepts within ``modelopt.torch.opt`` and how to use them to implement a model optimization algorithm. """ diff --git a/modelopt/torch/opt/conversion.py b/modelopt/torch/opt/conversion.py index bb3fe0ce4a5..a480b7d940c 100644 --- a/modelopt/torch/opt/conversion.py +++ b/modelopt/torch/opt/conversion.py @@ -19,8 +19,8 @@ (mode) of the algorithm. This module provides the state manager, which is a standardized interface (class) to record and store state information in the model. -Op top of the state manager, this module provides utilities to save a history of these modifications -("modelopt state dict") and restoring a unmodified model to the state indicated in the state dict. +On top of the state manager, this module provides utilities to save a history of these modifications +("modelopt state dict") and restoring an unmodified model to the state indicated in the state dict. """ import copy diff --git a/modelopt/torch/opt/dynamic.py b/modelopt/torch/opt/dynamic.py index 6c0e9dd3c5c..18baa449e9b 100644 --- a/modelopt/torch/opt/dynamic.py +++ b/modelopt/torch/opt/dynamic.py @@ -608,7 +608,7 @@ def convert(cls, module: nn.Module, **setup_kwargs: Any) -> "DynamicModule": Patching is achieved by updating the ``__class__`` attribute of the module to its dynamic counterpart. The dynamic counterpart is a subclass of the original class, hence, we ensure the module is fully compatible with the original class. Simultaneously, we can inject the - corresponding dynamic behavior in a standardized and rigoruos fashion. + corresponding dynamic behavior in a standardized and rigorous fashion. """ def bind_forward_method_if_needed(self): diff --git a/modelopt/torch/opt/mode.py b/modelopt/torch/opt/mode.py index d9ac1bf07bf..4ccac5473b0 100644 --- a/modelopt/torch/opt/mode.py +++ b/modelopt/torch/opt/mode.py @@ -94,7 +94,7 @@ def next_prohibited_modes(self) -> set[str] | None: """Modes that should not be applied after this mode. This is used to specify that certain modes should not be applied after this mode. - A ``None`` value indicates that incompatibitly check is skipped for this mode. + A ``None`` value indicates that the incompatibility check is skipped for this mode. This is useful if there are many allowed modes but only a few prohibited modes. Returns: diff --git a/modelopt/torch/opt/plugins/diffusers.py b/modelopt/torch/opt/plugins/diffusers.py index 3942507522b..58ff9735f3d 100644 --- a/modelopt/torch/opt/plugins/diffusers.py +++ b/modelopt/torch/opt/plugins/diffusers.py @@ -26,7 +26,9 @@ def _new_from_pretrained(cls, /, pretrained_model_name_or_path, *args, **kwargs): """Patch for `cls.from_pretrained` method to restore ModelOpt state.""" - with _patch_model_init_for_modelopt(cls, pretrained_model_name_or_path): + with _patch_model_init_for_modelopt( + cls, pretrained_model_name_or_path, subfolder=kwargs.get("subfolder") or "" + ): model = types.MethodType(cls._modelopt_cache["from_pretrained"].__func__, cls)( pretrained_model_name_or_path, *args, **kwargs ) diff --git a/modelopt/torch/opt/plugins/huggingface.py b/modelopt/torch/opt/plugins/huggingface.py index 19b2e357631..dc383c56566 100644 --- a/modelopt/torch/opt/plugins/huggingface.py +++ b/modelopt/torch/opt/plugins/huggingface.py @@ -20,11 +20,10 @@ import threading import warnings from contextlib import contextmanager +from pathlib import Path from typing import Any import torch -from huggingface_hub import try_to_load_from_cache -from huggingface_hub.errors import HFValidationError from modelopt.torch.utils import print_rank_0 @@ -58,21 +57,29 @@ def register_for_patching(name: str, cls: type, patch_methods: list[tuple[str, A methods_list.append(patch_methods) -def _get_modelopt_state_path(model_name_or_path: str) -> str: +def _get_modelopt_state_path(model_name_or_path: str, subfolder: str = "") -> str: """Get the path to the ModelOpt state file or empty string if not found. Also handles HF model card as input path. However for hf hub models, we dont have modelopt_state at the moment. """ if os.path.isdir(model_name_or_path): - return os.path.join(model_name_or_path, _MODELOPT_STATE_SAVE_NAME) + return os.path.join(model_name_or_path, subfolder, _MODELOPT_STATE_SAVE_NAME) + # Optional [hf] dependency: this module is imported even when huggingface_hub is not installed. + from huggingface_hub import try_to_load_from_cache + from huggingface_hub.errors import HFValidationError + + filename = ( + f"{subfolder}/{_MODELOPT_STATE_SAVE_NAME}" if subfolder else _MODELOPT_STATE_SAVE_NAME + ) try: - return try_to_load_from_cache(model_name_or_path, _MODELOPT_STATE_SAVE_NAME) or "" + cache_path = try_to_load_from_cache(model_name_or_path, filename) + return cache_path if isinstance(cache_path, str) else "" except HFValidationError: return "" @contextmanager -def _patch_model_init_for_modelopt(cls, model_path): +def _patch_model_init_for_modelopt(cls, model_path, subfolder=""): """Patch for `cls.init` method to restore ModelOpt state after `init`.""" # Note: Keeping original config in local as the package will be shared among threads added_original_init = False @@ -87,7 +94,7 @@ def _patch_model_init_for_modelopt(cls, model_path): @functools.wraps(_original__init__) def new_init_fn(self, *args, **kwargs): - modelopt_state_path = _get_modelopt_state_path(model_path) + modelopt_state_path = _get_modelopt_state_path(model_path, subfolder=subfolder) _original__init__(self, *args, **kwargs) if os.path.isfile(modelopt_state_path): restore_from_modelopt_state(self, modelopt_state_path=modelopt_state_path) @@ -111,6 +118,9 @@ def _new_save_pretrained(self, save_directory, *args, **kwargs): path = _get_modelopt_state_path(save_directory) torch.save(modelopt_state(self), path) print_rank_0(f"Saved ModelOpt state to {path}") + else: + # A checkpoint saved without ModelOpt state must not inherit an older sidecar. + Path(save_directory, _MODELOPT_STATE_SAVE_NAME).unlink(missing_ok=True) return outputs diff --git a/modelopt/torch/prune/fastnas.py b/modelopt/torch/prune/fastnas.py index 4852efdad2b..2385f5b19b2 100644 --- a/modelopt/torch/prune/fastnas.py +++ b/modelopt/torch/prune/fastnas.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Module implementing ``fasnas`` pruning algorithm for search.""" +"""Module implementing ``fastnas`` pruning algorithm for search.""" import warnings from collections.abc import Callable diff --git a/modelopt/torch/puzzletron/anymodel/README.md b/modelopt/torch/puzzletron/anymodel/README.md index 1c8c68b60bf..4c4ae11641a 100644 --- a/modelopt/torch/puzzletron/anymodel/README.md +++ b/modelopt/torch/puzzletron/anymodel/README.md @@ -84,7 +84,7 @@ Step 3: Configure YAML files Update the main model config YAML: - Set `descriptor` to match the name used in `@ModelDescriptorFactory.register_decorator("your_model_name")` -- See example: [llama_3_1_8b_instruct.yaml](../../../../tests/gpu/torch/puzzletron/resources/configs/llama_3_1_8b_instruct/llama_3_1_8b_instruct.yaml) +- See example: [Llama-3.1-8B-Instruct.yaml](../../../../tests/gpu/torch/puzzletron/resources/configs/meta-llama/Llama-3.1-8B-Instruct/Llama-3.1-8B-Instruct.yaml) Update pruning YAML files (`ffn_pruning.yaml`, `expert_pruning.yaml`, etc.): @@ -92,7 +92,7 @@ Update pruning YAML files (`ffn_pruning.yaml`, `expert_pruning.yaml`, etc.): - Set `layer_descriptor._target_` to your layer descriptor class - Set `hook_class` to the activation hook for scoring - Set `target_layer` in `activation_hooks_kwargs` to the layer name for hook attachment -- See examples in [configs/llama_3_1_8b_instruct/pruning/](../../../../tests/gpu/torch/puzzletron/resources/configs/llama_3_1_8b_instruct/pruning/) +- See examples in [configs/meta-llama/Llama-3.1-8B-Instruct/pruning/](../../../../tests/gpu/torch/puzzletron/resources/configs/meta-llama/Llama-3.1-8B-Instruct/pruning/) ## End-to-end example diff --git a/modelopt/torch/puzzletron/puzzletron_nas_plugin.py b/modelopt/torch/puzzletron/puzzletron_nas_plugin.py index 253674f97af..1295fcf1f4f 100644 --- a/modelopt/torch/puzzletron/puzzletron_nas_plugin.py +++ b/modelopt/torch/puzzletron/puzzletron_nas_plugin.py @@ -105,7 +105,7 @@ def convert_puzzletron_model(model: nn.Module, config: PuzzletronConfig) -> Conv 2. Score the pruning activations. 3. Prune the model and save pruned checkpoints - The output of this step will be used by mnt.search() to perform the NAS search. + The output of this step will be used by mtn.search() to perform the NAS search. """ # Required for mtn.search() to read NAS configuration model.hydra_config_dir = config.hydra_config_dir diff --git a/modelopt/torch/puzzletron/tools/post_init_sparse.py b/modelopt/torch/puzzletron/tools/post_init_sparse.py index ab939e3bb4b..0ff8f841dab 100644 --- a/modelopt/torch/puzzletron/tools/post_init_sparse.py +++ b/modelopt/torch/puzzletron/tools/post_init_sparse.py @@ -82,9 +82,9 @@ def do_sparsity(self, model: nn.Module, mask_dict=None): if mask_dict is None: state_dict_for_sparsifying = { - k.rstrip(".weight"): v + k.removesuffix(".weight"): v for k, v in model.state_dict().items() - if k.rstrip(".weight") in full_name_layers + if k.removesuffix(".weight") in full_name_layers } mask_dict = self.calculate_masks(state_dict_for_sparsifying) # print('Apply sparsity') diff --git a/modelopt/torch/quantization/calib/histogram.py b/modelopt/torch/quantization/calib/histogram.py index e27a5471e30..73ebc327892 100644 --- a/modelopt/torch/quantization/calib/histogram.py +++ b/modelopt/torch/quantization/calib/histogram.py @@ -218,7 +218,7 @@ def _normalize_distr(distr): if summ != 0: distr = distr / summ - bins = calib_hist[:] + bins = calib_hist.copy() bins[0] = bins[1] total_data = np.sum(bins) diff --git a/modelopt/torch/quantization/nn/modules/quant_conv.py b/modelopt/torch/quantization/nn/modules/quant_conv.py index 375fdc96e0c..8842540c87e 100644 --- a/modelopt/torch/quantization/nn/modules/quant_conv.py +++ b/modelopt/torch/quantization/nn/modules/quant_conv.py @@ -118,7 +118,7 @@ def _implicit_gemm_forward(self, input): padding=self.padding, dilation=self.dilation, act_amax=act_amax, - quant_act=self.input_quantizer.is_enabled, + quant_act=self.input_quantizer.is_enabled and self.input_quantizer._if_quant, fp4_block_size=fp4_block_size, ) return self.output_quantizer(output) diff --git a/modelopt/torch/quantization/plugins/accelerate.py b/modelopt/torch/quantization/plugins/accelerate.py index 35f6777fc8c..eb206db2c5c 100644 --- a/modelopt/torch/quantization/plugins/accelerate.py +++ b/modelopt/torch/quantization/plugins/accelerate.py @@ -233,6 +233,9 @@ def patched_from_pretrained(cls, /, pretrained_model_name_or_path, *args, **kwar "dtype", kwargs.get("torch_dtype", getattr(config, "torch_dtype", torch.float16)) ) model = cls.from_config(config, dtype=torch_dtype) + # init_empty_weights re-wraps every registered parameter, so the tie made inside + # from_config is lost; without it the absent tied checkpoint key (e.g. lm_head) stays meta. + model.tie_weights() mtq.quantize(model, quant_cfg) mtq.compress(model, config=mtq.CompressConfig(quant_gemm=quant_gemm)) diff --git a/modelopt/torch/quantization/plugins/attention.py b/modelopt/torch/quantization/plugins/attention.py index 2113edea8a7..e4843a3a2ec 100644 --- a/modelopt/torch/quantization/plugins/attention.py +++ b/modelopt/torch/quantization/plugins/attention.py @@ -52,7 +52,9 @@ def register_attention_for_kv_quant(attention_cls: type) -> bool: model_module = inspect.getmodule(attention_cls) head = ast.parse(source_code) - bmm_ops = ("matmul", "bmm", "baddbmm") + # No baddbmm: patch() wires args[0] and args[1], which miss its batch1 and batch2 operands + # (torch.baddbmm passes the additive input first, and Bloom passes them as keywords) + bmm_ops = ("matmul", "bmm") sdpa_ops = ("scaled_dot_product_attention",) def is_bmm(node): @@ -72,6 +74,12 @@ def is_sdpa(node): def is_bin_matmul(node): return isinstance(node, ast.BinOp) and isinstance(node.op, ast.MatMult) + def walk_post_order(node): + # Source order, children first; expects the q/k score matmul before the matmul consuming it + for child in ast.iter_child_nodes(node): + yield from walk_post_order(child) + yield node + def patch(node, quantizer_names, transpose=False): for index, quantizer_name in enumerate(quantizer_names): if quantizer_name is None: @@ -165,7 +173,7 @@ def patch_binop(node, quantizer_names, transpose=False): bmm_nodes = [] sdpa_nodes = [] bin_matmul_nodes = [] - for node in ast.walk(head): + for node in walk_post_order(head): if is_bmm(node): bmm_nodes.append(node) if is_sdpa(node): @@ -186,16 +194,16 @@ def patch_binop(node, quantizer_names, transpose=False): # after transpose, the quantization will be per-token, i.e., # self.k_bmm_quantizer(key_states.transpose(-1, -2).transpose(-1, -2)).transpose(-1, -2) # removing the additional transpose is doable but not trivial - patch(bmm_nodes[0], quantizer_names=(None, "v_bmm_quantizer")) - patch(bmm_nodes[1], quantizer_names=("q_bmm_quantizer", "k_bmm_quantizer"), transpose=True) + patch(bmm_nodes[0], quantizer_names=("q_bmm_quantizer", "k_bmm_quantizer"), transpose=True) + patch(bmm_nodes[1], quantizer_names=(None, "v_bmm_quantizer")) print("Patching 2 BMM/Matmul operators with quantizers") if len(bin_matmul_nodes) == 2: patch_binop( - bin_matmul_nodes[1], + bin_matmul_nodes[0], quantizer_names=("q_bmm_quantizer", "k_bmm_quantizer"), transpose=True, ) - patch_binop(bin_matmul_nodes[0], quantizer_names=(None, "v_bmm_quantizer")) + patch_binop(bin_matmul_nodes[1], quantizer_names=(None, "v_bmm_quantizer")) print("Patching 2 @ operators with quantizers") if len(sdpa_nodes) == 1: diff --git a/modelopt/torch/quantization/utils/core_utils.py b/modelopt/torch/quantization/utils/core_utils.py index 63f349dd10a..1c488c459e8 100644 --- a/modelopt/torch/quantization/utils/core_utils.py +++ b/modelopt/torch/quantization/utils/core_utils.py @@ -985,6 +985,9 @@ def fsdp2_aware_weight_update( Returns: None """ + # Bound only once setup succeeds; the finally skips FSDP bookkeeping until then so an + # early error (e.g. an OOM from the unshard below) propagates and FSDP state stays intact. + fsdp_param_group = None try: if isinstance(root_model, FSDPModule): # Get FSDP root module, if none is returned, then the update is not made to a submodule of an FSDPModule @@ -1006,8 +1009,9 @@ def fsdp2_aware_weight_update( root_module.unshard() # Get FSDPParam list - fsdp_param_group = fully_shard.state(root_module)._fsdp_param_group - fsdp_param_mapping = create_fsdp_param_mapping(fsdp_param_group.fsdp_params, root_model) + param_group = fully_shard.state(root_module)._fsdp_param_group + fsdp_param_mapping = create_fsdp_param_mapping(param_group.fsdp_params, root_model) + fsdp_param_group = param_group # Assert that all the modules in the module list are present in this fsdp_param_group if len(modules_to_update) > 1: @@ -1027,7 +1031,7 @@ def fsdp2_aware_weight_update( finally: from modelopt.torch.quantization.qtensor.base_qtensor import QFSDPParam, QTensorWrapper - if isinstance(root_model, FSDPModule): + if fsdp_param_group is not None: # Update FSDPParam list for module in modules_to_update: for param_name, param in module.named_parameters(): diff --git a/modelopt/torch/sparsity/attention_sparsity/conversion.py b/modelopt/torch/sparsity/attention_sparsity/conversion.py index b2acd137930..38944517740 100644 --- a/modelopt/torch/sparsity/attention_sparsity/conversion.py +++ b/modelopt/torch/sparsity/attention_sparsity/conversion.py @@ -332,18 +332,7 @@ def restore_sparse_attention_state(model: nn.Module, state_dict: dict[str, Any]) for name, module in get_named_sparse_attention_modules(model): module_name = get_unwrapped_name(name, model) if module_name in state_dict: - module_state = state_dict[module_name] - - # Restore method and config - if "method" in module_state: - module._method = module_state["method"] - if "method_config" in module_state: - # Restore config attributes - for key, val in module_state["method_config"].items(): - setattr(module, f"_{key}", val) - - # Re-setup with restored config - module._setup() + module.set_from_modelopt_state(state_dict[module_name]) def update_sparse_attention_metadata( @@ -356,21 +345,10 @@ def update_sparse_attention_metadata( config: Configuration used metadata: Metadata dict to update """ - sparse_state = {} - - for name, module in get_named_sparse_attention_modules(model): - module_name = get_unwrapped_name(name, model) - - # Save the method configuration that was used - # _method_config already contains the validated config dict - module_state = { - "method": module._sparse_method_instance.name, - "method_config": module._method_config.copy(), - } - - sparse_state[module_name] = module_state - - metadata["sparse_attention_state"] = sparse_state + metadata["sparse_attention_state"] = { + get_unwrapped_name(name, model): module.get_modelopt_state() + for name, module in get_named_sparse_attention_modules(model) + } metadata["sparse_attention_config"] = ( config.model_dump() if hasattr(config, "model_dump") else vars(config) ) @@ -394,7 +372,9 @@ def _get_sparse_softmax_export_config(module: SparseAttentionModule) -> dict[str return sparse_softmax_config -def export_sparse_attention_config(model: nn.Module) -> dict[str, Any] | None: +def export_sparse_attention_config( + model: nn.Module, *, for_diffusion: bool = False +) -> dict[str, Any] | None: """Extract sparse attention config for export to config.json. Extracts calibrated skip-softmax parameters and N:M sparse-softmax metadata @@ -406,6 +386,10 @@ def export_sparse_attention_config(model: nn.Module) -> dict[str, Any] | None: Args: model: Model with sparse attention applied + for_diffusion: Emit the single-phase schema read by the TRT-LLM VisualGen + runtime: one ``threshold_scale_factor.coefficients`` block and a scalar + ``target_sparsity``. The default is the LLM schema, with per-phase + (``prefill``/``decode``) blocks and a ``{phase: ratio}`` dict. Returns: Dictionary with sparse attention config for HuggingFace config.json export. @@ -495,6 +479,21 @@ def export_sparse_attention_config(model: nn.Module) -> dict[str, Any] | None: skip_group["threshold_scale_factor"] = export_threshold_scale_factor(calibration_params) if target_sparse_ratio is not None: skip_group["target_sparsity"] = target_sparse_ratio + if for_diffusion: + # Diffusion calibrates a single (prefill) phase; VisualGen rejects phase-keyed dicts. + tsf = skip_group["threshold_scale_factor"] + phases = {p: tsf.pop(p) for p in ("prefill", "decode") if p in tsf} + if phases: + tsf["coefficients"] = phases.get("prefill") or next(iter(phases.values())) + ratio = skip_group.get("target_sparsity") + if isinstance(ratio, dict): + scalar = ( + ratio["prefill"] if "prefill" in ratio else next(iter(ratio.values()), None) + ) + if scalar is None: + del skip_group["target_sparsity"] # VisualGen needs a scalar + else: + skip_group["target_sparsity"] = scalar config_groups[f"group_{group_idx}"] = skip_group group_idx += 1 diff --git a/modelopt/torch/sparsity/attention_sparsity/sparse_attention.py b/modelopt/torch/sparsity/attention_sparsity/sparse_attention.py index 4f29f382d2f..1979539e726 100644 --- a/modelopt/torch/sparsity/attention_sparsity/sparse_attention.py +++ b/modelopt/torch/sparsity/attention_sparsity/sparse_attention.py @@ -15,6 +15,7 @@ """Extensible sparse attention module.""" +import copy from typing import Any from modelopt.torch.opt.dynamic import DynamicModule, _DMRegistryCls @@ -23,6 +24,9 @@ from .methods import get_sparse_method from .stats_manager import SparseAttentionStatsManager +# Method-instance attributes set by calibration (see ``SparseAttentionMethod``). +_CALIBRATION_STATE_KEYS = ("calibration_params", "target_sparse_ratio") + class SparseAttentionModule(DynamicModule): """Generic sparse attention module wrapper for applying sparsity to attention layers. @@ -167,6 +171,37 @@ def get_threshold_info(self) -> dict[str, Any]: return self._sparse_method_instance.get_threshold_info() return {"type": "none", "value": None} + def get_modelopt_state(self) -> dict[str, Any]: + """Get the module state to save in the ModelOpt checkpoint metadata. + + Similar to ``TensorQuantizer.get_modelopt_state``. Calibration results live on the + method instance rather than in the state dict, so they are saved here. + """ + method = self._sparse_method_instance + modelopt_state = { + "method": method.name, + "method_config": self._method_config.copy(), + "enabled": self.is_enabled, + } + for key in _CALIBRATION_STATE_KEYS: + val = getattr(method, key, None) + if val is not None: + modelopt_state[key] = copy.deepcopy(val) + return modelopt_state + + def set_from_modelopt_state(self, modelopt_state: dict[str, Any]): + """Restore the state saved by :meth:`get_modelopt_state`. + + Similar to ``TensorQuantizer.set_from_modelopt_state``. The method and its config are + rebuilt from the mode config before this is called; this restores the runtime state on + top of them. + """ + if "enabled" in modelopt_state: + self._enabled = bool(modelopt_state["enabled"]) + for key in _CALIBRATION_STATE_KEYS: + if key in modelopt_state: + setattr(self._sparse_method_instance, key, copy.deepcopy(modelopt_state[key])) + def _setup(self): """Setup called by DynamicModule.""" # Apply default configuration if not yet configured diff --git a/modelopt/torch/speculative/plugins/megatron_eagle.py b/modelopt/torch/speculative/plugins/megatron_eagle.py index 22c337ed41d..65e68d8f280 100644 --- a/modelopt/torch/speculative/plugins/megatron_eagle.py +++ b/modelopt/torch/speculative/plugins/megatron_eagle.py @@ -368,11 +368,11 @@ def sharded_state_dict( ) if num_layers > 1: - # For multi-layer EAGLE, we use heterogenuous distrubted checkpoint + # For multi-layer EAGLE, we use heterogeneous distributed checkpoint sharded_prefix = state_dict_prefix sharded_pp_offset = [] else: - # For single EAGLE, we use homogenous distrubted checkpoint for backward compatibility + # For single EAGLE, we use homogeneous distributed checkpoint for backward compatibility sharded_prefix = layer_prefix sharded_pp_offset = [ (0, global_layer_offset, num_layers) @@ -592,7 +592,7 @@ def forward( """Forward function.""" # NOTE: Even if sequence_parallel is used, the rotary_seq_len must be in the original # length. Since we get the seq_len from hidden_states.shape[0], we need to - # multiply the the tp back. + # multiply the tp back. # Similarly, if context parallel is used, the rotary_seq_len must also be # multiplied by context parallel size. rotary_seq_len = hidden_states.shape[0] @@ -830,7 +830,7 @@ def _get_eagle_module_inputs( ttt_step: int = 0, ): """Getting EAGLE module inputs.""" - # gather_from_sequence_parallel_region gathers from the first dimention + # gather_from_sequence_parallel_region gathers from the first dimension # so we need to transpose input_ids first # [b,s] -> [s,b] input_ids = input_ids.clone().transpose(0, 1).contiguous() diff --git a/modelopt/torch/speculative/utils.py b/modelopt/torch/speculative/utils.py index 4db1f4b0f76..910536fdb10 100644 --- a/modelopt/torch/speculative/utils.py +++ b/modelopt/torch/speculative/utils.py @@ -17,9 +17,10 @@ import contextlib import copy +import hashlib +import importlib.abc import importlib.util import json -import os import sys import warnings from collections import Counter, defaultdict, deque @@ -28,12 +29,19 @@ import torch import torch.distributed import transformers -from huggingface_hub import snapshot_download +from huggingface_hub import hf_hub_download from torch import nn from torch.nn.attention import SDPBackend, sdpa_kernel from transformers.cache_utils import DynamicCache KIMI_K2_REPO_ID = "moonshotai/Kimi-K2-Thinking" +# These third-party files are executed, so pin them; bump the revision and hashes together after +# reviewing the upstream diff. +KIMI_K2_REVISION = "a51ccc050d73dab088bf7b0e2dd9b30ae85a4e55" +KIMI_K2_FILE_SHA256 = { + "configuration_deepseek.py": "54f29dd0d56e9a13f639c703e524daa71794e360050012aea3788ad7cfc38b17", + "modeling_deepseek.py": "ce753fc9076ce13ce9557e680e2f5bd68ae8929a8225772968a9a22a1d17b7d4", +} KIMI_K2_PACKAGE_NAME = "kimi_k2_temp" @@ -82,7 +90,7 @@ def calibrate_frequent_vocab(tokenizer, text, target_vocab_size, output_file=Non def get_default_attention_mask_and_position_ids(input_ids: torch.Tensor): - """Compute default attention_mask ans position_ids given input_ids.""" + """Compute default attention_mask and position_ids given input_ids.""" seq_len = input_ids.shape[-1] position_ids = torch.arange(seq_len, dtype=torch.long, device=input_ids.device) @@ -499,39 +507,81 @@ def _patch_dynamic_cache_compatibility() -> None: ) -def _import_module_from_path(module_path: str, module_name: str, package_name: str): - """Dynamically import a module from file path.""" - spec = importlib.util.spec_from_file_location(module_name, module_path) - if spec is None or spec.loader is None: +class _VerifiedSourceLoader(importlib.abc.InspectLoader): + """Execute already-verified source bytes, never re-reading the file or its ``__pycache__``.""" + + def __init__(self, path: str, source: bytes): + self.path = path + self.source = source + + def get_source(self, fullname): + return importlib.util.decode_source(self.source) + + def get_code(self, fullname): + return self.source_to_code(self.source, self.path) + + +def _download_verified_kimi_k2_file(filename: str) -> tuple[str, bytes]: + """Download ``filename`` at the pinned revision and return its path and hash-checked content.""" + # Not snapshot_download: for a commit revision it calls the Hub even with HF_HUB_OFFLINE set, + # so cached runs would fail offline. + path = hf_hub_download(KIMI_K2_REPO_ID, filename, revision=KIMI_K2_REVISION) + with open(path, "rb") as f: + source = f.read() + digest = hashlib.sha256(source).hexdigest() + expected = KIMI_K2_FILE_SHA256[filename] + if digest != expected: + raise RuntimeError( + f"sha256 of {path} is {digest}, expected {expected}; refusing to execute it." + ) + return path, source + + +def _import_module_from_source( + module_path: str, source: bytes, module_name: str, package_name: str +): + """Import already-verified ``source`` as ``module_name``, with ``module_path`` as its file.""" + loader = _VerifiedSourceLoader(module_path, source) + spec = importlib.util.spec_from_file_location(module_name, module_path, loader=loader) + if spec is None: raise ImportError(f"Failed to load module {module_name} from {module_path}") module = importlib.util.module_from_spec(spec) module.__package__ = package_name sys.modules[module_name] = module - spec.loader.exec_module(module) + try: + loader.exec_module(module) + except BaseException: + # Like importlib, don't leave a half-initialized module registered + sys.modules.pop(module_name, None) + raise return module def _setup_kimi_k2_decoder(): """Setup Kimi-K2 decoder with proper imports and patches.""" - # Download repository files - kimi_k2_path = snapshot_download(repo_id=KIMI_K2_REPO_ID, allow_patterns="*.py") - - # Patch deprecated usage in Kimi implementation - _patch_dynamic_cache_compatibility() + # Verify both files before executing either, so a refusal runs none of the downloaded code + config_module_path, config_source = _download_verified_kimi_k2_file("configuration_deepseek.py") + model_module_path, model_source = _download_verified_kimi_k2_file("modeling_deepseek.py") # Import required modules - config_module_path = os.path.join(kimi_k2_path, "configuration_deepseek.py") - model_module_path = os.path.join(kimi_k2_path, "modeling_deepseek.py") - - _import_module_from_path( - config_module_path, f"{KIMI_K2_PACKAGE_NAME}.configuration_deepseek", KIMI_K2_PACKAGE_NAME + _import_module_from_source( + config_module_path, + config_source, + f"{KIMI_K2_PACKAGE_NAME}.configuration_deepseek", + KIMI_K2_PACKAGE_NAME, ) - kimi_k2_module = _import_module_from_path( - model_module_path, f"{KIMI_K2_PACKAGE_NAME}.modeling_deepseek", KIMI_K2_PACKAGE_NAME + kimi_k2_module = _import_module_from_source( + model_module_path, + model_source, + f"{KIMI_K2_PACKAGE_NAME}.modeling_deepseek", + KIMI_K2_PACKAGE_NAME, ) + # Patch deprecated usage in Kimi implementation + _patch_dynamic_cache_compatibility() + # Patch Kimi Attention to init rope lazily to avoid save/load meta tensor error original_init_rope = kimi_k2_module.DeepseekV3Attention._init_rope original_forward = kimi_k2_module.DeepseekV3Attention.forward diff --git a/modelopt/torch/utils/dataset_utils.py b/modelopt/torch/utils/dataset_utils.py index f882a6e7364..76260c216e3 100644 --- a/modelopt/torch/utils/dataset_utils.py +++ b/modelopt/torch/utils/dataset_utils.py @@ -25,7 +25,6 @@ from pathlib import Path from typing import TYPE_CHECKING, Any, Literal -import requests import torch from torch.utils.data import DataLoader from torch.utils.data.distributed import DistributedSampler @@ -1318,6 +1317,7 @@ def download_hf_dataset_as_jsonl( Returns: List of paths to downloaded JSONL files. """ + import requests from datasets import load_dataset from huggingface_hub.utils import build_hf_headers diff --git a/modelopt/torch/utils/network.py b/modelopt/torch/utils/network.py index d9038faae86..101c6518bd0 100644 --- a/modelopt/torch/utils/network.py +++ b/modelopt/torch/utils/network.py @@ -278,7 +278,7 @@ def standardize_named_model_args( Returns: A tuple (args_normalized, args_with_default) where args_normalized is a dictionary of ordered model args where the key represents a unique - serialized string based on the the argument's name in the function signature and the + serialized string based on the argument's name in the function signature and the value contains the actual argument, args_with_default is a set indicating whether the argument was retrieved from the default value in the function signature of the model's ``forward()`` method or whether the @@ -520,7 +520,7 @@ def run_forward_loop( Args: model: The model with which we run forward. data_loader: An iterator with data samples. - max_iters: Number of batches to run; by default it is infiinite or until ``data_loader`` + max_iters: Number of batches to run; by default it is infinite or until ``data_loader`` is exhausted. collect_func: A ``Callable`` that takes a batch of data from the ``data_loader`` as input and returns the input to ``model.forward()`` such that the return value @@ -585,7 +585,7 @@ def create_param_grad_clear_hook(param): """Create a hook to clear gradients for a parameter. The hook will be fired after the gradient is accumulated for the parameter. - Important: For this to work, ``accum_grad`` should be kept alive as longs as this utility is needed. + Important: For this to work, ``accum_grad`` should be kept alive as long as this utility is needed. """ # For methods such as AutoQuantize, involving backward - # We want to clear parameter gradients as soon as they are computed to save memory diff --git a/modelopt/torch/utils/serialization.py b/modelopt/torch/utils/serialization.py index dc880b86b80..6ccae6ee8fd 100644 --- a/modelopt/torch/utils/serialization.py +++ b/modelopt/torch/utils/serialization.py @@ -56,7 +56,7 @@ def safe_save(obj: Any, f: str | os.PathLike | BinaryIO, **kwargs) -> None: def safe_load(f: str | os.PathLike | BinaryIO | bytes, **kwargs) -> Any: """Load a checkpoint securely using ``weights_only=True`` by default. - NOTE: We dont set default ``weights_only`` (interpret as True for torch>=2.6) so you can override it with + NOTE: We don't set default ``weights_only`` (interpret as True for torch>=2.6) so you can override it with ``export TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1`` if you see ``pickle.UnpicklingError`` and trust the checkpoint. """ if isinstance(f, (bytes, bytearray)): diff --git a/modelopt_recipes/configs/ptq/presets/README.md b/modelopt_recipes/configs/ptq/presets/README.md index 2aeec2d2b33..ed12bd66fa6 100644 --- a/modelopt_recipes/configs/ptq/presets/README.md +++ b/modelopt_recipes/configs/ptq/presets/README.md @@ -6,14 +6,9 @@ from `modelopt.torch.quantization.config` (e.g., `FP8_DEFAULT_CFG`, `FP8_KV_CFG`). Presets compose from the reusable snippets in `configs/numerics/` and -`configs/ptq/units/` via the `$import` system, and are split into two +`configs/ptq/units/` via the `$import` system, and are split into three kinds: -Preset files are also reusable snippets when imported by recipes or other -configs, so each preset must declare a `# modelopt-schema: ...` preamble. -Current preset files use `QuantizeConfig` so they are validated after their -own imports have been resolved. - - **`model/`** — *full* quantization presets. Each file is a complete, self-contained config (it sets `algorithm` and a full `quant_cfg` with a base-disable-all prefix + standard exclusions) and can be passed @@ -30,6 +25,11 @@ own imports have been resolved. attention and softmax quantizer choices that differ from the generic `model/` presets. +Preset files are also reusable snippets when imported by recipes or other +configs, so each preset must declare a `# modelopt-schema: ...` preamble. +Current preset files use `QuantizeConfig` so they are validated after their +own imports have been resolved. + **Note:** The main purpose of these presets is to support the existing `hf_ptq.py` script's `--qformat` / `--kv_cache_qformat` flags and other code paths that reference diff --git a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml index cac4f6d80aa..cd5097dbded 100644 --- a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml +++ b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml @@ -46,7 +46,7 @@ enable: false - quantizer_name: '*router*' enable: false - - quantizer_name: 'mtp.*' + - quantizer_name: '*mtp.*' enable: false - quantizer_name: 'output.*' enable: false diff --git a/modelopt_recipes/models/MiniMaxAI/MiniMax-M2.7/ptq/nvfp4_experts_only-kv_fp8_cast.yaml b/modelopt_recipes/models/MiniMaxAI/MiniMax-M2.7/ptq/nvfp4_experts_only-kv_fp8_cast.yaml index 1cf8b612d84..cdabe12919b 100644 --- a/modelopt_recipes/models/MiniMaxAI/MiniMax-M2.7/ptq/nvfp4_experts_only-kv_fp8_cast.yaml +++ b/modelopt_recipes/models/MiniMaxAI/MiniMax-M2.7/ptq/nvfp4_experts_only-kv_fp8_cast.yaml @@ -26,7 +26,9 @@ imports: metadata: description: >- MiniMax M2.7 expert-only NVFP4 weights and activations with max calibration; - FP8 KV-cache cast mode; attention projections, router, and LM head stay BF16. + FP8 KV-cache cast mode; attention projections, router, and LM head are + excluded from additional quantization. The source stores attention weights + in FP8, router weights in FP32, and LM-head weights in BF16. Reproduces the published nvidia/MiniMax-M2.7-NVFP4 checkpoint from MiniMaxAI/MiniMax-M2.7. @@ -40,5 +42,9 @@ quantize: cfg: {$import: nvfp4} - quantizer_name: '*block_sparse_moe.experts.*input_quantizer' cfg: {$import: nvfp4} + - quantizer_name: '*mlp.experts.*weight_quantizer' + cfg: {$import: nvfp4} + - quantizer_name: '*mlp.experts.*input_quantizer' + cfg: {$import: nvfp4} - $import: kv_fp8_cast - $import: default_disabled_quantizers diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index af3c10a117d..1d4b754bfb6 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -476,7 +476,10 @@ checkpoint's** quant config verbatim: and the KV cache use FP8. It uses max calibration. - **`models/MiniMaxAI/MiniMax-M2.7/ptq/nvfp4_experts_only-kv_fp8_cast`** reproduces `nvidia/MiniMax-M2.7-NVFP4` with max-calibrated NVFP4 W4A4 experts (block size - 16) and FP8 KV-cache cast mode; attention projections, router, and LM head stay BF16. + 16) and FP8 KV-cache cast mode, supporting both `block_sparse_moe.experts` and + `mlp.experts` layouts, including fused experts. Attention projections, router, + and LM head are excluded from additional quantization: the source checkpoint + stores attention weights in FP8, router weights in FP32, and LM-head weights in BF16. - **`models/Qwen/Qwen3.8-27B/ptq/nvfp4_w4a4_mlp_fp8_attn_local_hessian`** reproduces NVFP4 W4A4 mixed precision quantization recipe used for `nvidia/Qwen3.8-27B-NVFP4`: MLP projections and `lm_head` use NVFP4 W4A4, self-attention and diff --git a/plugins/modelopt/skills/evaluation/SKILL.md b/plugins/modelopt/skills/evaluation/SKILL.md index 342177ea29b..9f5aab58643 100644 --- a/plugins/modelopt/skills/evaluation/SKILL.md +++ b/plugins/modelopt/skills/evaluation/SKILL.md @@ -6,21 +6,21 @@ license: Apache-2.0 # https://github.com/NVIDIA-NeMo/Evaluator/tree/f1fa073/packages/nemo-evaluator-launcher/.claude/skills/nel-assistant --- -## NeMo Evaluator Launcher Assistant +# NeMo Evaluator Launcher Assistant Guide the user through creating NEL YAML configs, running evaluations, and monitoring progress. NEL manages serving for canaries and full runs; see [Step 3's serving-handoff check](references/launcher-workflow.md#step-3--configure-deployment) for external endpoints. -### Completed-run analysis +## Completed-run analysis For completed runs, skip launch/config; use launching-evals for artifacts/analysis and apply [run-validation.md](references/run-validation.md) before scores. Its policy overrides blanket uncapping and fixed token/context limits. -### Workspace integration +## Workspace integration If `MODELOPT_WORKSPACE_ROOT` is set, use the common skill's `workspace-management.md` and reuse existing workspaces (PTQ → checkpoint smoke test → NEL evaluation; carry compatibility requirements into NEL's deployment config). -### Workflow +## Workflow ```text - [ ] Step 0: Check workspace (if MODELOPT_WORKSPACE_ROOT set) @@ -38,7 +38,7 @@ If `MODELOPT_WORKSPACE_ROOT` is set, use the common skill's `workspace-managemen --- -### nel-next path (Terminal-Bench 2.x, SWE-bench, …) — branch here FIRST +## nel-next path (Terminal-Bench 2.x, SWE-bench, …) — branch here FIRST A few **agentic** AA benchmarks do **not** run on the currently validated `nemo-evaluator-launcher` 0.2.6 path (Steps 1–9 don't apply). They run on **nel-next** @@ -56,7 +56,7 @@ Steps 1–9 below are currently validated with 0.2.6 — use them for everything --- -### MRCR (NeMo Gym `simple_agent`) path — branch here too +## MRCR (NeMo Gym `simple_agent`) path — branch here too MRCR **does** run on the currently validated 0.2.6 `nel` launcher (as a `nemo_gym` task, not nel-next), so Steps 1–9 apply — but it is mechanically special and @@ -99,43 +99,43 @@ Read only the sections needed for the current stage; retain the dry-run → cana full-run gates. Existing configs must pass Step 3's serving-handoff check before Step 8. Paths in that reference are relative to this skill directory. -### Step 1 — Prerequisites +## Step 1 — Prerequisites Read [Step 1](references/launcher-workflow.md#step-1--prerequisites). -### Step 2 — Build base config (when not using shortcut) +## Step 2 — Build base config (when not using shortcut) Read [Step 2](references/launcher-workflow.md#step-2--build-base-config-when-not-using-shortcut). -### Step 3 — Configure deployment +## Step 3 — Configure deployment Read [Step 3](references/launcher-workflow.md#step-3--configure-deployment). -### Step 4 — Fill remaining ??? values +## Step 4 — Fill remaining ??? values Read [Step 4](references/launcher-workflow.md#step-4--fill-remaining--values). -### Step 5 — Confirm tasks (iterative) +## Step 5 — Confirm tasks (iterative) Read [Step 5](references/launcher-workflow.md#step-5--confirm-tasks-iterative). -### Step 6 — Multi-node +## Step 6 — Multi-node Read [Step 6](references/launcher-workflow.md#step-6--multi-node). -### Step 7 — Interceptors +## Step 7 — Interceptors Read [Step 7](references/launcher-workflow.md#step-7--interceptors). -### Step 7.5 — Container registry auth (SLURM private images only) +## Step 7.5 — Container registry auth (SLURM private images only) Read [Step 7.5](references/launcher-workflow.md#step-75--container-registry-auth-slurm-private-images-only). -### Step 8 — Run evaluation (gated dry-run → canary → full) +## Step 8 — Run evaluation (gated dry-run → canary → full) Read [Step 8](references/launcher-workflow.md#step-8--run-evaluation-gated-dry-run--canary--full). -### Step 9 — Verify completed run +## Step 9 — Verify completed run Apply [run-validation.md](references/run-validation.md)'s **Bounded Evaluation-Failure Policy** and **Timeout and Output-Limit Accounting** before reporting scores. For diff --git a/plugins/modelopt/skills/monitor/SKILL.md b/plugins/modelopt/skills/monitor/SKILL.md index 14ce4e14c32..01d619e11ad 100644 --- a/plugins/modelopt/skills/monitor/SKILL.md +++ b/plugins/modelopt/skills/monitor/SKILL.md @@ -1,6 +1,6 @@ --- name: monitor -description: Monitor submitted jobs (PTQ, evaluation, deployment) on SLURM clusters. Use when the user asks "check job status", "is my job done", "monitor my evaluation", "what's the status of the PTQ", "check on job ", or after any skill submits a long-running job. Also triggers on "nel status", "squeue", or any request to check progress of a previously submitted job. +description: Monitor submitted jobs (PTQ, evaluation, deployment) on SLURM clusters. Use when the user asks "check job status", "is my job done", "monitor my evaluation", "what's the status of the PTQ", "check on job" followed by a SLURM job ID, or after any skill submits a long-running job. Also triggers on "nel status", "squeue", or any request to check progress of a previously submitted job. --- # Job Monitor diff --git a/pyproject.toml b/pyproject.toml index 9a8d7cc4a76..648afc7b5c4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -59,7 +59,8 @@ onnx = [ "lief", "ml_dtypes", "onnx-graphsurgeon>=0.6.1", - "onnx~=1.21.0", + "onnx~=1.21.0; platform_system != 'Windows'", + "onnx~=1.23.2; platform_system == 'Windows'", "onnxconverter-common~=1.16.0", # Python 3.10 support for ModelOpt ONNX is deprecated and will be removed in a future release. # Native Windows ARM64 does not support ModelOpt ONNX on Python 3.10. diff --git a/tests/examples/diffusers/sparsity/test_sparsity.py b/tests/examples/diffusers/sparsity/test_sparsity.py index bca94e3dafb..4067251cc01 100644 --- a/tests/examples/diffusers/sparsity/test_sparsity.py +++ b/tests/examples/diffusers/sparsity/test_sparsity.py @@ -229,8 +229,9 @@ def test_wan22_calibrated_export(tiny_wan22_path, tmp_path): # threshold_scale_factor lives inside the skip_softmax group. tsf = group_0["threshold_scale_factor"] assert tsf["formula"] == "a * exp(b * target_sparsity)" - assert tsf["prefill"]["a"] == pytest.approx(math.exp(test_log_a)) - assert tsf["prefill"]["b"] == pytest.approx(test_b) + assert tsf["coefficients"]["a"] == pytest.approx(math.exp(test_log_a)) + assert tsf["coefficients"]["b"] == pytest.approx(test_b) + assert isinstance(group_0["target_sparsity"], float) # scalar, as VisualGen expects # Calibrated mode — no raw_threshold. assert "raw_threshold" not in group_0 diff --git a/tests/examples/diffusers/test_distillation_trainer.py b/tests/examples/diffusers/test_distillation_trainer.py new file mode 100644 index 00000000000..e9753469799 --- /dev/null +++ b/tests/examples/diffusers/test_distillation_trainer.py @@ -0,0 +1,161 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Exercise LTX checkpoint finalization without LTX model assets or a GPU.""" + +import importlib.util +import itertools +import json +import logging +import sys +from contextlib import nullcontext +from pathlib import Path +from types import ModuleType, SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +from accelerate import Accelerator +from pydantic import BaseModel +from safetensors.torch import load_file + + +@pytest.fixture +def trainer_module(monkeypatch): + # LTX is an optional, separately licensed dependency. Substitute its model setup + # and UI, but import the complete example and execute its train/save/export methods. + progress = nullcontext() + progress.update_training = lambda **kwargs: None + for name, attributes in { + "ltx_trainer": {"logger": logging.getLogger(__name__)}, + "ltx_trainer.config": {"ConfigBaseModel": BaseModel, "LtxTrainerConfig": BaseModel}, + "ltx_trainer.model_loader": {"load_transformer": None}, + "ltx_trainer.trainer": { + "IS_MAIN_PROCESS": True, + "LtxvTrainer": object, + "TrainingStats": SimpleNamespace, + }, + "ltx_trainer.gpu_utils": {"get_gpu_memory_gb": lambda device: 0}, + "ltx_trainer.hf_hub_utils": {"push_to_hub": None}, + "ltx_trainer.progress": {"TrainingProgress": lambda **kwargs: progress}, + }.items(): + module = ModuleType(name) + module.__dict__.update(attributes) + monkeypatch.setitem(sys.modules, name, module) + + path = ( + Path(__file__).resolve().parents[3] + / "examples/diffusers/distillation/distillation_trainer.py" + ) + spec = importlib.util.spec_from_file_location("ltx_checkpoint_example", path) + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, module) + spec.loader.exec_module(module) + monkeypatch.setattr(module, "is_global_rank0", lambda: True) + # Deterministically cross the one-minute deadline during the first optimizer step. + clock = itertools.count(step=60) + monkeypatch.setattr(module, "time", SimpleNamespace(time=lambda: next(clock))) + return module + + +@pytest.mark.parametrize("grad_accum", [1, 2]) +@pytest.mark.parametrize( + ("steps", "interval", "deadline", "saved_steps"), + [ + pytest.param(2, 1, None, [1, 2], id="interval-aligned"), + pytest.param(3, 2, None, [2, 3], id="nonaligned-final"), + pytest.param(3, 1, 1, [1], id="time-limit-with-interval"), + pytest.param(3, 2, 1, [1], id="time-limit-without-interval"), + ], +) +def test_train_final_checkpoint( + trainer_module, tmp_path, grad_accum, steps, interval, deadline, saved_steps +): + class TinyTrainer(trainer_module.DistillationTrainer): + def __init__(self): + self._accelerator = Accelerator(cpu=True, gradient_accumulation_steps=grad_accum) + self._config = SimpleNamespace( + seed=42, + output_dir=str(tmp_path), + model=SimpleNamespace(training_mode="full"), + optimization=SimpleNamespace( + steps=steps, + gradient_accumulation_steps=grad_accum, + max_grad_norm=1.0, + batch_size=1, + ), + checkpoints=SimpleNamespace(interval=interval, keep_last_n=1, precision="float32"), + validation=SimpleNamespace(interval=None, skip_initial_validation=True), + wandb=SimpleNamespace(log_validation_videos=False), + hub=SimpleNamespace(push_to_hub=False), + ) + self._distillation_config = trainer_module.DistillationConfig(must_save_by=deadline) + self._transformer = torch.nn.Linear(2, 1) + self._lr_scheduler = None + self._wandb_run = None + self._checkpoint_paths = [] + + def _init_optimizer(self): + optimizer = torch.optim.AdamW(self._transformer.parameters(), lr=0.01) + self._transformer, self._optimizer = self._accelerator.prepare( + self._transformer, optimizer + ) + self._trainable_params = list(self._transformer.parameters()) + + def _init_dataloader(self): + self._dataloader = [torch.ones(1, 2)] * (steps * grad_accum) + + def _training_step(self, batch): + return self._transformer(batch).square().mean() + + def _init_timestep_sampler(self): + pass + + def _save_config(self): + pass + + def _log_metrics(self, *args): + pass + + def _log_training_stats(self, *args): + pass + + trainer = TinyTrainer() + save_state = trainer._save_training_state + observed_steps = [] + + def record_save(): + observed_steps.append(trainer._global_step) + return save_state() + + trainer._save_training_state = record_save + trainer._save_checkpoint = Mock(wraps=trainer._save_checkpoint) + try: + exported_path, _ = trainer.train(disable_progress_bars=True) + assert observed_steps == saved_steps + trainer._save_checkpoint.assert_called_once() + assert trainer._checkpoint_paths == [exported_path] + exported = load_file(exported_path) + assert all( + torch.equal(exported[k], v) for k, v in trainer._transformer.state_dict().items() + ) + checkpoint_dir = tmp_path / "checkpoints" + assert sorted(p.name for p in checkpoint_dir.glob("step_*")) == [ + f"step_{saved_steps[-1]:06d}" + ] + metadata = checkpoint_dir / f"step_{saved_steps[-1]:06d}" / "distillation_metadata.json" + assert json.loads(metadata.read_text())["global_step"] == saved_steps[-1] + finally: + trainer._accelerator.free_memory() diff --git a/tests/examples/hf_ptq/test_example_utils.py b/tests/examples/hf_ptq/test_example_utils.py index a6e44a0c363..0ed9f41c730 100644 --- a/tests/examples/hf_ptq/test_example_utils.py +++ b/tests/examples/hf_ptq/test_example_utils.py @@ -371,21 +371,55 @@ def test_is_diffusion_gemma(hf_config, expected): @pytest.mark.parametrize( - ("trust_remote_code", "expect_bundled_code"), - [(True, True), (False, False)], + ( + "architecture", + "remote_config", + "builtin_config_error", + "trust_remote_code", + "expect_bundled", + ), + [ + ("DeepseekV3ForCausalLM", False, None, True, True), + ("DeepseekV3ForCausalLM", False, None, False, False), + # Dense Nemotron-H: the built-in class cannot parse the remote config -> bundled code. + ("NemotronHForCausalLM", True, KeyError("-"), True, True), + # MoE Nemotron-H (or a newer Transformers): the built-in config parses -> built-in class. + ("NemotronHForCausalLM", True, None, True, False), + # An unrelated failure (e.g. Hub retrieval) is not an incompatibility: it propagates. + ("NemotronHForCausalLM", True, OSError("hub down"), True, None), + ], ) -def test_get_model_deepseek_honors_trust_remote_code( - monkeypatch, trust_remote_code, expect_bundled_code +def test_get_model_selects_bundled_or_builtin_class( + monkeypatch, + tmp_path, + *, + architecture, + remote_config, + builtin_config_error, + trust_remote_code, + expect_bundled, ): - """DeepSeek ships bundled modeling code; --trust_remote_code selects it, else built-in.""" + """--trust_remote_code selects bundled code for DeepSeek and for a config the built-in + class cannot represent; otherwise the built-in class is used.""" used = {} - hf_config = SimpleNamespace( - architectures=["DeepseekV3ForCausalLM"], + config_cls = type( + "_Config", + (SimpleNamespace,), + {"__module__": "transformers_modules.ckpt.config" if remote_config else __name__}, + ) + hf_config = config_cls( + architectures=[architecture], dtype=torch.bfloat16, - model_type="deepseek_v3", + model_type="nemotron_h" if "NemotronH" in architecture else "deepseek_v3", torch_dtype=torch.bfloat16, ) + def _load_config(*args, **kwargs): + # The re-derivation of a built-in config is the call made without trust_remote_code. + if builtin_config_error and not kwargs.get("trust_remote_code"): + raise builtin_config_error + return hf_config + class FakeModel: def eval(self): return None @@ -407,20 +441,43 @@ def from_pretrained(*args, output_loading_info=False, **kwargs): return Fake - monkeypatch.setattr(example_utils.AutoConfig, "from_pretrained", lambda *a, **k: hf_config) + monkeypatch.setattr(example_utils.AutoConfig, "from_pretrained", _load_config) monkeypatch.setattr(example_utils, "AutoModelForCausalLM", _record("bundled")) - monkeypatch.setattr( - example_utils.transformers, "DeepseekV3ForCausalLM", _record("builtin"), raising=False - ) + monkeypatch.setattr(example_utils.transformers, architecture, _record("builtin"), raising=False) + monkeypatch.setattr(example_utils, "prepare_model_for_loading", lambda *a, **k: nullcontext()) monkeypatch.setattr(example_utils, "is_nemotron_vl", lambda config: False) monkeypatch.setattr(example_utils, "is_speculative", lambda config: False) monkeypatch.setattr(example_utils, "init_empty_weights", lambda include_buffers: nullcontext()) monkeypatch.setattr(example_utils, "get_max_memory", lambda: {0: 1024}) monkeypatch.setattr(example_utils, "infer_auto_device_map", lambda model, max_memory: {"": 0}) - example_utils.get_model("checkpoint", device="cpu", trust_remote_code=trust_remote_code) + if expect_bundled is None: + with pytest.raises(OSError): + example_utils.get_model( + str(tmp_path), device="cpu", trust_remote_code=trust_remote_code + ) + return + example_utils.get_model(str(tmp_path), device="cpu", trust_remote_code=trust_remote_code) + + assert used["path"] == ("bundled" if expect_bundled else "builtin") + + +def test_nemotron_h_prefers_bundled_code_for_unsupported_block_type(monkeypatch): + """A built-in config that parses but names a block type the built-in model lacks needs the + checkpoint's bundled code.""" + modeling = pytest.importorskip("transformers.models.nemotron_h.modeling_nemotron_h") + remote_cls = type("_Config", (SimpleNamespace,), {"__module__": "transformers_modules.c.cfg"}) + hf_config = remote_cls(model_type="nemotron_h") + monkeypatch.setattr( + example_utils.AutoConfig, + "from_pretrained", + lambda *a, **k: SimpleNamespace(layers_block_type=["mamba", "mlp"]), + ) - assert used["path"] == ("bundled" if expect_bundled_code else "builtin") + monkeypatch.setattr(modeling, "MIXER_TYPES", {"mamba": object}) + assert example_utils.prefers_bundled_code("nemotron_h", hf_config, "ckpt", {}) + monkeypatch.setattr(modeling, "MIXER_TYPES", {"mamba": object, "mlp": object}) + assert not example_utils.prefers_bundled_code("nemotron_h", hf_config, "ckpt", {}) # ---------- get_tokenizer ---------------------------------------------------- diff --git a/tests/examples/llm_sparsity/weight_sparsity/test_llama_sparsify.py b/tests/examples/llm_sparsity/weight_sparsity/test_llama_sparsify.py index 6277b0fa5b4..f6a7e31d7ec 100644 --- a/tests/examples/llm_sparsity/weight_sparsity/test_llama_sparsify.py +++ b/tests/examples/llm_sparsity/weight_sparsity/test_llama_sparsify.py @@ -112,6 +112,19 @@ def _test_llama_sparsity_finetune(tiny_llama_path, tmp_path, data_path, sparsity max_length=64, ) + # Finally restore the finetuned (FSDP-saved) state into an unsharded model and export it + run_example_command( + extend_cmd_parts( + ["python", "export_hf_ckpt.py"], + model_name_or_path=tiny_llama_path, + modelopt_restore_path=finetune_output / "finetuned_modelopt_state.pth", + dtype=dtype, + model_max_length=64, + output_dir=tmp_path / "export_output", + ), + "llm_sparsity/weight_sparsity", + ) + @pytest.mark.parametrize( ("sparsity_fmt", "dtype"), diff --git a/tests/gpu/torch/export/test_fsdp2_export.py b/tests/gpu/torch/export/test_fsdp2_export.py index 939fe3e581f..bbd1195d164 100644 --- a/tests/gpu/torch/export/test_fsdp2_export.py +++ b/tests/gpu/torch/export/test_fsdp2_export.py @@ -16,6 +16,7 @@ import json from functools import partial from pathlib import Path +from unittest.mock import patch import pytest import torch @@ -40,6 +41,7 @@ ) from modelopt.torch.export.unified_export_hf_streaming import _export_fsdp2_checkpoint_streaming from modelopt.torch.quantization.utils import ( + core_utils, enable_weight_access_and_writeback, fsdp2_aware_weight_update, module_name_maps, @@ -83,6 +85,35 @@ def _update_weight_test(rank, size): assert torch.allclose(torch.zeros(4, 4).to(output.device).to(output.dtype), output) +def _update_weight_error_propagates_test(rank, size): + """A setup error must propagate unchanged and leave the FSDP param group intact.""" + with patch_fsdp_mp_dtypes(): + model = ToyModel(dims=[4, 4], bias=False).to("cuda") + fully_shard(model.linears) + fully_shard(model) + + torch.distributed.barrier() + + class _SentinelError(RuntimeError): + pass + + def _raise(*args, **kwargs): + raise _SentinelError + + param_group = fully_shard.state(model.linears)._fsdp_param_group + fsdp_params = list(param_group.fsdp_params) + + # The unshard failure stands in for an OOM. + for target, attr in [(model.linears, "unshard"), (core_utils, "create_fsdp_param_mapping")]: + with ( + patch.object(target, attr, _raise), + pytest.raises(_SentinelError), + fsdp2_aware_weight_update(model, model.linears), + ): + pass + assert param_group.fsdp_params == fsdp_params + + def _compress_weight_test(rank, size): """Test fsdp2 weight update context for weight compression -> only value,shape and dtype changed""" with patch_fsdp_mp_dtypes(): @@ -262,6 +293,10 @@ def test_fsdp2_weight_update_context_for_export(dist_workers): dist_workers.run(_update_weight_test) +def test_fsdp2_weight_update_context_error_propagates(dist_workers): + dist_workers.run(_update_weight_error_propagates_test) + + @pytest.mark.parametrize( "quant_config", [ diff --git a/tests/gpu/torch/export/test_layerwise_export.py b/tests/gpu/torch/export/test_layerwise_export.py index 84244a8fec1..e442fe11c39 100644 --- a/tests/gpu/torch/export/test_layerwise_export.py +++ b/tests/gpu/torch/export/test_layerwise_export.py @@ -27,7 +27,7 @@ get_tiny_llama, get_tiny_qwen3_moe, ) -from safetensors.torch import load_file +from safetensors.torch import load_file, save_file import modelopt.torch.quantization as mtq from modelopt.torch.export.layerwise_export import ( @@ -610,20 +610,38 @@ def test_tied_embeddings_are_refused(tmp_path, make_cfg): mtq.quantize(model, cfg, _calib) -def test_moe_export_matches(tmp_path): - """MoE layers take a different path: fused expert inputs and gate/up amax sync.""" +@pytest.mark.parametrize("carry_auxiliary_weight", [False, True]) +def test_moe_export_matches(tmp_path, carry_auxiliary_weight): + """Packing fused experts must not reclassify their source tensors as unplaced.""" + source = tmp_path / "source" + source.mkdir() + baseline, layerwise = _build_moe_model(), _build_moe_model() + source_state = {k: v.cpu().contiguous() for k, v in baseline.state_dict().items()} + auxiliary_key = "mtp.eh_proj.weight" + if carry_auxiliary_weight: + source_state[auxiliary_key] = torch.ones(2, 2, dtype=torch.bfloat16) + save_file(source_state, source / "model.safetensors") + for model in (baseline, layerwise): + model._modelopt_source_checkpoint = str(source) + # Transformers can hide auxiliary weights from its unexpected_keys report. + model._modelopt_unplaced_source_keys = [] + baseline_dir = tmp_path / "baseline" base = _nvfp4_cfg() base["algorithm"] = {"method": "max", "layerwise": {"enable": True}} - export_hf_checkpoint(mtq.quantize(_build_moe_model(), base, _calib), export_dir=baseline_dir) + export_hf_checkpoint(mtq.quantize(baseline, base, _calib), export_dir=baseline_dir) export_dir = tmp_path / "fused" - _layerwise_quantize( - _build_moe_model(), _layerwise_cfg(export_dir, tmp_path / "ckpt", base=_nvfp4_cfg()) + mtq.quantize( + layerwise, _layerwise_cfg(export_dir, tmp_path / "ckpt", base=_nvfp4_cfg()), _calib ) + export_hf_checkpoint(layerwise, export_dir=export_dir) - _assert_same_checkpoint(_load_checkpoint(baseline_dir), _load_checkpoint(export_dir)) + exported = _load_checkpoint(export_dir) + _assert_same_checkpoint(_load_checkpoint(baseline_dir), exported) _assert_same_quant_config(baseline_dir, export_dir) + if carry_auxiliary_weight: + assert torch.equal(exported[auxiliary_key], source_state[auxiliary_key]) def test_export_consumes_the_model_without_affecting_the_checkpoint(tmp_path): diff --git a/tests/gpu/torch/puzzletron/test_nas_search.py b/tests/gpu/torch/puzzletron/test_nas_search.py index ae0dff47be5..2cdc3e26d44 100644 --- a/tests/gpu/torch/puzzletron/test_nas_search.py +++ b/tests/gpu/torch/puzzletron/test_nas_search.py @@ -66,7 +66,7 @@ def _test_nas_search_multiprocess_job( ) # - # Run the mnt.search() step + # Run the mtn.search() step # mtn.search( converted_model, diff --git a/tests/gpu/torch/quantization/test_gptq.py b/tests/gpu/torch/quantization/test_gptq.py index 496db08578a..b4b965b3154 100644 --- a/tests/gpu/torch/quantization/test_gptq.py +++ b/tests/gpu/torch/quantization/test_gptq.py @@ -208,6 +208,20 @@ def _run_fused_gptq_nvfp4( return w +@requires_triton +def test_fused_nvfp4_gptq_runs_on_the_weights_gpu(need_2_gpus): + """A layer device_map placed on another GPU is updated there, not from the current GPU.""" + torch.manual_seed(42) + weight, _, block_amax, global_scale, h_inv = _make_nvfp4_test_data(16, 64, 256) + expected = _run_fused_gptq_nvfp4(weight, block_amax, global_scale, h_inv, 128, 16) + + weight, block_amax, h_inv = (t.to("cuda:1") for t in (weight, block_amax, h_inv)) + got = _run_fused_gptq_nvfp4(weight, block_amax, global_scale, h_inv, 128, 16) + + assert torch.cuda.current_device() == 0 + torch.testing.assert_close(got.cpu(), expected.cpu(), rtol=0, atol=0) + + _NVFP4_QUANT_BLOCK_SIZES = [16, 128] _NVFP4_GPTQ_BLOCK_SIZES = [16, 128] diff --git a/tests/gpu/torch/quantization/test_quant_conv_gpu.py b/tests/gpu/torch/quantization/test_quant_conv_gpu.py new file mode 100644 index 00000000000..845db583c77 --- /dev/null +++ b/tests/gpu/torch/quantization/test_quant_conv_gpu.py @@ -0,0 +1,65 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""GPU regression tests for quantized convolutions.""" + +import pytest +import torch + +import modelopt.torch.quantization as mtq + + +@pytest.mark.parametrize("conv_kind", ["conv3d", "wan", "wan_cached"]) +@pytest.mark.parametrize("quantize_weight", [False, True]) +@torch.no_grad() +def test_nvfp4_conv3d_disable_quant(cuda_capability, conv_kind, quantize_weight): + """Disabling activation QDQ must match bypassing its quantizer, even in the fused kernel.""" + if cuda_capability < (8, 9): + pytest.skip("NVFP4 quantization requires compute capability >= 8.9") + + torch.manual_seed(1234) + conv_cls = torch.nn.Conv3d + if conv_kind != "conv3d": + # Diffusers and its quantization plugin are optional for the regular Conv3d case. + pytest.importorskip("onnx") + pytest.importorskip("diffusers") + from diffusers.models.autoencoders.autoencoder_kl_wan import WanCausalConv3d + + import modelopt.torch.quantization.plugins.diffusion.diffusers # noqa: F401 + + conv_cls = WanCausalConv3d + + model = conv_cls(16, 16, kernel_size=3, padding=1).to("cuda", torch.bfloat16).eval() + x = torch.randn(1, 16, 3, 6, 6, device="cuda", dtype=torch.bfloat16) + kwargs = {} + if conv_kind == "wan_cached": + kwargs["cache_x"] = torch.randn(1, 16, 1, 6, 6, device="cuda", dtype=torch.bfloat16) + model = mtq.quantize(model, mtq.NVFP4_DEFAULT_CFG, lambda m: m(x, **kwargs)) + if not quantize_weight: + model.weight_quantizer.disable_quant() + + enabled = model(x, **kwargs) + model.input_quantizer.disable_quant() + disabled = model(x, **kwargs) + + # Bypass the activation quantizer entirely as a control, keeping the weight state fixed. + model.input_quantizer.disable() + bypassed = model(x, **kwargs) + torch.testing.assert_close(disabled, bypassed, rtol=0, atol=0) + assert not torch.equal(enabled, disabled) + + model.input_quantizer.enable() + model.input_quantizer.enable_quant() + torch.testing.assert_close(model(x, **kwargs), enabled, rtol=0, atol=0) diff --git a/tests/gpu/torch/quantization/test_real_quantize_cuda.py b/tests/gpu/torch/quantization/test_real_quantize_cuda.py index 74ff7e93048..5af7cee0c8f 100644 --- a/tests/gpu/torch/quantization/test_real_quantize_cuda.py +++ b/tests/gpu/torch/quantization/test_real_quantize_cuda.py @@ -197,10 +197,12 @@ def forward_loop(model, run_backward=False): @pytest.mark.parametrize("quant_config", [mtq.FP8_DEFAULT_CFG]) -def test_real_quantize_linear(quant_config, tmp_path): - tiny_llama_dir = create_tiny_llama_dir(tmp_path) +@pytest.mark.parametrize("tie_word_embeddings", [False, True]) +def test_real_quantize_linear(quant_config, tie_word_embeddings, tmp_path): + tiny_llama_dir = create_tiny_llama_dir(tmp_path, tie_word_embeddings=tie_word_embeddings) with init_quantized_weights(quant_config): model = AutoModelForCausalLM.from_pretrained(tiny_llama_dir) + assert (model.lm_head.weight is model.model.embed_tokens.weight) == tie_word_embeddings for _, module in model.named_modules(): if ( diff --git a/tests/gpu/torch/quantization/test_tensor_quant_cuda.py b/tests/gpu/torch/quantization/test_tensor_quant_cuda.py index c1cbd6a09fa..3ab5edd0db4 100644 --- a/tests/gpu/torch/quantization/test_tensor_quant_cuda.py +++ b/tests/gpu/torch/quantization/test_tensor_quant_cuda.py @@ -219,6 +219,15 @@ def test_zero_block_scale_zeroes_block(self): assert torch.equal(output, torch.zeros_like(output)) + @pytest.mark.skipif(not triton_kernel.IS_AVAILABLE, reason="triton kernel is not available") + def test_static_blockwise_fp4_over_int32_numel(self): + x = torch.randn(2**31 + 2**20, device="cuda", dtype=torch.bfloat16) + amax = x.view(-1, 16).abs().amax(-1).float() + output = triton_kernel.static_blockwise_fp4_fake_quant(x, amax, amax.max()) + tail = triton_kernel.static_blockwise_fp4_fake_quant(x[-4096:], amax[-256:], amax.max()) + + assert torch.equal(output[-4096:], tail) + @pytest.mark.skipif(get_cuda_ext_mx() is None, reason="cuda_ext_mx is not available") @pytest.mark.parametrize( "set_torch_dtype", [torch.float, torch.float16, torch.bfloat16], indirect=True diff --git a/tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py b/tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py index 2bbcc108d69..18cc2807a8e 100644 --- a/tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py +++ b/tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py @@ -25,6 +25,7 @@ from _test_utils.torch.megatron.models import get_mcore_gpt_model from _test_utils.torch.megatron.utils import run_mcore_inference_with_dummy_input from _test_utils.torch.misc import set_seed +from megatron.core import parallel_state import modelopt.torch.distill as mtd from modelopt.torch.distill.plugins.megatron import ( @@ -465,26 +466,64 @@ def test_topk_logits_kl_loss_numerics_full_vocab_matches_dense(): assert torch.allclose(dense, topk, atol=1e-6) -def test_topk_logits_kl_loss_numerics_ghost_token_reference(): - """Top-K + ghost token matches a hand-written reference on the K+1 bucketed distributions.""" - cfg = SimpleNamespace(tensor_model_parallel_size=1) - student, teacher = _make_loss_inputs() +def _test_topk_logits_kl_loss_numerics_ghost_token_reference( + top_p, reverse, temperature, rank, size +): + parallel_state.initialize_model_parallel(tensor_model_parallel_size=size) + cfg = SimpleNamespace(tensor_model_parallel_size=size) + device = torch.device("cuda", rank) + # Equal vocab shards (torch.chunk is uneven when size does not divide the vocab) + student, teacher = _make_loss_inputs(vocab=16 * size) + student = student.detach().to(device).requires_grad_() + teacher = teacher.to(device) + student_shard = student.chunk(size, dim=-1)[rank].detach().contiguous().requires_grad_() + teacher_shard = teacher.chunk(size, dim=-1)[rank].contiguous() k = 4 - loss = TopLogitsKLLoss(cfg, top_k=k)(student, teacher)[0] + loss = TopLogitsKLLoss(cfg, top_k=k, top_p=top_p, reverse=reverse, temperature=temperature)( + student_shard, teacher_shard + )[0] - q_full = F.log_softmax(teacher, dim=-1) - p_full = F.log_softmax(student, dim=-1) + q_full = F.log_softmax(teacher / temperature, dim=-1) + p_full = F.log_softmax(student / temperature, dim=-1) _, idx = torch.topk(teacher, k, dim=-1) q_k, p_k = q_full.gather(-1, idx), p_full.gather(-1, idx) - q_rest = torch.log1p(-q_k.exp().sum(-1, keepdim=True)) - p_rest = torch.log1p(-p_k.exp().sum(-1, keepdim=True)) + keep = torch.ones_like(q_k, dtype=torch.bool) + if top_p is not None: + q_probs = q_k.exp() + keep = (q_probs.cumsum(-1) - q_probs) < top_p + assert not keep.all(), "test inputs should produce some truncation" + q_rest = torch.log1p(-(q_k.exp() * keep).sum(-1, keepdim=True)) + p_rest = torch.log1p(-(p_k.exp() * keep).sum(-1, keepdim=True)) q = torch.cat([q_k, q_rest], -1) p = torch.cat([p_k, p_rest], -1) - ref = (q.exp() * (q - p)).sum(-1).transpose(0, 1) - assert torch.allclose(loss, ref, atol=1e-6) + keep = torch.cat([keep, torch.ones_like(keep[..., :1])], -1) # Sanity: total mass within the K+1 buckets is 1 for both distributions. - assert torch.allclose(q.exp().sum(-1), torch.ones_like(q[..., 0]), atol=1e-5) - assert torch.allclose(p.exp().sum(-1), torch.ones_like(p[..., 0]), atol=1e-5) + assert torch.allclose((q.exp() * keep).sum(-1), torch.ones_like(q[..., 0]), atol=1e-5) + assert torch.allclose((p.exp() * keep).sum(-1), torch.ones_like(p[..., 0]), atol=1e-5) + if reverse: + p, q = q, p + ref = (keep * q.exp() * (q - p)).sum(-1).transpose(0, 1) + torch.testing.assert_close(loss, ref, atol=1e-6, rtol=1e-5) + loss.mean().backward() + ref.mean().backward() + torch.testing.assert_close( + student_shard.grad, student.grad.chunk(size, dim=-1)[rank], atol=1e-6, rtol=1e-5 + ) + + +@pytest.mark.parametrize( + ("top_p", "reverse", "temperature"), + [(None, False, 1.0), (0.5, False, 2.0), (0.5, True, 1.0)], +) +def test_topk_logits_kl_loss_numerics_ghost_token_reference( + dist_workers_size_2, top_p, reverse, temperature +): + """TP loss and gradients match a full-vocab reference on the K+1 bucketed distributions.""" + dist_workers_size_2.run( + partial( + _test_topk_logits_kl_loss_numerics_ghost_token_reference, top_p, reverse, temperature + ) + ) @pytest.mark.parametrize("temperature", [0.5, 2.0, 3.7]) diff --git a/tests/unit/examples/test_plugin_calib.py b/tests/unit/examples/test_plugin_calib.py new file mode 100644 index 00000000000..01eee88e4ec --- /dev/null +++ b/tests/unit/examples/test_plugin_calib.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +from examples.diffusers.quantization.calib.plugin_calib import PercentileCalibrator + + +@pytest.mark.parametrize( + ("observations", "expected_amax"), + [ + pytest.param([2.0, 4.0, 6.0, 8.0], 2.0, id="one-cycle"), + pytest.param([2.0, 4.0, 6.0, 8.0] * 3, 2.0, id="repeated-cycles"), + pytest.param([2.0, 4.0, 6.0, 8.0, 6.0, 8.0, 10.0, 12.0], 4.0, id="varying-cycles"), + pytest.param([2.0, 8.0, 6.0, 10.0, 6.0], 4.0, id="partial-final-cycle"), + ], +) +def test_min_mean_averages_observations_per_step(observations, expected_amax): + calibrator = PercentileCalibrator( + num_bits=8, axis=None, percentile=1.0, total_step=4, collect_method="min-mean" + ) + for value in observations: + calibrator.collect(torch.tensor([value])) + + torch.testing.assert_close(calibrator.compute_amax(), torch.tensor([expected_amax])) diff --git a/tests/unit/onnx/quantization/autotune/test_benchmark.py b/tests/unit/onnx/quantization/autotune/test_benchmark.py new file mode 100644 index 00000000000..c6d98a164c4 --- /dev/null +++ b/tests/unit/onnx/quantization/autotune/test_benchmark.py @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the default timing-cache file in autotune/benchmark.py.""" + +import stat +import sys +import tempfile +from pathlib import Path + +from modelopt.onnx.quantization.autotune.benchmark import Benchmark, default_timing_cache_file + + +class _NoOpBenchmark(Benchmark): + def run(self, path_or_bytes, log_file=None): + return 0.0 + + +def test_default_timing_cache_is_private_and_shared_within_a_process(tmp_path, monkeypatch): + monkeypatch.setattr(tempfile, "tempdir", str(tmp_path)) + default_timing_cache_file.cache_clear() + first = Path(_NoOpBenchmark().timing_cache_file) + second = Path(_NoOpBenchmark().timing_cache_file) + default_timing_cache_file.cache_clear() + + # A fixed shared path would let another local user pre-create the file or plant a symlink. + assert first == second + assert first.name == "trtexec_timing.cache" + assert first.parent.parent == tmp_path + if sys.platform != "win32": # POSIX file modes + assert stat.S_IMODE(first.parent.stat().st_mode) == 0o700 diff --git a/tests/unit/onnx/quantization/test_extensions.py b/tests/unit/onnx/quantization/test_extensions.py new file mode 100644 index 00000000000..faab210ba11 --- /dev/null +++ b/tests/unit/onnx/quantization/test_extensions.py @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import importlib + +import cppimport + +import modelopt.onnx.quantization.extensions as ext + + +def test_failed_build_falls_back_to_python(monkeypatch): + def build_fails(*args, **kwargs): + # What setuptools raises when, e.g., there is no C++ compiler. + raise SystemExit("error: command 'c++' failed: No such file or directory") + + monkeypatch.setattr(cppimport, "imp", build_fails) + try: + importlib.reload(ext) + assert ext.round_and_pack_ext is None + finally: + monkeypatch.undo() + importlib.reload(ext) diff --git a/tests/unit/onnx/test_external_data.py b/tests/unit/onnx/test_external_data.py new file mode 100644 index 00000000000..ddb77e997f8 --- /dev/null +++ b/tests/unit/onnx/test_external_data.py @@ -0,0 +1,42 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import numpy as np +import onnx +import pytest + +from modelopt.onnx.utils import save_onnx + + +@pytest.mark.parametrize("last_byte", [0x00, 0x0A, 0x1A]) +def test_save_external_data_preserves_tensor_boundaries(tmp_path, last_byte): + # Windows text-mode file opening can truncate a trailing Ctrl-Z before append. + first = np.full(1024, 0x41, dtype=np.uint8) + first[-1] = last_byte + second = np.full(1024, 0x42, dtype=np.uint8) + tensors = [ + onnx.numpy_helper.from_array(value, name) + for name, value in [("first", first), ("second", second)] + ] + graph = onnx.helper.make_graph([], "external_data", [], [], initializer=tensors) + model = onnx.helper.make_model(graph) + path = tmp_path / "model.onnx" + + save_onnx(model, str(path), save_as_external_data=True) + + assert (tmp_path / "model.onnx_data").read_bytes() == first.tobytes() + second.tobytes() + restored = onnx.load(path) + for actual, expected in zip(restored.graph.initializer, [first, second], strict=True): + np.testing.assert_array_equal(onnx.numpy_helper.to_array(actual), expected) diff --git a/tests/unit/onnx/test_onnx_utils.py b/tests/unit/onnx/test_onnx_utils.py index 590a9f24186..209e9c46036 100644 --- a/tests/unit/onnx/test_onnx_utils.py +++ b/tests/unit/onnx/test_onnx_utils.py @@ -304,7 +304,7 @@ def test_random_onnx_weights(): # Removed model weights should be greater than 18 MB assert original_model_size - len(onnx_model_wo_weights) > 18e6 - # After assigning random weights, model size should be slightly greater than the the original + # After assigning random weights, model size should be slightly greater than the original # size due to some extra metadata onnx_model_randomized = randomize_weights_onnx_bytes(onnx_model_wo_weights) assert len(onnx_model_randomized) > original_model_size diff --git a/tests/unit/recipe/test_minimax_m2_7_recipe.py b/tests/unit/recipe/test_minimax_m2_7_recipe.py new file mode 100644 index 00000000000..fd54dcd7888 --- /dev/null +++ b/tests/unit/recipe/test_minimax_m2_7_recipe.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch +import torch.nn as nn + +import modelopt.torch.quantization as mtq +from modelopt.recipe import load_recipe +from modelopt.torch.quantization.nn import TensorQuantizer +from modelopt.torch.quantization.plugins.huggingface import register_fused_experts_on_the_fly + + +class _FusedExperts(nn.Module): + def __init__(self): + super().__init__() + self.num_experts = 2 + self.intermediate_dim = 32 + self.gate_up_proj = nn.Parameter(torch.randn(2, 64, 32)) + self.down_proj = nn.Parameter(torch.randn(2, 32, 32)) + + def forward(self, hidden_states): + return hidden_states + + +@pytest.mark.parametrize("layout", ["block_sparse_moe", "mlp"]) +@pytest.mark.parametrize("fused", [False, True]) +def test_minimax_m2_7_expert_selection(layout, fused): + model = nn.Module() + model.model = nn.Module() + layer = nn.Module() + model.model.layers = nn.ModuleList([layer]) + moe = nn.Module() + layer.add_module(layout, moe) + moe.experts = ( + _FusedExperts() + if fused + else nn.ModuleList( + [nn.ModuleDict({p: nn.Linear(32, 32) for p in ("gate_proj", "up_proj", "down_proj")})] + ) + ) + moe.gate = nn.Linear(32, 2) + moe.shared_experts = nn.Linear(32, 32) + layer.self_attn = nn.ModuleDict( + {p: nn.Linear(32, 32) for p in ("q_proj", "k_proj", "v_proj", "o_proj")} + ) + model.lm_head = nn.Linear(32, 32) + preserved = {name: p.detach().clone() for name, p in model.named_parameters()} + register_fused_experts_on_the_fly(model) + + config = load_recipe( + "models/MiniMaxAI/MiniMax-M2.7/ptq/nvfp4_experts_only-kv_fp8_cast" + ).quantize.model_dump() + assert config["algorithm"]["method"] == "max" + config["algorithm"] = None + mtq.quantize(model, config) + + enabled = { + name: q + for name, q in model.named_modules() + if isinstance(q, TensorQuantizer) and q.is_enabled + } + assert len(enabled) == 6 # Three W4A4 projections, or four fused weights and two inputs. + for name, q in enabled.items(): + assert f".{layout}.experts." in name + assert q.num_bits == (2, 1) + assert q.block_sizes[-1] == 16 + assert q.block_sizes["type"] == "dynamic" + for name, p in model.named_parameters(): + torch.testing.assert_close(p, preserved[name], rtol=0, atol=0) diff --git a/tests/unit/recipe/test_presets.py b/tests/unit/recipe/test_presets.py index 3474cd51a43..167aa46e5a0 100644 --- a/tests/unit/recipe/test_presets.py +++ b/tests/unit/recipe/test_presets.py @@ -24,6 +24,7 @@ import argparse import pytest +import torch import modelopt.torch.quantization as mtq from modelopt.recipe import load_recipe, presets @@ -95,6 +96,25 @@ def test_w4a16_nvfp4_preset_disables_vllm_marlin_incompatible_projections(): } <= disabled_quantizers +def test_default_exclusions_disable_mtp_at_any_depth(): + # Megatron VLMs nest the MTP block under ``language_model.``; HF and Megatron LLMs keep it at the root. + model = torch.nn.Module() + for path in ("mtp", "language_model.mtp", "language_model.decoder"): + parent = model + for part in path.split("."): + if not hasattr(parent, part): + parent.add_module(part, torch.nn.Module()) + parent = getattr(parent, part) + parent.add_module("proj", torch.nn.Linear(16, 16)) + + quant_cfg = load_recipe("general/ptq/fp8_default-kv_fp8").quantize.model_dump() + mtq.quantize(model, {**quant_cfg, "algorithm": None}) + + assert model.language_model.decoder.proj.input_quantizer.is_enabled + assert not model.mtp.proj.input_quantizer.is_enabled + assert not model.language_model.mtp.proj.input_quantizer.is_enabled + + def test_local_hessian_layerwise_is_scoped_to_the_model_recipe(): # Layerwise calibration needs identifiable decoder layers, so it belongs to the model # recipe the published numbers come from -- not to the shared preset, which also backs diff --git a/tests/unit/torch/deploy/_runtime/tensorrt/test_engine_builder.py b/tests/unit/torch/deploy/_runtime/tensorrt/test_engine_builder.py index ff7f77cf617..4385ef133bf 100755 --- a/tests/unit/torch/deploy/_runtime/tensorrt/test_engine_builder.py +++ b/tests/unit/torch/deploy/_runtime/tensorrt/test_engine_builder.py @@ -15,9 +15,10 @@ import json import os +import stat import sys +import tempfile from pathlib import Path -from tempfile import gettempdir from types import ModuleType from unittest import mock @@ -30,12 +31,15 @@ tensorrt.__version__ = "10.10" sys.modules["tensorrt"] = tensorrt -from modelopt.torch._deploy._runtime.tensorrt.engine_builder import build_engine, profile_engine +from modelopt.torch._deploy._runtime.tensorrt.engine_builder import ( + _default_artifact_dir, + build_engine, + profile_engine, +) @pytest.fixture -def setup_mocks(): - tmp_path = Path(gettempdir()) / "modelopt_build/trt_artifacts" +def setup_mocks(tmp_path): engine_bytes = b"engine_bytes" dummy_hash = "997930f2" expected_engine_bytes = b"\x99y0\xf2" + b"A" * 28 + b",engine_bytes" @@ -125,3 +129,30 @@ def test_profile_engine_variants(setup_mocks, enable_layerwise, profiling_runs): assert profiled_data == {} assert out == "success" + + +def _write_engine_like_trtexec(cmd): + for arg in cmd: + if arg.startswith("--saveEngine="): + Path(arg.split("=", 1)[1]).write_bytes(b"engine_bytes") + return 0, "success" + + +def test_default_artifact_dir_is_private_and_reused_within_a_process(setup_mocks, monkeypatch): + tmp_path, mock_run, mock_onnx, _ = setup_mocks + shared_tmp = tmp_path / "shared_tmp" + shared_tmp.mkdir() + monkeypatch.setattr(tempfile, "tempdir", str(shared_tmp)) + _default_artifact_dir.cache_clear() + mock_run.side_effect = _write_engine_like_trtexec + + build_engine(onnx_bytes=mock_onnx) + build_engine(onnx_bytes=mock_onnx) + profile_engine(engine_bytes=b"\x99y0\xf2" + b"A" * 28 + b",engine_bytes") + _default_artifact_dir.cache_clear() + + # A fixed shared path would let another local user pre-create it or plant symlinks in it. One + # unpredictable private directory per process keeps the timing cache shared across calls. + (artifact_dir,) = shared_tmp.iterdir() + if sys.platform != "win32": # POSIX file modes + assert stat.S_IMODE(artifact_dir.stat().st_mode) == 0o700 diff --git a/tests/unit/torch/distill/test_distill.py b/tests/unit/torch/distill/test_distill.py index 7c938aee90d..721ff4ef9a6 100644 --- a/tests/unit/torch/distill/test_distill.py +++ b/tests/unit/torch/distill/test_distill.py @@ -15,6 +15,7 @@ import inspect import warnings +from contextlib import nullcontext import pytest import torch @@ -26,6 +27,17 @@ import modelopt.torch.opt as mto +class LinearModel(nn.Module): + def __init__(self, scale): + super().__init__() + self.linear = nn.Linear(2, 2, bias=False) # Wx + with torch.no_grad(): + self.linear.weight.copy_(scale * torch.eye(2)) + + def forward(self, inputs): + return self.linear(inputs) + + def get_input_tensor(): """Dummy input tensor.""" return torch.rand(2, 3, 112, 112) @@ -310,6 +322,63 @@ def test_student_fwd_only(distillation_model): assert distillation_model._intermediate_output is not None +@pytest.mark.parametrize("context_name", ["only_student_forward", "only_teacher_forward"]) +@pytest.mark.parametrize("inner_enabled", [False, True]) +@pytest.mark.parametrize("raise_in_inner", [False, True]) +def test_nested_forward_contexts(context_name, inner_enabled, raise_in_inner): + student, teacher = LinearModel(1), LinearModel(2) + model = mtd.convert( + student, mode=[("kd_loss", {"teacher_model": teacher, "criterion": nn.MSELoss()})] + ) + context = getattr(model, context_name) + inputs = torch.ones(1, 2) + calls = {"student": 0, "teacher": 0} + + def count_student(module, inputs): + calls["student"] += 1 + + def count_teacher(module, inputs): + calls["teacher"] += 1 + + model.linear.register_forward_pre_hook(count_student) + teacher.linear.register_forward_pre_hook(count_teacher) + expected = inputs if context_name == "only_student_forward" else 2 * inputs + + with context(): + torch.testing.assert_close(model(inputs), expected) + with ( + pytest.raises(RuntimeError, match="inner failure") if raise_in_inner else nullcontext(), + context(enable=inner_enabled), + ): + torch.testing.assert_close(model(inputs), expected) + if raise_in_inner: + raise RuntimeError("inner failure") + torch.testing.assert_close(model(inputs), expected) + + assert calls["teacher"] == (0 if context_name == "only_student_forward" else 3) + assert calls["student"] == (3 if context_name == "only_student_forward" else 0) + torch.testing.assert_close(model(inputs), inputs) + assert calls["teacher"] == (1 if context_name == "only_student_forward" else 4) + assert calls["student"] == (4 if context_name == "only_student_forward" else 1) + + with pytest.raises(RuntimeError, match="outer failure"), context(): + raise RuntimeError("outer failure") + with context(enable=False): + torch.testing.assert_close(model(inputs), inputs) + assert calls["teacher"] == (2 if context_name == "only_student_forward" else 5) + assert calls["student"] == (5 if context_name == "only_student_forward" else 2) + + other_context = ( + model.only_teacher_forward + if context_name == "only_student_forward" + else model.only_student_forward + ) + with context(): + with other_context(): + torch.testing.assert_close(model(inputs), inputs) + torch.testing.assert_close(model(inputs), expected) + + def test_train_eval_mode_switch(distillation_model): distillation_model.train() assert distillation_model._intermediate_output is None diff --git a/tests/unit/torch/export/test_export_diffusers.py b/tests/unit/torch/export/test_export_diffusers.py index 126550a978c..001b4245b72 100644 --- a/tests/unit/torch/export/test_export_diffusers.py +++ b/tests/unit/torch/export/test_export_diffusers.py @@ -17,6 +17,7 @@ import inspect import json from pathlib import Path +from types import SimpleNamespace import numpy as np import pytest @@ -26,6 +27,7 @@ get_tiny_dit, get_tiny_flux, get_tiny_flux2, + get_tiny_qwen_image_transformer, get_tiny_unet, get_tiny_wan22_transformer, get_tiny_wan22_vae, @@ -36,11 +38,13 @@ import diffusers from safetensors import safe_open -from safetensors.torch import save_file +from safetensors.torch import load_file, save_file from transformers import ( CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, + LlamaConfig, + LlamaModel, UMT5Config, UMT5EncoderModel, ) @@ -53,6 +57,7 @@ hide_quantizers_from_state_dict, ) from modelopt.torch.export.unified_export_hf import _postprocess_safetensors, export_hf_checkpoint +from modelopt.torch.quantization.nn import TensorQuantizer def _load_config(config_path): @@ -102,6 +107,25 @@ def test_export_diffusers_models_non_quantized(tmp_path, model_factory): assert "quantization_config" not in config_data +@pytest.mark.parametrize("quantized", [False, True]) +@pytest.mark.parametrize("save_modelopt_state", [False, True]) +@torch.no_grad() +def test_export_diffusers_save_modelopt_state( + tmp_path, enable_hf_checkpointing, quantized, save_modelopt_state +): + model = get_tiny_unet().eval() + inputs = generate_diffusion_dummy_inputs(model, torch.device("cpu"), torch.float32) + mtq.quantize(model, mtq.FP8_DEFAULT_CFG, lambda m: m(**inputs)) + if not quantized: + # Keep ModelOpt state while exercising the non-quantized save branch. + mtq.disable_quantizer(model, "*") + export_dir = tmp_path / "export_unet" + + export_hf_checkpoint(model, export_dir=export_dir, save_modelopt_state=save_modelopt_state) + + assert (export_dir / "modelopt_state.pth").is_file() == save_modelopt_state + + @pytest.fixture def tiny_pipeline(request, tmp_path, tiny_tokenizer): """Build and reload a tiny pipeline without accessing the Hugging Face Hub.""" @@ -281,6 +305,177 @@ def test_flux2_dummy_inputs_shape(): assert "guidance" in inputs +@pytest.mark.parametrize("pos_embed_seq_len", [None, 514], ids=["i2v", "flf2v"]) +@torch.no_grad() +def test_wan_dummy_inputs_support_qkv_fusion(pos_embed_seq_len): + torch.manual_seed(0) + model = get_tiny_wan22_transformer( + image_dim=32, added_kv_proj_dim=24, pos_embed_seq_len=pos_embed_seq_len + ).eval() + inputs = generate_diffusion_dummy_inputs(model, torch.device("cpu"), torch.float32) + assert inputs is not None + image_batch_size = 1 if pos_embed_seq_len is None else 2 + assert inputs["encoder_hidden_states_image"].shape == (image_batch_size, 257, 32) + assert inputs["encoder_hidden_states"].shape == (1, 512, 32) + + # Per-tensor INT8 exercises both input and weight amax fusion on CPU. + quant_cfg = copy.deepcopy(mtq.INT8_DEFAULT_CFG) + quant_cfg["quant_cfg"].append( + {"quantizer_name": "*weight_quantizer", "cfg": {"num_bits": 8, "axis": None}} + ) + mtq.quantize(model, quant_cfg, lambda m: m(**inputs)) + + groups = [] + for block in model.blocks: + groups.extend( + [ + (block.attn1.to_q, block.attn1.to_k, block.attn1.to_v), + (block.attn2.to_k, block.attn2.to_v), + (block.attn2.add_k_proj, block.attn2.add_v_proj), + ] + ) + # Distinct values ensure a skipped fusion cannot pass on already-equal scales. + for group in groups: + for index, module in enumerate(group, start=1): + module.input_quantizer.amax = torch.tensor(float(index)) + module.weight_quantizer.amax = torch.tensor(float(index)) + + unified_export_hf._fuse_qkv_linears_diffusion(model, strict=True) + + for group in groups: + expected_amax = torch.tensor(float(len(group))) + for module in group: + torch.testing.assert_close(module.input_quantizer.amax, expected_amax) + torch.testing.assert_close(module.weight_quantizer.amax, expected_amax) + + +@pytest.mark.parametrize( + ("model_factory", "export_attention_quantizers", "expect_attention_quantizers"), + [ + (get_tiny_qwen_image_transformer, None, False), + (get_tiny_flux, None, True), + (get_tiny_flux, False, False), + (get_tiny_flux, True, True), + ], +) +def test_export_diffusers_attention_quantizers( + tmp_path, model_factory, export_attention_quantizers, expect_attention_quantizers +): + model = model_factory().eval() + quant_cfg = copy.deepcopy(mtq.FP8_DEFAULT_CFG) + quant_cfg["quant_cfg"].append( + { + "quantizer_name": "*[qkv]_bmm_quantizer", + "cfg": {"num_bits": (4, 3), "axis": None}, + "enable": True, + } + ) + + def calibrate(m): + inputs = generate_diffusion_dummy_inputs(m, torch.device("cpu"), torch.float32) + assert inputs is not None + m(**inputs) + + mtq.quantize(model, quant_cfg, calibrate) + # The example disables sensitive blocks after calibration, leaving their amax buffers. + mtq.disable_quantizer(model, "transformer_blocks.0.*") + quantizers = { + name: module + for name, module in model.named_modules() + if isinstance(module, TensorQuantizer) + } + attention_amax = { + f"{name}._amax": module.amax.clone() + for name, module in quantizers.items() + if name.endswith(("q_bmm_quantizer", "k_bmm_quantizer", "v_bmm_quantizer")) + and module.is_enabled + } + assert attention_amax + assert any(not module.is_enabled and hasattr(module, "_amax") for module in quantizers.values()) + + kwargs = ( + {"export_attention_quantizers": export_attention_quantizers} + if export_attention_quantizers is not None + else {} + ) + export_hf_checkpoint(model, export_dir=tmp_path, max_shard_size="1KB", **kwargs) + + exported = {} + for shard in tmp_path.glob("*.safetensors"): + exported.update(load_file(str(shard))) + expected_attention = attention_amax if expect_attention_quantizers else {} + assert {key for key in exported if "quantizer" in key} == set(expected_attention) + for key, value in expected_attention.items(): + torch.testing.assert_close(exported[key], value) + assert any(value.dtype == torch.float8_e4m3fn for value in exported.values()) + assert any(key.endswith(".weight_scale") for key in exported) + config = _load_config(tmp_path / "config.json")["quantization_config"] + assert not any(key.startswith("kv_cache") for key in config) + for name, module in quantizers.items(): + assert model.get_submodule(name) is module + for key, value in attention_amax.items(): + torch.testing.assert_close(model.state_dict()[key], value) + + +@pytest.mark.parametrize("is_transformers_component", [False, True]) +@pytest.mark.parametrize("export_attention_quantizers", [None, False, True]) +def test_export_diffusion_component_kv_metadata( + tmp_path, is_transformers_component, export_attention_quantizers +): + if is_transformers_component: + model = LlamaModel( + LlamaConfig( + vocab_size=32, + hidden_size=16, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + ) + ).eval() + + def calibrate(m): + m(input_ids=torch.ones(1, 4, dtype=torch.long)) + + else: + # Exercise the native nn.Module save path used by LTX-2, with real attention. + model = nn.Sequential(get_tiny_flux()).eval() + + def calibrate(m): + inputs = generate_diffusion_dummy_inputs(m[0], torch.device("cpu"), torch.float32) + m[0](**inputs) + + quant_cfg = copy.deepcopy(mtq.FP8_DEFAULT_CFG) + quant_cfg["quant_cfg"].append( + {"quantizer_name": "*[kv]_bmm_quantizer", "cfg": {"num_bits": (4, 3), "axis": None}} + ) + mtq.quantize(model, quant_cfg, calibrate) + kv_amax = { + f"{name}._amax": module.amax.clone() + for name, module in model.named_modules() + if name.endswith(("k_bmm_quantizer", "v_bmm_quantizer")) and module.is_enabled + } + assert kv_amax + # LTX-2 exposes its native transformer through the stage-1 ledger. + pipe = SimpleNamespace(stage_1_model_ledger=SimpleNamespace(transformer=lambda: model)) + unified_export_hf._export_diffusers_checkpoint( + pipe, None, tmp_path, None, export_attention_quantizers=export_attention_quantizers + ) + + config = _load_config(tmp_path / "config.json")["quantization_config"] + assert ("kv_cache_scheme" in config) == is_transformers_component + if is_transformers_component: + assert config["kv_cache_scheme"] == {"dynamic": False, "num_bits": 8, "type": "float"} + else: + assert not any(key.startswith("kv_cache") for key in config) + + exported = load_file(str(tmp_path / "model.safetensors")) + expected_kv = kv_amax if is_transformers_component or export_attention_quantizers else {} + assert {key for key in exported if "quantizer" in key} == set(expected_kv) + for key, value in expected_kv.items(): + torch.testing.assert_close(exported[key], value) + + def test_svdquant_diffusers_export_promotes_clean_keys(): """Fast CPU check of the diffusers SVDQuant export promotion. diff --git a/tests/unit/torch/export/test_quantized_weight_state.py b/tests/unit/torch/export/test_quantized_weight_state.py index 682bc0eb229..ef0734aff96 100644 --- a/tests/unit/torch/export/test_quantized_weight_state.py +++ b/tests/unit/torch/export/test_quantized_weight_state.py @@ -156,6 +156,23 @@ def test_functional_fp8_export_is_repeatable(): torch.testing.assert_close(base[name], repeated[name]) +def test_functional_fp8_export_matches_the_fake_quantized_bf16_weight(): + torch.manual_seed(0) + module = nn.Linear(64, 64, bias=False).to(torch.bfloat16) + mtq.quantize( + module, + copy.deepcopy(mtq.FP8_DEFAULT_CFG), + lambda m: m(torch.randn(2, 64, dtype=torch.bfloat16)), + ) + state = capture_quantized_weight_export_state(module) + + exported = export_quantized_weight_tensors(module.weight, state, torch.bfloat16) + + dequantized = exported["weight"].float() * exported["weight_scale"] + expected = module.weight_quantizer(module.weight) + torch.testing.assert_close(dequantized.to(torch.bfloat16), expected, rtol=0, atol=0) + + _CAPTURE_FORMAT_CONFIGS = { "fp8_pc_pt": {"num_bits": (4, 3), "axis": 0}, "fp8_pb_wo": {"num_bits": (4, 3), "block_sizes": {-2: 128, -1: 128}}, diff --git a/tests/unit/torch/opt/plugins/test_diffusers_save_load.py b/tests/unit/torch/opt/plugins/test_diffusers_save_load.py index cee5dbefbf8..5a6561fd28c 100644 --- a/tests/unit/torch/opt/plugins/test_diffusers_save_load.py +++ b/tests/unit/torch/opt/plugins/test_diffusers_save_load.py @@ -14,23 +14,61 @@ # limitations under the License. import pytest +import torch from _test_utils.torch.diffusers_models import ( create_tiny_unet_dir, df_modelopt_state_and_output_tester, + get_tiny_unet, + get_unet_dummy_inputs, ) from _test_utils.torch.opt.utils import apply_mode_with_sampling from diffusers import UNet2DConditionModel +import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.nn import TensorQuantizer + @pytest.mark.parametrize("model_cls", [UNet2DConditionModel]) -def test_unet_save_restore(tmp_path, model_cls): +@pytest.mark.parametrize("use_subfolder", [False, True], ids=["direct", "subfolder"]) +def test_unet_save_restore(tmp_path, model_cls, use_subfolder): tiny_unet_dir = create_tiny_unet_dir(tmp_path) model_ref = model_cls.from_pretrained(tiny_unet_dir) # TODO: Add calibrate, compress mode to the test model_ref = apply_mode_with_sampling( model_ref, ["sparse_magnitude", "export_sparse", "quantize"] ) - model_ref.save_pretrained(tiny_unet_dir / "modelopt_model") + component_dir = tmp_path / "checkpoint" / "unet" + model_ref.save_pretrained(component_dir) - model_test = model_cls.from_pretrained(tiny_unet_dir / "modelopt_model") + load_path = component_dir.parent if use_subfolder else component_dir + load_kwargs = {"subfolder": "unet"} if use_subfolder else {} + model_test = model_cls.from_pretrained(load_path, local_files_only=True, **load_kwargs) df_modelopt_state_and_output_tester(model_ref, model_test) + + +@pytest.mark.parametrize("save_modelopt_state", [True, False]) +@torch.no_grad() +def test_plain_save_replaces_quantized_checkpoint(tmp_path, save_modelopt_state): + quantized = get_tiny_unet().eval() + inputs = get_unet_dummy_inputs(quantized) + mtq.quantize(quantized, mtq.FP8_DEFAULT_CFG, lambda m: m(**inputs)) + quantized.save_pretrained(tmp_path) + assert (tmp_path / "modelopt_state.pth").is_file() + + plain = get_tiny_unet().eval() + original_state = (tmp_path / "modelopt_state.pth").read_bytes() + with pytest.raises(ValueError): + plain.save_pretrained(tmp_path, max_shard_size="invalid") + assert (tmp_path / "modelopt_state.pth").read_bytes() == original_state + + plain.save_pretrained(tmp_path / "fresh") + reference = UNet2DConditionModel.from_pretrained(tmp_path / "fresh", local_files_only=True) + plain.save_pretrained(tmp_path, save_modelopt_state=save_modelopt_state) + restored = UNet2DConditionModel.from_pretrained(tmp_path, local_files_only=True) + + assert not any(isinstance(module, TensorQuantizer) for module in restored.modules()) + assert not (tmp_path / "modelopt_state.pth").exists() + torch.testing.assert_close(restored.state_dict(), plain.state_dict(), rtol=0, atol=0) + torch.testing.assert_close( + restored(**inputs).sample, reference(**inputs).sample, rtol=0, atol=0 + ) diff --git a/tests/unit/torch/opt/plugins/test_hf_patching.py b/tests/unit/torch/opt/plugins/test_hf_patching.py index 719550f7c78..ba568bb0790 100644 --- a/tests/unit/torch/opt/plugins/test_hf_patching.py +++ b/tests/unit/torch/opt/plugins/test_hf_patching.py @@ -13,6 +13,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +from unittest.mock import Mock + +import huggingface_hub import pytest import torch import torch.nn as nn @@ -26,10 +29,31 @@ import modelopt.torch.distill as mtd import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq +from modelopt.torch.opt.plugins.huggingface import _get_modelopt_state_path from modelopt.torch.opt.plugins.transformers import _restore_qtensor_wrappers from modelopt.torch.quantization.qtensor import QTensorWrapper +@pytest.mark.parametrize( + ("load_kwargs", "cache_result", "filename", "expected"), + [ + ({}, "cached-state.pth", "modelopt_state.pth", "cached-state.pth"), + ({"subfolder": "unet"}, "cached-state.pth", "unet/modelopt_state.pth", "cached-state.pth"), + ({"subfolder": "unet"}, None, "unet/modelopt_state.pth", ""), + ({"subfolder": "unet"}, huggingface_hub._CACHED_NO_EXIST, "unet/modelopt_state.pth", ""), + ], + ids=["default-cached", "subfolder-cached", "subfolder-missing", "subfolder-known-missing"], +) +def test_get_modelopt_state_path_cached(monkeypatch, load_kwargs, cache_result, filename, expected): + cache_lookup = Mock(return_value=cache_result) + monkeypatch.setattr(huggingface_hub, "try_to_load_from_cache", cache_lookup) + + result = _get_modelopt_state_path("test-org/test-model", **load_kwargs) + + cache_lookup.assert_called_once_with("test-org/test-model", filename) + assert result == expected + + @pytest.mark.parametrize( ("model_cls", "teacher_model_type"), [ diff --git a/tests/unit/torch/opt/plugins/test_transformers_save_load.py b/tests/unit/torch/opt/plugins/test_transformers_save_load.py index 1e526fc09c0..3faf3de79c1 100644 --- a/tests/unit/torch/opt/plugins/test_transformers_save_load.py +++ b/tests/unit/torch/opt/plugins/test_transformers_save_load.py @@ -22,11 +22,15 @@ from _test_utils.torch.opt.utils import apply_mode_with_sampling from _test_utils.torch.transformers_models import ( create_tiny_llama_dir, + get_tiny_qwen3vl, tf_modelopt_state_and_output_tester, ) from safetensors.torch import load_file from transformers import AutoConfig, AutoModelForCausalLM, LlamaForCausalLM +import modelopt.torch.opt as mto +import modelopt.torch.quantization as mtq +from modelopt.torch.export import export_hf_checkpoint from modelopt.torch.opt.plugins.transformers import _legacy_tied_weights_keys_as_dict @@ -44,6 +48,32 @@ def test_causal_lm_save_restore(tmp_path, model_cls): tf_modelopt_state_and_output_tester(model_ref, model_test) +@pytest.mark.parametrize( + "quant_config", [mtq.FP8_DEFAULT_CFG, mtq.NVFP4_DEFAULT_CFG], ids=["fp8", "nvfp4"] +) +@torch.no_grad() +def test_vlm_export_without_modelopt_state(tmp_path, quant_config): + """Export separately converted VLM components without requiring a single ModelOpt state.""" + model = get_tiny_qwen3vl().eval() + disabled_config = {"quant_cfg": [{"quantizer_name": "*", "enable": False}], "algorithm": "max"} + mtq.quantize(model.model.visual, disabled_config) + mtq.quantize(model.lm_head, disabled_config) + mtq.quantize( + model.model.language_model, + quant_config, + lambda _: model(input_ids=torch.ones(1, 4, dtype=torch.long)), + ) + state_path = tmp_path / "modelopt_state.pth" + torch.save(mto.modelopt_state(model.model.language_model), state_path) + + export_hf_checkpoint(model, export_dir=tmp_path, save_modelopt_state=False) + + assert not state_path.exists() + saved = load_file(tmp_path / "model.safetensors") + assert any(key.endswith(".weight_scale") for key in saved) + assert any("visual" in key for key in saved) + + @pytest.mark.parametrize("tie_word_embeddings", [False, True]) def test_save_pretrained_with_legacy_tied_weights_keys(tmp_path, tie_word_embeddings): """A model declaring 4.x list-style `_tied_weights_keys` must still save (nvbug 6518665). diff --git a/tests/unit/torch/puzzletron/test_post_init_sparse.py b/tests/unit/torch/puzzletron/test_post_init_sparse.py new file mode 100644 index 00000000000..5cab9b5869c --- /dev/null +++ b/tests/unit/torch/puzzletron/test_post_init_sparse.py @@ -0,0 +1,52 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for post-initialization structured sparsity.""" + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from modelopt.torch.puzzletron.tools.post_init_sparse import SparsityMethod2o4 + + +@pytest.mark.parametrize("layer_name", ["dense", "gate", "proj"]) +def test_do_sparsity_preserves_layer_name(layer_name): + block = nn.Module() + block.mlp = nn.Module() + linear = nn.Linear(4, 2, bias=False) + block.mlp.add_module(layer_name, linear) + + model = nn.Module() + model.model = nn.Module() + model.model.layers = nn.ModuleList([block]) + model.config = SimpleNamespace( + block_configs=[ + SimpleNamespace( + ffn=SimpleNamespace(sparsify=[layer_name]), + attention=SimpleNamespace(sparsify=[]), + ) + ] + ) + with torch.no_grad(): + linear.weight.copy_(torch.arange(1, 9, dtype=torch.float32).reshape(2, 4)) + + SparsityMethod2o4().do_sparsity(model) + + expected_mask = torch.tensor([[0, 0, 1, 1], [0, 0, 1, 1]], dtype=torch.float32) + torch.testing.assert_close(linear.weight_mask, expected_mask) + torch.testing.assert_close(linear.weight, linear.weight_orig * expected_mask) diff --git a/tests/unit/torch/quantization/plugins/test_attention_quant.py b/tests/unit/torch/quantization/plugins/test_attention_quant.py index 22a86971a6d..ff9e9404ea4 100644 --- a/tests/unit/torch/quantization/plugins/test_attention_quant.py +++ b/tests/unit/torch/quantization/plugins/test_attention_quant.py @@ -52,6 +52,54 @@ def forward(self, hidden_states, **kwargs): return F.scaled_dot_product_attention(q, k, v), None +class FlatMatmulAttention(nn.Module): + def forward(self, q, k, v): + scores = torch.matmul(q, k.transpose(-2, -1)) + probs = torch.softmax(scores, dim=-1) + return torch.matmul(probs, v) + + +class FlatBMMAttention(nn.Module): + def forward(self, q, k, v): + scores = torch.bmm(q, k.transpose(-2, -1)) + probs = torch.softmax(scores, dim=-1) + return torch.bmm(probs, v) + + +class FlatBinMatmulAttention(nn.Module): + def forward(self, q, k, v): + scores = q @ k.transpose(-2, -1) + probs = torch.softmax(scores, dim=-1) + return probs @ v + + +class NestedMatmulAttention(nn.Module): + def forward(self, q, k, v): + return torch.matmul(torch.softmax(torch.matmul(q, k.transpose(-2, -1)), dim=-1), v) + + +class NestedBMMAttention(nn.Module): + def forward(self, q, k, v): + return torch.bmm(torch.softmax(torch.bmm(q, k.transpose(-2, -1)), dim=-1), v) + + +class NestedBinMatmulAttention(nn.Module): + def forward(self, q, k, v): + return torch.softmax(q @ k.transpose(-2, -1), dim=-1) @ v + + +class BAddBMMAttention(nn.Module): + def forward(self, q, k, v, bias): + scores = torch.baddbmm(bias, q, k.transpose(-2, -1)) + return torch.bmm(torch.softmax(scores, dim=-1), v) + + +class BAddBMMKeywordAttention(nn.Module): + def forward(self, q, k, v, bias): + scores = torch.baddbmm(input=bias, batch1=q, batch2=k.transpose(-2, -1)) + return torch.bmm(torch.softmax(scores, dim=-1), v) + + kv_cache_config = { "quant_cfg": [ {"quantizer_name": "*[kv]_bmm_quantizer", "cfg": {"num_bits": 4}, "enable": True}, @@ -115,6 +163,37 @@ def test_kv_quant_hf(model_getter, attn_cls): mtq.unregister(attn_cls) +@pytest.mark.parametrize( + "attn_cls", + [ + FlatMatmulAttention, + FlatBMMAttention, + FlatBinMatmulAttention, + NestedMatmulAttention, + NestedBMMAttention, + NestedBinMatmulAttention, + ], +) +def test_kv_quant_operand_wiring(attn_cls): + """Each q/k/v_bmm_quantizer must calibrate on its own operand, whatever the layout.""" + q, k, v = torch.randn(1, 4, 8), 10 * torch.randn(1, 4, 8), 100 * torch.randn(1, 4, 8) + assert mtq.plugins.register_attention_for_kv_quant(attn_cls) + try: + model = attn_cls() + mtq.quantize(model, kv_cache_config, lambda model: model(q, k, v)) + for name, operand in zip("qkv", (q, k, v)): + quantizer = getattr(model, f"{name}_bmm_quantizer") + assert quantizer.amax == operand.abs().max(), f"{name}_bmm_quantizer saw another tensor" + finally: + mtq.unregister(attn_cls) + + +@pytest.mark.parametrize("attn_cls", [BAddBMMAttention, BAddBMMKeywordAttention]) +def test_kv_quant_rejects_baddbmm(attn_cls): + """baddbmm takes the additive input first, so its operands cannot be wired by position.""" + assert not mtq.plugins.register_attention_for_kv_quant(attn_cls) + + def test_kv_quant_bert(): """Test KV cache quantization on BERT model with decorated attention.""" model_test = get_tiny_bert() diff --git a/tests/unit/torch/quantization/test_calibrator.py b/tests/unit/torch/quantization/test_calibrator.py index 9f7d77ce6f7..2d489b14543 100644 --- a/tests/unit/torch/quantization/test_calibrator.py +++ b/tests/unit/torch/quantization/test_calibrator.py @@ -89,6 +89,32 @@ def test_track_amax_raises(self): class TestHistogramCalibrator: + @pytest.mark.parametrize("torch_hist", [False, True]) + def test_entropy_preserves_collected_histogram(self, torch_hist): + calibrator = calib.HistogramCalibrator(4, None, False, num_bins=32, torch_hist=torch_hist) + reference = calib.HistogramCalibrator(4, None, False, num_bins=32, torch_hist=torch_hist) + first_batch = torch.cat((torch.zeros(100), torch.arange(1, 33).float())) + calibrator.collect(first_batch) + reference.collect(first_batch) + expected_hist = np.asarray(calibrator._calib_hist).copy() + expected_edges = np.asarray(calibrator._calib_bin_edges).copy() + percentile = calibrator.compute_amax("percentile", percentile=50) + + first_entropy = calibrator.compute_amax("entropy", start_bin=16) + second_entropy = calibrator.compute_amax("entropy", start_bin=16) + + np.testing.assert_array_equal(calibrator._calib_hist, expected_hist) + np.testing.assert_array_equal(calibrator._calib_bin_edges, expected_edges) + assert calibrator.compute_amax("percentile", percentile=50) == percentile + assert first_entropy == second_entropy + next_batch = torch.tensor([0.0, 2.0, 4.0, 64.0]) + calibrator.collect(next_batch) + reference.collect(next_batch) + np.testing.assert_array_equal(calibrator._calib_hist, reference._calib_hist) + assert calibrator.compute_amax("percentile", percentile=50) == reference.compute_amax( + "percentile", percentile=50 + ) + @pytest.mark.skip(reason="TODO: Fix assertions in test_grow") def test_grow(self, verbose): x_1 = torch.tensor([0, 255, 255, 255, 255, 255]) diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_conversion.py b/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_conversion.py index 6cfdd2775b9..e70f2b7de65 100644 --- a/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_conversion.py +++ b/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_conversion.py @@ -188,6 +188,60 @@ def test_restore_sparse_attention_model(self): assert hasattr(module, "_method") assert module._method == "flash_skip_softmax" + def test_restore_keeps_calibration_params(self, tmp_path): + """Calibrated thresholds survive an mto.save/mto.restore roundtrip.""" + calibration_params = {"prefill": {"a": 3.14, "b": 7.5}, "decode": {"a": 0.5, "b": 9.0}} + target_sparse_ratio = {"prefill": 0.4, "decode": 0.6} + + model = sparse_attn.sparsify(SimpleAttentionModel(), FLASH_SKIP_SOFTMAX_DEFAULT_CFG) + for module in model.modules(): + if isinstance(module, SparseAttentionModule): + module._sparse_method_instance.calibration_params = calibration_params + module._sparse_method_instance.target_sparse_ratio = target_sparse_ratio + + mto.save(model, tmp_path / "model.pth") + model_restored = mto.restore(SimpleAttentionModel(), tmp_path / "model.pth") + + sparse_modules = [ + m for m in model_restored.modules() if isinstance(m, SparseAttentionModule) + ] + assert sparse_modules + for module in sparse_modules: + method = module._sparse_method_instance + assert method.calibration_params == calibration_params + assert method.target_sparse_ratio == target_sparse_ratio + assert module.get_threshold_info()["type"] == "dynamic_calibrated" + + def test_restore_keeps_disabled_modules(self): + """Modules disabled after sparsify stay disabled after restore.""" + model = sparse_attn.sparsify(SimpleAttentionModel(), FLASH_SKIP_SOFTMAX_DEFAULT_CFG) + disable_sparse_attention(model, "*attention*") + + model_restored = mto.restore_from_modelopt_state( + SimpleAttentionModel(), mto.modelopt_state(model) + ) + + sparse_modules = [ + m for m in model_restored.modules() if isinstance(m, SparseAttentionModule) + ] + assert sparse_modules + assert not any(m.is_enabled for m in sparse_modules) + + def test_restore_legacy_state_without_calibration_keys(self): + """Metadata saved before calibration state was persisted still restores.""" + model = sparse_attn.sparsify(SimpleAttentionModel(), FLASH_SKIP_SOFTMAX_DEFAULT_CFG) + modelopt_state = mto.modelopt_state(model) + for _, mode_state in modelopt_state["modelopt_state_dict"]: + for module_state in mode_state["metadata"]["sparse_attention_state"].values(): + module_state.pop("enabled", None) + + model_restored = mto.restore_from_modelopt_state(SimpleAttentionModel(), modelopt_state) + + for module in model_restored.modules(): + if isinstance(module, SparseAttentionModule): + assert module.is_enabled + assert module._sparse_method_instance.calibration_params is None + class TestSparseAttentionModuleMethods: """Test SparseAttentionModule methods.""" @@ -387,6 +441,15 @@ def test_exports_when_calibration_present(self): assert group_0["target_sparsity"] == {"prefill": 0.4, "decode": 0.6} assert out["producer"]["name"] == "modelopt" + # Diffusion (VisualGen) schema: one coefficient block, scalar target_sparsity. + group_0 = export_sparse_attention_config(model, for_diffusion=True)["config_groups"][ + "group_0" + ] + tsf = group_0["threshold_scale_factor"] + assert tsf["coefficients"] == {"a": 3.14, "b": 7.5} + assert "prefill" not in tsf and "decode" not in tsf + assert group_0["target_sparsity"] == 0.4 + def test_exports_sparse_softmax_metadata(self): """N:M sparse-softmax params are exported without calibration metadata.""" model = SimpleAttentionModel() diff --git a/tests/unit/torch/speculative/test_speculative_utils.py b/tests/unit/torch/speculative/test_speculative_utils.py index 3de4999a69e..3ac7e51bbbd 100644 --- a/tests/unit/torch/speculative/test_speculative_utils.py +++ b/tests/unit/torch/speculative/test_speculative_utils.py @@ -19,12 +19,22 @@ was 2 (field count) and every conversation got dropped by the num_input_tokens <= 10 filter. """ +import hashlib +import sys +from unittest.mock import patch + import pytest import torch from _test_utils.torch.transformers_models import get_tiny_tokenizer from transformers import BatchEncoding -from modelopt.torch.speculative.utils import get_conversation_input_ids +import modelopt.torch.speculative.utils as spec_utils +from modelopt.torch.speculative.utils import ( + KIMI_K2_FILE_SHA256, + KIMI_K2_PACKAGE_NAME, + _setup_kimi_k2_decoder, + get_conversation_input_ids, +) CONVERSATIONS = [ {"role": "user", "content": "Explain why the sky is blue in a few sentences."}, @@ -73,3 +83,65 @@ def apply_chat_template(self, conversations, **kwargs): assert input_ids == expected # The pre-fix code saw len() in {1, 2} here and silently dropped the conversation. assert len(input_ids) > 10 + + +def _serve_kimi_k2_files(tmp_path, monkeypatch, sources): + """Serve ``sources`` as the downloaded Kimi-K2 files, with their hashes pinned.""" + for filename, source in sources.items(): + (tmp_path / filename).write_text(source) + digest = hashlib.sha256(source.encode()).hexdigest() + monkeypatch.setitem(KIMI_K2_FILE_SHA256, filename, digest) + monkeypatch.setattr( + spec_utils, "hf_hub_download", lambda repo_id, filename, **kwargs: str(tmp_path / filename) + ) + # Keep the global DynamicCache patch out of the rest of the test session. + monkeypatch.setattr(spec_utils, "_patch_dynamic_cache_compatibility", lambda: None) + + +def test_kimi_k2_decoder_refuses_unverified_code(tmp_path, monkeypatch): + """No downloaded Kimi-K2 code runs unless every file matches its pinned hash.""" + marker = tmp_path / "executed" + run_marker = f"open({str(marker)!r}, 'w').close()\n" + _serve_kimi_k2_files(tmp_path, monkeypatch, dict.fromkeys(KIMI_K2_FILE_SHA256, run_marker)) + # Only the modeling file fails its check, so the valid config file must not run either. + (tmp_path / "modeling_deepseek.py").write_text(run_marker + "# tampered\n") + + with pytest.raises(RuntimeError): + _setup_kimi_k2_decoder() + assert not marker.exists() + + +def test_kimi_k2_decoder_imports_verified_code(tmp_path, monkeypatch): + """Verified files load as one package, so the modeling file's relative import resolves.""" + modeling = ( + "from .configuration_deepseek import DeepseekV3Config\n" + "class DeepseekV3Attention:\n" + " _init_rope = forward = None\n" + "class DeepseekV3DecoderLayer:\n" + " config_class = DeepseekV3Config\n" + " forward = None\n" + ) + config = "class DeepseekV3Config:\n pass\n" + _serve_kimi_k2_files( + tmp_path, + monkeypatch, + {"configuration_deepseek.py": config, "modeling_deepseek.py": modeling}, + ) + + with patch.dict(sys.modules): + decoder_cls = _setup_kimi_k2_decoder() + assert decoder_cls.config_class.__module__ == f"{KIMI_K2_PACKAGE_NAME}.configuration_deepseek" + + +def test_kimi_k2_decoder_drops_failed_module(tmp_path, monkeypatch): + """A module that raises while importing is not left half-initialized in sys.modules.""" + _serve_kimi_k2_files( + tmp_path, + monkeypatch, + {"configuration_deepseek.py": "", "modeling_deepseek.py": "raise ImportError\n"}, + ) + + with patch.dict(sys.modules): + with pytest.raises(ImportError): + _setup_kimi_k2_decoder() + assert f"{KIMI_K2_PACKAGE_NAME}.modeling_deepseek" not in sys.modules diff --git a/tests/unit/torch/test_import.py b/tests/unit/torch/test_import.py new file mode 100644 index 00000000000..3ceb933c8b3 --- /dev/null +++ b/tests/unit/torch/test_import.py @@ -0,0 +1,41 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import subprocess +import sys + +# Top-level modules of the optional ``[hf]`` extra, plus ``requests`` which it only pulls in +# transitively. A base ``pip install nvidia-modelopt`` provides none of them. +_HF_EXTRA_MODULES = [ + "accelerate", + "datasets", + "deepspeed", + "diffusers", + "httpx", + "huggingface_hub", + "nltk", + "peft", + "requests", + "sentencepiece", + "tiktoken", + "transformers", + "wonderwords", +] + + +def test_modelopt_torch_imports_without_hf_extra(): + # A None entry in sys.modules makes importing that module raise ModuleNotFoundError. + code = f"import sys; sys.modules.update(dict.fromkeys({_HF_EXTRA_MODULES!r})); import modelopt.torch" + subprocess.run([sys.executable, "-c", code], check=True) diff --git a/tools/launcher/pyproject.toml b/tools/launcher/pyproject.toml index 8ba825776d4..02adf5991fe 100644 --- a/tools/launcher/pyproject.toml +++ b/tools/launcher/pyproject.toml @@ -4,6 +4,7 @@ version = "0.1.0" description = "ModelOpt Launcher — submit quantization, training, and evaluation jobs to Slurm clusters" requires-python = ">=3.10" dependencies = [ + "leptonai<0.47", # 0.47 removed leptonai.api.v1, which nemo-run (<=0.11.1) imports "nemo-run>=0.8.0", "pyyaml", ]