diff --git a/.agents/skills/deployment/references/support-matrix.md b/.agents/skills/deployment/references/support-matrix.md index 265f1b58560..e07d1b115a5 100644 --- a/.agents/skills/deployment/references/support-matrix.md +++ b/.agents/skills/deployment/references/support-matrix.md @@ -2,26 +2,18 @@ ## Unified HF Checkpoint — Framework Compatibility -| Model | Quant Format | TRT-LLM | vLLM | SGLang | -|-------|-------------|---------|------|--------| -| Llama 3.x | FP8 | yes | yes | yes | -| Llama 3.x | FP4 | yes | yes | yes | -| Llama 4 | FP8 | yes | — | yes | -| Llama 4 | FP4 | yes | — | — | -| DeepSeek R1 | FP8 | yes | yes | yes | -| DeepSeek R1 | FP4 | yes | yes | yes | -| DeepSeek V3 | FP8 | yes | yes | yes | -| DeepSeek V3 | FP4 | yes | yes | yes | -| Qwen 3 | FP8 | yes | yes | yes | -| Qwen 3 | FP4 | yes | yes | — | -| Qwen 3 MoE | FP8 | yes | yes | yes | -| Qwen 3 MoE | FP4 | yes | — | — | -| Qwen 2.5 | FP8 | yes | yes | yes | -| Qwen 2.5 | FP4 | yes | yes | — | -| QwQ-32B | FP8 | yes | yes | yes | -| QwQ-32B | FP4 | yes | yes | — | -| Mixtral 8x7B | FP8 | yes | yes | yes | -| Mixtral 8x7B | FP4 | yes | — | — | +**Do not maintain a copy of the matrix here.** The single source of truth is +`docs/source/deployment/3_unified_hf.rst` ("Model Support Matrix"), and every entry in it is drawn +from `tests/examples/hf_ptq/test_deploy.py`. + +Read that doc's legend before reporting a model as supported: the cases are marked `release` and do +not run on PR CI, and each is a load-and-generate smoke check on the text path — so an entry is +declared coverage, not proof the combination serves correctly. + +To answer "is model X supported on framework Y", read one of those two files — `test_deploy.py` is +the more precise answer, since it also carries the exact checkpoint, tensor-parallel size, and +minimum SM version per entry. It covers language models, VLMs (Qwen2.5-VL, Qwen3-VL, +Nemotron Omni), EAGLE3/Medusa drafters, and diffusion models. ## Supported Quantization Formats @@ -50,13 +42,13 @@ | SGLang | `quantization="modelopt"` | `quantization="modelopt_fp4"` | | TRT-LLM | auto-detected from checkpoint | auto-detected from checkpoint | -## Models not in this list +## Models not in the matrix -This matrix covers officially validated combinations. For unlisted models: +The matrix covers the combinations modelopt tracks, not the full set of what will run. For unlisted models: 1. **Check the framework's own docs** — vLLM and SGLang support many HuggingFace models natively. Use WebSearch to check `vllm supported models` or `sglang supported models`. 2. **Try it** — if the model uses standard `nn.Linear` layers and has `hf_quant_config.json`, vLLM/SGLang will likely work with `--quantization modelopt`. -3. **Ask the user** — if unsure, ask: "This model isn't in the validated support matrix. Would you like to try deploying it anyway?" +3. **Ask the user** — if unsure, ask: "This model isn't in the support matrix. Would you like to try deploying it anyway?" ## Notes @@ -64,4 +56,5 @@ This matrix covers officially validated combinations. For unlisted models: - **B300/GB300 are `sm_103`** and need a **CUDA-13** serving image — from v0.20.0 the unsuffixed tag is CUDA-13 (`-cu129` opts back to CUDA 12); `cu12` images lack the `sm_103` FP4 kernel and serve NVFP4 as gibberish or error out. See the CUDA-13 note in the deployment `SKILL.md`. - **Verify the GPU with `nvidia-smi`** before choosing the image — cluster GPU labels can be stale. - INT4_AWQ and W4A8_AWQ are only supported by TRT-LLM (not vLLM or SGLang). -- Source: `examples/hf_ptq/README.md` and `docs/source/deployment/3_unified_hf.rst` +- For VLMs, only the language model is quantized; the vision encoder stays in high precision, so multimodal serving depends on the framework's own support for that architecture. +- Source: `docs/source/deployment/3_unified_hf.rst` and `tests/examples/hf_ptq/test_deploy.py` diff --git a/CHANGELOG.rst b/CHANGELOG.rst index eff4cea75bf..ca49e30738a 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -88,6 +88,9 @@ Changelog - Fix HF checkpoint export failing with ``AttributeError: 'list' object has no attribute 'keys'`` for models whose modeling code still declares tied weights in the ``transformers<5`` list format (NVBug 6518665, observed on ``stepfun-ai/Step-3.7-Flash``). transformers 5.0 changed ``_tied_weights_keys`` to a ``{target: source}`` dict and ``save_pretrained`` calls ``.keys()`` on every submodule's declaration without a type check, so such models — common among ``trust_remote_code`` checkpoints — load fine but die at the end of PTQ, after calibration. ModelOpt's ``save_pretrained`` patch now normalizes a list-style declaration to the equivalent dict for the duration of the save (each entry mapped to itself, which is what the legacy list meant) and restores the original attribute afterwards. - Fix unified HF export of multimodal models whose vision tower carries its own ``PrefixChange`` conversion (``LlavaForConditionalGeneration`` on ``transformers>=5.12`` — NVBug 6525511). transformers collects conversion mappings recursively and scopes each sub-model's transforms to that sub-module via ``scope_prefix``, matching only keys under that prefix. ModelOpt's quant-aware reverse conversion read the raw patterns and ignored ``scope_prefix``, so the vision tower's "add a ``vision_model.`` prefix" rule was applied to *every* key in the state dict: an exported llava-1.5-13b checkpoint had all 758 tensors moved under a bogus top-level ``vision_model.`` namespace (``vision_model.language_model.*``, ``vision_model.lm_head.*``), and vLLM rejected it with ``ValueError: There is no module or parameter named 'vision_model' in LlavaForConditionalGeneration``. Reverse rename rules now carry their scope and are applied only to keys under it, matching transformers' own ``WeightTransform._scoped_match`` semantics. ``Gemma3ForConditionalGeneration`` was affected identically and is fixed by the same change. - Fix QLoRA export in ``examples/llm_qat/export.py`` failing with ``AssertionError: Model already has modelopt state!`` (NVBug 6542481). The QLoRA training output is an adapter-only checkpoint, so ``from_pretrained`` resolves the quantized base model from ``adapter_config.json`` and ``enable_huggingface_checkpointing`` already restores its ModelOpt state; the export then restored a second time. It now restores only when the loaded model is not already converted. Two further breakages on the same path are also fixed: ``_restore_qtensor_wrappers`` matched no modules because PEFT re-parents the quantized linear as ``.base_layer`` while ``q_tensor_state`` is keyed by the name it was saved with (the packed NVFP4 weight then reached ``F.linear`` and raised a shape error), and ``postprocess_state_dict`` silently dropped every ``base_layer.*`` key missing from a hand-maintained rename map — losing the NVFP4 ``weight_scale_2`` global scale and any linear ``bias`` (Qwen2-style q/k/v biases), and leaving ``base_layer`` in the exported AWQ ``pre_quant_scale`` key. The rename is now a generic ``.base_layer.`` strip. +- Fix ``--use_fsdp2`` PTQ (``examples/hf_ptq``) failing on models that hold a few parameters in a dtype other than the model's own, with ``AssertionError: FSDP expects uniform original parameter dtype`` on the first calibration forward. Nemotron-3-Nano is one such model: its MoE router gates are declared ``float32`` while the rest of the checkpoint is bfloat16, so each decoder layer's FSDP2 shard group mixed dtypes. ``fsdp2_wrap`` now passes those off-dtype parameters to ``fully_shard(ignored_params=...)``, leaving them replicated in their original dtype instead of casting them, and warns with their names and their share of the model. +- Fix ``--use_fsdp2`` HF export making no progress for hours on large MoE checkpoints. ``create_fsdp_param_mapping`` resolved each ``FSDPParam``'s module by scanning every ``model.named_parameters()``, and export calls it once per quantized module, so the cost was quadratic in (parameters x modules): harmless for dense models, intractable for a MoE with many experts. Exporting Nemotron-3-Nano-30B-A3B (6,243 parameter tensors, 6,004 quantized modules) spent an estimated 1.9 hours there with every GPU idle. The parameter index is now built once per mapping instead of once per ``FSDPParam`` (1151 ms -> 5.1 ms per call), preserving the previous ``named_parameters()``-order resolution for tied weights. +- Update HuggingFace checkpoint export to use name-based tied-weight deduplication instead of the previous address-based approach. The address-based deduplication could incorrectly drop an untied weight that happened to share memory with a tied one, producing an incomplete checkpoint (observed as a false positive on MiniMax-M2.7). 0.45 (2026-07-02) ^^^^^^^^^^^^^^^^^ diff --git a/docs/source/deployment/3_unified_hf.rst b/docs/source/deployment/3_unified_hf.rst index 59a2782c5e7..ccef639d00e 100644 --- a/docs/source/deployment/3_unified_hf.rst +++ b/docs/source/deployment/3_unified_hf.rst @@ -51,48 +51,175 @@ The unified HF export API supports the following quantization formats: 5. INT4_AWQ - 4-bit integer with AWQ optimization 6. W4A8_AWQ - 4-bit weights and 8-bit activations with AWQ optimization -Framework-Specific Support +Minimum Framework Versions -------------------------- -TensorRT-LLM -~~~~~~~~~~~~ +=============== ================= +Framework Minimum version +=============== ================= +TensorRT-LLM v1.2.0 +vLLM v0.10.1 +SGLang v0.4.10 +=============== ================= + +These are the oldest versions expected to load a unified HF checkpoint. The deployment suite itself +targets newer ones — TensorRT-LLM containers in ``.github/workflows/`` are on the 1.3.x line. Older +TensorRT-LLM releases may still serve FP8 checkpoints; that is simply not exercised, so v1.2.0 is +the oldest version stated here rather than the oldest that works. + +.. _unified-hf-support-matrix: + +Model Support Matrix +-------------------- + +What this matrix is based on +~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Entries are drawn from the release deployment suite, +`tests/examples/hf_ptq/test_deploy.py `_. +For each entry it loads the exported checkpoint in the framework and generates from four short text +prompts, asserting that each returns non-empty output. + +Two limits are worth stating plainly, because they bound what any ✅ below can mean: + +* **These are declared cases, not PR-gated coverage.** The suite is marked ``release`` and collects + only when pytest is given ``--run-release``, which no workflow in ``.github/workflows/`` currently + passes. A green check on a pull request does not mean these cases ran. +* **Each case is a load-and-generate smoke check on the text path.** It does not verify accuracy, + image or audio inputs, diffusion output, or that speculative decoding actually engages. + +Legend: + +* ✅ — declared in the release deployment suite, subject to the two limits above. +* ⚠ — expected to work, but not a suite entry: either carried over from earlier documentation, or + present as a case that does not exercise the feature the row names. +* ``-`` — not in the suite. It may still work; see `Models not listed here`_. + +Language models +~~~~~~~~~~~~~~~ + +============================================ ============== ============ ====== ======== +Model Quant format TensorRT-LLM vLLM SGLang +============================================ ============== ============ ====== ======== +Llama 3.1, 3.3 FP8, NVFP4 ✅ ✅ ✅ +Llama 4 Scout, Maverick FP8 ✅ ✅ ✅ +Llama 4 Scout NVFP4 ✅ ✅ ✅ +Llama 4 Maverick NVFP4 ⚠ \- \- +Llama Nemotron Super 49B v1, v1.5 FP8 ✅ ✅ ✅ +Llama Nemotron Ultra 253B v1 FP8 ✅ ✅ ✅ +Nemotron 3 Nano 30B-A3B FP8, NVFP4 ✅ ✅ ✅ +Nemotron 3 Super 120B-A12B FP8, NVFP4 ✅ ✅ ✅ +Nemotron 3 Ultra 550B-A55B NVFP4 ✅ ✅ ✅ +DeepSeek R1, R1-0528 NVFP4 ✅ ✅ ✅ +DeepSeek R1, V3 FP8 ⚠ ⚠ ⚠ +DeepSeek V3, V3.1, V3.2 NVFP4 ✅ ✅ ✅ +DeepSeek V4 Flash NVFP4 ✅ ✅ ✅ +DeepSeek V4 Pro NVFP4 \- ✅ ✅ +Qwen 3 8B, 14B FP8, NVFP4 ✅ ✅ ✅ +Qwen 3 32B NVFP4 ✅ ✅ ✅ +Qwen 3 MoE 235B-A22B FP8, NVFP4 ✅ ✅ ✅ +Qwen 3 MoE 30B-A3B NVFP4 ✅ ✅ ✅ +Qwen 3 Coder 480B-A35B NVFP4 ✅ ✅ ✅ +Qwen 3-Next 80B-A3B NVFP4 ✅ ✅ ✅ +Qwen 3.5 397B-A17B NVFP4 ✅ ✅ ✅ +Qwen 3.5 122B-A10B, Qwen 3.6 35B-A3B NVFP4 \- ✅ \- +Qwen 2.5 FP8 ⚠ ⚠ ⚠ +Qwen 2.5 NVFP4 ⚠ ⚠ \- +QwQ-32B FP8 ⚠ ⚠ ⚠ +QwQ-32B NVFP4 ⚠ ⚠ \- +Gemma 4 31B NVFP4 ✅ ✅ ✅ +Gemma 4 26B-A4B NVFP4 \- ✅ \- +GLM-4.7, GLM-5, GLM-5.2 NVFP4 ✅ ✅ ✅ +GLM-5.1 NVFP4 \- ✅ ✅ +Kimi K2-Thinking, K2.5 NVFP4 ✅ ✅ ✅ +Kimi K2.6 NVFP4 \- ✅ \- +MiniMax M2.5, M3 NVFP4 ✅ ✅ ✅ +Mixtral 8x7B FP8 ⚠ ⚠ ⚠ +Mixtral 8x7B NVFP4 ⚠ \- \- +============================================ ============== ============ ====== ======== + +Vision-language and multimodal models +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +For VLMs, modelopt quantizes the language model only; the vision encoder is kept in high precision. +The exported checkpoint therefore relies on the serving framework's own multimodal support for that +architecture — see the +`TensorRT-LLM multimodal support matrix `_. + +.. important:: + ✅ in this table is **text-only smoke coverage**. The suite sends the same plain-text prompts it + uses for language models, so no image or audio input reaches the processor or vision encoder. + These entries show that the quantized checkpoint loads and that its language path generates — + they do not demonstrate multimodal serving. + +============================================ ============== ============ ====== ======== +Model Quant format TensorRT-LLM vLLM SGLang +============================================ ============== ============ ====== ======== +Qwen 2.5-VL 7B FP8, NVFP4 ✅ ✅ ✅ +Qwen 3-VL 235B-A22B NVFP4 ✅ ✅ ✅ +Nemotron 3 Nano Omni 30B-A3B FP8, NVFP4 ✅ ✅ ✅ +============================================ ============== ============ ====== ======== + +Speculative decoding drafters +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Drafters are deployed on top of their base checkpoint. + +Two caveats specific to this table: + +* **Most entries are doubly conditional.** Beyond the ``--run-release`` gate, the drafter cases in + ``test_eagle`` also require ``MODELOPT_LOCAL_EAGLE_MODEL`` to point at a directory containing the + drafter, and skip otherwise. The exception is EAGLE3 for Kimi K2.6, which is declared in + ``test_kimi`` without that gate — which is also why it is the one row with vLLM coverage. +* **Medusa is marked ⚠ because the case does not exercise Medusa.** The shared harness builds a + speculative-decoding configuration only when the model ID contains ``eagle``, so the Medusa entry + performs ordinary generation. It shows the checkpoint loads and serves; it does not validate + Medusa decoding. + +============================================================ ============ ============ ====== ======== +Drafter Quant format TensorRT-LLM vLLM SGLang +============================================================ ============ ============ ====== ======== +EAGLE3 for Llama 3.3 70B, Llama 4 Maverick FP8 ✅ \- ✅ +EAGLE3 for Qwen 3 235B-A22B (incl. Thinking-2507, FP4) BF16, NVFP4 ✅ \- ✅ +EAGLE3 for Qwen 3 30B-A3B-Thinking-2507 BF16 ✅ \- ✅ +EAGLE3 for Kimi K2-Thinking, K2.5 NVFP4 ✅ \- ✅ +EAGLE3 for Kimi K2.6 NVFP4 ✅ ✅ ✅ +EAGLE3 for gpt-oss-120b BF16 ✅ \- ✅ +Medusa for Llama 3.1 8B FP8 ⚠ \- ⚠ +============================================================ ============ ============ ====== ======== + +Diffusion models +~~~~~~~~~~~~~~~~ + +============================================ ============== ============ ====== ======== +Model Quant format TensorRT-LLM vLLM SGLang +============================================ ============== ============ ====== ======== +Wan 2.2 T2V A14B FP8, NVFP4 ⚠ \- ⚠ +DiffusionGemma 26B-A4B NVFP4 ✅ ✅ ✅ +============================================ ============== ============ ====== ======== + +Wan 2.2 is marked ⚠ because its cases run through the same autoregressive text helper as the +language models and assert on generated text. They never call a diffusion or video serving API, so +they do not substantiate text-to-video deployment. -Models: - * Llama 4, 3.x (FP8, NVFP4) - * Qwen 3, 2.5 (FP8, NVFP4) - * Qwen 3 MoE (FP8, NVFP4) - * Qwen 3-VL (FP8, NVFP4) - * Deepseek R1/V3 (NVFP4) - * Mixtral 8x7B (FP8, NVFP4) - * Medusa (FP8) - * Eagle (FP8) - -Requirements: TensorRT-LLM v0.17.0 or later - -vLLM -~~~~ - -Models: - * Llama 4, 3.x (FP8, NVFP4) - * Qwen 3, 2.5 (FP8, NVFP4) - * Qwen 3 MoE (FP8, NVFP4) - * Mixtral 8x7B (FP8) - * Deepseek R1/V3 (NVFP4) - -Requirements: vLLM v0.10.1 or later - -SGLang -~~~~~~ +.. note:: + NVFP4 inference requires Blackwell GPUs. Hopper can produce an NVFP4 checkpoint but cannot serve + it. On B300/GB300 (``sm_103``) use a CUDA-13 build of the serving framework; CUDA-12 builds lack + the ``sm_103`` FP4 kernels. -Models: - * Llama 4, 3.x (FP8, NVFP4) - * Qwen 3, 2.5 (FP8, NVFP4) - * Qwen 3 MoE (FP8, NVFP4) - * Deepseek R1/V3 (NVFP4) +Models not listed here +~~~~~~~~~~~~~~~~~~~~~~ -Requirements: SGLang v0.4.10 or later +This matrix records the combinations modelopt validates. It is not an exhaustive list of what will +run: vLLM, SGLang, and TensorRT-LLM load unified HF checkpoints generically, so a model built from +standard ``nn.Linear`` layers with an ``hf_quant_config.json`` will often deploy without any modelopt +change. Check the serving framework's own model support list first, then try it. -Note: While other models and quantization formats may work, they have not been thoroughly tested and validated. +The exact checkpoints behind every ✅ above, including tensor-parallel size and minimum SM +version, are listed in +`tests/examples/hf_ptq/test_deploy.py `__; +most are published under the +`NVIDIA Hugging Face organization `_. Deployment with Selected Inference Frameworks @@ -102,7 +229,7 @@ Deployment with Selected Inference Frameworks Follow the `TensorRT-LLM installation instructions. `_ - Currently we support fp8 and nvfp4 quantized models for TensorRT-LLM deployment, you need v0.17.0 or later version of TensorRT-LLM. + FP8 and NVFP4 quantized models are supported; you need v1.2.0 or later version of TensorRT-LLM. To run modelopt quantized model from Huggingface model hub, e.g., `nvidia/Llama-3.1-8B-Instruct-FP8`_, refer to the sample code below: @@ -136,7 +263,8 @@ Deployment with Selected Inference Frameworks Follow `vLLM installation instructions. `_ - Currently we support fp8 quantized models (without fp8 kv cache) for vLLM deployment, you need v0.6.5 or later version of vLLM. + FP8 and NVFP4 quantized models are supported; you need v0.10.1 or later version of vLLM. Pass + ``quantization="modelopt"`` for FP8 and ``quantization="modelopt_fp4"`` for NVFP4. To run modelopt quantized model from Huggingface model hub, e.g., `nvidia/Llama-3.1-8B-Instruct-FP8`_, refer to the sample code below: @@ -171,7 +299,8 @@ Deployment with Selected Inference Frameworks Follow the `SGLang installation instructions. `_ - Currently we support fp8 quantized models (without fp8 kv cache) for SGLang deployment, you need to use the main branch of SGLang (since Jan 6, 2025) and build it from source. + FP8 and NVFP4 quantized models are supported; you need v0.4.10 or later version of SGLang. Pass + ``quantization="modelopt"`` for FP8 and ``quantization="modelopt_fp4"`` for NVFP4. To run modelopt quantized model from Huggingface model hub, e.g., `nvidia/Llama-3.1-8B-Instruct-FP8`_, refer to the sample code below: diff --git a/examples/diffusers/README.md b/examples/diffusers/README.md index a9efb5fc3a3..3c8e5c80876 100644 --- a/examples/diffusers/README.md +++ b/examples/diffusers/README.md @@ -78,7 +78,7 @@ mtq.quantize(model=transformer, config=quant_config, forward_func=forward_pass) > *1.The w4a8_awq is an experimental quantization scheme that may result in a higher accuracy penalty.* -> *2.A selective set of the popular models are internally tested. The actual model support list may be longer. NVFP4 inference requires Blackwell GPUs and TensorRT-LLM v0.17 or later* +> *2.A selective set of the popular models are internally tested. The actual model support list may be longer. NVFP4 inference requires Blackwell GPUs and TensorRT-LLM v1.2 or later* > *3.The SVDQuant Perf in TRT might not good as the [Nunchaku: MIT-Nvidia](https://github.com/nunchaku-tech/nunchaku) at this moment.* diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 8d6be9c3845..90c94c67948 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -128,7 +128,7 @@ Please reference our [framework scripts](#framework-scripts) and our [docs](http > *2.For some models, there is only support for exporting quantized checkpoints.* \ > *3.W4A8_AWQ is only available on some models but not all* \ > *4.For some models, KV cache quantization may result in a higher accuracy penalty.* \ -> *5.A selective set of the popular models are internally tested. The actual model support list may be longer. NVFP4 inference requires Blackwell GPUs and TensorRT-LLM v0.17 or later* \ +> *5.A selective set of the popular models are internally tested. The actual model support list may be longer. NVFP4 inference requires Blackwell GPUs and TensorRT-LLM v1.2 or later* \ > *6.Some models currently support export to HF format only.* \ > *7.[PTQ for DeepSeek](../deepseek/README.md)* \ > *8.GLM-4.7 has MTP (Multi-Token Prediction) layers that are automatically loaded and excluded from quantization.* \ @@ -586,27 +586,24 @@ print(llm_fp8.generate(["What's the age of the earth? "])) ### Unified HF Checkpoint Deployment Model Support Matrix -| Model | Quant format | TRT-LLM | vLLM | SGLang | -| :---: | :---: | :---: | :---: | :---: | -| LLAMA 3.x | FP8 | ✅ | ✅ | ✅ | -| LLAMA 3.x | FP4 | ✅ | ✅ | ✅ | -| LLAMA 4 | FP8 | ✅ | - | ✅ | -| LLAMA 4 | FP4 | ✅ | - | - | -| DS-R1 | FP8 | ✅ | ✅ | ✅ | -| DS-R1 | FP4 | ✅ | ✅ | ✅ | -| DS-V3 | FP8 | ✅ | ✅ | ✅ | -| DS-V3 | FP4 | ✅ | ✅ | ✅ | -| QWen3 | FP8 | ✅ | ✅ | ✅ | -| QWen3 | FP4 | ✅ | ✅ | - | -| QWen3 MoE | FP8 | ✅ | ✅ | ✅ | -| QWen3 MoE | FP4 | ✅ | - | - | -| QWen3.5 MoE | FP4 | - | - | ✅ | -| QWen2.5 | FP8 | ✅ | ✅ | ✅ | -| QWen2.5 | FP4 | ✅ | ✅ | - | -| QwQ-32B | FP8 | ✅ | ✅ | ✅ | -| QwQ-32B | FP4 | ✅ | ✅ | - | -| Mixtral 8x7B | FP8 | ✅ | ✅ | ✅ | -| Mixtral 8x7B | FP4 | ✅ | - | - | +The deployment support matrix — which model families and quantization formats are covered on +TRT-LLM, vLLM, and SGLang, including vision-language models, speculative decoding drafters, and +diffusion models — lives in the documentation so there is a single copy to keep current: + +**[Unified HF Checkpoint → Model Support Matrix](https://nvidia.github.io/Model-Optimizer/deployment/3_unified_hf.html#model-support-matrix)** + +Each entry there is drawn from [`tests/examples/hf_ptq/test_deploy.py`](../../tests/examples/hf_ptq/test_deploy.py), +which loads the exported checkpoint in each framework and generates from short text prompts. That +file is also the place to look for the exact checkpoint, tensor-parallel size, and minimum SM +version behind each entry. + +> *Note: those cases are marked `release` and run out-of-band — no workflow currently passes +> `--run-release` — and each is a load-and-generate smoke check on the text path. Read the legend in +> the docs before treating an entry as verified support.* + +> *Note: the matrix records what modelopt validates, not the full set of what will run. vLLM, SGLang, +> and TRT-LLM load unified HF checkpoints generically, so unlisted models frequently deploy without +> any modelopt change — check the serving framework's own model support list and try it.* ### (Legacy) TensorRT-LLM Checkpoints diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index b8eef827d95..636aeffe99d 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -20,7 +20,6 @@ import json import logging import os -import shutil import warnings from collections.abc import Callable, Iterable from dataclasses import dataclass @@ -44,6 +43,7 @@ ) from modelopt.torch.export.model_utils import is_multimodal_model +from modelopt.torch.export.plugins.hf_checkpoint_utils import copy_non_safetensor_files_from_ckpt try: from huggingface_hub import snapshot_download @@ -56,6 +56,52 @@ SPECULATIVE_MODEL_LIST = ["Eagle", "Medusa"] +_HF_SIDECAR_DOWNLOAD_ALLOW_PATTERNS = [ + "*.jinja", + "*.json", + "*.md", + "*.model", + "*.py", + "*.tiktoken", + "*.txt", + "LICENSE*", + "NOTICE*", +] +_HF_PTQ_WEIGHT_FILE_PATTERNS = ( + "*.safetensors", + "*.safetensors.index.json", + "*.bin", + "*.bin.index.json", + "*.ckpt", + "*.gguf", + "*.h5", + "*.msgpack", + "*.npy", + "*.npz", + "*.onnx", + "*.pb", + "*.pickle", + "*.pkl", + "*.pt", + "*.pth", + "*.tar", + "*.tar.bz2", + "*.tar.gz", + "*.tar.xz", + "*.tflite", + "*.tgz", + "*.zip", +) +_HF_PTQ_EXPORT_OWNED_FILES = { + "config.json", + "hf_quant_config.json", + "quant_config.json", + "quantization_config.json", + "quantize_config.json", + "recipe.yaml", + "recipe.yml", +} + @dataclass class DistributedState: @@ -667,6 +713,17 @@ def _apply_dtype_to_config(model_kwargs, config_dtype, architecture, apply_confi return model_kwargs +def _fmt_max_memory(max_memory: dict) -> str: + """Format a ``{device: bytes}`` budget dict into a human-readable string.""" + parts = [] + for key in sorted(max_memory.keys(), key=lambda k: (isinstance(k, str), k)): + val = max_memory[key] + label = f"{val / 1024**3:.1f} GiB" if isinstance(val, int) else str(val) + key_str = f"GPU {key}" if isinstance(key, int) else str(key) + parts.append(f" {key_str}: {label}") + return "\n".join(parts) + + def get_model( ckpt_path, device="cuda", @@ -674,9 +731,21 @@ def get_model( trust_remote_code=False, use_seq_device_map=False, attn_implementation=None, + offload_folder=None, + max_cpu_memory_gb=None, + max_gpu_memory_gb=None, ): print(f"Initializing model from {ckpt_path}") + _disk_offload = offload_folder is not None + if _disk_offload and max_cpu_memory_gb is None: + warnings.warn( + "offload_folder is set but max_cpu_memory_gb is not specified. " + "CPU memory usage during model load will be unbounded. " + "Pass max_cpu_memory_gb to cap CPU usage.", + UserWarning, + ) + device_map = "auto" if device == "cpu": device_map = "cpu" @@ -749,6 +818,20 @@ def has_pack_quantized_config(config): return True return False + # Only the general load path below threads max_memory/offload_folder into + # from_pretrained; the specialized loaders build their own calls. + if _disk_offload and ( + is_speculative(hf_config) + or has_pack_quantized_config(hf_config) + or get_original_hf_quant_method(hf_config) == "mxfp4" + ): + warnings.warn( + "offload_folder is ignored for speculative, pack-quantized, and MXFP4 " + "checkpoints: these use dedicated load paths that cannot offload. The model " + "will be loaded fully resident.", + UserWarning, + ) + if is_speculative(hf_config): model = AutoModelForCausalLM.from_pretrained( ckpt_path, @@ -793,7 +876,11 @@ def has_pack_quantized_config(config): raise ValueError(f"Model config at {ckpt_path} has no architectures defined") architecture = hf_config.architectures[0] - if not hasattr(transformers, architecture) or "Deepseek" in architecture: + # 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 + + if not hasattr(transformers, architecture) or use_bundled_code: if not hasattr(transformers, architecture): warnings.warn( f"Architecture {architecture} not found in transformers: {transformers.__version__}. " @@ -830,24 +917,41 @@ def has_pack_quantized_config(config): model = from_config(config_for_init, **model_kwargs2) max_memory = get_max_memory() - inferred_device_map = infer_auto_device_map(model, max_memory=max_memory) - - on_cpu = "cpu" in inferred_device_map.values() - - if on_cpu: - for _device in max_memory: - if isinstance(_device, int): - max_memory[_device] *= gpu_mem_percentage + if _disk_offload: + for _k in max_memory: + if isinstance(_k, int): + if max_gpu_memory_gb is not None: + max_memory[_k] = int(max_gpu_memory_gb * 1024**3) + else: + max_memory[_k] = int(max_memory[_k] * gpu_mem_percentage) + if max_cpu_memory_gb is not None: + max_memory["cpu"] = int(max_cpu_memory_gb * 1024**3) + model_kwargs["max_memory"] = max_memory print( - "Model does not fit to the GPU mem. " - f"We apply the following memory limit for calibration: \n{max_memory}\n" - "If you hit GPU OOM issue, please adjust `gpu_mem_percentage` or " - "reduce the calibration `batch_size` manually." + "Disk-offload mode enabled. " + f"Memory budgets: {_fmt_max_memory(max_memory)}\n" + f"Offload folder: {offload_folder}\n" + "Weights exceeding GPU+CPU budgets will be streamed from disk." ) - model_kwargs["max_memory"] = max_memory + else: + inferred_device_map = infer_auto_device_map(model, max_memory=max_memory) + if "cpu" in inferred_device_map.values(): + for _device in max_memory: + if isinstance(_device, int): + max_memory[_device] *= gpu_mem_percentage + + print( + "Model does not fit to the GPU mem. " + f"We apply the following memory limit for calibration: \n{max_memory}\n" + "If you hit GPU OOM issue, please adjust `gpu_mem_percentage` or " + "reduce the calibration `batch_size` manually." + ) + model_kwargs["max_memory"] = max_memory model_kwargs2 = _apply_dtype_to_config(model_kwargs, config_dtype, architecture) + if _disk_offload: + model_kwargs2["offload_folder"] = offload_folder model = auto_model_module.from_pretrained( ckpt_path, device_map=device_map, @@ -916,11 +1020,13 @@ def _resolve_model_path(model_name_or_path: str, trust_remote_code: bool = False try: local_path = snapshot_download( repo_id=model_name_or_path, - allow_patterns=["*.py", "*.json"], # Only download Python files and config + allow_patterns=_HF_SIDECAR_DOWNLOAD_ALLOW_PATTERNS, ) return local_path except Exception as e: - print(f"Warning: Could not download model files using snapshot_download: {e}") + print( + f"Warning: Could not download checkpoint sidecars using snapshot_download: {e}" + ) # Fallback: try to find in HuggingFace cache from transformers.utils import TRANSFORMERS_CACHE @@ -955,49 +1061,31 @@ def _resolve_model_path(model_name_or_path: str, trust_remote_code: bool = False return model_name_or_path -def copy_custom_model_files(source_path: str, export_path: str, trust_remote_code: bool = False): - """Copy processor/tokenizer artifacts (and, with trust_remote_code, custom code) to export. - - Processor and tokenizer *data* artifacts -- e.g. a VLM's ``preprocessor_config.json``, - ``merges.txt``/``vocab.json``, and the processor helper modules -- are needed by the - deployment stack (vLLM/SGLang) even when the model itself runs on native (non-remote) - transformers code. transformers 5.x restructured many VLM configs and no longer - re-saves these on ``save_pretrained`` for models loaded natively, so without copying - them a native-path export is missing e.g. ``preprocessor_config.json`` and fails to - load (``Can't load image processor``). These are copied regardless of - ``trust_remote_code``. Executable model/config code (``modeling*.py``, - ``configuration_*.py``, ``tokenization_*.py``, and other custom JSON) is only meaningful - with ``trust_remote_code`` and is copied only then. ``config.json`` and - ``model.safetensors.index.json`` are always skipped (handled by the export itself). +def copy_custom_model_files( + source_path: str, + export_path: str, + trust_remote_code: bool = False, + exclude_files: Iterable[str] | None = None, +): + """Copy source checkpoint sidecar files to an HF PTQ export. + + The HF PTQ script writes ModelOpt-owned metadata and quantized weights first, then + copies source checkpoint sidecars so tokenizer/processor files, remote-code modules, + README assets, parser plugins, and similar deployment files are preserved for both + native and ``trust_remote_code`` loads. Weight and weight-index files are skipped + to avoid copying the unquantized source weights. Export-owned metadata (``config.json``, + ``hf_quant_config.json``) and stale source quantization metadata are also skipped. + Source tokenizer and processor files intentionally still win because Transformers may + not regenerate all metadata in the source format. The exported ``tokenizer_config.json`` + wins when it has a separate chat template. Callers that write a generation config can + exclude it; the TensorRT-LLM export retains the source generation config. Args: source_path: Path to the original model directory or HuggingFace model ID export_path: Path to the exported model directory - trust_remote_code: Whether trust_remote_code was used (gates the executable code files) + trust_remote_code: Passed to HuggingFace model-ID resolution; does not control copying. + exclude_files: Additional source file names to skip. """ - # Deployment-critical processor/tokenizer artifacts: safe to copy regardless of - # trust_remote_code (data + processor helpers, not model code). - always_copy_patterns = [ - "preprocessor_config.json", - "processor_config.json", - "image_processing*.py", - "processing_*.py", - "video_processing*.py", - "feature_extraction_*.py", - "added_tokens.json", - "special_tokens_map.json", - "vocab.json", - "merges.txt", - "tokenizer.model", - ] - # Executable custom model/config code + other custom JSON: only used with trust_remote_code. - code_patterns = [ - "configuration_*.py", - "modeling*.py", - "tokenization_*.py", - "*.json", - ] - # Resolve the source path (handles both local paths and HF model IDs) resolved_source_path = _resolve_model_path(source_path, trust_remote_code) @@ -1018,29 +1106,23 @@ def copy_custom_model_files(source_path: str, export_path: str, trust_remote_cod print(f"Warning: Export directory {export_path} does not exist") return - patterns = [*always_copy_patterns, *(code_patterns if trust_remote_code else [])] - - copied_files: list[str] = [] - for pattern in patterns: - for file_path in source_dir.glob(pattern): - if file_path.is_file(): - # Skip config.json and model.safetensors.index.json as they're handled separately - if file_path.name in ["config.json", "model.safetensors.index.json"]: - continue - if file_path.name in copied_files: # e.g. matched by both pattern lists - continue - dest_path = export_dir / file_path.name - try: - shutil.copy2(file_path, dest_path) - copied_files.append(file_path.name) - print(f"Copied custom model file: {file_path.name}") - except Exception as e: - print(f"Warning: Failed to copy {file_path.name}: {e}") + exclude_files = _HF_PTQ_EXPORT_OWNED_FILES | set(exclude_files or ()) + if (export_dir / "chat_template.jinja").is_file(): + exclude_files.add("tokenizer_config.json") + + copied_files = copy_non_safetensor_files_from_ckpt( + source_dir, + export_dir, + exclude_files=exclude_files, + exclude_patterns=_HF_PTQ_WEIGHT_FILE_PATTERNS, + ) if copied_files: - print(f"Successfully copied {len(copied_files)} custom model files to {export_path}") + for file_name in copied_files: + print(f"Copied checkpoint sidecar file: {file_name}") + print(f"Successfully copied {len(copied_files)} checkpoint sidecar files to {export_path}") else: - print("No custom model files found to copy") + print("No checkpoint sidecar files found to copy") def _layerwise_checkpoint_dir_location(algorithm) -> tuple[str, str] | None: diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 23265a47e70..a2f8e986f10 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -578,6 +578,9 @@ def load_model(args: argparse.Namespace): trust_remote_code=args.trust_remote_code, use_seq_device_map=args.use_seq_device_map, attn_implementation=args.attn_implementation, + offload_folder=args.offload_folder, + max_cpu_memory_gb=args.max_cpu_memory_gb, + max_gpu_memory_gb=args.max_gpu_memory_gb, ) else: assert args.qformat in QUANT_CFG_CHOICES, ( @@ -855,11 +858,12 @@ def export_quantized( print("This is normal for some VLM architectures that don't use AutoProcessor") start_time = time.time() - if ( + is_tensorrt_llm_export = ( model_type in ["t5", "bart", "whisper"] or args.sparsity_fmt != "dense" or "int8_sq" in args.qformat - ): + ) + if is_tensorrt_llm_export: if ( args.inference_tensor_parallel != 1 or args.inference_pipeline_parallel != 1 ) and args.qformat == "nvfp4_svdquant": @@ -931,7 +935,13 @@ def export_quantized( # from the source checkpoint take precedence over regenerated ones (which may # differ in format due to newer transformers versions). if args.dist_state.is_main: - copy_custom_model_files(args.pyt_ckpt_path, export_path, args.trust_remote_code) + exclude_files = None if is_tensorrt_llm_export else {"generation_config.json"} + copy_custom_model_files( + args.pyt_ckpt_path, + export_path, + args.trust_remote_code, + exclude_files=exclude_files, + ) end_time = time.time() print_rank_0( @@ -1620,6 +1630,38 @@ def parse_args() -> argparse.Namespace: "openai/gpt-oss-20b) and the target qformat is NVFP4-family." ), ) + parser.add_argument( + "--offload_folder", + type=str, + default=None, + help=( + "Path to a local folder for disk-offloaded model weights. " + "When set, activates disk-offload mode: model weights that exceed the GPU+CPU " + "budgets are streamed from disk during calibration and export. " + "Pair with --max_cpu_memory_gb to cap CPU RAM usage. " + "Incompatible with --low_memory_mode and --use_seq_device_map." + ), + ) + parser.add_argument( + "--max_cpu_memory_gb", + type=float, + default=None, + help=( + "Maximum CPU RAM budget in GiB for disk-offload model loading. " + "Only effective when --offload_folder is set. " + "Weights beyond this limit are streamed from disk." + ), + ) + parser.add_argument( + "--max_gpu_memory_gb", + type=float, + default=None, + help=( + "Maximum GPU memory budget per device in GiB for disk-offload model loading. " + "Only effective when --offload_folder is set. " + "Defaults to 80%% of available GPU memory when not specified." + ), + ) args = parser.parse_args() if args.moe_calib_experts_ratio is not None and not (0.0 < args.moe_calib_experts_ratio <= 1.0): @@ -1654,6 +1696,31 @@ def parse_args() -> argparse.Namespace: if args.use_fsdp2 and args.cast_mxfp4_to_nvfp4: parser.error("--use_fsdp2 does not support --cast_mxfp4_to_nvfp4.") + if args.offload_folder is not None and args.low_memory_mode: + parser.error("--offload_folder (disk-offload) is not compatible with --low_memory_mode.") + + if args.offload_folder is not None and args.use_seq_device_map: + parser.error( + "--offload_folder (disk-offload) is not compatible with --use_seq_device_map; " + "device_map=auto is used for disk-offload to let accelerate place layers across " + "GPU, CPU, and disk." + ) + + if args.offload_folder is not None and args.device == "cpu": + parser.error( + "--offload_folder (disk-offload) is not compatible with --device cpu; " + "device_map=cpu makes accelerate ignore the memory budgets and offload folder, " + "loading the whole model into RAM." + ) + + if args.offload_folder is None and ( + args.max_cpu_memory_gb is not None or args.max_gpu_memory_gb is not None + ): + parser.error( + "--max_cpu_memory_gb/--max_gpu_memory_gb only apply to disk-offload loading; " + "pass --offload_folder to enable it." + ) + return args diff --git a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_trtllm.py b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_trtllm.py index df9688cc3c3..54e8e0a837d 100644 --- a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_trtllm.py +++ b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_trtllm.py @@ -30,6 +30,8 @@ from tqdm import tqdm as tqdm from transformers import AutoConfig, AutoTokenizer +from modelopt.torch.speculative.utils import get_conversation_input_ids + REMOVE_THINK_CHAT_TEMPLATE = ( "{% if '' in content %}{% set content = content.split('')[-1] %}{% endif %}" ) @@ -263,10 +265,8 @@ async def submit_generates(): num_invalid += 1 continue - input_ids = tokenizer.apply_chat_template(conversations, add_generation_template=False) - num_input_tokens = ( - input_ids.shape[1] if isinstance(input_ids, torch.Tensor) else len(input_ids) - ) + input_ids = get_conversation_input_ids(tokenizer, conversations) + num_input_tokens = len(input_ids) if num_input_tokens <= 10 or num_input_tokens > args.max_seq_len: num_skipped_too_long += 1 continue diff --git a/examples/speculative_decoding/collect_hidden_states/send_conversations_for_hiddens.py b/examples/speculative_decoding/collect_hidden_states/send_conversations_for_hiddens.py index e664e6d46fd..f0bbe4f951e 100644 --- a/examples/speculative_decoding/collect_hidden_states/send_conversations_for_hiddens.py +++ b/examples/speculative_decoding/collect_hidden_states/send_conversations_for_hiddens.py @@ -26,6 +26,8 @@ from tqdm import tqdm from transformers import AutoTokenizer +from modelopt.torch.speculative.utils import get_conversation_input_ids + def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( @@ -149,9 +151,7 @@ async def main(args: argparse.Namespace) -> None: f, ) - input_ids = tokenizer.apply_chat_template( - conversations, return_tensors=None, add_generation_template=False, tokenize=True - ) + input_ids = get_conversation_input_ids(tokenizer, conversations) num_input_tokens = len(input_ids) if num_input_tokens <= 10 or num_input_tokens > args.max_seq_len: num_too_long += 1 diff --git a/examples/speculative_decoding/scripts/send_conversation_vllm.py b/examples/speculative_decoding/scripts/send_conversation_vllm.py index b4e4cdf9f5c..9271af121bc 100644 --- a/examples/speculative_decoding/scripts/send_conversation_vllm.py +++ b/examples/speculative_decoding/scripts/send_conversation_vllm.py @@ -26,6 +26,8 @@ from tqdm import tqdm from transformers import AutoTokenizer +from modelopt.torch.speculative.utils import get_conversation_input_ids + def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( @@ -199,9 +201,7 @@ async def main(args: argparse.Namespace) -> None: f, ) - input_ids = tokenizer.apply_chat_template( - conversations, return_tensors=None, add_generation_template=False, tokenize=True - ) + input_ids = get_conversation_input_ids(tokenizer, conversations) num_input_tokens = len(input_ids) if num_input_tokens <= 10 or num_input_tokens > args.max_seq_len: num_too_long += 1 diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index 800b51daca9..21a8a3fe246 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/modelopt/torch/export/hf_export_handlers.py @@ -45,7 +45,7 @@ def _export_weight( # install the built-in handlers while retaining this legacy helper's import path. from .unified_export_hf import _export_quantized_weight - _export_quantized_weight(module, ctx.dtype, weight_name, _tied_cache=ctx.tied_cache) + _export_quantized_weight(module, ctx.dtype, weight_name) # Preparation handlers are registered in the same precedence as the legacy MoE prepass. @@ -128,14 +128,13 @@ def _export_moe_linear(name: str, module: nn.Module, ctx: ExportContext) -> None @ExportModuleRegistry.register(predicate=_has_fused_experts_quantizers) def _export_fused_experts_module(name: str, module: nn.Module, ctx: ExportContext) -> None: - """Split and quantize a fused-experts module with plural weight quantizers.""" + """Split and quantize a fused-experts module with plural weight quantizers. + + Tied experts are packed independently and their duplicate keys are dropped by name + in postprocess_state_dict; no per-module dedup cache is used. + """ with fsdp2_aware_weight_update(ctx.model, module, reshard=False): - _export_fused_experts( - module, - ctx.dtype, - _moe_tied_cache=ctx.moe_tied_cache, - _tied_cache=ctx.tied_cache, - ) + _export_fused_experts(module, ctx.dtype) @ExportModuleRegistry.register(predicate=is_quantlinear) diff --git a/modelopt/torch/export/model_utils.py b/modelopt/torch/export/model_utils.py index 1729dbfffcf..f27405ec83f 100755 --- a/modelopt/torch/export/model_utils.py +++ b/modelopt/torch/export/model_utils.py @@ -14,7 +14,7 @@ # limitations under the License. """Utility functions for model type detection and classification.""" -import re +import warnings import torch.nn as nn @@ -70,7 +70,12 @@ {MODEL_NAME_TO_TYPE=} """ -__all__ = ["get_language_model_from_vl", "get_model_type", "is_multimodal_model"] +__all__ = [ + "TiedWeightMap", + "get_language_model_from_vl", + "get_model_type", + "is_multimodal_model", +] def get_model_type(model): @@ -149,82 +154,62 @@ def get_language_model_from_vl(model) -> list[nn.Module] | None: return None -def _collect_canonical_tied_patterns( - model: nn.Module, -) -> tuple[list[re.Pattern], list[str]]: - """Walk the model and collect canonical-side tied-weight matchers. +class TiedWeightMap: + """Name-based lookups over HF's ``{alias: canonical}`` tie map (``model.all_tied_weights_keys``). - Patterns are submodule-prefixed regexes from each module's - ``_tied_weights_keys`` dict-style declaration (the prefix matters - for nested models where the dict lives on an inner submodule). - Side substrings are dot-separated tokens that appear only on the - canonical side of those declarations — needed because modelopt's - per-expert unpacking creates post-export keys (e.g. - ``…experts.Y.gate_proj.input_scale``) that HF's regexes never knew - about. List-style (legacy) declarations are skipped. + Export sites ask for a *group key*: both sides of a tie share one key, an untied parameter + returns ``None``. The key is a name, so it survives packing / FSDP / offload, where a + ``data_ptr`` would not. """ - patterns: list[re.Pattern] = [] - alias_token_set: set[str] = set() - canonical_token_set: set[str] = set() - - def _tokens(s: str) -> set[str]: - """Identifiers in a regex string, with regex specials as separators.""" - return {tok for tok in re.split(r"[^A-Za-z0-9_]+", s) if tok} - - for name, submodule in model.named_modules(): - tied = getattr(submodule, "_tied_weights_keys", None) - if not isinstance(tied, dict) or not tied: - continue - prefix = f"{name}." if name else "" - for alias_pat, canonical_pat in tied.items(): - patterns.append(re.compile(prefix + canonical_pat)) - alias_token_set.update(_tokens(prefix + alias_pat)) - canonical_token_set.update(_tokens(prefix + canonical_pat)) - - # Tokens unique to the canonical side become substring matchers. - side_substrings = sorted(canonical_token_set - alias_token_set) - return patterns, side_substrings - - -def _reorder_canonical_first(state_dict: dict, model: nn.Module) -> dict: - r"""Reorder ``state_dict`` so canonical-side tied keys iterate first. - - Lets the downstream first-wins data_ptr dedup keep canonical names. - Uses both regex patterns and substring matchers from - :func:`_collect_canonical_tied_patterns`. Gated on the model class - name to scope the reorder to DiffusionGemma; other tied - encoder-decoder models that ship dict-style ``_tied_weights_keys`` - can be added to the allowlist here. Mirrors the ``model_type`` - dispatch used for the Whisper and Nemotron-VL branches elsewhere - in ``unified_export_hf.py``. - """ - model_type = type(model).__name__.lower() - if "diffusiongemma" not in model_type and "diffusion_gemma" not in model_type: - return state_dict - - canonical_patterns, side_substrings = _collect_canonical_tied_patterns(model) - if not canonical_patterns and not side_substrings: - return state_dict - - def _has_side_substring(key: str) -> bool: - # Require the token to appear as a proper dot-separated path - # component, not just as a substring of an unrelated identifier. - for tok in side_substrings: - if ( - f".{tok}." in key - or key.startswith(f"{tok}.") - or key.endswith(f".{tok}") - or key == tok - ): - return True - return False - - head: dict = {} - tail: dict = {} - for k, v in state_dict.items(): - if any(p.search(k) for p in canonical_patterns) or _has_side_substring(k): - head[k] = v - else: - tail[k] = v - head.update(tail) - return head + + def __init__(self, model: nn.Module) -> None: + """Source the tie map from HF's ``all_tied_weights_keys`` (transformers >=5.0). + + HF's ``{target: source}`` == our ``{alias: canonical}``, resolved at load, config-gated, + ``torch.equal``-pruned, and name-based so it survives FSDP shard / offload. Absent on + transformers <5.0 -> empty map (the ``data_ptr`` backstop in postprocess is the net). + """ + all_tied = getattr(model, "all_tied_weights_keys", None) + # Warn whenever a tie is declared (embedding tie or any ``_tied_weights_keys`` entry, e.g. + # encoder/decoder or fused-MoE) but the name-based map is missing, not just for embeddings. + declares_tie = bool( + getattr(getattr(model, "config", None), "tie_word_embeddings", False) + ) or bool(getattr(model, "_tied_weights_keys", None)) + if all_tied is None and declares_tie: + warnings.warn( + "This model may contain tied/shared weights, but deduplicating them on export " + "requires transformers>=5.0 (it uses model.all_tied_weights_keys, which is only " + "supported in newer versions). On older versions the exported checkpoint may keep " + "duplicate copies of the tied weights (larger files), and tied weights may not be " + "deduplicated correctly during export. Upgrade to transformers>=5.0 for correct " + "tied-weight export." + ) + # Drop any self-entry (alias == canonical): HF should not emit one, but a target==source + # pair would schedule the kept canonical for deletion, so filter it out defensively. + self.alias_to_canonical: dict[str, str] = { + alias: canonical for alias, canonical in (all_tied or {}).items() if alias != canonical + } + self.canonical_names: set[str] = set(self.alias_to_canonical.values()) + + def group_key(self, param_full_name: str) -> str | None: + """Canonical group key for a parameter name, or ``None`` if untied. + + Both sides of a tie return the same key, so it does not matter which side export + visits first. + """ + if param_full_name in self.alias_to_canonical: + return self.alias_to_canonical[param_full_name] + if param_full_name in self.canonical_names: + return param_full_name + return None + + def container_group_key(self, container_name: str, first_proj_attr: str) -> str | None: + """Group key for a fused-experts container, or ``None`` if untied. + + The tie lives on the container's 3-D projection (e.g. ``…experts.gate_up_proj``); + stripping that suffix gives one key shared by all the container's projections. + """ + gk = self.group_key(f"{container_name}.{first_proj_attr}") + if gk is None: + return None + return gk.removesuffix(f".{first_proj_attr}") diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index 787e173959e..4ce60b192ee 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -23,35 +23,6 @@ import torch.nn as nn -def _alias_per_expert_subtree_from_prior(module: nn.Module, prior: nn.Module, n: int) -> None: - """Build per-expert subtree on ``module`` by aliasing ``prior``'s packed buffers. - - For each expert ``idx`` in ``0..n-1``, creates ``module.{idx}.{gate,up,down}_proj`` - sub-modules whose ``weight`` / ``weight_scale`` / ``weight_scale_2`` / - ``input_scale`` are aliased to the prior side's already-packed tensors. - data_ptr equality is preserved so the downstream - ``postprocess_state_dict`` dedup collapses the duplicates at write time. - Called by ``_export_fused_experts`` on the tied-experts cache-hit fast path. - """ - for _idx in range(n): - _prior_expert = getattr(prior, str(_idx), None) - if _prior_expert is None: - continue - _cur_expert = nn.Module() - for _proj_name in ("gate_proj", "up_proj", "down_proj"): - _prior_proj = getattr(_prior_expert, _proj_name, None) - if _prior_proj is None: - continue - _cur_proj = nn.Module() - if hasattr(_prior_proj, "weight"): - _cur_proj.weight = _prior_proj.weight - for _attr in ("weight_scale", "weight_scale_2", "input_scale"): - if hasattr(_prior_proj, _attr): - _cur_proj.register_buffer(_attr, getattr(_prior_proj, _attr)) - _cur_expert.add_module(_proj_name, _cur_proj) - module.add_module(str(_idx), _cur_expert) - - def _delete_fused_moe_source_attrs(module: nn.Module) -> None: """Remove the 3-D fused source params and per-expert quantizer ModuleLists. @@ -77,8 +48,6 @@ def _delete_fused_moe_source_attrs(module: nn.Module) -> None: def _export_fused_experts( module: nn.Module, dtype: torch.dtype, - _moe_tied_cache: dict[tuple[int, int], nn.Module] | None = None, - _tied_cache: dict[int, nn.Module] | None = None, ) -> None: """Split fused MoE expert weights and export per-expert quantization scales. @@ -100,19 +69,10 @@ def _export_fused_experts( {E}.up_proj.weight, {E}.up_proj.weight_scale, ... {E}.down_proj.weight, {E}.down_proj.weight_scale, ... - Tied-experts dedup is opt-in via ``_moe_tied_cache``: when multiple - fused-expert modules share their 3-D source params via HF - ``_tied_weights_keys``, the unpacking creates fresh per-expert tensors - that break the tie. With ``_moe_tied_cache`` provided (tuple-keyed by - ``(.data_ptr(), down_proj.data_ptr())``), the alias step - at the end re-points the per-expert ``weight`` / ``weight_scale`` / - ``weight_scale_2`` / ``input_scale`` buffers at a previously-processed - module sharing the same source memory. ``_tied_cache`` (int-keyed) is - threaded through to the per-projection ``_export_quantized_weight`` - calls so wrapper-level dedup uses the same scope as standalone Linears. - Both caches are owned by the caller (typically - ``_export_transformers_checkpoint``) and scoped to one export - invocation; when ``None`` the corresponding alias step is skipped. + Tied experts are not deduped here: when multiple fused-expert modules share their + 3-D source params via HF ``_tied_weights_keys``, each is split and packed + independently to byte-identical per-expert tensors, and the duplicate keys are + dropped by name in ``postprocess_state_dict`` (the single dedup authority). """ from modelopt.torch.export.unified_export_hf import _export_quantized_weight from modelopt.torch.quantization.plugins.huggingface import _get_fused_expert_intermediate_dim @@ -124,25 +84,6 @@ def _export_fused_experts( # Only the gated split needs the per-expert intermediate dim (gate|up boundary). expert_dim = _get_fused_expert_intermediate_dim(module) if is_gated else None - # Capture source tensor identities BEFORE unpacking (the source - # attrs are deleted at the end of this function). - _source_key = ( - getattr(module, first_proj_attr).data_ptr(), - module.down_proj.data_ptr(), - ) - - # Tied-experts fast path: if this exact (first_proj, down) source-tensor pair - # has been processed before, alias all per-expert buffers directly from the - # prior module — no unpacking, no per-expert packing, no transient buffers - # thrown away. Cache miss falls through to the full unpack/pack below and - # registers this module as the prior for any later tied module. - if _moe_tied_cache is not None: - _prior = _moe_tied_cache.get(_source_key) - if _prior is not None and _prior is not module: - _alias_per_expert_subtree_from_prior(module, _prior, n) - _delete_fused_moe_source_attrs(module) - return - # 1. Shared input quantizers — one per projection type, shared across all experts. first_proj_input_q = getattr(module, f"{first_proj_attr}_input_quantizer") first_proj_weight_quantizers = getattr(module, f"{first_proj_attr}_weight_quantizers") @@ -271,7 +212,7 @@ def _export_fused_experts( wrapper.weight_quantizer = w_quantizer wrapper.input_quantizer = i_quantizer - _export_quantized_weight(wrapper, dtype, _tied_cache=_tied_cache) + _export_quantized_weight(wrapper, dtype) proj = nn.Module() proj.weight = wrapper.weight @@ -286,13 +227,6 @@ def _export_fused_experts( # 4. Remove fused params and quantizer lists — replaced by per-expert submodules _delete_fused_moe_source_attrs(module) - # 5. Register this module in the dedup cache so any later tied module - # (same source data_ptr pair) takes the fast path at the top of this - # function. Reached only on cache miss; cache-hit modules early-exited - # above before any unpack work. - if _moe_tied_cache is not None: - _moe_tied_cache[_source_key] = module - def save_expert_token_count_table(model: nn.Module, output_dir: str | Path | None = None): """Collect expert_token_count from all quantized MoE layers and save as an HTML table. diff --git a/modelopt/torch/export/plugins/hf_checkpoint_utils.py b/modelopt/torch/export/plugins/hf_checkpoint_utils.py index ddae8e2b409..54058c742db 100644 --- a/modelopt/torch/export/plugins/hf_checkpoint_utils.py +++ b/modelopt/torch/export/plugins/hf_checkpoint_utils.py @@ -15,10 +15,12 @@ """Hugging Face checkpoint utility.""" +import fnmatch import json import os import shutil import warnings +from collections.abc import Iterable from pathlib import Path from typing import Any @@ -253,25 +255,50 @@ def load_multimodal_components( return multimodal_state_dict -def copy_non_safetensor_files_from_ckpt(src: str | os.PathLike, dst: str | os.PathLike): +def _matches_any_pattern(file_name: str, patterns: tuple[str, ...]) -> bool: + return any(fnmatch.fnmatchcase(file_name, pattern) for pattern in patterns) + + +def copy_non_safetensor_files_from_ckpt( + src: str | os.PathLike, + dst: str | os.PathLike, + *, + exclude_files: Iterable[str] | None = None, + exclude_patterns: Iterable[str] | None = None, +) -> list[str]: """Copy every non-safetensors file from a local HF checkpoint dir verbatim. Use as a baseline so tokenizer files, remote_code ``*.py``, README, LICENSE, etc. - are preserved from the source. The caller is expected to overwrite the files - modelopt owns (``config.json``, ``generation_config.json``, ``hf_quant_config.json``, - ``preprocessor_config.json``) after this step. + are preserved from the source. Callers can exclude additional files or patterns when + copying after export-owned metadata has already been written. Args: src: Source HF checkpoint directory. Must be a local path. dst: Destination directory; created if missing. + exclude_files: Exact file names to skip. + exclude_patterns: Glob patterns for additional files to skip. + + Returns: + File names copied into ``dst``. """ if not os.path.isdir(src): raise ValueError(f"Invalid source path: {src}. It should be a directory.") + exclude_files = set(exclude_files or ()) + exclude_patterns = tuple(exclude_patterns or ()) + copied_files = [] os.makedirs(dst, exist_ok=True) - for entry in os.listdir(src): + for entry in sorted(os.listdir(src)): + if entry in exclude_files or _matches_any_pattern(entry, exclude_patterns): + continue sp = os.path.join(src, entry) if not os.path.isfile(sp): continue if entry.endswith(".safetensors") or entry == "model.safetensors.index.json": continue - shutil.copy2(sp, dst) + try: + shutil.copy2(sp, dst) + except OSError as error: + warnings.warn(f"Failed to copy checkpoint sidecar {entry}: {error}") + continue + copied_files.append(entry) + return copied_files diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index cc894d0ffd5..c86af3aa9f5 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -16,6 +16,7 @@ """Utils for quantization including scaling factors adjustments.""" import logging +from collections import defaultdict from collections.abc import Generator from types import SimpleNamespace from typing import Any @@ -70,6 +71,7 @@ QUANTIZATION_W4A8_NVFP4_FP8, QUANTIZATION_W4A16_NVFP4, ) +from .model_utils import TiedWeightMap logger = logging.getLogger(__name__) @@ -959,6 +961,22 @@ def from_quantized_weight( raise NotImplementedError(f"quantization format {quantization} not supported") +_KV_CACHE_REPLACEMENTS: dict[str, str] = { + "k_bmm_quantizer._amax": "k_proj.k_scale", + "v_bmm_quantizer._amax": "v_proj.v_scale", + "k_bmm_quantizer._bias_value": "k_proj.k_bias", + "v_bmm_quantizer._bias_value": "v_proj.v_bias", + "input_quantizer._pre_quant_scale": "pre_quant_scale", +} +_BASE_SKIP_KEYS: tuple[str, ...] = ( + "output_quantizer", + "_amax", + "_bias_value", + "input_quantizer._pre_quant_scale", + "weight_shape", +) + + def _strip_base_layer(key: str, is_modelopt_qlora: bool) -> str: """Drop the `base_layer` component PEFT inserts, which deployment does not expect. @@ -967,11 +985,78 @@ def _strip_base_layer(key: str, is_modelopt_qlora: bool) -> str: return key.replace(".base_layer.", ".") if is_modelopt_qlora else key +def _maybe_squeeze_scale(key: str, value: Any) -> Any: + """Squeeze a leading dim=1 from 3-D scale tensors of shape (1, n, m).""" + if ( + "scale" in key + and isinstance(value, torch.Tensor) + and value.dim() == 3 + and value.shape[0] == 1 + ): + return value.squeeze(0) + return value + + +def _postprocess_single_tensor( + key: str, + value: torch.Tensor, + kv_cache_max_bound: float, + kv_cache_format: str | None, + is_modelopt_qlora: bool = False, +) -> tuple[str | None, torch.Tensor | None]: + """Per-tensor subset of :func:`postprocess_state_dict`, for streaming export. + + Returns ``(new_key, new_value)`` to emit, or ``(None, None)`` to skip. + Tied-weight dedup is NOT performed here; callers should pre-compute alias + keys from ``model._tied_weights_keys`` and filter them at the call site. + """ + replacements = _KV_CACHE_REPLACEMENTS + skip_keys = _BASE_SKIP_KEYS + + # Skip problematic VL model parameters + if key == "vision_model.radio_model.summary_idxs": + return None, None + + # Skip real quant parameters + if any(key.endswith("weight_quantizer." + q) for q in RealQuantLinear.list_of_scale_tensors): + return None, None + + # Skip LoRA adapters for QLoRA models + if is_modelopt_qlora and "lora" in key: + return None, None + + # Keys not related to quantizers: keep as-is + if all(sk not in key for sk in skip_keys): + new_key = _strip_base_layer(key, is_modelopt_qlora) + return new_key, _maybe_squeeze_scale(new_key, value) + + # Apply replacements if the key matches any suffix in the replacements dict + for old_suffix, new_suffix in replacements.items(): + if key.endswith(old_suffix): + prefix = key[: -len(old_suffix)] + if "_amax" in key: + assert kv_cache_format in [KV_CACHE_FP8, KV_CACHE_NVFP4, KV_CACHE_NVFP4_AFFINE], ( + "Invalid KV cache quantization format." + ) + assert kv_cache_max_bound > 0, "Maxbound must be greater than zero." + value = value.float() / kv_cache_max_bound + if kv_cache_format == KV_CACHE_FP8 and value.item() > 0.5: + logger.warning( + "Large KV activations detected. Quantized KV cache may lead to higher accuracy drop." + ) + new_key = _strip_base_layer(prefix + new_suffix, is_modelopt_qlora) + return new_key, _maybe_squeeze_scale(new_key, value) + + # Key has a skip_key but no replacement matched — drop it + return None, None + + def postprocess_state_dict( state_dict: dict, maxbound: float, quantization: str | None, is_modelopt_qlora: bool = False, + tied_map: "TiedWeightMap | None" = None, ) -> dict: """Filters out keys related to weight quantizers and updates KV cache related keys. @@ -980,24 +1065,19 @@ def postprocess_state_dict( maxbound: The maximum bound value for the output quantizer. quantization: The KV cache quantization format. is_modelopt_qlora: Whether the model is a modelopt-trained QLoRA model. + tied_map: Optional :class:`TiedWeightMap`. When provided, tied-weight + dedup is authoritative and name-based: a declared alias key whose canonical + counterpart is present is dropped, independent of tensor address. This is + what makes dedup correct under the FSDP full-state-dict gather (and offload), + where tied tensors are materialized at distinct addresses and the address + pass below cannot see the tie. The address pass is retained as a backstop for + undeclared genuine shares and coincidental collisions. Returns: The filtered state_dict without unnecessary keys like '_amax' and non KV cache output quantizers. """ - replacements = { - "k_bmm_quantizer._amax": "k_proj.k_scale", - "v_bmm_quantizer._amax": "v_proj.v_scale", - "k_bmm_quantizer._bias_value": "k_proj.k_bias", - "v_bmm_quantizer._bias_value": "v_proj.v_bias", - "input_quantizer._pre_quant_scale": "pre_quant_scale", - } - skip_keys = [ - "output_quantizer", - "_amax", - "_bias_value", - "input_quantizer._pre_quant_scale", - "weight_shape", - ] + replacements = _KV_CACHE_REPLACEMENTS + skip_keys = _BASE_SKIP_KEYS def _export_key(key: str) -> str: return _strip_base_layer(key, is_modelopt_qlora) @@ -1036,15 +1116,7 @@ def _export_key(key: str) -> str: post_state_dict[_export_key(prefix + new_suffix)] = value break - # Squeeze scales with a leading dimension of 1 - for key, value in post_state_dict.items(): - if ( - "scale" in key - and isinstance(value, torch.Tensor) - and value.dim() == 3 - and value.shape[0] == 1 - ): - post_state_dict[key] = value.squeeze(0) + post_state_dict = {k: _maybe_squeeze_scale(k, v) for k, v in post_state_dict.items()} # remove real quant parameters from the state dict keys_to_delete = [] @@ -1060,27 +1132,148 @@ def _export_key(key: str) -> str: for key in post_state_dict: if "lora" in key and key not in keys_to_delete: keys_to_delete.append(key) - # Check for tied weights and remove duplicates - seen_tensors = {} + # Name-based tied-weight dedup (authoritative, address-independent): for each declared + # {alias: canonical} tie, drop the alias's own exported keys when the canonical is present. + # Per-parameter (not per-prefix), so an untied sibling (a bias, an untied projection) is kept. + # This is what makes it correct under the FSDP gather / offload, where tied tensors land at + # distinct addresses. + dropped_dense_prefixes: set[str] = set() + if tied_map is not None and tied_map.alias_to_canonical: + # Expand each tied *pre-pack* parameter name into the concrete keys export produced. If + # export starts emitting a new scale companion or fused-projection name, extend + # `weight_suffixes` / `proj_splits` here -- nothing else needs to change. + # (pre_quant_scale is the AWQ / NVFP4_AWQ / SVDQuant companion, renamed in the KV-cache pass.) + weight_suffixes = ( + "weight", + "weight_scale", + "weight_scale_2", + "input_scale", + "pre_quant_scale", + ) + # A tied 3-D fused projection splits into these per-expert 2-D projection names. + proj_splits = { + "gate_up_proj": ("gate_proj", "up_proj"), + "up_proj": ("up_proj",), + "down_proj": ("down_proj",), + } + + # group name -> [(alias_key, canonical_key), ...], dropped all-or-none. + alias_groups: dict[str, list[tuple[str, str]]] = {} + # dense group name -> module prefix, for the leftover-companion check below. + dense_prefixes: dict[str, str] = {} + # fused-MoE container (alias_prefix, canonical_prefix) -> its tied per-expert projection names. + moe_containers: dict[tuple[str, str], set[str]] = {} + + for alias, canonical in tied_map.alias_to_canonical.items(): + a_pre, _, a_name = alias.rpartition(".") + c_pre, _, c_name = canonical.rpartition(".") + if a_name == c_name == "weight": + # Dense tie: the weight and its own scale companions only (never a sibling bias). + members = [ + (f"{a_pre}.{s}" if a_pre else s, f"{c_pre}.{s}" if c_pre else s) + for s in weight_suffixes + ] + members = [(ak, ck) for ak, ck in members if ak in post_state_dict] + if members: + alias_groups[alias] = members + dense_prefixes[alias] = a_pre + elif a_name in proj_splits and a_name == c_name: + # Fused-MoE: export splits the 3-D container into per-expert keys; record those names. + moe_containers.setdefault((a_pre, c_pre), set()).update(proj_splits[a_name]) + elif alias in post_state_dict: + alias_groups[alias] = [(alias, canonical)] + + for (a_pre, c_pre), tied_proj_names in moe_containers.items(): + # Rewrite only the tied projections' per-expert keys, so an untied projection (e.g. + # down_proj when only gate_up_proj is tied) or a router/bias child is left alone. + prefix = f"{a_pre}." if a_pre else "" + members = [ + (key, c_pre + key[len(a_pre) :]) + for key in post_state_dict + if key.startswith(prefix) + and any(part in tied_proj_names for part in key[len(prefix) :].split(".")) + ] + if members: + # Key by the full (alias, canonical) pair: one alias container's projections may + # tie to different canonical containers, and keying by a_pre alone would overwrite. + alias_groups[f"{a_pre} -> {c_pre}"] = members + + # Atomic drop: remove a group only when every alias key has its canonical twin present, + # so mixed quant state (one side quantized, the other not) never orphans a scale. + for gk, members in alias_groups.items(): + missing = [ak for ak, ck in members if ck not in post_state_dict] + if missing: + logger.warning( + f"Skipping name-based dedup of '{gk}': tied sides have mismatched keys " + f"(e.g. '{missing[0]}' has no canonical counterpart, likely differing " + f"quantization state); keeping both sides to avoid orphaned tensors." + ) + continue + # Safety: a declared tie must export identical bytes on both sides. If quantization + # diverged them, dropping the alias would silently corrupt it (HF re-ties canonical over + # it on load) -- so raise. Mirrors HF's own torch.equal decline-to-tie. + for ak, ck in members: + av, cv = post_state_dict[ak], post_state_dict[ck] + if ( + isinstance(av, torch.Tensor) + and isinstance(cv, torch.Tensor) + and not av.is_meta + and not cv.is_meta + and not torch.equal(av, cv) + ): + raise RuntimeError( + f"Tied-weight export mismatch: '{ak}' differs from its canonical '{ck}'. " + f"The tie is declared but the two sides quantized to different values, so " + f"deduplicating would corrupt '{ak}'. Ensure both tied modules use the same " + f"quantization format/config." + ) + for ak, _ in members: + keys_to_delete.append(ak) + if gk in dense_prefixes: + dropped_dense_prefixes.add(dense_prefixes[gk]) + logger.warning( + f"Tied weight (declared): dropping {len(members)} alias key(s) for '{gk}'; " + f"canonical kept." + ) - # Remove any tied weights if found. + # Apply the name drops first so the address pass below runs on the reduced dict and never + # re-processes a declared alias. dict.fromkeys dedups the delete list while preserving order. + for key in dict.fromkeys(keys_to_delete): + post_state_dict.pop(key, None) + + # Address backstop (pre-existing, unchanged): drop any remaining keys that still share a + # data_ptr. Declared ties are already gone above, so this only catches an undeclared + # same-address share (an unquantized/unpacked tie); a quantized tie packs to distinct storage + # and never reaches here. Zero-pointer (meta) tensors are left to serialization. + seen_tensors: dict = {} + backstop_delete = [] for key, value in post_state_dict.items(): - if isinstance(value, torch.Tensor): - # Use tensor data pointer to identify tied weights + if isinstance(value, torch.Tensor) and value.data_ptr() != 0: tensor_id = value.data_ptr() if tensor_id in seen_tensors: - # This is a tied weight, mark for deletion and warn - keys_to_delete.append(key) + backstop_delete.append(key) logger.warning( f"Found tied weight: '{key}' is tied to '{seen_tensors[tensor_id]}'. " f"Removing duplicate '{key}' from the exported state dict." ) else: seen_tensors[tensor_id] = key - - for key in keys_to_delete: + for key in backstop_delete: del post_state_dict[key] + # If a dense tie was dropped but a non-`bias` key still remains under its prefix, it is likely + # a new quantizer companion missing from `weight_suffixes` -- warn so it gets added. + for pre in dropped_dense_prefixes: + leftovers = [ + k for k in post_state_dict if k.startswith(f"{pre}.") and k.rsplit(".", 1)[-1] != "bias" + ] + if leftovers: + logger.warning( + f"Tied weight '{pre}.weight' was deduped but companion key(s) {leftovers} remain " + f"under '{pre}.'; likely an un-enumerated quantizer companion -- add its suffix to " + f"weight_suffixes in postprocess_state_dict so it is dropped with the tie." + ) + return post_state_dict @@ -1527,27 +1720,24 @@ def has_quantized_modules(model: nn.Module) -> bool: ) -def sync_tied_input_amax(model: nn.Module) -> int: - """Max-merge input_quantizer amaxes across modules sharing a weight ``data_ptr``. +def sync_tied_input_amax(model: nn.Module, tied_map: "TiedWeightMap | None" = None) -> int: + """Max-merge ``input_quantizer`` amaxes across modules that share a weight, in place. - Mutates ``model`` in place: overwrites the ``.amax`` buffer on every - affected ``input_quantizer`` with the per-group maximum. Intended to - run as part of an export pipeline that already replaces weights with - packed bytes downstream — i.e. the model is not expected to be reused - after this helper runs. + Tied modules whose forward paths see different activation ranges (encoder vs decoder in + YOCO-style models) must end up with one ``input_scale`` covering every side. Run BEFORE + per-module export so the merged amax flows into ``input_scale`` derivation; the model is + not expected to be reused afterward. - Closes the loop on ``input_scale`` for HF-tied modules whose forward - paths see different activation distributions (encoder vs decoder in - YOCO-style models). Must run BEFORE per-module export so the merged - amax flows into ``input_scale`` derivation. Handles both dense - Linears (keyed by ``weight.data_ptr()``) and fused MoE (keyed by - ``(, down_proj)`` data_ptr tuple). Returns the number of - tied groups merged. + Declared ties are grouped by name via :class:`TiedWeightMap` (dense Linears and fused-MoE + containers). A physically shared but *undeclared* weight is grouped by ``id(weight)`` as a + fallback, so the side the address backstop later drops still had its amax merged in here. + Returns the number of groups merged; pass ``tied_map`` to reuse one, else it is built here. """ - from collections import defaultdict + if tied_map is None: + tied_map = TiedWeightMap(model) - by_dp: dict = defaultdict(list) - for _, m in model.named_modules(): + by_group: dict = defaultdict(list) + for name, m in model.named_modules(): # Fused MoE: 3-D source tensors with shared input quantizers first_proj_attr = getattr(m, "_first_proj_attr", "gate_up_proj") first_proj = getattr(m, first_proj_attr, None) @@ -1558,15 +1748,25 @@ def sync_tied_input_amax(model: nn.Module) -> int: and hasattr(m, "down_proj") and first_proj.dim() == 3 ): - key = ("moe", first_proj.data_ptr(), m.down_proj.data_ptr()) - by_dp[key].append(m) + gk = tied_map.container_group_key(name, first_proj_attr) + if gk is not None: + by_group[("moe", gk)].append(m) + else: + # Undeclared share: group by projection identity so its amaxes still merge + # (symmetric with the dense fallback below). + by_group[("moe_shared", id(first_proj))].append(m) # Dense quantized Linear with an input_quantizer elif ( hasattr(m, "input_quantizer") and hasattr(m, "weight") and isinstance(m.weight, torch.nn.Parameter) ): - by_dp[("dense", m.weight.data_ptr())].append(m) + gk = tied_map.group_key(f"{name}.weight" if name else "weight") + if gk is not None: + by_group[("dense", gk)].append(m) + else: + # Undeclared share: group by object identity so its amaxes still merge. + by_group[("dense_shared", id(m.weight))].append(m) def _merge(quantizers: list) -> bool: """Max-merge amaxes across the quantizer list. Returns True on merge.""" @@ -1594,7 +1794,7 @@ def _merge(quantizers: list) -> bool: return True synced = 0 - for key, modules in by_dp.items(): + for key, modules in by_group.items(): if len(modules) < 2: continue if key[0] == "moe": diff --git a/modelopt/torch/export/registry.py b/modelopt/torch/export/registry.py index 260cb32eea3..5af2a8c2c0e 100644 --- a/modelopt/torch/export/registry.py +++ b/modelopt/torch/export/registry.py @@ -27,13 +27,11 @@ """ from collections.abc import Callable -from dataclasses import dataclass, field +from dataclasses import dataclass import torch import torch.nn as nn -from modelopt.torch.utils.distributed import is_fsdp2_model - __all__ = [ "ExportContext", "ExportHandler", @@ -46,27 +44,16 @@ class ExportContext: """Shared state for a single export invocation, passed to every handler call. - The tied-weight dedup caches must be scoped to one export invocation: a - process-global cache would carry stale entries whose ``data_ptr`` keys can be - recycled by PyTorch's allocator across exports, causing silent false-positive - aliasing. ``tied_cache`` (int keys) holds dense Linear / per-expert wrapper - dedup; ``moe_tied_cache`` (tuple keys) holds MoE fused-experts module dedup. + Tied-weight dedup is not a handler concern: the driver builds one name-based + :class:`TiedWeightMap` and feeds it to ``sync_tied_input_amax`` and + ``postprocess_state_dict`` directly. Both dense and fused-MoE tied weights are packed + independently and their duplicate keys are dropped by name there, so the context + carries no tied-weight map (handlers never consulted it). """ model: nn.Module dtype: torch.dtype is_modelopt_qlora: bool = False - tied_cache: dict[int, nn.Module] | None = field(default_factory=dict) - moe_tied_cache: dict[tuple[int, int], nn.Module] | None = field(default_factory=dict) - - def __post_init__(self) -> None: - # FSDP2 may recycle data_ptr() values as modules are resharded, so pointer-keyed dedup can - # falsely alias distinct weights. Disable it for FSDP2; consequently, legitimately tied - # packed weights and scale buffers are not re-aliased and may be stored as duplicates. - # TODO: replace this with stable, name-based tied-group deduplication. - if is_fsdp2_model(self.model): - self.tied_cache = None - self.moe_tied_cache = None ExportHandler = Callable[[str, nn.Module, ExportContext], None] diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index ecb69a3f906..77429b1cfaf 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -60,6 +60,7 @@ from modelopt.torch.quantization.qtensor.base_qtensor import QTensorWrapper from modelopt.torch.quantization.qtensor.nvfp4_tensor import _cast_per_block_scale_to_fp8 from modelopt.torch.quantization.utils import fsdp2_aware_weight_update, quantizer_attr_names +from modelopt.torch.quantization.utils.core_utils import has_accelerate_offload from modelopt.torch.utils.dataset_utils import _disable_use_cache from modelopt.torch.utils.distributed import is_fsdp2_model @@ -91,7 +92,7 @@ QUANTIZATION_W4A8_NVFP4_FP8, QUANTIZATION_W4A16_NVFP4, ) -from .model_utils import _reorder_canonical_first, get_language_model_from_vl, is_multimodal_model +from .model_utils import TiedWeightMap, get_language_model_from_vl, is_multimodal_model from .plugins import SpeculativeDecodingExporter, has_spec_opt, sanitize_hf_config_for_deployment from .quant_aware_conversion import ( build_reverse_name_mapper, @@ -569,24 +570,17 @@ def _export_quantized_weight( sub_module: nn.Module, dtype: torch.dtype, weight_name: str = "weight", - _tied_cache: dict[int, nn.Module] | None = None, ): """For the given weight attr of the sub_module, export the quantization info of it. The export includes converting weight tensor to correct quantized values and quantized dtype, and registering scaling factors. - Tied-weight dedup is opt-in via ``_tied_cache``: the setattr below replaces - ``.weight`` with a fresh ``nn.Parameter`` wrapping packed bytes, breaking - any HF-level tie. When the caller passes a ``_tied_cache`` dict (keyed by - the pre-pack ``weight.data_ptr()``), the alias step at the end re-points - ``weight`` / ``weight_scale`` / ``weight_scale_2`` at a previously-processed - module sharing the same source memory so the downstream data_ptr dedup can - collapse them. The cache is owned by the caller (typically - ``_export_transformers_checkpoint``) and scoped to one export invocation; - when ``_tied_cache`` is ``None`` (the default) the alias step is skipped - entirely. Uses memory identity only — no ``_tied_weights_keys`` lookup, - no-op for non-tied modules. + Tied-weight dedup is not handled here: both sides of a tie are packed independently + (identically, once ``sync_tied_input_amax`` has equalized their scales), and the + duplicate is dropped by name in :func:`postprocess_state_dict`. Deduping per-module at + pack time only ever made the packed tensors share an address for the address-based + drop; the name-based drop needs no such aliasing. """ quantization_format = get_quantization_format(sub_module) if quantization_format == QUANTIZATION_NONE: @@ -596,12 +590,13 @@ def _export_quantized_weight( quantizer_attrs = quantizer_attr_names(weight_name) weight: nn.Parameter = getattr(sub_module, weight_name) - # Capture source identity BEFORE any tensor-creating operation below. - # For HF-tied weights this matches across all modules sharing the - # underlying Parameter; the cache lookup at the end of this function - # uses it to detect ties whose Python identity is about to be broken - # by the setattr on `weight_name` further down. - _tied_source_data_ptr = weight.data_ptr() + if weight.is_meta: + raise RuntimeError( + f"Weight '{weight_name}' of {type(sub_module).__name__} is a meta tensor during " + "export. If the model was loaded with disk/CPU offload, use export_hf_checkpoint() " + "which dispatches to the streaming writer that materialises weights layer-by-layer." + ) + weight_quantizer: TensorQuantizer | SequentialQuantizer = getattr( sub_module, quantizer_attrs.weight_quantizer ) @@ -812,35 +807,90 @@ def _export_quantized_weight( if weight_scale is not None: sub_module.register_buffer(quantizer_attrs.weight_scale, weight_scale) - # Tied-weight dedup: if a previously-processed module shared the same - # source weight memory, alias the packed weight + scale buffers so the - # downstream data_ptr dedup in postprocess_state_dict can collapse them. - # input_scale is safe to alias because sync_tied_input_amax (earlier in - # this export) already max-merged the per-side amaxes. Gated on the - # caller-owned _tied_cache so the dedup state is scoped to one export. - if _tied_cache is not None: - _prior = _tied_cache.get(_tied_source_data_ptr) - if _prior is not None and _prior is not sub_module: - if hasattr(_prior, weight_name): - setattr(sub_module, weight_name, getattr(_prior, weight_name)) - for _attr in ( - quantizer_attrs.weight_scale, - quantizer_attrs.weight_scale_2, - quantizer_attrs.input_scale, - ): - if not hasattr(_prior, _attr): - continue - if _attr in sub_module._buffers: - del sub_module._buffers[_attr] - elif hasattr(sub_module, _attr): - delattr(sub_module, _attr) - sub_module.register_buffer(_attr, getattr(_prior, _attr)) - else: - _tied_cache[_tied_source_data_ptr] = sub_module - torch.cuda.empty_cache() +def _dispatch_export_handler(name: str, sub_module: nn.Module, ctx: ExportContext) -> None: + """QLoRA skip, unpack-weight preprocessing, and handler dispatch for one module.""" + if ctx.is_modelopt_qlora and hasattr(sub_module, "base_layer"): + return + # Restore unpacked weight so the export path can read the live quantizer state. + if hasattr(sub_module, "weight_packed") or ( + "QuantFP8Linear" in type(sub_module).__name__ and sub_module.weight.element_size() <= 1 + ): + sub_module.unpack_weight() + handler = ExportModuleRegistry.match(sub_module) + if handler is not None: + handler(name, sub_module, ctx) + + +def _resolve_export_dtype(model: nn.Module, dtype: torch.dtype | None) -> torch.dtype: + """Return the export dtype, defaulting to the model's own and warning on a mismatch.""" + if dtype is None: + return model.config.torch_dtype + if dtype != model.config.torch_dtype: + warnings.warn( + f"Model's original dtype ({model.config.torch_dtype}) differs from target dtype " + f"({dtype}), which may lead to numerical errors." + ) + return dtype + + +def _prepare_moe_inputs( + model: nn.Module, + dtype: torch.dtype, + is_modelopt_qlora: bool, +) -> None: + """Handle input quantizers of experts that are not calibrated. + + Each MoE block is dispatched by its experts container to the matching preparation + handler. + """ + prepare_ctx = ExportContext(model=model, dtype=dtype, is_modelopt_qlora=is_modelopt_qlora) + for name, sub_module in model.named_modules(): + if is_moe(sub_module) and hasattr(sub_module, "experts"): + handler = PrepareMoEInputsRegistry.match(sub_module.experts) + if handler is None: + # Unsupported MoE model structure + raise NotImplementedError( + f"MoE model with experts type '{type(sub_module.experts).__name__}' is not supported in export." + f"Please file an issue or add support for this model architecture." + ) + handler(name, sub_module, prepare_ctx) + + +def _add_mtp_exclusions(model: nn.Module, quant_config: dict) -> None: + """Add MTP layer prefixes to exclude_modules if they were excluded from quantization. + + This ensures they appear in ``quantization_config["ignore"]`` in ``config.json``. + """ + mtp_layer_prefixes = getattr(model, "_mtp_layer_prefixes", None) + if mtp_layer_prefixes: + exclude_modules = quant_config["quantization"].setdefault("exclude_modules", []) + for prefix in mtp_layer_prefixes: + # Add wildcard pattern to exclude all submodules under this MTP layer + pattern = f"{prefix}*" + if pattern not in exclude_modules: + exclude_modules.append(pattern) + print(f"Adding MTP layer to quantization_config ignore: {pattern}") + + +def _warn_on_unsynced_moe_gate_up(model: nn.Module) -> None: + """Safety net for gate/up weight quantizer amaxes that resmoothing did not reach. + + ``requantize_resmooth_fused_llm_layers`` can miss experts that the dummy forward + never activated, or that use non-standard expert naming. + """ + synced = sync_moe_gate_up_amax(model) + if synced: + warnings.warn( + f"Found {synced} MoE expert gate/up projection pair(s) with mismatched " + f"weight_scale_2 after requantize_resmooth_fused_llm_layers. " + f"This typically means the dummy forward did not activate these experts. " + f"Taking element-wise max of amaxes for serving-engine fusion." + ) + + def _process_quantized_modules( model: nn.Module, dtype: torch.dtype, @@ -857,9 +907,7 @@ def _process_quantized_modules( is_modelopt_qlora: Whether the model is a modelopt-trained QLoRA model. If True, modules with base_layer attribute are skipped. """ - # Per-call tied-weight dedup caches inside the context. Created fresh on - # every invocation so cache state is scoped to one export and cannot leak - # into a later call (see ExportContext). + # No per-module dedup cache: tied duplicates are dropped by name in postprocess_state_dict. ctx = ExportContext(model=model, dtype=dtype, is_modelopt_qlora=is_modelopt_qlora) fsdp_module_to_reshard = None @@ -874,20 +922,7 @@ def _process_quantized_modules( fsdp_module_to_reshard = sub_module - # We skip QuantLoraLinear module for modelopt QLoRA - if ctx.is_modelopt_qlora and hasattr(sub_module, "base_layer"): - continue - - # Preprocessing: restore unpacked weight so the export path can read - # the live quantizer state. Falls through to the handler dispatch below. - if hasattr(sub_module, "weight_packed") or ( - "QuantFP8Linear" in type(sub_module).__name__ and sub_module.weight.element_size() <= 1 - ): - sub_module.unpack_weight() - - handler = ExportModuleRegistry.match(sub_module) - if handler is not None: - handler(name, sub_module, ctx) + _dispatch_export_handler(name, sub_module, ctx) def _export_transformers_checkpoint( @@ -900,6 +935,11 @@ def _export_transformers_checkpoint( The packed checkpoint will be consumed by the TensorRT-LLM unified converter. + Builds the whole quantized state dict in memory, so it requires every weight to be + resident. Models with accelerate CPU/disk offload are rejected here and handled by + :func:`_export_transformers_checkpoint_streaming`, which materializes one layer at a + time; :func:`export_hf_checkpoint` picks between the two. + Args: model: the full torch model to export. The actual quantized model may be a submodule. dtype: the weights data type to export the unquantized layers or the default model data type if None. @@ -907,70 +947,45 @@ def _export_transformers_checkpoint( Returns: post_state_dict: Dict containing quantized weights quant_config: config information to export hf_quant_cfg.json - """ - if dtype is None: - dtype = model.config.torch_dtype - elif dtype != model.config.torch_dtype: - warnings.warn( - f"Model's original dtype ({model.config.torch_dtype}) differs from target dtype " - f"({dtype}), which may lead to numerical errors." - ) - # Handle input quantizers of experts that are not calibrated. Each MoE block is - # dispatched by its experts container to the matching preparation handler. - prepare_ctx = ExportContext(model=model, dtype=dtype, is_modelopt_qlora=is_modelopt_qlora) - for name, sub_module in model.named_modules(): - if is_moe(sub_module) and hasattr(sub_module, "experts"): - handler = PrepareMoEInputsRegistry.match(sub_module.experts) - if handler is None: - # Unsupported MoE model structure - raise NotImplementedError( - f"MoE model with experts type '{type(sub_module.experts).__name__}' is not supported in export." - f"Please file an issue or add support for this model architecture." - ) - handler(name, sub_module, prepare_ctx) + Raises: + NotImplementedError: if the model has accelerate offload hooks. + """ + dtype = _resolve_export_dtype(model, dtype) + # One tied-weight map for the whole export (amax sync + final dedup in postprocess_state_dict). + # Sourced from HF's name-based all_tied_weights_keys, so it is correct even under FSDP/offload. + tied_map = TiedWeightMap(model) + _prepare_moe_inputs(model, dtype, is_modelopt_qlora) # Resmooth and requantize fused layers # TODO: Handle mixed precision requantize_resmooth_fused_llm_layers(model) + # Offloaded models need their weights materialized layer-by-layer, which this + # whole-state-dict path cannot do; export_hf_checkpoint() streams them instead. + if has_accelerate_offload(model): + raise NotImplementedError( + "_export_transformers_checkpoint does not support disk/CPU-offloaded models. " + "Use export_hf_checkpoint() which dispatches to _export_transformers_checkpoint_streaming." + ) + # Remove all hooks from the model try: from accelerate.hooks import remove_hook_from_module remove_hook_from_module(model, recurse=True) except ImportError: - warnings.warn("accelerate is not installed, hooks will not be removed") + pass # no accelerate installed → no offload hooks exist to remove quant_config = get_quant_config(model, is_modelopt_qlora=is_modelopt_qlora) - # Add MTP layer prefixes to exclude_modules if they were excluded from quantization - # This ensures they appear in quantization_config["ignore"] in config.json - mtp_layer_prefixes = getattr(model, "_mtp_layer_prefixes", None) - if mtp_layer_prefixes: - exclude_modules = quant_config["quantization"].setdefault("exclude_modules", []) - for prefix in mtp_layer_prefixes: - # Add wildcard pattern to exclude all submodules under this MTP layer - pattern = f"{prefix}*" - if pattern not in exclude_modules: - exclude_modules.append(pattern) - print(f"Adding MTP layer to quantization_config ignore: {pattern}") + _add_mtp_exclusions(model, quant_config) - # Safety net: sync any gate/up weight quantizer amaxes that - # requantize_resmooth_fused_llm_layers did not reach (e.g. experts not - # activated during the dummy forward, or non-standard expert naming). - synced = sync_moe_gate_up_amax(model) - if synced: - warnings.warn( - f"Found {synced} MoE expert gate/up projection pair(s) with mismatched " - f"weight_scale_2 after requantize_resmooth_fused_llm_layers. " - f"This typically means the dummy forward did not activate these experts. " - f"Taking element-wise max of amaxes for serving-engine fusion." - ) + _warn_on_unsynced_moe_gate_up(model) - # Merge per-side input_quantizer amaxes BEFORE _process_quantized_modules, - # so the merged value flows into input_scale derivation downstream. - synced_input = sync_tied_input_amax(model) + # Merge per-side input_quantizer amaxes BEFORE export, so the retained tied weight's single + # input_scale covers every side's activation range (else the dropped side clips at inference). + synced_input = sync_tied_input_amax(model, tied_map) if synced_input: print( f"sync_tied_input_amax: max-merged input_quantizer amaxes across " @@ -978,11 +993,9 @@ def _export_transformers_checkpoint( ) # Process all quantized modules and export weights - _process_quantized_modules(model, dtype, is_modelopt_qlora) - - # Reconstruct fused MoELinear: per-expert _QuantLinear weights → original 3D format from modelopt.torch.quantization.plugins.huggingface import _reconstruct_fused_moe_linear + _process_quantized_modules(model, dtype, is_modelopt_qlora) _reconstruct_fused_moe_linear(model) if is_fsdp2_model(model): @@ -999,14 +1012,12 @@ def _export_transformers_checkpoint( kv_cache_max_bound = 448 kv_cache_format = quant_config["quantization"]["kv_cache_quant_algo"] - # Reorder so canonical-side tied keys (per HF's _tied_weights_keys) - # iterate first into postprocess_state_dict's first-wins data_ptr dedup. - # Self-gated to DiffusionGemma inside _reorder_canonical_first; no-op - # for every other model. - quantized_state_dict = _reorder_canonical_first(quantized_state_dict, model) - quantized_state_dict = postprocess_state_dict( - quantized_state_dict, kv_cache_max_bound, kv_cache_format, is_modelopt_qlora + quantized_state_dict, + kv_cache_max_bound, + kv_cache_format, + is_modelopt_qlora, + tied_map=tied_map, ) return quantized_state_dict, quant_config @@ -1446,6 +1457,37 @@ def export_speculative_decoding( exporter.export(export_dir, dtype) +def _write_hf_export_config( + model: nn.Module, + hf_quant_config: dict | None, + export_dir: Path, +) -> None: + """Write hf_quant_config.json (if quantized) and embed quantization_config into config.json.""" + quantization_details = (hf_quant_config or {}).get("quantization", {}) + is_quantized_export = ( + quantization_details.get("quant_algo") is not None + or quantization_details.get("kv_cache_quant_algo") is not None + ) + quantization_config = None + if hf_quant_config is not None and is_quantized_export: + with open(f"{export_dir}/hf_quant_config.json", "w") as file: + json.dump(hf_quant_config, file, indent=4) + quantization_config = convert_hf_quant_config_format(hf_quant_config) + + original_config = f"{export_dir}/config.json" + with open(original_config) as file: + config_data = json.load(file) + sanitize_hf_config_for_deployment(config_data, model) + if quantization_config is not None: + config_data["quantization_config"] = quantization_config + if export_sparse_attention_config is not None: + sparse_attn_config = export_sparse_attention_config(model) + if sparse_attn_config is not None: + config_data["sparse_attention_config"] = sparse_attn_config + with open(original_config, "w") as file: + json.dump(config_data, file, indent=4) + + def export_hf_checkpoint( model: Any, dtype: torch.dtype | None = None, @@ -1503,7 +1545,44 @@ def export_hf_checkpoint( and torch.distributed.is_initialized() and is_fsdp2_model(model) ) + # Offloaded models take the streaming path: it materializes one layer at a time and + # writes each straight to a shard file, so peak memory is one layer plus one shard + # buffer instead of the whole quantized state dict. + _offloaded = has_accelerate_offload(model) + try: + if _offloaded: + # Imported here rather than at module scope: the streaming exporter imports the + # shared prep helpers from this module, so a top-level import would be circular. + from .unified_export_hf_streaming import _export_transformers_checkpoint_streaming + + if save_modelopt_state: + warnings.warn( + "save_modelopt_state=True is not supported in the streaming offload export " + "path and will be ignored." + ) + _, hf_quant_config = _export_transformers_checkpoint_streaming( + model, + dtype, + export_dir=export_dir, + max_shard_size=max_shard_size, + extra_state_dict=extra_state_dict, + **kwargs, + ) + if getattr(model, "hf_quantizer", None) is not None: + model.hf_quantizer = None + try: + name_mapper = build_reverse_name_mapper(model) + if name_mapper is not None and hf_quant_config: + revert_quant_config_names(hf_quant_config.get("quantization", {}), name_mapper) + except Exception as exc: + warnings.warn( + f"Quant-aware reverse weight conversion skipped ({exc}); exported tensor " + "names may not match the original HF hub checkpoint." + ) + _write_hf_export_config(model, hf_quant_config, export_dir) + return + post_state_dict, hf_quant_config = _export_transformers_checkpoint(model, dtype, **kwargs) # Remove hf_quantizer from model so post_state_dict can be exported. @@ -1538,26 +1617,6 @@ def export_hf_checkpoint( if is_distributed and torch.distributed.get_rank() != 0: return - # Only treat the export as quantized when at least one quant_algo field is set. - # get_quant_config always returns a dict (even for sparsity-only or unmodified models), - # so emitting hf_quant_config.json unconditionally produces a file with - # "quant_algo": null that downstream loaders (e.g. TensorRT-LLM) reject as a - # malformed pre-quantized checkpoint. - quantization_details = (hf_quant_config or {}).get("quantization", {}) - is_quantized_export = ( - quantization_details.get("quant_algo") is not None - or quantization_details.get("kv_cache_quant_algo") is not None - ) - - if is_quantized_export: - # Save hf_quant_config.json for backward compatibility - with open(f"{export_dir}/hf_quant_config.json", "w") as file: - json.dump(hf_quant_config, file, indent=4) - - hf_quant_config = convert_hf_quant_config_format(hf_quant_config) - else: - hf_quant_config = None - # Keep transformers' own revert_weight_conversion disabled (the quant-aware reverse # above replaces it): it can't handle quantized state dicts (RuntimeError on 0-d scalar # scale tensors). Patch both the source and importing module since modeling_utils does @@ -1577,25 +1636,7 @@ def export_hf_checkpoint( finally: _unpatch_revert_weight_conversion(_patches) - original_config = f"{export_dir}/config.json" - config_data = {} - - with open(original_config) as file: - config_data = json.load(file) - - sanitize_hf_config_for_deployment(config_data, model) - - if hf_quant_config is not None: - config_data["quantization_config"] = hf_quant_config - - # Add sparse attention config if available - if export_sparse_attention_config is not None: - sparse_attn_config = export_sparse_attention_config(model) - if sparse_attn_config is not None: - config_data["sparse_attention_config"] = sparse_attn_config - - with open(original_config, "w") as file: - json.dump(config_data, file, indent=4) + _write_hf_export_config(model, hf_quant_config, export_dir) except Exception as e: warnings.warn( diff --git a/modelopt/torch/export/unified_export_hf_streaming.py b/modelopt/torch/export/unified_export_hf_streaming.py new file mode 100644 index 00000000000..50b5d797c6a --- /dev/null +++ b/modelopt/torch/export/unified_export_hf_streaming.py @@ -0,0 +1,449 @@ +# 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. + +"""Streaming HF checkpoint export for disk/CPU-offloaded models. + +Kept apart from :mod:`unified_export_hf` so the resident exporter cannot drift back into +being offload-aware: the only edge between them is the dispatch in +``export_hf_checkpoint``, which imports :func:`_export_transformers_checkpoint_streaming` +lazily to keep the dependency acyclic. +""" + +import contextlib +import itertools +import json +import shutil +import warnings +from pathlib import Path +from typing import Any + +import torch +import torch.nn as nn +from safetensors.torch import save_file + +from .quant_aware_conversion import build_reverse_name_mapper +from .quant_utils import _postprocess_single_tensor, get_quant_config +from .registry import ExportContext +from .unified_export_hf import ( + _add_mtp_exclusions, + _dispatch_export_handler, + _patch_revert_weight_conversion, + _prepare_moe_inputs, + _resolve_export_dtype, + _sanitize_generation_config_for_save, + _unpatch_revert_weight_conversion, + _warn_on_unsynced_moe_gate_up, + requantize_resmooth_fused_llm_layers, +) + +__all__ = ["_export_transformers_checkpoint_streaming"] + + +class _StreamingShardWriter: + """Write tensors to safetensors shard files without accumulating the full state dict. + + Buffers tensors up to ``max_shard_size`` bytes, flushes to a numbered temp file, then + at :meth:`finalize` renames temp files to canonical shard names once the total shard + count is known. + + Peak memory = 1 layer (being materialized) + 1 shard buffer, not the full checkpoint. + """ + + def __init__(self, export_dir: Path | str, max_shard_size: int) -> None: + self._export_dir = Path(export_dir) + self._max_shard_size = max_shard_size + self._buffer: dict[str, torch.Tensor] = {} + self._buffer_bytes: int = 0 + self._part_files: list[Path] = [] + self._total_bytes: int = 0 + # Maps tensor key → part-file index (recorded at flush time) + self._key_to_part: dict[str, int] = {} + # data_ptr of every buffered tensor, so aliases never reach save_file. + self._buffer_storage: set[int] = set() + + def _flush(self) -> None: + if not self._buffer: + return + part_idx = len(self._part_files) + part_path = self._export_dir / f"__shard_part_{part_idx:05d}.safetensors" + save_file(self._buffer, str(part_path)) + for key in self._buffer: + self._key_to_part[key] = part_idx + self._part_files.append(part_path) + self._total_bytes += self._buffer_bytes + self._buffer = {} + self._buffer_storage = set() + self._buffer_bytes = 0 + + def add(self, key: str, tensor: torch.Tensor) -> None: + """Buffer a tensor, flushing the current shard to disk when it is full. + + ``save_file`` rejects tensors sharing storage, which two keys can still do here + when the tensor reaches us already on CPU (so ``_stream_tensor``'s ``.cpu()`` was + a no-op rather than a copy). Copy on collision rather than dropping one of them: + offloaded export writes tied weights as separate entries, so every key must + survive. ``data_ptr()`` only has to hold within one buffer, whose entries stay + alive until :meth:`_flush`. + """ + if tensor.data_ptr() in self._buffer_storage: + tensor = tensor.clone() + + self._buffer_storage.add(tensor.data_ptr()) + self._buffer[key] = tensor + self._buffer_bytes += tensor.nbytes + if self._buffer_bytes >= self._max_shard_size: + self._flush() + + def finalize(self) -> dict[str, str]: + """Flush remaining buffer, rename part files, write model.safetensors.index.json. + + Returns the weight_map ``{key: shard_filename}`` written to the index. + Single-shard exports use ``model.safetensors`` without an index file. + """ + self._flush() + n_shards = len(self._part_files) + if n_shards == 0: + return {} + + if n_shards == 1: + final_name = "model.safetensors" + self._part_files[0].rename(self._export_dir / final_name) + return dict.fromkeys(self._key_to_part, final_name) + + for i, part_path in enumerate(self._part_files): + part_path.rename(self._export_dir / f"model-{i + 1:05d}-of-{n_shards:05d}.safetensors") + + weight_map = { + key: f"model-{part_idx + 1:05d}-of-{n_shards:05d}.safetensors" + for key, part_idx in self._key_to_part.items() + } + total_size = self._total_bytes + index_path = self._export_dir / "model.safetensors.index.json" + with open(index_path, "w") as f: + json.dump({"metadata": {"total_size": total_size}, "weight_map": weight_map}, f) + return weight_map + + +def _parse_shard_size(size: int | str) -> int: + """Convert a shard-size string (e.g. ``"10GB"``, ``"500MB"``) to bytes. + + Mirrors transformers' ``convert_file_size_to_int``, which reads ``GB``/``MB``/``KB`` + as decimal and only ``GiB``/``MiB``/``KiB`` as binary. That helper was removed from + ``transformers.utils`` in transformers 5.x, so the fallback below is the live path + there, not a rarely-taken branch. + """ + try: + from transformers.utils import convert_file_size_to_int + + return convert_file_size_to_int(size) + except ImportError: + pass + if isinstance(size, int): + return size + s = size.strip().upper() + for suffix, multiplier in ( + ("GIB", 1024**3), + ("MIB", 1024**2), + ("KIB", 1024), + ("GB", 1000**3), + ("MB", 1000**2), + ("KB", 1000), + ): + if s.endswith(suffix): + return int(float(s[: -len(suffix)]) * multiplier) + return int(s) + + +def _assert_no_split_rules(model: nn.Module) -> None: + """Refuse to stream a model whose conversion mapping needs tensor-level splits.""" + from .quant_aware_conversion import _build_reverse_rules + + try: + split_rules, _, _ = _build_reverse_rules(model) + except Exception: + return # build_reverse_name_mapper reports the failure with a warning + if split_rules: + raise NotImplementedError( + "Disk/CPU-offloaded export cannot reverse tensor-level split rules in this " + "model's transformers conversion mapping: the streaming path reverses names " + "per tensor, while splits need the full state dict. Export without offloading." + ) + + +def _export_transformers_checkpoint_streaming( + model: nn.Module, + dtype: torch.dtype | None = None, + is_modelopt_qlora: bool = False, + export_dir: Path | str = ".", + max_shard_size: int | str = "10GB", + extra_state_dict: dict[str, torch.Tensor] | None = None, + **kwargs, +) -> tuple[None, dict[str, Any]]: + """Export a disk/CPU-offloaded model by streaming tensors layer-by-layer to shard files. + + The offloaded counterpart of :func:`_export_transformers_checkpoint`, which builds the + whole quantized state dict at once and so needs every weight resident. Here each + decoder layer is materialized, exported, and written to a shard file before the next + one is touched, bounding peak memory at one layer plus one shard buffer. + + Model-level preparation (MoE input handling, resmooth/requantize, quant config) matches + the resident path. The per-tensor work does not: instead of ``postprocess_state_dict`` + over a finished dict, each tensor goes through :func:`_postprocess_single_tensor` as it + is produced. Two consequences follow from having no whole-dict view: + + - Tied weights are dropped by *name* from ``_tied_weights_keys`` (data_ptr is meaningless + once weights move host<->device); see the TODO below on adopting ``all_tied_weights_keys``. + - Conversion mappings that need tensor-level splits cannot be reversed one tensor at a + time, so they are rejected up front rather than exported incorrectly. + + Args: + model: the full torch model to export, carrying accelerate offload hooks. + dtype: weight dtype for unquantized layers, or the model's dtype if None. + is_modelopt_qlora: whether the model is a ModelOpt QLoRA model. + export_dir: directory to write shards and config artifacts into. + max_shard_size: shard size limit, as bytes or a string such as ``"10GB"``. + extra_state_dict: tensors the model itself never holds (e.g. MTP weights, which HF + leaves orphaned) and which would otherwise be missing from the export. + + Returns: + ``(None, quant_config)``. No state dict is returned because none is ever + assembled; shards, ``config.json``, and ``generation_config.json`` are written to + ``export_dir`` directly. The caller writes ``hf_quant_config.json`` and merges + ``quantization_config`` into ``config.json``. + + Raises: + NotImplementedError: if the model's conversion mapping contains split rules. + RuntimeError: if decoder layers cannot be discovered for layer-wise materialization. + """ + from modelopt.torch.quantization.plugins.huggingface import _reconstruct_fused_moe_linear + from modelopt.torch.quantization.utils.core_utils import ( + enable_weight_access_and_writeback, + requires_weight_materialization, + ) + from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector + + export_dir = Path(export_dir) + # Materialization dispatch walks the module tree from the root; without this cache each + # call re-derives it, which is O(N^2) over a MoE model's expert modules. + name_to_module = dict(model.named_modules()) + + # --- Same model-level setup as _export_transformers_checkpoint --- + dtype = _resolve_export_dtype(model, dtype) + _prepare_moe_inputs(model, dtype, is_modelopt_qlora) + + requantize_resmooth_fused_llm_layers(model) + + quant_config = get_quant_config(model, is_modelopt_qlora=is_modelopt_qlora) + + _add_mtp_exclusions(model, quant_config) + + _warn_on_unsynced_moe_gate_up(model) + + # --- Per-tensor constants --- + kv_cache_max_bound = 448 + kv_cache_format = quant_config["quantization"]["kv_cache_quant_algo"] + + # --- Tied alias keys to skip --- + # data_ptr() is unreliable for disk-offloaded weights, so we use _tied_weights_keys. + # Only apply when tie_word_embeddings=True: _tied_weights_keys can list keys whose + # weights are not actually shared (e.g. if the model was saved with tie_word_embeddings=False + # but the attribute was never cleared), which would incorrectly drop lm_head.weight. + # + # TODO(tied-map): the resident path reads HF's ``all_tied_weights_keys`` (covers dict-style/MoE + # ties); this path could too, to close the streaming gap for offloaded 5.x models -- but that + # swap needs offload-specific validation (meta tensors, per-tensor order, disk round-trip) first. + raw_tied_keys: set[str] = ( + set(getattr(model, "_tied_weights_keys", None) or []) + if getattr(model.config, "tie_word_embeddings", False) + else set() + ) + + # --- Name mapper for per-tensor key reversal --- + # Tensor names are applied inline; quant config names are handled by the caller. + # Renames are all a per-tensor pass can reverse. The batch path additionally runs + # revert_weight_conversion_quant_aware() for split rules, which need the whole state + # dict to regroup tensors, so refuse rather than emit fused tensors under unfused + # hub keys. + _assert_no_split_rules(model) + name_mapper = None + try: + name_mapper = build_reverse_name_mapper(model) + except Exception as exc: + warnings.warn( + f"Reverse name mapper unavailable ({exc}); exported tensor names may not match " + "the original HF hub checkpoint." + ) + + tied_alias_keys: set[str] = ( + {name_mapper(k) for k in raw_tied_keys} if name_mapper is not None else raw_tied_keys + ) + + # --- Decoder layers --- + decoder_layers = LayerActivationCollector.get_decoder_layers(model) + if decoder_layers is None: + raise RuntimeError( + "Streaming export requires discoverable decoder layers. " + "The model architecture is not supported by LayerActivationCollector." + ) + decoder_layer_ids = {id(m) for m in decoder_layers} + # Descendants too, not just the layers: an offloaded layer's children return to meta + # when its window closes, so a child-level check would re-enter and re-export weights + # this pass already packed. + decoder_owned_ids = {id(m) for layer in decoder_layers for m in layer.modules()} + + # --- Persistent-buffer predicate (mirrors state_dict() which excludes non-persistent) --- + def _is_persistent_buffer(name: str) -> bool: + parts = name.split(".") + mod: nn.Module = model + for part in parts[:-1]: + mod = getattr(mod, part, mod) + return parts[-1] not in getattr(mod, "_non_persistent_buffers_set", frozenset()) + + # --- Stream tensors to shard files --- + shard_size_bytes = _parse_shard_size(max_shard_size) + writer = _StreamingShardWriter(export_dir, shard_size_bytes) + ctx = ExportContext(model=model, dtype=dtype, is_modelopt_qlora=is_modelopt_qlora) + seen_keys: set[str] = set() + + def _stream_tensor(full_key: str, tensor: torch.Tensor) -> None: + new_key, new_value = _postprocess_single_tensor( + full_key, tensor, kv_cache_max_bound, kv_cache_format, is_modelopt_qlora + ) + if new_key is None or new_value is None: + return + if name_mapper is not None: + new_key = name_mapper(new_key) + if new_key in tied_alias_keys: + return + writer.add(new_key, new_value.detach().contiguous().cpu()) + + # Decoder layers: materialize one at a time + for layer_name, layer_module in model.named_modules(): + if id(layer_module) not in decoder_layer_ids: + continue + with enable_weight_access_and_writeback( + layer_module, model, name_to_module, writeback=False + ): + for sub_name, sub_mod in layer_module.named_modules(): + full_name = f"{layer_name}.{sub_name}" if sub_name else layer_name + _dispatch_export_handler(full_name, sub_mod, ctx) + _reconstruct_fused_moe_linear(layer_module) + prefix = f"{layer_name}." if layer_name else "" + for key, tensor in layer_module.state_dict().items(): + full_key = prefix + key + if full_key in seen_keys: + continue + seen_keys.add(full_key) + _stream_tensor(full_key, tensor) + # Release GPU tensors added by export handlers before hook.post_forward + # runs, to prevent cross-layer accumulation on disk-offloaded models. + # + # Two categories accumulate without explicit cleanup: + # + # 1. CUDA *buffers* on any sub-module (weight_scale, weight_scale_2, + # input_scale): AlignDevicesHook.post_forward uses offload_buffers=False + # by default, so it never offloads buffers. Pre-existing buffers in + # disk-offloaded layers live on CPU, so any CUDA buffer encountered here + # was registered by the export handlers and is safe to drop. + # + # 2. CUDA *parameters* on modules WITHOUT _hf_hook: _export_fused_experts + # creates fresh nn.Module objects (one per expert x projection) and adds + # them to the layer via add_module() *after* weight_access_and_writeback + # captured its materialized list. hook.post_forward never visits these + # new modules, so their packed NVFP4 weight parameters (~5 GB per MoE + # layer) stay live on GPU. Modules WITH _hf_hook are original model + # modules whose parameters hook.post_forward will meta-ify; leave those + # alone. + for sub_mod in layer_module.modules(): + for buf_name in list(sub_mod._buffers): + buf = sub_mod._buffers[buf_name] + if buf is not None and buf.device.type == "cuda": + sub_mod._buffers[buf_name] = None + if not hasattr(sub_mod, "_hf_hook"): + for param_name, param in list(sub_mod._parameters.items()): + if param is not None and param.device.type == "cuda": + sub_mod._parameters[param_name] = None + torch.cuda.empty_cache() + + # Non-decoder modules whose weights are not directly readable (embed_tokens, norm, + # lm_head, ...). Containers are skipped: their children get their own window. + for name, module in model.named_modules(): + if id(module) in decoder_owned_ids: + continue + if not requires_weight_materialization(module, model, name_to_module): + continue + with enable_weight_access_and_writeback(module, model, name_to_module, writeback=False): + for sub_name, sub_mod in module.named_modules(): + full_name = f"{name}.{sub_name}" if sub_name else name + _dispatch_export_handler(full_name, sub_mod, ctx) + prefix = f"{name}." if name else "" + for key, tensor in module.state_dict().items(): + full_key = prefix + key + if full_key in seen_keys or tensor.is_meta: + continue + seen_keys.add(full_key) + _stream_tensor(full_key, tensor) + + # GPU-resident parameters and persistent buffers (not covered by the above loops). + # named_buffers() includes non-persistent buffers that state_dict() excludes; filter them. + for name, tensor in itertools.chain( + model.named_parameters(), + ((n, b) for n, b in model.named_buffers() if _is_persistent_buffer(n)), + ): + if name in seen_keys or tensor is None or tensor.is_meta: + continue + seen_keys.add(name) + _stream_tensor(name, tensor) + + # Tensors the model never held — e.g. MTP weights, which HF leaves orphaned because it + # builds only num_hidden_layers decoders. They are already materialized and skip the + # per-tensor postprocessing, matching how the batch path merges them after + # postprocess_state_dict; only the hub-name reversal applies. + for name, tensor in (extra_state_dict or {}).items(): + if name in seen_keys: + continue + seen_keys.add(name) + writer.add( + name_mapper(name) if name_mapper is not None else name, + tensor.detach().contiguous().cpu(), + ) + + writer.finalize() + + # Write non-weight artifacts: config.json, generation_config.json, and the custom + # modeling *.py files that trust_remote_code models (e.g. NemotronH) need. + # We avoid model.save_pretrained(state_dict={}) here because MoE models (e.g. DSR1) + # have expert weights that share underlying storage across layers; safetensors' shared- + # tensor check fires even when saving an empty state dict, crashing the export after + # all shards are already written correctly. + _sanitize_generation_config_for_save(model) + _patches = _patch_revert_weight_conversion() + try: + model.config.save_pretrained(str(export_dir)) + finally: + _unpatch_revert_weight_conversion(_patches) + if hasattr(model, "generation_config") and model.generation_config is not None: + with contextlib.suppress(Exception): + model.generation_config.save_pretrained(str(export_dir)) + + # Copy custom modeling *.py files for trust_remote_code checkpoints. + _src_dir = Path(getattr(model.config, "_name_or_path", "") or "") + if _src_dir.is_dir(): + for _py in _src_dir.glob("*.py"): + _dst = export_dir / _py.name + if not _dst.exists(): + shutil.copy2(_py, _dst) + + return None, quant_config diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 153bc42ce3f..dde24513086 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -1772,10 +1772,14 @@ def get_nemotron_h_decoder_layers(model: nn.Module) -> nn.ModuleList | None: if not _is_supported_hf_model(model): return None - if hasattr(model, "backbone") and hasattr(model.backbone, "layers"): - layers = model.backbone.layers - if len(layers) > 0 and hasattr(layers[0], "block_type"): - return layers + # Custom remote-code checkpoint uses model.backbone.layers; + # native transformers NemotronHModel uses model.model.layers. + for container_attr in ("backbone", "model"): + container = getattr(model, container_attr, None) + if container is not None and hasattr(container, "layers"): + layers = container.layers + if layers and hasattr(layers[0], "block_type"): + return layers return None diff --git a/modelopt/torch/quantization/utils/core_utils.py b/modelopt/torch/quantization/utils/core_utils.py index 1bdf23da64a..cff04351f9b 100644 --- a/modelopt/torch/quantization/utils/core_utils.py +++ b/modelopt/torch/quantization/utils/core_utils.py @@ -16,6 +16,7 @@ """Quantization utilities.""" import copy +import itertools from collections import namedtuple from contextlib import ExitStack, contextmanager, nullcontext from typing import TYPE_CHECKING, Any @@ -25,7 +26,7 @@ import torch.nn.functional as F from torch.distributed.fsdp import FSDPModule, MixedPrecisionPolicy, fully_shard from torch.distributed.fsdp._fully_shard._fsdp_param import FSDPParam -from torch.distributed.tensor import Replicate +from torch.distributed.tensor import DTensor, Replicate from modelopt.torch.quantization.config import QuantizerCfgEntry from modelopt.torch.utils import get_unwrapped_name, print_rank_0 @@ -619,6 +620,45 @@ def enable_weight_access_and_writeback( yield +def requires_weight_materialization(module, root_model, name_to_module: dict | None = None) -> bool: + """Whether ``module``'s own weights are currently unreadable and need a window. + + Mirrors the dispatch in :func:`enable_weight_access_and_writeback`, so callers + deciding *whether* to open a window agree with what opening one would do. Two things + must hold: the module owns tensors that are not directly readable right now + (offloaded to meta, or a sharded ``DTensor``), and a context exists that can + materialize them. Modules already materialized are excluded -- re-entering a window + would re-run export handlers over already-packed weights. + """ + if not any( + t is not None and (t.is_meta or isinstance(t, DTensor)) + for t in itertools.chain(module._parameters.values(), module._buffers.values()) + ): + return False + if _get_enclosing_fsdp_module(module, root_model, name_to_module) is not None: + return True + if is_quantized_parallel_linear(module) and hasattr(module, "_hf_tp_plan"): + return True + hook = getattr(module, "_hf_hook", None) + if hook is None: + return False + from ..plugins.accelerate import _get_offload_hook + + return _get_offload_hook(hook) is not None + + +def has_accelerate_offload(module: nn.Module) -> bool: + """Return True if any module in ``module`` has a CPU- or disk-offload accelerate hook.""" + try: + from ..plugins.accelerate import _get_offload_hook + except ImportError: + return False + + return any( + _get_offload_hook(getattr(m, "_hf_hook", None)) is not None for m in module.modules() + ) + + @contextmanager def persistent_materialization(layer, writeback: bool = True): """Keep all layer weights materialized on GPU for the duration. @@ -770,22 +810,30 @@ def _disable_fsdp_unshard_reshard(layer): yield -def get_prefixed_param_names(parent_model, target_module): +def build_param_index(model): + """Map ``id(param)`` to its ``(position, name)`` in ``model.named_parameters()``. + + Lets callers resolve many modules against one walk of the parameters instead of one walk + each; the position keeps "first in ``named_parameters()`` order" resolvable. + """ + return {id(param): (i, name) for i, (name, param) in enumerate(model.named_parameters())} + + +def get_prefixed_param_names(parent_model, target_module, param_index=None): """Get parameter names for a target module prefixed with the parent model name. This function is used to get full parameter name from FSDPParam module_info which stores the unprefixed parameter name. + Pass ``param_index`` (see :func:`build_param_index`) when resolving many target modules + against the same parent, so the parent's parameters are walked once rather than per module. """ + if param_index is None: + param_index = build_param_index(parent_model) target_ids = {id(p) for p in target_module.parameters()} - return next( - ( - name.rsplit(".", 1)[0] - for name, param in parent_model.named_parameters() - if id(param) in target_ids - ), - None, # default value if no match - ) + # Lowest position == first in named_parameters() order, matching a linear scan's result. + match = min((param_index[pid] for pid in target_ids if pid in param_index), default=None) + return match[1].rsplit(".", 1)[0] if match is not None else None def create_fsdp_param_mapping(fsdp_param_list, model): @@ -798,10 +846,14 @@ def create_fsdp_param_mapping(fsdp_param_list, model): Returns: dict: Full parameter name → FSDP parameter. """ + # Built once per call, not once per FSDPParam: export resolves every quantized module, so the + # per-param walk made this quadratic in (params x modules) and stalled MoE exports for hours. + # It cannot be cached across calls -- callers swap in quantized params between them. + param_index = build_param_index(model) mapping = {} for param in fsdp_param_list: # Get the module name - module_name = get_prefixed_param_names(model, param._module_info.module) + module_name = get_prefixed_param_names(model, param._module_info.module, param_index) if module_name is not None: # Get the parameter name from _module_info and construct full param name param_name = param._module_info.param_name diff --git a/modelopt/torch/speculative/utils.py b/modelopt/torch/speculative/utils.py index 8c5418bb6b8..50a08482876 100644 --- a/modelopt/torch/speculative/utils.py +++ b/modelopt/torch/speculative/utils.py @@ -22,6 +22,7 @@ import sys import warnings from collections import Counter, defaultdict, deque +from collections.abc import Mapping import torch import torch.distributed @@ -40,13 +41,27 @@ ) +def get_conversation_input_ids(tokenizer, conversations): + """Return the chat-template token ids as a flat list[int].""" + input_ids = tokenizer.apply_chat_template( + conversations, add_generation_prompt=False, tokenize=True, return_tensors=None + ) + if isinstance(input_ids, Mapping): + input_ids = input_ids["input_ids"] # transformers>=5 returns a BatchEncoding + if isinstance(input_ids, torch.Tensor): + input_ids = input_ids.flatten().tolist() + if input_ids and isinstance(input_ids[0], list | tuple): + assert len(input_ids) == 1, f"expected a single conversation, got {len(input_ids)}" + input_ids = list(input_ids[0]) # unwrap a single-conversation batch dim + assert not input_ids or isinstance(input_ids[0], int), ( + f"expected flat token ids, got {input_ids[:1]}" + ) + return input_ids + + def calibrate_frequent_vocab(tokenizer, text, target_vocab_size, output_file=None): """Given a calibration text, find the most common vocabs and return the mapping.""" - conversations = tokenizer.apply_chat_template(text) - # Transformers5.x returns a BatchEncoding from apply_chat_template - if hasattr(conversations, "input_ids"): - conversations = conversations.input_ids - counter = Counter(conversations) + counter = Counter(get_conversation_input_ids(tokenizer, text)) vocab = counter.most_common(target_vocab_size) mapping = torch.zeros(target_vocab_size, dtype=torch.int64) assert len(vocab) == target_vocab_size, ( diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 245ca81de79..c3cea6acaf4 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -252,6 +252,68 @@ def is_fsdp2_model(model) -> bool: return any(isinstance(m, FSDPModule) for m in model.modules()) +def _off_dtype_params(model) -> set[torch.nn.Parameter]: + """Params whose dtype differs from the model's dominant (by element count) param dtype. + + FSDP2 needs one dtype per shard group, but HF models routinely keep a few params in fp32 for + stability (e.g. MoE router gates). Pass these to ``fully_shard(ignored_params=...)``. + + TODO: Drop this and shard the off-dtype params once a stable PyTorch release includes FSDP2 + mixed-precision parameter dtype support (already on nightly). + """ + numel_by_dtype: dict[torch.dtype, int] = {} + for param in model.parameters(): + numel_by_dtype[param.dtype] = numel_by_dtype.get(param.dtype, 0) + param.numel() + if len(numel_by_dtype) <= 1: + return set() + + # Lazy import: logging imports this module at top level (circular). + from modelopt.torch.utils.logging import warn_rank_0 + + dominant = max(numel_by_dtype, key=lambda d: numel_by_dtype[d]) + off_dtype = {n: p for n, p in model.named_parameters() if p.dtype != dominant} + off_numel = sum(numel_by_dtype[d] for d in numel_by_dtype if d != dominant) + names = sorted(off_dtype) + warn_rank_0( + f"Model has mixed parameter dtypes {set(numel_by_dtype)}; FSDP2 needs one dtype per shard " + f"group, so {len(names)} non-{dominant} parameter(s) " + f"({100 * off_numel / sum(numel_by_dtype.values()):.2f}% of elements) will stay replicated " + f"rather than sharded: {names[:3]}{' ...' if len(names) > 3 else ''}" + ) + return set(off_dtype.values()) + + +def _move_to_fsdp_device(model, params: set[torch.nn.Parameter]) -> None: + """Move ``params`` onto the device FSDP2 computes on for ``model``. + + ``fully_shard`` only moves the params it manages, so ignored ones would be stranded on + whatever device the caller built the model on. Meta params are left alone for deferred init. + + The device comes from a sharded param's mesh rather than its local shard: under + ``cpu_offload`` the shard rests on CPU while compute still happens on the accelerator. + """ + # Lazy import: logging imports this module at top level (circular). + from modelopt.torch.utils.logging import warn_rank_0 + + mesh = next((p.device_mesh for p in model.parameters() if isinstance(p, DTensor)), None) + if mesh is None: + warn_rank_0( + f"FSDP2 sharded no parameter of {type(model).__name__}, so the compute device for " + f"{len(params)} unsharded off-dtype parameter(s) cannot be determined; leaving them " + "where they are. Move them to the compute device or the forward will fail." + ) + return + + device = ( + torch.device("cpu") + if mesh.device_type == "cpu" + else torch.device(mesh.device_type, getattr(torch, mesh.device_type).current_device()) + ) + for param in params: + if not param.is_meta and param.device != device: + param.data = param.data.to(device) + + def fsdp2_wrap(model, shard_root=True, mp_policy=None, cpu_offload: bool = False): """Auto-detect a HF causal-LM's decoder layers and FSDP2 ``fully_shard`` each one. @@ -259,6 +321,11 @@ def fsdp2_wrap(model, shard_root=True, mp_policy=None, cpu_offload: bool = False sharded instead of replicated per rank; pass ``shard_root=False`` to leave the root replicated (only decoder layers sharded). Returns the detected decoder layers so callers can reuse the detection result. + + Parameters whose dtype differs from the model's dominant one are excluded from the wrap (see + :func:`_off_dtype_params`), since FSDP2 rejects a shard group that mixes dtypes. They stay + replicated and are moved onto the shards' device, which ``fully_shard`` does not do for the + params it ignores. """ # Lazy import: layerwise_calib imports this module at top level (circular). from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector @@ -274,6 +341,9 @@ def fsdp2_wrap(model, shard_root=True, mp_policy=None, cpu_offload: bool = False fsdp_kwargs["mp_policy"] = mp_policy if cpu_offload: fsdp_kwargs["offload_policy"] = CPUOffloadPolicy() + ignored_params = _off_dtype_params(model) + if ignored_params: + fsdp_kwargs["ignored_params"] = ignored_params # Snapshot/restore config.architectures: some HF builders mutate it during fully_shard. config = getattr(model, "config", None) @@ -282,6 +352,8 @@ def fsdp2_wrap(model, shard_root=True, mp_policy=None, cpu_offload: bool = False fully_shard(layer, **fsdp_kwargs) if shard_root: fully_shard(model, **fsdp_kwargs) + if ignored_params: + _move_to_fsdp_device(model, ignored_params) if config is not None and architectures: config.architectures = architectures diff --git a/modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_offload.yaml b/modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_offload.yaml new file mode 100644 index 00000000000..aa525cb4188 --- /dev/null +++ b/modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_offload.yaml @@ -0,0 +1,49 @@ +# 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. + +imports: + base_disable_all: configs/ptq/units/base_disable_all + default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + nvfp4: configs/numerics/nvfp4 + kv_fp8: configs/ptq/units/kv_fp8 + +metadata: + recipe_type: ptq + description: > + NVFP4 static weight and dynamic activation for expert layers only (W4A4), FP8 KV cache, + max layerwise calibration with calib_mutates_weights=False for disk-offloaded single-GPU + PTQ. Weights stay as meta tensors between layers; export_hf_checkpoint materializes them. +quantize: + algorithm: + method: max + layerwise: + enable: true + calib_mutates_weights: false + quant_cfg: + - $import: base_disable_all + - quantizer_name: '*.experts.*weight_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*.experts.*input_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*block_sparse_moe*weight_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*block_sparse_moe*input_quantizer' + cfg: + $import: nvfp4 + - $import: kv_fp8 + - $import: default_disabled_quantizers diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index 847117acc0c..cb3d8d70257 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -29,7 +29,7 @@ supported combinations. ### The shipped recipes
-All 20 general/ptq/ recipes (click to expand) +All 21 general/ptq/ recipes (click to expand) | Recipe | Model body | KV cache | Calibration | |--------|-----------|----------|-------------| @@ -46,6 +46,7 @@ supported combinations. | `nvfp4_experts_only-kv_fp8` | NVFP4 W4A4, MoE experts only | FP8 (calibrated) | max | | `nvfp4_experts_only-kv_fp8_cast` | NVFP4 W4A4, MoE experts only | FP8 (constant amax) | max | | `nvfp4_experts_only-kv_fp8_layerwise` | NVFP4 W4A4, MoE experts only | FP8 (calibrated) | max, layerwise | +| `nvfp4_experts_only-kv_fp8_layerwise_offload` | NVFP4 W4A4, MoE experts only | FP8 (calibrated) | max, layerwise (non-mutating, for disk offload) | | `nvfp4_experts_only_mse-kv_fp8_cast` | NVFP4 W4A4, MoE experts only | FP8 (constant amax) | MSE + FP8 sweep | | `nvfp4_experts_only_input_scale1-kv_fp8_cast` | NVFP4 W4A4, MoE experts only, expert `input_scale` pinned to 1.0 | FP8 (constant amax) | max (weights); expert activations uncalibrated | | `nvfp4_omlp_only-kv_fp8` | NVFP4 W4A4, o_proj + MLP/MoE | FP8 (calibrated) | max | diff --git a/tests/_test_utils/torch/quantization/tied_modules.py b/tests/_test_utils/torch/quantization/tied_modules.py index 32afc213386..7a98cc30b1d 100644 --- a/tests/_test_utils/torch/quantization/tied_modules.py +++ b/tests/_test_utils/torch/quantization/tied_modules.py @@ -94,15 +94,8 @@ def wrap_in_parent_with_tied_keys( when ``decoder_canonical=True``, list-style (legacy, no canonical/alias distinction) when ``decoder_canonical=False``. - Used by tests for :func:`_collect_canonical_tied_patterns` and - :func:`_reorder_canonical_first`. The legacy list-style branch exercises - the "no patterns extracted" negative case. - - The parent's class name contains ``DiffusionGemma`` so the model_type - gate inside :func:`_reorder_canonical_first` (mirrors the existing - whisper / nemotron-vl dispatch in ``unified_export_hf.py``) passes for - test parents — without this, the function early-returns before - reaching the patterns step. + Used by tests for :class:`TiedWeightMap`. The list-style branch exercises the + "no canonical/alias info" negative case (an empty tie map). """ parent_cls = type("DiffusionGemmaTestParent", (nn.Module,), {}) parent = parent_cls() @@ -115,8 +108,12 @@ def wrap_in_parent_with_tied_keys( parent._tied_weights_keys = { rf"^encoder\.{re.escape(weight_attr)}$": f"decoder.{weight_attr}", } + # transformers >=5.0 resolves the declaration into concrete {alias: canonical} names on + # the model (``all_tied_weights_keys``); TiedWeightMap reads that attribute. + parent.all_tied_weights_keys = {f"encoder.{weight_attr}": f"decoder.{weight_attr}"} else: - # Legacy list-style: just a list of tied paths, no canonical info. + # Legacy list-style: just a list of tied paths, no canonical info -> empty resolved map. parent._tied_weights_keys = [f"encoder.{weight_attr}"] + parent.all_tied_weights_keys = {} return parent diff --git a/tests/examples/hf_ptq/test_example_utils.py b/tests/examples/hf_ptq/test_example_utils.py index 57b9676ee92..b904143dcd7 100644 --- a/tests/examples/hf_ptq/test_example_utils.py +++ b/tests/examples/hf_ptq/test_example_utils.py @@ -49,6 +49,95 @@ def _write_safetensors(path, tensors): save_file(tensors, str(path), metadata={"format": "pt"}) +def test_copy_custom_model_files_preserves_non_weight_sidecars(tmp_path): + source_dir = tmp_path / "source" + export_dir = tmp_path / "export" + source_dir.mkdir() + export_dir.mkdir() + + source_files = { + "super_v3_reasoning_parser.py": "class Parser: pass\n", + "modeling_custom.py": "class Model: pass\n", + "README.md": "# Source model\n", + "LICENSE": "license text\n", + "chat_template.jinja": "{{ messages }}\n", + "tokenizer_config.json": '{"chat_template": "source"}\n', + "generation_config.json": '{"source": "generation"}\n', + "config.json": '{"source": "config"}\n', + "hf_quant_config.json": '{"source": "quant"}\n', + "quant_config.json": '{"source": "stale quant"}\n', + "quantize_config.json": '{"source": "stale quant"}\n', + "recipe.yaml": "quantize: {}\n", + "model.safetensors.index.json": '{"weight_map": {}}\n', + "model-00001-of-00001.safetensors": "source weights\n", + "model.gguf": "source weights\n", + } + for file_name, contents in source_files.items(): + (source_dir / file_name).write_text(contents) + + (export_dir / "config.json").write_text('{"export": "config"}\n') + (export_dir / "generation_config.json").write_text('{"export": "generation"}\n') + (export_dir / "hf_quant_config.json").write_text('{"export": "quant"}\n') + (export_dir / "chat_template.jinja").write_text("{{ exported_messages }}\n") + (export_dir / "tokenizer_config.json").write_text('{"chat_template": "export"}\n') + + example_utils.copy_custom_model_files(str(source_dir), str(export_dir), trust_remote_code=False) + + for file_name in [ + "super_v3_reasoning_parser.py", + "modeling_custom.py", + "README.md", + "LICENSE", + "chat_template.jinja", + "generation_config.json", + ]: + assert (export_dir / file_name).read_text() == source_files[file_name] + + assert (export_dir / "config.json").read_text() == '{"export": "config"}\n' + assert (export_dir / "hf_quant_config.json").read_text() == '{"export": "quant"}\n' + assert (export_dir / "tokenizer_config.json").read_text() == '{"chat_template": "export"}\n' + assert not (export_dir / "quant_config.json").exists() + assert not (export_dir / "quantize_config.json").exists() + assert not (export_dir / "recipe.yaml").exists() + assert not (export_dir / "model.safetensors.index.json").exists() + assert not (export_dir / "model-00001-of-00001.safetensors").exists() + assert not (export_dir / "model.gguf").exists() + + (export_dir / "generation_config.json").write_text('{"export": "generation"}\n') + example_utils.copy_custom_model_files( + str(source_dir), + str(export_dir), + exclude_files={"generation_config.json"}, + ) + assert (export_dir / "generation_config.json").read_text() == '{"export": "generation"}\n' + + +def test_resolve_model_path_snapshot_download_stays_allowlisted(monkeypatch, tmp_path): + snapshot_dir = tmp_path / "snapshot" + + def fake_snapshot_download(**kwargs): + assert kwargs == { + "repo_id": "org/model", + "allow_patterns": example_utils._HF_SIDECAR_DOWNLOAD_ALLOW_PATTERNS, + } + return str(snapshot_dir) + + def fake_from_pretrained(*args, **kwargs): + assert (args, kwargs) == (("org/model",), {"trust_remote_code": False}) + return SimpleNamespace(_name_or_path="org/model") + + monkeypatch.setattr( + example_utils.AutoConfig, + "from_pretrained", + fake_from_pretrained, + ) + monkeypatch.setattr(example_utils, "snapshot_download", fake_snapshot_download) + + assert example_utils._resolve_model_path("org/model", trust_remote_code=False) == str( + snapshot_dir + ) + + def test_load_mtp_weights_inlined_orphaned(tmp_path): # GLM-5.1: HF builds only num_hidden decoders → MTP keys orphaned. main_keys = ["model.embed_tokens.weight", "model.layers.0.x.weight"] @@ -419,3 +508,55 @@ def from_pretrained(*args, **kwargs): ) def test_is_diffusion_gemma(hf_config, expected): assert example_utils.is_diffusion_gemma(hf_config) is expected + + +@pytest.mark.parametrize( + ("trust_remote_code", "expect_bundled_code"), + [(True, True), (False, False)], +) +def test_get_model_deepseek_honors_trust_remote_code( + monkeypatch, trust_remote_code, expect_bundled_code +): + """DeepSeek ships bundled modeling code; --trust_remote_code selects it, else built-in.""" + used = {} + hf_config = SimpleNamespace( + architectures=["DeepseekV3ForCausalLM"], + dtype=torch.bfloat16, + model_type="deepseek_v3", + torch_dtype=torch.bfloat16, + ) + + class FakeModel: + def eval(self): + return None + + def _record(tag): + class Fake: + @staticmethod + def from_config(config, **kwargs): + used["path"] = tag + return FakeModel() + + _from_config = from_config + + @staticmethod + def from_pretrained(*args, **kwargs): + used["path"] = tag + return FakeModel() + + return Fake + + monkeypatch.setattr(example_utils.AutoConfig, "from_pretrained", lambda *a, **k: hf_config) + monkeypatch.setattr(example_utils, "AutoModelForCausalLM", _record("bundled")) + monkeypatch.setattr( + example_utils.transformers, "DeepseekV3ForCausalLM", _record("builtin"), raising=False + ) + 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) + + assert used["path"] == ("bundled" if expect_bundled_code else "builtin") diff --git a/tests/gpu/torch/export/test_fsdp2_export.py b/tests/gpu/torch/export/test_fsdp2_export.py index 1897f3d6db3..5f7c186c0ed 100644 --- a/tests/gpu/torch/export/test_fsdp2_export.py +++ b/tests/gpu/torch/export/test_fsdp2_export.py @@ -19,10 +19,15 @@ import torch from _test_utils.torch.export.utils import SmallQKVModel, ToyModel from _test_utils.torch.misc import minimum_sm +from _test_utils.torch.quantization.tied_modules import ( + make_tied_linear_pair, + wrap_in_parent_with_tied_keys, +) from torch.distributed._composable.fsdp import fully_shard import modelopt.torch.quantization as mtq from modelopt.torch.export.layer_utils import is_quantlinear +from modelopt.torch.export.model_utils import TiedWeightMap from modelopt.torch.export.unified_export_hf import ( _export_quantized_weight, requantize_resmooth_fused_llm_layers, @@ -204,6 +209,38 @@ def calib_fn(x): _compare_parameters_and_buffers(model, non_fsdp_model) +def _tied_map_survives_fsdp2_test(rank, size): + """TiedWeightMap (from HF's name-based all_tied_weights_keys) survives fully_shard. + + ``fully_shard`` splits the shared ``nn.Parameter`` into distinct per-module shards, but the HF + map is a plain name dict on the model, unaffected by sharding, so TiedWeightMap still resolves + the tie both before and after -- no pre-shard capture needed. + """ + with patch_fsdp_mp_dtypes(): + enc, dec = make_tied_linear_pair(in_features=32, out_features=32) + model = wrap_in_parent_with_tied_keys(enc, dec, decoder_canonical=True).to("cuda") + + # Pre-shard: the HF map resolves the tie. + assert TiedWeightMap(model).alias_to_canonical == {"encoder.weight": "decoder.weight"} + + fully_shard(model.encoder) + fully_shard(model.decoder) + fully_shard(model) + torch.distributed.barrier() + + # Post-shard: the HF name-based map is untouched -> TiedWeightMap still resolves the tie. + tied_map = TiedWeightMap(model) + assert tied_map.alias_to_canonical == {"encoder.weight": "decoder.weight"} + assert tied_map.group_key("encoder.weight") == "decoder.weight" + assert tied_map.group_key("decoder.weight") == "decoder.weight" + + +def test_fsdp2_tied_map_survives_shard(dist_workers): + if torch.cuda.device_count() < 2: + pytest.skip("needs >=2 GPUs to shard a tied weight into distinct params") + dist_workers.run(_tied_map_survives_fsdp2_test) + + @minimum_sm(90) def test_fsdp2_weight_compress_context_for_export(dist_workers): dist_workers.run(_compress_weight_test) diff --git a/tests/gpu/torch/export/test_offload_export.py b/tests/gpu/torch/export/test_offload_export.py new file mode 100644 index 00000000000..9ee6139077e --- /dev/null +++ b/tests/gpu/torch/export/test_offload_export.py @@ -0,0 +1,194 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""GPU integration tests for offload-aware unified HF export. + +Tests the full round-trip: + tiny LLaMA (CPU-offloaded via accelerate) + → FP8 layerwise calibration (calib_mutates_weights=False) + → export_hf_checkpoint + → assert no meta tensors in output safetensors + → assert hf_quant_config.json present with fp8 format +""" + +import copy +import json + +import pytest +import torch +from _test_utils.torch.transformers_models import create_tiny_llama_dir +from accelerate import init_empty_weights, load_checkpoint_and_dispatch +from safetensors import safe_open +from transformers import AutoConfig, AutoModelForCausalLM + +import modelopt.torch.quantization as mtq +from modelopt.torch.export import export_hf_checkpoint + + +def _make_cpu_offloaded_model(tmp_path, num_hidden_layers=3, tiny_llama_dir=None): + """Tiny LLaMA with first decoder layer offloaded to CPU, rest on GPU.""" + if tiny_llama_dir is None: + tiny_llama_dir = create_tiny_llama_dir(tmp_path, num_hidden_layers=num_hidden_layers) + config = AutoConfig.from_pretrained(tiny_llama_dir) + + with init_empty_weights(): + model = AutoModelForCausalLM.from_config(config) + + # First layer on CPU to exercise the offload path; lm_head / embed on GPU. + device_map = {} + for n, _m in model.named_modules(): + if "layers" not in n or n.split("layers.")[-1].isdigit(): + device_map[n] = 0 + device_map["model.layers.0"] = "cpu" + + model = load_checkpoint_and_dispatch(model, tiny_llama_dir, device_map=device_map) + return model, config, tiny_llama_dir + + +def _layerwise_fp8_cfg(): + cfg = copy.deepcopy(mtq.FP8_DEFAULT_CFG) + algo = cfg.get("algorithm", "max") + method = algo if isinstance(algo, str) else algo.get("method", "max") + # calib_mutates_weights is a field of LayerwiseConfig (nested), not of the algorithm. + cfg["algorithm"] = {"method": method, "layerwise": {"calib_mutates_weights": False}} + return cfg + + +@pytest.mark.parametrize("quant_cfg", [mtq.FP8_DEFAULT_CFG, _layerwise_fp8_cfg()]) +def test_export_hf_checkpoint_cpu_offloaded(tmp_path, quant_cfg): + """export_hf_checkpoint must succeed on a CPU-offloaded model and produce valid weights. + + Regression guard against the pre-fix bug where remove_hook_from_module was called + before weight materialization, causing meta tensors to be serialized as empty safetensors. + """ + num_hidden_layers = 3 + model, _config, _llama_dir = _make_cpu_offloaded_model( + tmp_path / "offloaded", num_hidden_layers=num_hidden_layers + ) + model.eval() + + def forward_loop(m): + ids = torch.randint(0, m.config.vocab_size, (1, 32)).cuda() + with torch.no_grad(): + m(ids) + + model = mtq.quantize(model, quant_cfg, forward_loop) + + export_dir = tmp_path / "hf_export" + export_dir.mkdir() + export_hf_checkpoint(model, export_dir=str(export_dir)) + + # --- Assertions --- + + # 1. hf_quant_config.json must exist and declare fp8 + quant_config_path = export_dir / "hf_quant_config.json" + assert quant_config_path.exists(), "hf_quant_config.json not written" + with open(quant_config_path) as f: + quant_config = json.load(f) + assert quant_config["quantization"]["quant_algo"] == "FP8", ( + f"Expected FP8, got {quant_config['quantization'].get('quant_algo')}" + ) + + # 2. All tensors in safetensors shards must be non-empty (no meta serialized as zeros) + safetensor_files = list(export_dir.glob("*.safetensors")) + assert safetensor_files, "No safetensors files written" + + for st_file in safetensor_files: + with safe_open(str(st_file), framework="pt") as st: + for key in list(st.keys()): + tensor = st.get_tensor(key) + assert tensor.numel() > 0, f"Zero-numel tensor for key '{key}' in {st_file.name}" + assert not tensor.is_meta, f"Meta tensor for key '{key}' in {st_file.name}" + # Weight tensors (not scales) must have non-zero norm — guards against all-zeros + # from meta serialization + if "weight" in key and "scale" not in key and "quantizer" not in key: + assert tensor.float().abs().sum() > 0, ( + f"All-zero weight tensor '{key}' in {st_file.name} — " + "possible meta tensor serialization bug" + ) + + +def _read_shards(export_dir): + """Map every exported tensor key to its (shape, dtype), unioned across shards.""" + tensors = {} + for st_file in sorted(export_dir.glob("*.safetensors")): + with safe_open(str(st_file), framework="pt") as st: + for key in st.keys(): # noqa: SIM118 + t = st.get_tensor(key) + assert key not in tensors, f"Duplicate key '{key}' across shards" + tensors[key] = (tuple(t.shape), t.dtype) + + index_path = export_dir / "model.safetensors.index.json" + if index_path.exists(): + with open(index_path) as f: + weight_map = json.load(f)["weight_map"] + assert set(weight_map) == set(tensors), "index weight_map disagrees with shard contents" + for shard_name in set(weight_map.values()): + assert (export_dir / shard_name).exists(), f"Indexed shard '{shard_name}' missing" + return tensors + + +@pytest.mark.parametrize("max_shard_size", ["10GB", "5KB"]) +def test_streaming_export_matches_batch_export(tmp_path, max_shard_size): + """The offloaded streaming path must emit the same tensors as the resident batch path. + + Every other assertion in this file iterates only the keys that were written, so none + of them can see a tensor going *missing* -- which is how the dropped ``mtp.*`` tensors + were found. The streaming exporter reimplements much of ``_export_transformers_checkpoint`` + plus ``postprocess_state_dict`` per tensor, so the two can drift silently. + + Both models load the same checkpoint, so key/shape/dtype parity is exact. Values are + not compared: layer 0 calibrates on CPU in the offloaded model and on GPU in the + resident one, and that alone can shift an amax in the last bits. + + The ``5KB`` case additionally drives multi-shard index generation through the real + export path (``_StreamingShardWriter`` is otherwise only tested in isolation). The + whole tiny export is ~28KB, so anything larger stays single-shard and would silently + skip that coverage. + """ + llama_dir = create_tiny_llama_dir(tmp_path / "src", num_hidden_layers=3) + + def forward_loop(m): + ids = torch.randint(0, m.config.vocab_size, (1, 32)).cuda() + with torch.no_grad(): + m(ids) + + def _export(model, subdir): + model.eval() + mtq.quantize(model, mtq.FP8_DEFAULT_CFG, forward_loop) + export_dir = tmp_path / subdir + export_dir.mkdir() + export_hf_checkpoint(model, export_dir=str(export_dir), max_shard_size=max_shard_size) + return _read_shards(export_dir) + + offloaded, _cfg, _dir = _make_cpu_offloaded_model( + tmp_path / "offloaded", tiny_llama_dir=llama_dir + ) + offloaded_tensors = _export(offloaded, "export_offloaded") + + resident = AutoModelForCausalLM.from_pretrained(llama_dir).cuda() + resident_tensors = _export(resident, "export_resident") + + missing = set(resident_tensors) - set(offloaded_tensors) + extra = set(offloaded_tensors) - set(resident_tensors) + assert not missing, f"Streaming export dropped {len(missing)} tensor(s): {sorted(missing)}" + assert not extra, f"Streaming export emitted {len(extra)} unexpected tensor(s): {sorted(extra)}" + + mismatched = { + k: (resident_tensors[k], offloaded_tensors[k]) + for k in resident_tensors + if resident_tensors[k] != offloaded_tensors[k] + } + assert not mismatched, f"shape/dtype drift between export paths: {mismatched}" diff --git a/tests/gpu/torch/utils/test_distributed.py b/tests/gpu/torch/utils/test_distributed.py new file mode 100644 index 00000000000..4091c7465dc --- /dev/null +++ b/tests/gpu/torch/utils/test_distributed.py @@ -0,0 +1,131 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""GPU/distributed tests for ``modelopt.torch.utils.distributed``.""" + +from functools import partial + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F +from _test_utils.torch.transformers_models import get_tiny_llama +from torch.distributed.checkpoint.state_dict import StateDictOptions, set_model_state_dict +from torch.distributed.tensor import DTensor + +from modelopt.torch.utils.distributed import fsdp2_wrap + +VOCAB_SIZE = 32 +N_EXPERTS = 4 + + +class _Fp32Router(nn.Module): + """An MoE router gate pinned to fp32, as Nemotron-3-Nano's modeling code declares it.""" + + def __init__(self, hidden_size: int): + super().__init__() + self.weight = nn.Parameter(torch.empty(N_EXPERTS, hidden_size, dtype=torch.float32)) + nn.init.normal_(self.weight, std=0.02) + + def forward(self, hidden_states): + return F.linear(hidden_states.float(), self.weight.float()) + + +class _RoutedMLP(nn.Module): + """Fronts a bf16 MLP with the fp32 router, so one decoder layer holds both dtypes.""" + + def __init__(self, mlp: nn.Module, hidden_size: int): + super().__init__() + self.mlp = mlp + self.gate = _Fp32Router(hidden_size) + + def forward(self, hidden_states): + scale = self.gate(hidden_states).softmax(-1)[..., :1].to(hidden_states.dtype) + return self.mlp(hidden_states) * scale + + +def _mixed_dtype_model(device): + model = get_tiny_llama(vocab_size=VOCAB_SIZE).to(device) + for layer in model.model.layers: + layer.mlp = _RoutedMLP(layer.mlp, model.config.hidden_size).to(device) + return model.eval() + + +def _test_fsdp2_wrap_mixed_dtypes(rank, size): + """A model with a few fp32 params must still wrap, forward, and load state dicts.""" + device = torch.device(f"cuda:{rank}") + model = _mixed_dtype_model(device) + assert {p.dtype for p in model.model.layers[0].parameters()} == { + torch.bfloat16, + torch.float32, + } + + fsdp2_wrap(model) + + # Raised "FSDP expects uniform original parameter dtype" before the ignored-param fix. + input_ids = torch.randint(0, VOCAB_SIZE, (1, 8), device=device) + with torch.no_grad(): + assert model(input_ids=input_ids).logits.shape == (1, 8, VOCAB_SIZE) + + # bf16 weights are sharded; the fp32 router is left replicated in its original dtype. + gate_weight = model.model.layers[0].mlp.gate.weight + sharded_weight = model.model.layers[0].mlp.mlp.up_proj.weight + assert isinstance(sharded_weight, DTensor) + assert not isinstance(gate_weight, DTensor) + assert gate_weight.dtype == torch.float32 + # Left out of the wrap, it still has to sit on the compute device alongside the shards. + assert gate_weight.device == sharded_weight.to_local().device + + # The FSDP2 loader pushes full tensors into each decoder layer; that must still reach the + # replicated fp32 param as well as the sharded bf16 ones. + layer = model.model.layers[0] + hidden_size = model.config.hidden_size + set_model_state_dict( + layer, + {"mlp.gate.weight": torch.full((N_EXPERTS, hidden_size), 3.0, device=device)}, + options=StateDictOptions(full_state_dict=True, broadcast_from_rank0=False, strict=False), + ) + assert torch.equal( + layer.mlp.gate.weight, torch.full((N_EXPERTS, hidden_size), 3.0, device=device) + ) + + +def test_fsdp2_wrap_mixed_dtypes(dist_workers): + dist_workers.run(_test_fsdp2_wrap_mixed_dtypes) + + +def _test_fsdp2_wrap_moves_ignored_params_to_device(rank, size, cpu_offload): + """A CPU-resident model must end up computing on GPU: fully_shard skips the params it ignores.""" + model = _mixed_dtype_model(torch.device("cpu")) + assert model.model.layers[0].mlp.gate.weight.device.type == "cpu" + + fsdp2_wrap(model, cpu_offload=cpu_offload) + + # Under cpu_offload the shard rests on CPU, but compute — and so the ignored params — is + # still on GPU, which is why the device is taken from the mesh and not from the local shard. + sharded_weight = model.model.layers[0].mlp.mlp.up_proj.weight + assert sharded_weight.to_local().device.type == ("cpu" if cpu_offload else "cuda") + assert model.model.layers[0].mlp.gate.weight.device.type == "cuda" + + input_ids = torch.randint(0, VOCAB_SIZE, (1, 8), device=torch.device(f"cuda:{rank}")) + with torch.no_grad(): + assert model(input_ids=input_ids).logits.shape == (1, 8, VOCAB_SIZE) + + +@pytest.mark.parametrize("cpu_offload", [False, True]) +def test_fsdp2_wrap_moves_ignored_params_to_device(dist_workers, cpu_offload): + dist_workers.run( + partial(_test_fsdp2_wrap_moves_ignored_params_to_device, cpu_offload=cpu_offload) + ) diff --git a/tests/unit/torch/export/test_export_registry.py b/tests/unit/torch/export/test_export_registry.py index 67647f9c31f..4c93388ca08 100644 --- a/tests/unit/torch/export/test_export_registry.py +++ b/tests/unit/torch/export/test_export_registry.py @@ -303,10 +303,10 @@ def test_process_quantized_modules_exports_via_registry(): assert weight.dtype == torch.float8_e4m3fn -def test_export_context_caches_are_per_instance(): - model = nn.Linear(2, 2) - ctx_a = ExportContext(model=model, dtype=torch.float16) - ctx_b = ExportContext(model=model, dtype=torch.float16) - ctx_a.tied_cache[123] = model - assert ctx_b.tied_cache == {} - assert ctx_b.moe_tied_cache == {} +def test_export_context_carries_no_resolver(): + # Tied-weight dedup is driven by the driver's resolver (fed to sync_tied_input_amax / + # postprocess_state_dict), not by handlers, so ExportContext no longer builds or holds + # a resolver. Building one per context would be dead work (handlers never read it) and, + # for large models, an avoidable O(#modules x #patterns x #params) alias-map pass. + ctx = ExportContext(model=nn.Linear(2, 2), dtype=torch.float16) + assert not hasattr(ctx, "resolver") diff --git a/tests/unit/torch/export/test_hf_checkpoint_utils.py b/tests/unit/torch/export/test_hf_checkpoint_utils.py index f3be2564312..08292c65f9f 100644 --- a/tests/unit/torch/export/test_hf_checkpoint_utils.py +++ b/tests/unit/torch/export/test_hf_checkpoint_utils.py @@ -24,7 +24,59 @@ hf_hub_errors = pytest.importorskip("huggingface_hub.errors") LocalEntryNotFoundError = hf_hub_errors.LocalEntryNotFoundError -from modelopt.torch.export import copy_hf_ckpt_remote_code, sanitize_hf_config_for_deployment +from modelopt.torch.export import ( + copy_hf_ckpt_remote_code, + copy_non_safetensor_files_from_ckpt, + sanitize_hf_config_for_deployment, +) +from modelopt.torch.export.plugins import hf_checkpoint_utils + + +def test_copy_non_safetensor_files_from_ckpt_supports_additional_exclusions(tmp_path): + src_dir = tmp_path / "src" + src_dir.mkdir() + (src_dir / "model.safetensors").write_text("weights") + (src_dir / "model.safetensors.index.json").write_text('{"weight_map": {}}') + (src_dir / "pytorch_model.bin").write_text("weights") + (src_dir / "stats.npy").write_text("stats") + (src_dir / "reasoning_parser.py").write_text("parser") + + default_dst = tmp_path / "default" + copy_non_safetensor_files_from_ckpt(src_dir, default_dst) + assert not (default_dst / "model.safetensors").exists() + assert not (default_dst / "model.safetensors.index.json").exists() + assert (default_dst / "pytorch_model.bin").exists() + assert (default_dst / "stats.npy").exists() + + filtered_dst = tmp_path / "filtered" + copy_non_safetensor_files_from_ckpt( + src_dir, + filtered_dst, + exclude_patterns=("*.bin", "*.npy"), + ) + assert (filtered_dst / "reasoning_parser.py").exists() + assert not (filtered_dst / "pytorch_model.bin").exists() + assert not (filtered_dst / "stats.npy").exists() + + +def test_copy_non_safetensor_files_from_ckpt_continues_after_copy_failure(tmp_path, monkeypatch): + src_dir = tmp_path / "src" + src_dir.mkdir() + (src_dir / "bad.py").write_text("bad") + (src_dir / "good.py").write_text("good") + + original_copy2 = hf_checkpoint_utils.shutil.copy2 + + def copy2(source, *args, **kwargs): + if source.endswith("bad.py"): + raise PermissionError("unreadable") + return original_copy2(source, *args, **kwargs) + + monkeypatch.setattr(hf_checkpoint_utils.shutil, "copy2", copy2) + with pytest.warns(UserWarning, match="bad.py"): + copied_files = copy_non_safetensor_files_from_ckpt(src_dir, tmp_path / "dst") + + assert copied_files == ["good.py"] def test_copy_hf_ckpt_remote_code_local_dir(tmp_path): diff --git a/tests/unit/torch/export/test_offload_export.py b/tests/unit/torch/export/test_offload_export.py new file mode 100644 index 00000000000..242a64f762b --- /dev/null +++ b/tests/unit/torch/export/test_offload_export.py @@ -0,0 +1,379 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for offload-aware unified HF export helpers (CPU-only, no GPU required).""" + +import json +import tempfile +from pathlib import Path + +import pytest +import torch +import torch.nn as nn +from safetensors import safe_open + +try: + from accelerate.hooks import AlignDevicesHook, add_hook_to_module + from accelerate.utils import set_module_tensor_to_device +except ImportError: + pytest.skip("accelerate not available", allow_module_level=True) + +from _test_utils.torch.quantization.tied_modules import ( + make_tied_linear_pair, + wrap_in_parent_with_tied_keys, +) + +import modelopt.torch.quantization as mtq +from modelopt.torch.export.model_config import KV_CACHE_FP8 +from modelopt.torch.export.model_utils import TiedWeightMap +from modelopt.torch.export.quant_utils import _postprocess_single_tensor +from modelopt.torch.export.unified_export_hf import _export_quantized_weight +from modelopt.torch.export.unified_export_hf_streaming import ( + _parse_shard_size, + _StreamingShardWriter, +) +from modelopt.torch.quantization.nn.modules.quant_linear import RealQuantLinear +from modelopt.torch.quantization.utils.core_utils import has_accelerate_offload + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_offloaded_linear(dim: int = 16): + """Return a Linear with a CPU-offload AlignDevicesHook attached and params on meta.""" + linear = nn.Linear(dim, dim, bias=False) + weights_map = {"weight": linear.weight.data.clone().cpu()} + hook = AlignDevicesHook(execution_device="cpu", offload=True, weights_map=weights_map) + add_hook_to_module(linear, hook) + set_module_tensor_to_device(linear, "weight", "meta") + return linear, weights_map + + +def _offload_module(module): + """Offload ``module`` like accelerate does: real weight to ``weights_map``, ``.weight`` to a meta Parameter.""" + weights_map = {"weight": module.weight.data.clone().cpu()} + hook = AlignDevicesHook(execution_device="cpu", offload=True, weights_map=weights_map) + add_hook_to_module(module, hook) + set_module_tensor_to_device(module, "weight", "meta") + + +# --------------------------------------------------------------------------- +# tied-weight alias map under offload +# --------------------------------------------------------------------------- + + +def test_tied_weight_map_from_hf_map_survives_offload(): + """TiedWeightMap reads HF's name-based all_tied_weights_keys, so offload does not change it.""" + enc, dec = make_tied_linear_pair(in_features=16, out_features=16) + model = wrap_in_parent_with_tied_keys(enc, dec, decoder_canonical=True) + + assert not has_accelerate_offload(model) + assert TiedWeightMap(model).alias_to_canonical == {"encoder.weight": "decoder.weight"} + + _offload_module(model.encoder) + _offload_module(model.decoder) + assert has_accelerate_offload(model) + + # The map is a plain name dict on the model; offload metas the weights but not the attribute. + tied_map = TiedWeightMap(model) + assert tied_map.alias_to_canonical == {"encoder.weight": "decoder.weight"} + assert tied_map.group_key("encoder.weight") == "decoder.weight" + + +# --------------------------------------------------------------------------- +# has_accelerate_offload +# --------------------------------------------------------------------------- + + +def test_has_accelerate_offload_true(): + linear, _ = _make_offloaded_linear() + assert has_accelerate_offload(linear) is True + + +def test_has_accelerate_offload_false_no_hooks(): + linear = nn.Linear(16, 16) + assert has_accelerate_offload(linear) is False + + +def test_has_accelerate_offload_false_non_offload_hook(): + """A hook with offload=False should not be detected as offloaded.""" + linear = nn.Linear(16, 16) + hook = AlignDevicesHook(execution_device="cpu", offload=False) + add_hook_to_module(linear, hook) + assert has_accelerate_offload(linear) is False + + +def test_has_accelerate_offload_detects_nested_module(): + """Offload hook on a child module should be detected when scanning the parent.""" + + class _Parent(nn.Module): + def __init__(self): + super().__init__() + self.child = nn.Linear(8, 8, bias=False) + + def forward(self, x): + return self.child(x) + + parent = _Parent() + weights_map = {"weight": parent.child.weight.data.clone().cpu()} + hook = AlignDevicesHook(execution_device="cpu", offload=True, weights_map=weights_map) + add_hook_to_module(parent.child, hook) + set_module_tensor_to_device(parent.child, "weight", "meta") + + assert has_accelerate_offload(parent) is True + + +# --------------------------------------------------------------------------- +# _export_quantized_weight meta guard +# --------------------------------------------------------------------------- + + +def test_meta_guard_raises_on_meta_weight(): + """_export_quantized_weight must raise RuntimeError when weight is a meta tensor.""" + linear = nn.Linear(16, 16, bias=False) + + mtq.quantize(linear, mtq.FP8_DEFAULT_CFG, lambda m: m(torch.randn(1, 16))) + + # Manually set weight to meta to simulate what happens after hooks are removed. + linear.weight = nn.Parameter(torch.empty(16, 16, device="meta")) + + with pytest.raises(RuntimeError, match="meta tensor"): + _export_quantized_weight(linear, torch.float32) + + +def test_meta_guard_not_raised_for_real_weight(): + """No RuntimeError when weight is a real (non-meta) tensor.""" + linear = nn.Linear(32, 32, bias=False) + mtq.quantize(linear, mtq.FP8_DEFAULT_CFG, lambda m: m(torch.randn(1, 32))) + # Should not raise + _export_quantized_weight(linear, torch.float32) + + +# --------------------------------------------------------------------------- +# _StreamingShardWriter +# --------------------------------------------------------------------------- + + +def test_streaming_shard_writer_single_shard(): + """Tensors fitting in one shard produce model.safetensors, no index, and round-trip.""" + with tempfile.TemporaryDirectory() as tmpdir: + a, b = torch.randn(4, 4), torch.zeros(2, 2) + writer = _StreamingShardWriter(tmpdir, max_shard_size=10 * 1024**3) + writer.add("a", a) + writer.add("b", b) + weight_map = writer.finalize() + + single = Path(tmpdir) / "model.safetensors" + index = Path(tmpdir) / "model.safetensors.index.json" + assert single.exists(), "model.safetensors not written" + assert not index.exists(), "index file must not exist for single-shard export" + assert set(weight_map.values()) == {"model.safetensors"} + assert set(weight_map.keys()) == {"a", "b"} + + with safe_open(str(single), framework="pt") as f: + assert torch.equal(f.get_tensor("a"), a) + assert torch.equal(f.get_tensor("b"), b) + + +def test_streaming_shard_writer_multi_shard(): + """Tensors exceeding max_shard_size produce an index whose shards all exist on disk. + + The file-existence half is a regression guard: an earlier code path called + model.save_pretrained(state_dict={}) after finalize(), triggering transformers' stale- + shard cleanup loop which matched and deleted every model-NNNNN-of-NNNNN.safetensors + file because filename_to_tensors was empty. + """ + with tempfile.TemporaryDirectory() as tmpdir: + # One float32 4x4 tensor = 64 bytes; set limit to 64 so each tensor goes to a new shard + writer = _StreamingShardWriter(tmpdir, max_shard_size=64) + writer.add("x", torch.ones(4, 4)) + writer.add("y", torch.ones(4, 4)) + weight_map = writer.finalize() + + index_path = Path(tmpdir) / "model.safetensors.index.json" + assert index_path.exists(), "model.safetensors.index.json not written" + assert weight_map["x"] != weight_map["y"], "keys must be in different shards" + + with open(index_path) as f: + index = json.load(f) + assert index["metadata"]["total_size"] > 0 + + for key, shard_name in index["weight_map"].items(): + shard_path = Path(tmpdir) / shard_name + assert shard_path.exists(), ( + f"Shard '{shard_name}' (for key '{key}') missing from disk after finalize()" + ) + assert shard_path.stat().st_size > 0, f"Shard file {shard_name} is empty" + + +def test_streaming_shard_writer_copies_tied_alias(): + """Two keys sharing storage must both survive; save_file rejects the alias itself. + + The name-based _tied_weights_keys filter misses ties transformers does not declare + (e.g. tie_word_embeddings=False but shared storage), so the writer needs its own + guard. It copies rather than drops: offloaded export writes tied weights as separate + entries, so losing a key here would leave the checkpoint short a tensor. + """ + with tempfile.TemporaryDirectory() as tmpdir: + shared = torch.ones(4, 4) + writer = _StreamingShardWriter(tmpdir, max_shard_size=10 * 1024**3) + writer.add("embed_tokens.weight", shared) + writer.add("lm_head.weight", shared) + weight_map = writer.finalize() + + assert set(weight_map) == {"embed_tokens.weight", "lm_head.weight"} + shard_file = Path(tmpdir) / weight_map["lm_head.weight"] + with safe_open(str(shard_file), framework="pt") as f: + assert torch.equal(f.get_tensor("embed_tokens.weight"), shared) + assert torch.equal(f.get_tensor("lm_head.weight"), shared) + + +def test_streaming_shard_writer_copies_aliased_view(): + """A distinct view onto shared storage must be copied, not dropped.""" + with tempfile.TemporaryDirectory() as tmpdir: + base = torch.arange(16, dtype=torch.float32).reshape(4, 4) + view = base.view(16) # same data_ptr, different shape + writer = _StreamingShardWriter(tmpdir, max_shard_size=10 * 1024**3) + writer.add("base", base) + writer.add("view", view) + weight_map = writer.finalize() + + assert set(weight_map) == {"base", "view"}, "aliased view must be kept, not dropped" + shard_file = Path(tmpdir) / weight_map["view"] + with safe_open(str(shard_file), framework="pt") as f: + assert torch.equal(f.get_tensor("view"), view) + assert torch.equal(f.get_tensor("base"), base) + + +def test_streaming_shard_writer_accepts_extra_tensors(): + """extra_state_dict tensors must land in the shards. + + MTP weights are orphaned — HF builds only num_hidden_layers decoders, so they are + never in model.state_dict() and reach export only via extra_state_dict. The streaming + path used to drop them, silently losing 19 tensors relative to the batch export. + """ + with tempfile.TemporaryDirectory() as tmpdir: + writer = _StreamingShardWriter(tmpdir, max_shard_size=10 * 1024**3) + writer.add("model.layers.0.weight", torch.ones(4, 4)) + writer.add("mtp.fc.weight", torch.full((2, 2), 7.0)) + weight_map = writer.finalize() + + assert "mtp.fc.weight" in weight_map + with safe_open(str(Path(tmpdir) / weight_map["mtp.fc.weight"]), framework="pt") as f: + assert torch.equal(f.get_tensor("mtp.fc.weight"), torch.full((2, 2), 7.0)) + + +# --------------------------------------------------------------------------- +# _postprocess_single_tensor +# --------------------------------------------------------------------------- + + +def test_postprocess_passthrough_normal_key(): + """Non-quantizer weights pass through unchanged.""" + key, val = _postprocess_single_tensor( + "model.layers.0.self_attn.q_proj.weight", torch.randn(4, 4), 448.0, None + ) + assert key == "model.layers.0.self_attn.q_proj.weight" + assert val is not None + assert val.shape == (4, 4) + + +@pytest.mark.parametrize( + "key", + [ + "model.layers.0.weight_quantizer._amax", # skip_keys hit, no replacement + "model.layers.0.output_quantizer._amax", # output_quantizer always dropped + "vision_model.radio_model.summary_idxs", # problematic VL parameter + *( # RealQuantLinear scale tensors + f"model.layers.0.weight_quantizer.{q}" for q in RealQuantLinear.list_of_scale_tensors + ), + ], +) +def test_postprocess_drops_key(key): + """Keys with no exported counterpart are dropped rather than renamed.""" + assert _postprocess_single_tensor(key, torch.tensor(1.0), 448.0, None) == (None, None) + + +def test_postprocess_kv_scale_renamed_and_divided(): + """k_bmm_quantizer._amax is renamed to k_proj.k_scale and divided by maxbound.""" + key, val = _postprocess_single_tensor( + "model.layers.0.self_attn.k_bmm_quantizer._amax", + torch.tensor(224.0), + 448.0, + KV_CACHE_FP8, + ) + assert key == "model.layers.0.self_attn.k_proj.k_scale" + assert abs(val.item() - 0.5) < 1e-5 + + +def test_postprocess_scale_squeezed(): + """3D scale tensors with shape[0]==1 are squeezed.""" + t = torch.ones(1, 4, 4) + key, val = _postprocess_single_tensor("model.weight_scale", t, 448.0, None) + assert key == "model.weight_scale" + assert val.shape == (4, 4), f"expected (4, 4), got {val.shape}" + + +# --------------------------------------------------------------------------- +# data_ptr identity under offload +# +# The tests below exist because ``data_ptr()`` only identifies a tensor while that +# tensor is resident. Getting this wrong silently exported wrong weights: meta +# tensors all report 0, and freed addresses are recycled by the allocator. +# --------------------------------------------------------------------------- + + +def test_tied_weights_exported_independently_without_cache(): + """Tied dense modules each pack their own weight instead of aliasing. + + Dense ties are no longer deduped at pack time (the duplicate is dropped by name in + postprocess_state_dict), so both sides pack independently to byte-identical tensors. + This checks the packing behavior only; it does not construct an offloaded model or + drive the streaming export path. + """ + shared = nn.Parameter(torch.randn(16, 16)) + first, second = nn.Linear(16, 16, bias=False), nn.Linear(16, 16, bias=False) + first.weight = second.weight = shared + + for linear in (first, second): + mtq.quantize(linear, mtq.FP8_DEFAULT_CFG, lambda m: m(torch.randn(1, 16))) + _export_quantized_weight(linear, torch.float16) + + assert first.weight.data_ptr() != second.weight.data_ptr() + assert torch.equal(first.weight, second.weight) + + +# --------------------------------------------------------------------------- +# _parse_shard_size +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("size", "expected"), + [ + (1234, 1234), + ("1234", 1234), + # transformers' convert_file_size_to_int reads GB/MB as decimal, GiB/MiB as binary + ("10GB", 10 * 1000**3), + ("500MB", 500 * 1000**2), + ("100KB", 100 * 1000), + ("10GiB", 10 * 1024**3), + ("500MiB", 500 * 1024**2), + ("100KiB", 100 * 1024), + ], +) +def test_parse_shard_size_units(size, expected): + assert _parse_shard_size(size) == expected diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py index 118331ce3d9..b9fa29238d7 100644 --- a/tests/unit/torch/export/test_unified_export_hf.py +++ b/tests/unit/torch/export/test_unified_export_hf.py @@ -15,8 +15,7 @@ """Tests for tied-weight helpers in unified_export_hf.""" -from collections import OrderedDict - +import pytest import torch from _test_utils.torch.quantization.tied_modules import ( make_tied_linear_pair, @@ -24,58 +23,163 @@ ) import modelopt.torch.quantization as mtq -from modelopt.torch.export.model_utils import ( - _collect_canonical_tied_patterns, - _reorder_canonical_first, +from modelopt.torch.export.model_utils import TiedWeightMap +from modelopt.torch.export.quant_utils import ( + fuse_prequant_layernorm, + postprocess_state_dict, + sync_tied_input_amax, ) -from modelopt.torch.export.quant_utils import fuse_prequant_layernorm, sync_tied_input_amax -from modelopt.torch.export.unified_export_hf import _export_quantized_weight from modelopt.torch.quantization.nn import TensorQuantizer -def test_collect_canonical_tied_patterns_dict_style(): - """Dict-style _tied_weights_keys yields regex patterns + canonical-side substrings.""" - enc, dec = make_tied_linear_pair() - parent = wrap_in_parent_with_tied_keys(enc, dec, decoder_canonical=True) - - patterns, side_substrings = _collect_canonical_tied_patterns(parent) +def test_hf_all_tied_weights_keys_contract(): + """Pin the transformers API we build tied_map from, so a version bump fails loud here. + + We rely on ``model.all_tied_weights_keys`` being a name-based ``{alias: canonical}`` dict + (``tie_word_embeddings=True`` -> lm_head aliases the embedding). If transformers renames it or + flips the direction, this breaks instead of silently skipping tied-weight dedup. + """ + pytest.importorskip( + "transformers", minversion="5.0" + ) # attribute only exists on transformers>=5.0 + from transformers import AutoModelForCausalLM, LlamaConfig + + cfg = LlamaConfig( + vocab_size=64, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=4, + tie_word_embeddings=True, + ) + cfg.architectures = ["LlamaForCausalLM"] + model = AutoModelForCausalLM.from_config(cfg) - assert len(patterns) >= 1 - # "decoder" is in the canonical RHS but not the alias LHS — must auto-derive. - # "encoder" is alias-only and must NOT be returned as canonical (would invert dedup). - assert "decoder" in side_substrings - assert "encoder" not in side_substrings + assert model.all_tied_weights_keys == {"lm_head.weight": "model.embed_tokens.weight"} + # TiedWeightMap consumes it verbatim. + assert TiedWeightMap(model).alias_to_canonical == { + "lm_head.weight": "model.embed_tokens.weight" + } -def test_collect_canonical_tied_patterns_list_style_yields_no_canonical_info(): - """Legacy list-style _tied_weights_keys carries no canonical/alias info — returns empty.""" - enc, dec = make_tied_linear_pair() - parent = wrap_in_parent_with_tied_keys(enc, dec, decoder_canonical=False) +def test_tied_weight_map_drops_self_entries(): + """A self-entry (alias == canonical) is filtered, so the kept canonical is never dropped.""" - patterns, side_substrings = _collect_canonical_tied_patterns(parent) + class _M(torch.nn.Module): + all_tied_weights_keys = {"a.weight": "b.weight", "b.weight": "b.weight"} - assert patterns == [] - assert side_substrings == [] + assert TiedWeightMap(_M()).alias_to_canonical == {"a.weight": "b.weight"} -def test_reorder_canonical_first_puts_decoder_keys_before_encoder_keys(): - """_reorder_canonical_first moves canonical-side state_dict keys ahead of alias-side keys.""" +def test_tied_group_resolver_group_key_is_shared_and_order_independent(): + """Both sides of a declared tie map to the same key; untied params map to None.""" enc, dec = make_tied_linear_pair() parent = wrap_in_parent_with_tied_keys(enc, dec, decoder_canonical=True) - sd = OrderedDict( - [ - ("encoder.weight", torch.zeros(1)), - ("unrelated.foo", torch.zeros(1)), - ("decoder.weight", torch.zeros(1)), - ] - ) + tied_map = TiedWeightMap(parent) + + assert tied_map.group_key("encoder.weight") == tied_map.group_key("decoder.weight") + assert tied_map.group_key("encoder.weight") == "decoder.weight" # canonical wins + assert tied_map.group_key("unrelated.weight") is None + + +def test_tied_group_resolver_per_layer_backreference(): + """container_group_key resolves each layer's tie independently (no cross-layer collapse). + + HF expands the per-layer regex/backreference into concrete ``all_tied_weights_keys`` names; + TiedWeightMap reads that and container_group_key must keep layers distinct. + """ + + class _Parent(torch.nn.Module): + all_tied_weights_keys = { + "encoder.layers.0.experts.gate_up_proj": "decoder.layers.0.experts.gate_up_proj", + "encoder.layers.1.experts.gate_up_proj": "decoder.layers.1.experts.gate_up_proj", + } - reordered = _reorder_canonical_first(sd, parent) - keys = list(reordered.keys()) + tied_map = TiedWeightMap(_Parent()) - assert keys.index("decoder.weight") < keys.index("encoder.weight") - assert set(reordered) == set(sd) # no drops or additions + assert ( + tied_map.container_group_key("encoder.layers.0.experts", "gate_up_proj") + == "decoder.layers.0.experts" + ) + assert ( + tied_map.container_group_key("encoder.layers.1.experts", "gate_up_proj") + == "decoder.layers.1.experts" + ) + # Encoder layer 0 must not collapse into decoder layer 1. + assert tied_map.container_group_key( + "encoder.layers.0.experts", "gate_up_proj" + ) != tied_map.container_group_key("encoder.layers.1.experts", "gate_up_proj") + + +def test_tied_group_resolver_parallel_pattern_declaration(): + """DiffusionGemma-style tie: container resolves to the decoder canonical; per-expert split keys drop by name. + + HF resolves DiffGemma's parallel-regex declaration into concrete ``all_tied_weights_keys`` + names (incl. the fused expert Parameters); TiedWeightMap reads that. + """ + + class _Root(torch.nn.Module): + all_tied_weights_keys = { + "model.encoder.language_model.layers.0.experts.gate_up_proj": ( + "model.decoder.layers.0.experts.gate_up_proj" + ), + "model.encoder.language_model.layers.0.experts.down_proj": ( + "model.decoder.layers.0.experts.down_proj" + ), + } + + tied_map = TiedWeightMap(_Root()) + + # container group key: encoder side resolves to the decoder canonical container + assert ( + tied_map.container_group_key( + "model.encoder.language_model.layers.0.experts", "gate_up_proj" + ) + == "model.decoder.layers.0.experts" + ) + # post-export per-expert split keys of the (fully tied) container are dropped by name. + enc = "model.encoder.language_model.layers.0.experts" + dec = "model.decoder.layers.0.experts" + shared = torch.randn(4, 4) # tied sides export identical bytes + sd = { + f"{enc}.3.gate_proj.weight": shared.clone(), + f"{dec}.3.gate_proj.weight": shared.clone(), + } + out = postprocess_state_dict(sd, maxbound=448, quantization=None, tied_map=tied_map) + assert f"{enc}.3.gate_proj.weight" not in out # alias split key dropped + assert f"{dec}.3.gate_proj.weight" in out # canonical kept + + +def test_postprocess_moe_alias_container_ties_to_two_canonicals(): + """One alias container whose projections tie to DIFFERENT canonical containers dedups both. + + Groups are keyed by the full (alias, canonical) pair, so gate_up_proj -> decA and + down_proj -> decB under the same 'enc.experts' don't overwrite each other. + """ + + class _M(torch.nn.Module): + all_tied_weights_keys = { + "enc.experts.gate_up_proj": "decA.experts.gate_up_proj", + "enc.experts.down_proj": "decB.experts.down_proj", + } + + tied_map = TiedWeightMap(_M()) + w = torch.randn(4, 4) # tied sides export identical bytes + sd = { + "enc.experts.0.gate_proj.weight": w.clone(), + "decA.experts.0.gate_proj.weight": w.clone(), + "enc.experts.0.down_proj.weight": w.clone(), + "decB.experts.0.down_proj.weight": w.clone(), + } + out = postprocess_state_dict(sd, maxbound=448, quantization=None, tied_map=tied_map) + assert "enc.experts.0.gate_proj.weight" not in out # tied to decA -> dropped + assert ( + "enc.experts.0.down_proj.weight" not in out + ) # tied to decB -> dropped (would leak w/o fix) + assert "decA.experts.0.gate_proj.weight" in out + assert "decB.experts.0.down_proj.weight" in out def _quantize_and_get_input_quantizers(parent): @@ -116,71 +220,236 @@ def test_sync_tied_input_amax_no_op_for_untied_modules(): assert torch.allclose(dec_q.amax, torch.tensor(5.0)) -def _calibrate_through_both_children(parent): - """Insert NVFP4 quantizers and run a one-shot forward through both children for calibration.""" +def test_sync_tied_input_amax_merges_undeclared_shared_weight(): + """Two Linears sharing a weight but declaring no tie still get their input amaxes merged by identity.""" + parent = torch.nn.Module() + parent.a = torch.nn.Linear(16, 32, bias=False) + parent.b = torch.nn.Linear(16, 32, bias=False) + parent.b.weight = parent.a.weight # undeclared physical share (same Parameter object) + + mtq.quantize(parent, mtq.FP8_DEFAULT_CFG, forward_loop=lambda m: None) + assert parent.a.weight is parent.b.weight # share survives quantize + # No _tied_weights_keys declared, so name-based grouping finds nothing to merge. + assert TiedWeightMap(parent).group_key("a.weight") is None + + parent.a.input_quantizer.amax = torch.tensor(2.0) + parent.b.input_quantizer.amax = torch.tensor(8.0) - def forward_loop(m): - x = torch.randn(2, 16) - m.encoder(x) - m.decoder(x) + sync_tied_input_amax(parent) - mtq.quantize(parent, mtq.NVFP4_DEFAULT_CFG, forward_loop=forward_loop) + expected = torch.tensor(8.0) + assert torch.allclose(parent.a.input_quantizer.amax, expected) + assert torch.allclose(parent.b.input_quantizer.amax, expected) -def test_export_quantized_weight_aliases_packed_weight_for_tied_linears(): - """Tied Linears share data_ptr for packed .weight and scale buffers after export.""" +def test_postprocess_name_based_drops_alias_across_distinct_addresses(): + """Declared alias dropped by name even when its tensor is at a different address (the FSDP full_state_dict case).""" enc, dec = make_tied_linear_pair() - parent = wrap_in_parent_with_tied_keys(enc, dec) - _calibrate_through_both_children(parent) + parent = wrap_in_parent_with_tied_keys(enc, dec, decoder_canonical=True) + tied_map = TiedWeightMap(parent) - # Per-call dedup cache (the production pattern: caller owns the cache, scoped - # to one export invocation). Threaded through both sides of the tied pair so - # the alias step at the end of _export_quantized_weight catches the dedup. - tied_cache: dict = {} - _export_quantized_weight(enc, torch.float16, "weight", _tied_cache=tied_cache) - _export_quantized_weight(dec, torch.float16, "weight", _tied_cache=tied_cache) + # Distinct storages (different data_ptr) but identical bytes (genuinely tied): the address + # pass could never collapse these, but the name pass does. + shared = torch.randn(4, 4) + sd = {"encoder.weight": shared.clone(), "decoder.weight": shared.clone()} + assert sd["encoder.weight"].data_ptr() != sd["decoder.weight"].data_ptr() - assert enc.weight.data_ptr() == dec.weight.data_ptr() - for scale_attr in ("weight_scale", "weight_scale_2"): - if hasattr(enc, scale_attr) and hasattr(dec, scale_attr): - assert getattr(enc, scale_attr).data_ptr() == getattr(dec, scale_attr).data_ptr() + out = postprocess_state_dict(sd, maxbound=448, quantization=None, tied_map=tied_map) + assert "decoder.weight" in out # canonical kept + assert "encoder.weight" not in out # alias dropped by name -def test_export_quantized_weight_no_alias_for_untied_linears(): - """Untied Linears keep independent data_ptrs after export — no false-positive aliasing.""" - parent = torch.nn.Module() - parent.encoder = torch.nn.Linear(16, 32, bias=False) - parent.decoder = torch.nn.Linear(16, 32, bias=False) - assert parent.encoder.weight.data_ptr() != parent.decoder.weight.data_ptr() - _calibrate_through_both_children(parent) - # Same fresh cache shape as the positive case — confirms that even with - # dedup enabled, untied modules with distinct source data_ptrs do not get - # falsely aliased. - tied_cache: dict = {} - _export_quantized_weight(parent.encoder, torch.float16, "weight", _tied_cache=tied_cache) - _export_quantized_weight(parent.decoder, torch.float16, "weight", _tied_cache=tied_cache) +def test_postprocess_name_based_keeps_alias_when_canonical_absent(): + """An alias is NOT dropped when its canonical counterpart is missing (no orphaning).""" + enc, dec = make_tied_linear_pair() + parent = wrap_in_parent_with_tied_keys(enc, dec, decoder_canonical=True) + tied_map = TiedWeightMap(parent) + + sd = {"encoder.weight": torch.randn(4, 4)} # canonical decoder.weight absent + out = postprocess_state_dict(sd, maxbound=448, quantization=None, tied_map=tied_map) + + assert "encoder.weight" in out + + +def test_postprocess_keeps_both_sides_when_tied_quant_state_differs(): + """Tied sides with differing quant state aren't deduped (atomic drop), so no scale is orphaned.""" + enc, dec = make_tied_linear_pair() + parent = wrap_in_parent_with_tied_keys(enc, dec, decoder_canonical=True) + tied_map = TiedWeightMap(parent) + + sd = { + # alias (encoder) exported as quantized: weight + companion scales + "encoder.weight": torch.randn(4, 4), + "encoder.weight_scale": torch.randn(4), + "encoder.input_scale": torch.randn(1), + # canonical (decoder) exported unquantized: weight only, no scales + "decoder.weight": torch.randn(4, 4), + } + + out = postprocess_state_dict(sd, maxbound=448, quantization=None, tied_map=tied_map) + + # Mismatched companion keys -> keep the whole alias group; no orphaned scales. + assert set(out) == set(sd) + + +def test_postprocess_name_based_drops_tied_expert_subtree_by_name(): + """A container-level declared expert tie drops every per-expert alias key by name, + keeping only the canonical subtree -- across distinct addresses (FSDP-safe).""" + + class _Parent(torch.nn.Module): + all_tied_weights_keys = { + "encoder.experts.gate_up_proj": "decoder.experts.gate_up_proj", + "encoder.experts.down_proj": "decoder.experts.down_proj", + } + + parent = _Parent() + tied_map = TiedWeightMap(parent) + assert tied_map.alias_to_canonical == { + "encoder.experts.gate_up_proj": "decoder.experts.gate_up_proj", + "encoder.experts.down_proj": "decoder.experts.down_proj", + } + + # Exported-style per-expert keys; tied sides carry identical bytes (distinct storage). + sd = {} + for e in range(2): + for proj in ("gate_proj", "up_proj", "down_proj"): + w, s = torch.randn(4, 4), torch.randn(4) + for side in ("encoder", "decoder"): + sd[f"{side}.experts.{e}.{proj}.weight"] = w.clone() + sd[f"{side}.experts.{e}.{proj}.weight_scale"] = s.clone() + + out = postprocess_state_dict(sd, maxbound=448, quantization=None, tied_map=tied_map) + + assert not any(k.startswith("encoder.experts.") for k in out) # all aliases dropped + assert all(k.startswith("decoder.experts.") for k in out) # only canonical remains + assert len(out) == 2 * 3 * 2 # 2 experts * 3 projections * (weight + weight_scale) + + +def test_postprocess_keeps_independent_bias_under_tied_weight(): + """A weight tie must not drop an independent bias sharing the module prefix (the NVBug 6525352 failure class).""" + + class _TwoLinear(torch.nn.Module): + all_tied_weights_keys = {"A.weight": "B.weight"} + + tied_map = TiedWeightMap(_TwoLinear()) + tied_w = torch.randn(4, 4) # A.weight is B.weight -> identical exported bytes + sd = { + "A.weight": tied_w.clone(), + "A.bias": torch.randn(4), # independent + "B.weight": tied_w.clone(), + "B.bias": torch.randn(4), # independent + } + out = postprocess_state_dict(sd, maxbound=448, quantization=None, tied_map=tied_map) + + assert "A.weight" not in out # tied weight dropped + assert "A.bias" in out # independent bias survives + assert "B.weight" in out and "B.bias" in out + + +def test_postprocess_partially_tied_container_dedups_only_tied_projections(): + """Only the tied projection's per-expert keys are deduped; an untied down_proj and a router child survive.""" + + class _Parent(torch.nn.Module): + all_tied_weights_keys = {"encoder.experts.gate_up_proj": "decoder.experts.gate_up_proj"} + + tied_map = TiedWeightMap(_Parent()) + assert tied_map.alias_to_canonical == { + "encoder.experts.gate_up_proj": "decoder.experts.gate_up_proj" + } + + # Tied projections (gate_proj/up_proj, from gate_up_proj) carry identical bytes across sides; + # untied down_proj and router differ. + sd = {} + for e in range(2): + for proj in ("gate_proj", "up_proj"): + w = torch.randn(4, 4) + sd[f"encoder.experts.{e}.{proj}.weight"] = w.clone() + sd[f"decoder.experts.{e}.{proj}.weight"] = w.clone() + for side in ("encoder", "decoder"): + sd[f"{side}.experts.{e}.down_proj.weight"] = torch.randn(4, 4) # untied + for side in ("encoder", "decoder"): + sd[f"{side}.experts.router.weight"] = torch.randn(4, 4) # non-projection child + + out = postprocess_state_dict(sd, maxbound=448, quantization=None, tied_map=tied_map) + + # Tied gate_up_proj (splits to gate_proj/up_proj) is deduped on the encoder (alias) side. + assert not any(".gate_proj." in k or ".up_proj." in k for k in out if k.startswith("encoder.")) + # Untied down_proj and the router survive on both sides. + assert all(f"encoder.experts.{e}.down_proj.weight" in out for e in range(2)) + assert "encoder.experts.router.weight" in out and "decoder.experts.router.weight" in out + # Decoder (canonical) side fully kept. + assert all( + f"decoder.experts.{e}.{p}.weight" in out + for e in range(2) + for p in ("gate_proj", "up_proj", "down_proj") + ) + + +def test_postprocess_backstop_collapses_keys_sharing_a_dataptr(): + """The address backstop drops a later key that shares a ``data_ptr`` with an earlier one (first-wins).""" + storage = torch.arange(4) + sd = {"short": storage[:2], "long": storage} # both start at offset 0 -> same data_ptr + assert sd["short"].data_ptr() == sd["long"].data_ptr() + + out = postprocess_state_dict(sd, maxbound=448, quantization=None) - assert parent.encoder.weight.data_ptr() != parent.decoder.weight.data_ptr() + assert len(out) == 1 and "short" in out # first-seen kept, later collision dropped -def test_export_quantized_weight_skips_alias_when_one_tied_side_is_unquantized(): - """Unquantized side early-returns; its .weight stays at the original shared Parameter.""" +def test_postprocess_backstop_keeps_keys_with_distinct_dataptrs(): + """Two slices at different offsets have distinct ``data_ptr``s, so the backstop leaves both.""" + base = torch.arange(4) + sd = {"first": base[:2], "second": base[2:]} # offsets 0 and 2 -> different data_ptr + assert sd["first"].data_ptr() != sd["second"].data_ptr() + + out = postprocess_state_dict(sd, maxbound=448, quantization=None) + + assert set(out) == {"first", "second"} # neither dropped + assert torch.equal(out["first"], torch.tensor([0, 1])) + assert torch.equal(out["second"], torch.tensor([2, 3])) + + +def test_postprocess_dense_tie_drops_pre_quant_scale_companion(): + """An AWQ-style dense tie drops ``pre_quant_scale`` with the weight (no orphaned companion).""" enc, dec = make_tied_linear_pair() - parent = wrap_in_parent_with_tied_keys(enc, dec) - original_shared_data_ptr = enc.weight.data_ptr() + parent = wrap_in_parent_with_tied_keys(enc, dec, decoder_canonical=True) + tied_map = TiedWeightMap(parent) + w, pqs = torch.randn(4, 4), torch.randn(4) # tied sides export identical bytes + sd = { + "encoder.weight": w.clone(), + "encoder.pre_quant_scale": pqs.clone(), + "decoder.weight": w.clone(), + "decoder.pre_quant_scale": pqs.clone(), + } + + out = postprocess_state_dict(sd, maxbound=448, quantization=None, tied_map=tied_map) + + assert "encoder.weight" not in out and "encoder.pre_quant_scale" not in out # both dropped + assert "decoder.weight" in out and "decoder.pre_quant_scale" in out # canonical kept + + +def test_postprocess_raises_when_tied_sides_export_different_values(): + """A declared tie whose two sides export different bytes must raise, not silently corrupt.""" + + class _TwoLinear(torch.nn.Module): + all_tied_weights_keys = {"A.weight": "B.weight"} + + tied_map = TiedWeightMap(_TwoLinear()) + sd = {"A.weight": torch.zeros(4, 4), "B.weight": torch.ones(4, 4)} # declared tie, but differ + with pytest.raises(RuntimeError, match="differs from its canonical"): + postprocess_state_dict(sd, maxbound=448, quantization=None, tied_map=tied_map) + - _calibrate_through_both_children(parent) - # is_enabled is a read-only property; .disable() is the canonical bypass. - dec.weight_quantizer.disable() +def test_postprocess_state_dict_preserves_zero_pointer_tensors(): + state_dict = { + "first": torch.empty(4, device="meta"), + "second": torch.empty(4, device="meta"), + } - tied_cache: dict = {} - _export_quantized_weight(enc, torch.float16, "weight", _tied_cache=tied_cache) - _export_quantized_weight(dec, torch.float16, "weight", _tied_cache=tied_cache) + processed = postprocess_state_dict(state_dict, maxbound=448, quantization=None) - assert enc.weight.data_ptr() != original_shared_data_ptr # encoder got fresh packed - assert dec.weight.data_ptr() == original_shared_data_ptr # decoder untouched - assert enc.weight.data_ptr() != dec.weight.data_ptr() + assert set(processed) == set(state_dict) def _linear_with_input_quantizer(): diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index c435b3698be..07ef2aaff24 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -641,12 +641,20 @@ def _spy_export(wrapper, dtype, **_kwargs): # Tests for tied-experts dedup in _export_fused_experts # --------------------------------------------------------------------------- def _build_two_moe_blocks(tie: bool) -> nn.Module: - """Build a parent with two _SyntheticSparseMoeBlock children, optionally with tied 3-D params.""" + """Build a parent with two _SyntheticSparseMoeBlock children, optionally with tied 3-D params. + + When ``tie`` is set the parent both shares the 3-D expert Parameters and declares the tie via + ``_tied_weights_keys``, so the name-based map resolves it (object sharing alone is not enough). + """ parent = nn.Module() parent.encoder = _SyntheticSparseMoeBlock() parent.decoder = _SyntheticSparseMoeBlock() if tie: tie_fused_experts_3d_params(parent.encoder.experts, parent.decoder.experts) + parent._tied_weights_keys = { + r"^encoder\.experts\.gate_up_proj$": "decoder.experts.gate_up_proj", + r"^encoder\.experts\.down_proj$": "decoder.experts.down_proj", + } return parent @@ -685,30 +693,16 @@ def _cleanup_registry(mod_type): if QuantModuleRegistry.get(mod_type) is not None: QuantModuleRegistry.unregister(mod_type) - def test_per_expert_buffers_share_data_ptr_for_tied_fused_experts(self): - """Two tied FusedExperts modules: every per-expert .weight + scale buffer shares data_ptr.""" + def test_tied_fused_experts_pack_independently_to_equal_values(self): + """Tied FusedExperts pack independently: distinct data_ptrs, equal bytes (dropped by name later).""" parent = _build_two_moe_blocks(tie=True) expert_type = type(parent.encoder.experts) self._cleanup_registry(expert_type) try: _calibrate_two_moe_blocks(parent) - # Per-call dedup caches threaded through both export calls; int keys - # for per-expert wrapper dedup, tuple keys for module-level dedup. - tied_cache: dict = {} - moe_tied_cache: dict = {} - _export_fused_experts( - parent.encoder.experts, - torch.float16, - _moe_tied_cache=moe_tied_cache, - _tied_cache=tied_cache, - ) - _export_fused_experts( - parent.decoder.experts, - torch.float16, - _moe_tied_cache=moe_tied_cache, - _tied_cache=tied_cache, - ) + _export_fused_experts(parent.encoder.experts, torch.float16) + _export_fused_experts(parent.decoder.experts, torch.float16) for idx in range(NUM_EXPERTS): enc_expert = getattr(parent.encoder.experts, str(idx)) @@ -716,41 +710,23 @@ def test_per_expert_buffers_share_data_ptr_for_tied_fused_experts(self): for proj_name in ("gate_proj", "up_proj", "down_proj"): enc_proj = getattr(enc_expert, proj_name) dec_proj = getattr(dec_expert, proj_name) - assert enc_proj.weight.data_ptr() == dec_proj.weight.data_ptr() - for scale_attr in ("weight_scale", "weight_scale_2"): - if hasattr(enc_proj, scale_attr) and hasattr(dec_proj, scale_attr): - assert ( - getattr(enc_proj, scale_attr).data_ptr() - == getattr(dec_proj, scale_attr).data_ptr() - ) + # independent storage (no aliasing) ... + assert enc_proj.weight.data_ptr() != dec_proj.weight.data_ptr() + # ... but byte-identical, so postprocess drops one by name + assert torch.equal(enc_proj.weight, dec_proj.weight) finally: self._cleanup_registry(expert_type) - def test_per_expert_buffers_have_independent_data_ptrs_for_untied_fused_experts(self): - """Two untied FusedExperts modules: per-expert buffers stay independent (no false-positive alias).""" + def test_untied_fused_experts_have_independent_buffers(self): + """Untied FusedExperts stay fully independent — no aliasing, distinct values.""" parent = _build_two_moe_blocks(tie=False) expert_type = type(parent.encoder.experts) self._cleanup_registry(expert_type) try: _calibrate_two_moe_blocks(parent) - # Same fresh caches as the positive case — confirms that even with - # dedup enabled, untied modules with distinct source data_ptrs do - # not get falsely aliased. - tied_cache: dict = {} - moe_tied_cache: dict = {} - _export_fused_experts( - parent.encoder.experts, - torch.float16, - _moe_tied_cache=moe_tied_cache, - _tied_cache=tied_cache, - ) - _export_fused_experts( - parent.decoder.experts, - torch.float16, - _moe_tied_cache=moe_tied_cache, - _tied_cache=tied_cache, - ) + _export_fused_experts(parent.encoder.experts, torch.float16) + _export_fused_experts(parent.decoder.experts, torch.float16) for idx in range(NUM_EXPERTS): enc_expert = getattr(parent.encoder.experts, str(idx)) diff --git a/tests/unit/torch/quantization/test_param_index.py b/tests/unit/torch/quantization/test_param_index.py new file mode 100644 index 00000000000..a78dd7b8cff --- /dev/null +++ b/tests/unit/torch/quantization/test_param_index.py @@ -0,0 +1,121 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the parameter-name index behind FSDP2 export's param mapping.""" + +from types import SimpleNamespace + +import torch.nn as nn + +from modelopt.torch.quantization.utils import core_utils +from modelopt.torch.quantization.utils.core_utils import build_param_index, get_prefixed_param_names + + +def _linear_scan(parent_model, target_module): + """The pre-index implementation, kept as the correctness oracle.""" + target_ids = {id(p) for p in target_module.parameters()} + return next( + ( + name.rsplit(".", 1)[0] + for name, param in parent_model.named_parameters() + if id(param) in target_ids + ), + None, + ) + + +def _moe_ish_model(n_layers=3, n_experts=8): + """Many small sibling modules, i.e. the shape that made the old scan quadratic.""" + return nn.Sequential( + *[ + nn.ModuleDict( + { + "attn": nn.Linear(8, 8), + "experts": nn.ModuleList( + [nn.Linear(8, 8, bias=False) for _ in range(n_experts)] + ), + } + ) + for _ in range(n_layers) + ] + ) + + +def test_matches_linear_scan_for_every_module(): + model = _moe_ish_model() + index = build_param_index(model) + for _, module in model.named_modules(): + if not any(True for _ in module.parameters()): + continue + assert get_prefixed_param_names(model, module, index) == _linear_scan(model, module) + + +def test_index_is_optional(): + """Callers that pass no index still get the same answer.""" + model = _moe_ish_model() + target = model[1]["experts"][3] + assert get_prefixed_param_names(model, target) == _linear_scan(model, target) + + +def test_returns_none_for_foreign_module(): + model = _moe_ish_model() + assert get_prefixed_param_names(model, nn.Linear(8, 8), build_param_index(model)) is None + + +def test_parameterless_module_returns_none(): + model = _moe_ish_model() + assert get_prefixed_param_names(model, nn.Identity(), build_param_index(model)) is None + + +def test_index_maps_every_parameter(): + model = _moe_ish_model() + index = build_param_index(model) + assert len(index) == len(list(model.parameters())) + for pos, (name, param) in enumerate(model.named_parameters()): + assert index[id(param)] == (pos, name) + + +def test_shared_parameter_resolves_to_first_occurrence(): + """A tied weight must resolve the same way the linear scan did: earliest name wins.""" + model = _moe_ish_model(n_layers=2, n_experts=2) + model[1]["experts"][0].weight = model[0]["experts"][0].weight + target = model[1]["experts"][0] + assert get_prefixed_param_names(model, target, build_param_index(model)) == _linear_scan( + model, target + ) + + +def test_mapping_walks_the_parameters_once(monkeypatch): + """The regression guard: one parameter walk per call, not one per FSDPParam. + + Walking per FSDPParam is what made MoE export quadratic and stalled it for hours. + """ + model = _moe_ish_model(n_layers=2, n_experts=8) + experts = [model[0]["experts"][i] for i in range(8)] + + calls = [] + real_build = core_utils.build_param_index + monkeypatch.setattr( + core_utils, "build_param_index", lambda m: (calls.append(m), real_build(m))[1] + ) + + fsdp_params = [ + SimpleNamespace(_module_info=SimpleNamespace(module=e, param_name="weight")) + for e in experts + ] + mapping = core_utils.create_fsdp_param_mapping(fsdp_params, model) + + assert len(calls) == 1, f"expected 1 parameter walk, got {len(calls)} (one per FSDPParam?)" + assert set(mapping) == {f"0.experts.{i}.weight" for i in range(8)} diff --git a/tests/unit/torch/speculative/test_speculative_utils.py b/tests/unit/torch/speculative/test_speculative_utils.py new file mode 100644 index 00000000000..3de4999a69e --- /dev/null +++ b/tests/unit/torch/speculative/test_speculative_utils.py @@ -0,0 +1,75 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for get_conversation_input_ids, the shared offline-dump tokenizer helper. + +apply_chat_template returns a BatchEncoding on transformers>=5, so the old len(input_ids) +was 2 (field count) and every conversation got dropped by the num_input_tokens <= 10 filter. +""" + +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 + +CONVERSATIONS = [ + {"role": "user", "content": "Explain why the sky is blue in a few sentences."}, + { + "role": "assistant", + "content": "Rayleigh scattering makes shorter blue wavelengths scatter more. " * 8, + }, +] + + +def _expected_ids(tokenizer): + rendered = tokenizer.apply_chat_template( + CONVERSATIONS, add_generation_prompt=False, tokenize=False + ) + return tokenizer(rendered, add_special_tokens=False)["input_ids"] + + +def test_matches_rendered_prompt(): + tokenizer = get_tiny_tokenizer() + input_ids = get_conversation_input_ids(tokenizer, CONVERSATIONS) + assert input_ids == _expected_ids(tokenizer) + assert len(input_ids) > 10 + + +@pytest.mark.parametrize( + "wrap", ["batch_encoding", "plain_dict", "tensor_2d", "batched_list", "plain_list"] +) +def test_normalizes_to_flat_list(wrap): + """Pin every apply_chat_template return shape to a flat list[int], regardless of version.""" + expected = _expected_ids(get_tiny_tokenizer()) + returns = { + "batch_encoding": BatchEncoding( + {"input_ids": expected, "attention_mask": [1] * len(expected)} + ), + "plain_dict": {"input_ids": expected, "attention_mask": [1] * len(expected)}, + "tensor_2d": torch.tensor([expected]), + "batched_list": [expected], + "plain_list": expected, + } + + class _Stub: + def apply_chat_template(self, conversations, **kwargs): + return returns[wrap] + + input_ids = get_conversation_input_ids(_Stub(), CONVERSATIONS) + assert input_ids == expected + # The pre-fix code saw len() in {1, 2} here and silently dropped the conversation. + assert len(input_ids) > 10 diff --git a/tests/unit/torch/utils/test_distributed.py b/tests/unit/torch/utils/test_distributed.py new file mode 100644 index 00000000000..7b630f93d50 --- /dev/null +++ b/tests/unit/torch/utils/test_distributed.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Pure-function tests for ``modelopt.torch.utils.distributed``.""" + +import pytest +import torch +import torch.nn as nn + +from modelopt.torch.utils.distributed import _off_dtype_params + + +def _model(*sizes_and_dtypes) -> nn.Module: + model = nn.Module() + for i, (numel, dtype) in enumerate(sizes_and_dtypes): + model.register_parameter(f"p{i}", nn.Parameter(torch.zeros(numel, dtype=dtype))) + return model + + +def test_off_dtype_params_uniform_is_empty(recwarn): + model = _model((8, torch.bfloat16), (4, torch.bfloat16)) + assert _off_dtype_params(model) == set() + assert not recwarn.list + + +def test_off_dtype_params_returns_only_the_minority(): + # An fp32 MoE router gate next to bf16 weights, as Nemotron-3-Nano ships it. + model = _model((100, torch.bfloat16), (5, torch.float32)) + with pytest.warns(UserWarning, match="mixed parameter dtypes"): + off = _off_dtype_params(model) + assert off == {model.p1} + + +@pytest.mark.parametrize( + ("params", "expected"), + [ + # Dominant dtype is by element count, not parameter count: three small fp32 params + # lose to one large bf16 param, and vice versa. + ([(100, torch.bfloat16), (5, torch.float32), (5, torch.float32), (5, torch.float32)], "p0"), + ( + [(100, torch.float32), (5, torch.bfloat16), (5, torch.bfloat16), (5, torch.bfloat16)], + "p0", + ), + ], +) +def test_off_dtype_params_dominant_is_by_numel(params, expected): + model = _model(*params) + with pytest.warns(UserWarning, match="mixed parameter dtypes"): + off = _off_dtype_params(model) + kept = {n for n, p in model.named_parameters() if p not in off} + assert kept == {expected}