diff --git a/.github/workflows/unit_tests.yml b/.github/workflows/unit_tests.yml index b258db24b69..6fbca185d9f 100644 --- a/.github/workflows/unit_tests.yml +++ b/.github/workflows/unit_tests.yml @@ -67,7 +67,9 @@ jobs: env: COVERAGE_PROCESS_START: ${{ github.workspace }}/pyproject.toml COVERAGE_FILE: ${{ github.workspace }}/.coverage - run: pip install nox uv && nox -s "unit-3.12(torch_214, tf_latest)" + # Everything but PRs also runs the slow tests marked ``manual`` (e.g. AutoNAS). + PYTEST_EXTRA: ${{ github.event_name != 'pull_request' && '--run-manual' || '' }} + run: pip install nox uv && nox -s "unit-3.12(torch_214, tf_latest)" -- $PYTEST_EXTRA - name: Upload coverage reports to Codecov uses: codecov/codecov-action@v7 with: diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 98bcfd6915d..38f8f6f5c50 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -13,8 +13,11 @@ Changelog *Quantization* +- Add metadata-only export specifications and non-mutating per-weight quantizer-state capture. +- HF export records canonical activation schemes, group/block sizes and weight-only MXFP4 metadata; mixed-format exports retain unquantized-layer exclusions, and non-square FP8 block layouts are rejected. - Add Q8_0 weight-only quantization with 32-value GGML blocks, packed unified HF and Megatron export, and a built-in ``q8_0`` PTQ recipe. - Add Hugging Face PTQ calibration and export support for Nemotron-H MTP modules, preserving calibrated expert input scales. +- Backfill checkpoint aliases for three more published NVFP4 releases: ``moonshotai/Kimi-K2.7-Code``, ``google/gemma-4-26B-A4B-it`` and ``google/diffusiongemma-26B-A4B-it``. Each imports an existing general recipe wholesale rather than copying its body. - Backfill checkpoint aliases for nine more published NVFP4 releases, so each is reachable from its source model's hub path: ``zai-org/GLM-5.1`` and ``GLM-5.2``, ``MiniMaxAI/MiniMax-M2.5`` and ``MiniMax-M3``, ``deepseek-ai/DeepSeek-V3.1`` and ``DeepSeek-V3.2``, ``Qwen/Qwen3-235B-A22B-Instruct-2507`` and ``-Thinking-2507``, and ``Qwen/Qwen3.6-27B``. Each imports an existing general or architecture recipe wholesale rather than copying its body. - Add composed Hugging Face AutoQuantize recipes that run fixed PTQ or weight AutoQuantize before a separate KV-cache AutoQuantize stage, with independent resumable checkpoints for the weight and @@ -47,6 +50,8 @@ Changelog *Megatron Framework (M-LM / M-Bridge)* +- Add ``freeze_mtp`` to the Megatron ``DistillationConfig`` (``--freeze_mtp`` in ``examples/megatron_bridge/distill.py``): the student's MTP head keeps its input weights and its loss is skipped, so distillation trains the rest of the model on the distillation loss alone. +- Add ``output_format: chat_jsonl`` to ``prepare_megatron_data_blend``: it writes a weighted blend as ``train.jsonl`` / ``validation.jsonl`` chat data for ``distill.py --sft --sft_hf_dataset``, so the distillation loss can be restricted to assistant tokens. See the `Megatron data preparation guide `_. - Add optional Top-P (nucleus) truncation to the Megatron top-k logits KD loss (now ``TopLogitsKLLoss``) via ``logit_kl_top_p`` and ``logit_kl_top_p_min_k`` in ``DistillationConfig``: after the global Top-K selection, only the smallest prefix whose cumulative teacher probability reaches ``top_p`` (with a floor of ``min_k`` entries) contributes to the KL, mirroring ``--logits-save-top-p`` / ``--logits-save-top-p-min-k`` in Megatron-LM's logits saver. - Add an end-to-end W4A4 NVFP4 PTQ and QAD tutorial for Qwen3.6-35B-A3B also covering evaluation and vLLM throughput benchmarking. See `examples/megatron_bridge/tutorials/Qwen3.6-35B-A3B/README.md `_ for details. - Add ``--mlflow `` to the ``examples/megatron_bridge`` scripts that write a checkpoint -- ``prune_minitron.py``, ``quantize.py``, ``distill.py``, ``export_quantized_megatron_to_hf.py`` and ``export_distilled_megatron_to_hf.py`` (MLflow's own ``MLFLOW_TRACKING_URI`` is honoured too). Each run records the invocation, its arguments as searchable params and its log -- a distillation records its training metrics instead of the rank-0 log -- and writes ``.experiment.json`` into the checkpoint it produced, so a pruning, a quantization, the distillation that refines its checkpoint and the export that deploys it can be traced to one another; uploading the checkpoints themselves stays off unless ``--mlflow_log_checkpoints`` is passed. @@ -63,6 +68,7 @@ Changelog **Backward Breaking Changes** +- ``modelopt.torch.utils.distributed.broadcast``/``allgather`` now reject objects of unregistered custom classes. Register them with ``torch.serialization.add_safe_globals``. - Bump minimum transformers version to ``5.5`` instead of ``4.57``; transformers 4.x is no longer supported. Upgrade with ``pip install -U "nvidia-modelopt[hf]"``. - ``modelopt.torch.distill.plugins.megatron.TopKLogitsKLLoss`` (``logit_kl_topk`` in ``DistillationConfig``) is renamed to ``TopLogitsKLLoss``, keeping the old name as a deprecated alias. It now normalizes both distributions over the full vocabulary instead of re-normalizing over the Top-K entries, and always appends a "ghost" token holding the probability mass outside the Top-K to both student and teacher (matching Megatron-LM's offline cached-logits KD loss). Loss values change for existing ``logit_kl_topk`` runs. - ``LogitsAndIntermediatesLossBalancer`` (Megatron distillation plugin) no longer rescales the distillation loss to the magnitude of the LM loss when the LM loss is included. The total is now the fixed convex combination ``(1 - alpha) * lm_loss + alpha * kd_loss`` with ``DistillationConfig.kd_loss_alpha`` (in [0, 1]), matching Megatron-LM's offline cached-logits KD. The default ``kd_loss_alpha=1.0`` skips the LM loss, as before. ``DistillationConfig.skip_lm_loss`` and ``kd_loss_scale`` are removed and raise ``ValueError`` if passed; set ``kd_loss_alpha`` instead. In ``examples/megatron_bridge/distill.py``, ``--kd_loss_alpha`` replaces ``--no_skip_lm_loss`` and ``--kd_loss_scale``. @@ -124,7 +130,16 @@ Changelog **Bug Fixes** +- Harden the distributed object broadcast/allgather helpers and the Megatron expert-parallel export gather so a compromised peer rank can no longer execute code through a crafted pickle payload: payloads load with ``weights_only=True`` and the export gathers raw bytes instead of pickled objects. +- Fix LSQ clamping small, representable FP8 block scales too aggressively, which could zero quantized values and suppress scale gradients. +- Fix ONNX AutoTune benchmark initialization with TensorRT 11 while preserving strongly typed networks on TensorRT 10. +- Keep vision-block Linears in high precision in the ``qwen3_vl`` and ``qwen3_5`` FP8 vision recipes so exported checkpoints avoid a vLLM vision-attention shape failure. The recipes still quantize primary and deepstack merger Linears; re-quantize affected checkpoints before deployment. +- Fix unified HF export writing a weight name the source checkpoint does not use when a model declares a backwards-compatibility *load* rename as a reversible ``WeightRenaming``. Reversing such a rule rewrote the canonical name to a legacy spelling, and a deployment runtime keyed on the canonical name then failed to load the quantized checkpoint. Each exported name is now resolved against the names the source checkpoint actually carries, for both the exported tensors and the quantization config's module references. +- Fix exports using the built-in Transformers NemotronH implementation retaining incompatible legacy remote-code config and model mappings. Re-export affected checkpoints before deployment. +- Fix ``export_hf_checkpoint`` dropping Diffusers pipeline settings, such as Wan 2.2 ``boundary_ratio``, from ``model_index.json`` when ``components`` is given. Re-export affected checkpoints to a new directory, since component exports do not overwrite an existing ``model_index.json``. +- Fix INT8 real quantization (``mtq.compress``) of FP16 weights zeroing or clipping channels with a small amax: their scales underflowed in FP16. Scales of newly compressed INT8 weights are now stored in FP32. - Fix Hugging Face exports classifying registered Qwen VLM parameters as unplaced source weights, which overwrote quantized tensors with original weights and incorrectly excluded their modules from quantization. Re-export affected checkpoints. +- Right-pad labeled calibration samples to avoid fully masked causal rows that can produce non-finite cuDNN attention gradients during AutoQuantize. Generation previews remove padding before decoding. - Fix calibration silently dropping samples when an out-of-memory error forces the batch to be split and a smaller working batch size is already known: the slices stopped tiling the batch, so the rows between the old and the new width never reached the model and the collected @@ -137,10 +152,17 @@ Changelog - Fix ``hf_ptq`` overwriting a model's existing ``pad_token`` with ``eos_token`` when the model already has a valid padding token. The exported tokenizer now preserves the source model's padding configuration. - Fix Hugging Face exports dropping off-index safetensors such as GLM-4.7's ``mtp.safetensors``. - Fix Megatron-Core checkpoint saving for quantized grouped MoE experts when tensor and expert parallelism are both enabled. +- Fix FP8 fake quantization treating a positive amax at or below 2**-24 as zero, which silently zeroed such tensors and channels and every block of a static NVFP4 weight (``NVFP4StaticQuantizer``, including LSQ) with a global amax below ~3.6e-7. Fake quantization now matches the exported checkpoint. +- Fix ``megatron_preprocess_data`` with ``workers > 1`` segfaulting when run in a process that has already initialized CUDA or other native threads (for example in-process after GPU work): the tokenization worker pool now uses the ``forkserver`` start method instead of ``fork``. - Fix unified HuggingFace export of RADIO-based VLMs retaining post-conversion vision and projector names instead of restoring the hub layout; deployment loaders could skip those weights. Streaming, offload, FSDP2, and layerwise export now reject reverse conversions that regroup tensors; use resident export for those models. +- Fix that same reverse rename double-prefixing tensors copied from the source checkpoint rather + than loaded by the model, when a conversion mapping renames a module to a suffix of its original + name. Those tensors already carry the hub name, so C-RADIO towers exported + ``vision_model.radio_model.radio_model.input_conditioner.norm_mean``. Re-export affected + checkpoints. - Fix shared ONNX export metadata and Diffusers attention policy: every ``NVFP4QuantExporter`` post-process now upgrades the default-domain opset to at least 23, all FP8 custom-op exports re-run ONNX shape/type inference after setting output metadata, and quantized SDPA derives FP8 MHA enablement from the live Q/K/V quantizers instead of honoring a caller-set ``_disable_fp8_mha`` attribute. - Fix ONNX FP16 conversion failing to preserve public output types when type inference changes a graph output declaration before output casts are inserted. - Fix ``examples/hf_ptq/hf_ptq.py`` discarding a completed PTQ run (no checkpoint exported) when the optional post-quantization sanity-check ``generate()`` call raised, for example because ``device_map="auto"`` placed part of the model on CPU. That failure is now caught and only skips the sanity check; export proceeds regardless. @@ -158,6 +180,10 @@ Changelog - Fix two issues in the vLLM offline hidden-state dump (``examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py``) that only surface on large runs. **Resume:** the filter that skips conversations whose ``.pt`` already exists now runs with ``load_from_cache_file=False``. It depends on on-disk state, which is not part of the fingerprint ``datasets`` computes from the function and the dataset, so with a persistent HF cache reused across a resumed or requeued run the cached "keep everything" result from an earlier run was replayed and the dump re-generated and overwrote conversations it had already finished (observed: tens of thousands of ``.pt`` rewritten while the output count stayed flat). **Staging:** generation is now chunked (``--save-chunk-size``, default 256), so each chunk is saved and its staged hidden states freed before the next chunk is generated. Previously the whole dataset was generated before anything was saved, which kept every conversation staged in the connector's ``shared_storage_path`` (``/dev/shm``, i.e. RAM, by default) at once and exhausted it partway through large dumps. Chunking also makes the dump incrementally durable, so an interrupted run keeps its finished conversations and resumes from them. The save path now also frees each conversation's staged hidden states in a ``finally``, so a conversation skipped mid-loop (e.g. a short ``loss_mask``) can no longer leak its staging file, and conversation ids are validated as plain filenames before being used to build output paths. - Fix YAML config I/O decoding with the locale codepage instead of UTF-8, which made a config containing any non-ASCII byte fail to load on a machine whose locale is not UTF-8 (notably Windows, where the default is cp1252). ``modelopt/recipe/loader.py``, the two ONNX autotune state files, the two transformers config readers, the distill config and the puzzletron profile now pass ``encoding="utf-8"`` explicitly. Only the YAML config paths are covered: these are the files most likely to carry non-ASCII text in comments, model names or paths, and the only ones read inside a user's process. - Hybrid (e.g. Nemotron-H) checkpoints saved by the ``examples/megatron_bridge`` scripts now record their layer spec in ``run_config.yaml`` in a form that reloads, so they can be converted to HuggingFace; a checkpoint saved by an earlier release still needs its ``model.hybrid_stack_spec`` block replaced by hand. +- Fix ONNX preprocessing with custom TensorRT plugin libraries on TensorRT 11. +- Fix GPTQ with Megatron data parallelism: each data-parallel replica of a module now builds its Hessian from the whole calibration set instead of only its own shard, so all replicas reach the same quantized weights. +- Fix ONNX quantization with the default ``concat_elimination`` pass crashing in calibration with ``KeyError: '_ReduceMin'`` when a ``Concat`` input is a constant, such as the class token of a ViT. +- Fix PTQ of native Transformers FP8 experts to apply source weight scales during calibration and export. Re-quantize affected checkpoints from the original FP8 source. 0.47.0 (2026-09-23) ^^^^^^^^^^^^^^^^^^^ diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 3f344153850..d778f1bb10b 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -183,6 +183,16 @@ nox -s "unit-3.12(torch_211, tf_latest)" keeping CI/CD fast and lean. Default to one test per behavior: use `@pytest.mark.parametrize` instead of near-duplicate test functions, and don't add a per-branch test for a helper that a higher-level test already exercises. More tests is not better coverage — every checked-in test is CI time and maintenance forever. + Before adding a test, check whether an existing unit, GPU, or example test already exercises the path; if so, + extend it instead. Do not add tests that only restate the implementation, assert on mocks, pin log/warning + wording, or re-cover a branch of a helper whose caller is tested. A PR that adds many tests should be able to + say what regression each one catches. +- **Keep `tests/unit` fast — the `linux` job has a hard 20-minute limit** and runs the whole suite serially. + A unit test should take well under 5s (check with `pytest --durations=20`). Tests that spawn processes or + subprocesses (`spawn_multiprocess_job`, `subprocess.run` importing torch), train for many steps, or build + real-size models belong in `tests/gpu*`; batch several cases into one spawn rather than one spawn per case, + or reuse the session-scoped `cpu_dist_workers` fixture for multi-rank CPU tests. + Slow tests for features without active maintenance (e.g. AutoNAS) should be `@pytest.mark.manual` or nightly-only. - **Exercise the behavior a test claims to validate.** Mocks are useful for focused interface and wiring coverage, but replacing the implementation under test does not validate its real behavior. Include an end-to-end test that runs the actual implementation whenever the test claims backend or runtime behavior. For example, a test of diff --git a/docs/source/guides/9_autotune.rst b/docs/source/guides/9_autotune.rst index 02f79395752..2e79e482692 100644 --- a/docs/source/guides/9_autotune.rst +++ b/docs/source/guides/9_autotune.rst @@ -569,10 +569,7 @@ Using TensorRT Python API # Create builder and logger logger = trt.Logger(trt.Logger.WARNING) builder = trt.Builder(logger) - network = builder.create_network( - 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) - | 1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED) - ) + network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)) parser = trt.OnnxParser(network, logger) # Parse optimized ONNX model diff --git a/examples/dataset/MEGATRON_DATA_PREP.md b/examples/dataset/MEGATRON_DATA_PREP.md index 91b89375907..c920ed02567 100644 --- a/examples/dataset/MEGATRON_DATA_PREP.md +++ b/examples/dataset/MEGATRON_DATA_PREP.md @@ -5,6 +5,7 @@ | From JSONL files | Tokenize local JSONL files | \[[Link](#from-jsonl-files)\] | | From Hugging Face Hub | Stream or download HF datasets and tokenize | \[[Link](#from-hugging-face-hub)\] | | Token-budgeted data blends | Prepare weighted subsets for fast experiments | \[[Link](#prepare-token-budgeted-data-blends)\] | +| Chat JSONL blends | Weighted chat data for assistant-only loss | \[[Link](#chat-jsonl-blends-for-assistant-only-loss)\] | | `reasoning_content` for Post-Training v3 | Control how chain-of-thought traces are handled | \[[Link](#reasoning_content-for-post-training-v3-datasets)\] | | Nemotron Pre/Post-Training Datasets | Ready-to-run commands for all Nemotron datasets | \[[Link](#ready-to-run-tokenization-commands)\] | @@ -172,6 +173,54 @@ The output contains tokenized Megatron `.bin`/`.idx` files, `data_blend.txt` wit and `config.yaml` recording how the blend was generated. The final token count can slightly exceed the target because the final document from each source is kept whole. +### Chat JSONL blends for assistant-only loss + +Tokenized `.bin`/`.idx` data carries no loss mask, so `--data_paths` training computes the loss on every token. +To compute it on assistant tokens only, set `output_format: chat_jsonl`. The same sources are then written, +untokenized, as `train.jsonl` and `validation.jsonl` with one `{"messages": [...]}` conversation per line, for +Megatron-Bridge's chat path (`examples/megatron_bridge/distill.py --sft --sft_hf_dataset /train.jsonl +--sft_hf_validation /validation.jsonl --sft_loss_mode assistant`), which applies the model's chat template +when it loads the data. Run it in an environment with Megatron-Bridge, whose text-only chat filter selects the rows. + +```yaml +output_dir: /path/to/chat_blend +output_format: chat_jsonl +target_samples: 36_800 # Total conversations; required, since the selected rows are mixed in memory. +shuffle: false # Optional; take the first rows of each split in order (default: buffer shuffle). +seed: 1234 # Optional; seeds the shuffle, the cross-source mix and the validation hold-out. +reasoning_content: native # Optional; native (default), inline or strip. +sources: + - hf_dataset: nvidia/Nemotron-Pretraining-SFT-v1 + config: Nemotron-SFT-General + split: train + content_field: text + weight: 20 + - hf_dataset: nvidia/Nemotron-SFT-Math-v3 + config: default + split: train + content_field: messages + weight: 17 +``` + +Each source contributes `round(target_samples * weight / sum(weights))` conversations that contain only text +and have an assistant turn; a source that rounds to zero is skipped. `content_field` must be `messages` or `text`: + +- A `text` row becomes one assistant turn, or is split into turns if it uses the `User` / + `Assistant` / `System` chat markers. A row with a marker for any other role (such as a SteerLM + label line) or with an empty turn other than `System` is dropped. +- `messages` and a row's `tools` are prepared as in the tokenizing mode (JSON-string `tool_calls` arguments + become dicts), and a top-level `system` prompt is kept for Megatron-Bridge to prepend. The one difference is + that `reasoning_content` stays on the message by default (`native`), because the chat template is applied + later, at training time. Use `inline` if your model's chat template ignores `reasoning_content`, so the + reasoning is kept as a `` block in `content` (see [below](#reasoning_content-for-post-training-v3-datasets)), + or `strip` to drop it. + +`files` may be `.jsonl` or `.jsonl.gz`; malformed records in them are skipped with a warning. The conversations are +mixed and `min(1000, n // 20)` of them are held out for validation. A source that fails to load, has no usable +rows, or runs out of them stops the run. Each source's selection is cached under `/parts`, keyed by +the settings that choose its rows, so an interrupted run resumes and a changed configuration selects afresh. +Nothing prunes `parts`; delete it once `train.jsonl` is written to reclaim the space. + ## `reasoning_content` for Post-Training v3 Datasets v3 datasets include a `reasoning_content` field in assistant messages (chain-of-thought separate from diff --git a/examples/gpt-oss/requirements.txt b/examples/gpt-oss/requirements.txt index e8618d6ca5e..7462574d4d0 100644 --- a/examples/gpt-oss/requirements.txt +++ b/examples/gpt-oss/requirements.txt @@ -1,5 +1,5 @@ kernels>=0.9.0,<0.13 -trackio +trackio>=0.41 # transformers>=5.3 avoids CVE-2026-4372 (RCE via the `kernels` Hub download path, which this example installs) transformers>=5.3 # TRL 1.15 enables a GPU-only fused LM head, breaking the CPU SFT smoke test. diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 1d4690bf00a..7e3c7d8fc44 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -516,6 +516,10 @@ formats in `kv_cache_quantized_layers`. `mtq.auto_quantize` returns the sensitiv selected recipe in its search state; `--kv_auto_quantize_checkpoint` stores that resumable state, including the candidate quantizer tensors needed for replay. +KV-only recipes leave weights in their original precision, so `hf_quant_config.json` contains +`quantization.quant_algo: null` and `quantized_layers: {}`. Validate their quantization through +`kv_cache_quant_algo`, `kv_cache_quantized_layers`, and the exported K/V scaling tensors. + KV sensitivity scoring runs one reference forward plus one forward per eligible-layer candidate for every scoring step. Choose the recipe's `auto_quantize.score_size` with the number of eligible layers and candidates in mind. Peak scoring memory also grows with the number of selected tokens diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index db29d5fbf61..d7a14a68192 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -286,11 +286,15 @@ def _validate_recipe_calibration(args: argparse.Namespace, recipe) -> None: if not isinstance(recipe, ModelOptPTQRecipe) or args.calib_with_images: return + visual_input_names = { + "*visual.*input_quantizer", + "*visual.*merger*.input_quantizer", + } visual_input_entry = next( ( entry for entry in reversed(recipe.quantize.quant_cfg) - if entry.quantizer_name == "*visual.*input_quantizer" + if entry.quantizer_name in visual_input_names ), None, ) @@ -757,6 +761,9 @@ def pre_quantize( preview_input_ids = batch[input_key][0:1] # Pass attention_mask to generate(): HF cannot infer it when pad_token == eos_token. preview_attention_mask = batch["attention_mask"][0:1] if "attention_mask" in batch else None + if input_key == "input_ids" and preview_attention_mask is not None: + preview_input_ids = preview_input_ids[:, preview_attention_mask[0].bool()] + preview_attention_mask = torch.ones_like(preview_input_ids) # Generate preview before quantization if args.skip_generate: diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index f71114cf1f3..aec3425295a 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -213,6 +213,11 @@ def get_args(): default=1.0, help="KD loss weight alpha in (1 - alpha) * lm_loss + alpha * kd_loss. 1.0 skips the LM loss entirely.", ) + parser.add_argument( + "--freeze_mtp", + action="store_true", + help="Keep the student's MTP head at its input weights: skip its loss and freeze it.", + ) parser.add_argument( "--no_async_save", action="store_true", @@ -493,6 +498,7 @@ def _build_model_provider(hf_path, load_weights=True, moe_grouped_gemm=True): logit_kl_topk=args.logit_kl_top_k, logit_kl_top_p=args.logit_kl_top_p, logit_kl_top_p_min_k=args.logit_kl_top_p_min_k, + freeze_mtp=args.freeze_mtp, ) # HF VLM configs expose ``vision_config``; Megatron-Bridge nests the text model under diff --git a/examples/onnx_ptq/custom_op_plugin/plugin/IdentityConvPlugin.cpp b/examples/onnx_ptq/custom_op_plugin/plugin/IdentityConvPlugin.cpp index 7b8a41d7b90..ecd7809b7c7 100644 --- a/examples/onnx_ptq/custom_op_plugin/plugin/IdentityConvPlugin.cpp +++ b/examples/onnx_ptq/custom_op_plugin/plugin/IdentityConvPlugin.cpp @@ -179,12 +179,14 @@ nvinfer1::DataType IdentityConv::getOutputDataType(int32_t index, return inputTypes[0]; } +#if NV_TENSORRT_MAJOR < 11 bool IdentityConv::isOutputBroadcastAcrossBatch(int32_t outputIndex, bool const *inputIsBroadcasted, int32_t nbInputs) const noexcept { return false; } bool IdentityConv::canBroadcastInputAcrossBatch(int32_t inputIndex) const noexcept { return false; } +#endif int32_t IdentityConv::enqueue(int32_t batchSize, void const *const *inputs, void *const *outputs, void *workspace, cudaStream_t stream) noexcept { diff --git a/examples/onnx_ptq/custom_op_plugin/plugin/IdentityConvPlugin.h b/examples/onnx_ptq/custom_op_plugin/plugin/IdentityConvPlugin.h index da6128424e0..1f111fb5a05 100644 --- a/examples/onnx_ptq/custom_op_plugin/plugin/IdentityConvPlugin.h +++ b/examples/onnx_ptq/custom_op_plugin/plugin/IdentityConvPlugin.h @@ -88,10 +88,12 @@ class IdentityConv : public nvinfer1::IPluginV2IOExt { char const *getPluginNamespace() const noexcept override; +#if NV_TENSORRT_MAJOR < 11 bool isOutputBroadcastAcrossBatch(int32_t outputIndex, bool const *inputIsBroadcasted, int32_t nbInputs) const noexcept override; bool canBroadcastInputAcrossBatch(int32_t inputIndex) const noexcept override; +#endif private: void deserialize(uint8_t const *data, size_t length); diff --git a/experimental/dms/README.md b/experimental/dms/README.md index 2107c80079a..edaa34bd3d6 100644 --- a/experimental/dms/README.md +++ b/experimental/dms/README.md @@ -48,7 +48,7 @@ This single install provides everything needed for training and evaluation (incl ### Train DMS adapters -**Note:** The number of GPUs determines the effective batch size. The configuration below was tested on a DGX H100 with 8× H100 80GB GPUs. For debugging with a smaller compute budget (e.g., a single RTX 5090), see [`scripts/train_small_debug.sh`](scripts/train_small_debug.sh). +**Note:** The number of GPUs determines the effective batch size. The configuration below was tested on a DGX H100 with 8× H100 80GB GPUs. For debugging with a smaller compute budget (e.g., a single RTX 5090), see [`scripts/train_small.sh`](scripts/train_small.sh). ```bash bash scripts/train.sh configs/qwen3_8b.yaml diff --git a/modelopt/onnx/export/_dynamo_adapter.py b/modelopt/onnx/export/_dynamo_adapter.py new file mode 100644 index 00000000000..193ddc33d4e --- /dev/null +++ b/modelopt/onnx/export/_dynamo_adapter.py @@ -0,0 +1,500 @@ +# 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. + +"""Dynamo-only normalization and validation for quantized ONNX graphs.""" + +from collections.abc import Iterable + +import numpy as np +import onnx +from onnx import numpy_helper + +from modelopt.onnx.quantization.graph_indexing import ( + get_tensor_consumer_nodes, + get_tensor_producer_nodes, +) + +__all__ = [] + +_DYNAMO_INT4_NODE_PREFIX = "__modelopt_dynamo_int4__" +_DYNAMO_MXFP8_NODE_PREFIX = "__modelopt_dynamo_mxfp8__" + + +def _attribute(node: onnx.NodeProto, name: str): + attribute = next((item for item in node.attribute if item.name == name), None) + return onnx.helper.get_attribute_value(attribute) if attribute is not None else None + + +def _single_consumer(consumers, name: str, weight_name: str) -> onnx.NodeProto: + matches = consumers.get(name, []) + if len(matches) != 1: + raise NotImplementedError( + f"Unsupported Dynamo quantized weight topology for '{weight_name}': " + f"expected one consumer; found {len(matches)}." + ) + return matches[0] + + +def _constant_tensor(producers, name: str) -> onnx.TensorProto: + producer = producers.get(name) + if isinstance(producer, onnx.TensorProto): + return producer + if isinstance(producer, onnx.NodeProto) and producer.op_type == "Constant": + value = _attribute(producer, "value") + if isinstance(value, onnx.TensorProto): + return value + raise NotImplementedError("Dynamo ONNX export requires constant weight shapes and scales.") + + +def _replace_initializer(graph: onnx.GraphProto, tensor: onnx.TensorProto) -> None: + existing = next((item for item in graph.initializer if item.name == tensor.name), None) + if existing is None: + graph.initializer.append(tensor) + else: + existing.CopyFrom(tensor) + + +def _remove_nodes(graph: onnx.GraphProto, nodes: Iterable[onnx.NodeProto]) -> None: + output_names = {output for node in nodes for output in node.output} + for index in reversed(range(len(graph.node))): + if not output_names.isdisjoint(graph.node[index].output): + del graph.node[index] + + +def _path_after_marker( + consumers, value_name: str, weight_name: str +) -> tuple[list[onnx.NodeProto], onnx.NodeProto]: + path = [] + passthrough_order = {"Identity": 0, "Reshape": 1, "Cast": 1, "Transpose": 2} + seen = set() + previous_order = -1 + node = _single_consumer(consumers, value_name, weight_name) + while node.op_type in passthrough_order: + current_order = passthrough_order[node.op_type] + if node.op_type in seen or current_order < previous_order: + raise NotImplementedError( + f"Unsupported Dynamo quantized weight topology for '{weight_name}': " + f"unexpected {node.op_type} ordering." + ) + seen.add(node.op_type) + previous_order = current_order + path.append(node) + value_name = node.output[0] + node = _single_consumer(consumers, value_name, weight_name) + if node.op_type not in {"MatMul", "Gemm"} or len(node.input) < 2 or node.input[1] != value_name: + raise NotImplementedError( + f"Unsupported Dynamo quantized weight topology for '{weight_name}': " + "expected terminal MatMul/Gemm at input 1." + ) + return path, node + + +def _weight_source(graph, producers, consumers, marker: onnx.NodeProto): + initializers = {item.name: item for item in graph.initializer} + marker_input = marker.input[0] + producer = producers.get(marker_input) + reshape = ( + producer if isinstance(producer, onnx.NodeProto) and producer.op_type == "Reshape" else None + ) + weight_name = reshape.input[0] if reshape is not None else marker_input + if weight_name not in initializers: + raise NotImplementedError( + f"Unsupported Dynamo quantized weight '{marker_input}': " + "data input must be a static initializer." + ) + _single_consumer(consumers, weight_name, weight_name) + if reshape is not None: + _single_consumer(consumers, marker_input, weight_name) + _constant_tensor(producers, reshape.input[1]) + return initializers[weight_name], reshape + + +def _is_static_value(producers, name: str, cache: dict[str, bool]) -> bool: + stack: list[tuple[str, list[str] | None]] = [(name, None)] + visiting = set() + while stack: + current, inputs = stack.pop() + if current in cache: + continue + if inputs is not None: + cache[current] = all(cache[input_name] for input_name in inputs) + visiting.remove(current) + continue + producer = producers.get(current) + if isinstance(producer, onnx.TensorProto): + cache[current] = True + elif not isinstance(producer, onnx.NodeProto): + cache[current] = False + elif producer.op_type == "Constant": + cache[current] = True + elif current in visiting or not any(producer.input): + cache[current] = False + else: + inputs = [input_name for input_name in producer.input if input_name] + visiting.add(current) + stack.append((current, inputs)) + stack.extend((input_name, None) for input_name in reversed(inputs)) + return cache[name] + + +def _has_static_weight_source(producers, name: str, cache: dict[str, bool]) -> bool: + producer = producers.get(name) + if isinstance(producer, onnx.TensorProto): + return True + if ( + isinstance(producer, onnx.NodeProto) + and producer.op_type == "Reshape" + and isinstance(producers.get(producer.input[0]), onnx.TensorProto) + ): + return True + # Shape and index inputs can be dynamic; only input 0 carries the weight data. + data_name = name + seen = set() + while ( + isinstance(producer, onnx.NodeProto) + and producer.op_type + in { + "Identity", + "Reshape", + "Cast", + "Transpose", + "Slice", + "Gather", + "GatherElements", + "GatherND", + "Expand", + "Squeeze", + "Unsqueeze", + "Flatten", + } + and producer.input[0] not in seen + ): + data_name = producer.input[0] + seen.add(data_name) + producer = producers.get(data_name) + if _is_static_value(producers, data_name, cache): + raise NotImplementedError( + f"Unsupported Dynamo quantized weight '{name}': unsupported static or dynamic transformation " + "before the quantization marker." + ) + return False + + +def _validate_restoring_reshape( + path, weight: onnx.TensorProto, producers, format_name: str +) -> None: + reshape = next((node for node in path if node.op_type == "Reshape"), None) + if reshape is None: + return + restored_shape = [ + int(dim) for dim in numpy_helper.to_array(_constant_tensor(producers, reshape.input[1])) + ] + if restored_shape != list(weight.dims): + raise NotImplementedError( + f"Unsupported Dynamo {format_name} weight '{weight.name}': restoring Reshape does " + "not match the source initializer shape." + ) + + +def _fold_prefix_reshape( + graph: onnx.GraphProto, + marker: onnx.NodeProto, + weight: onnx.TensorProto, + reshape: onnx.NodeProto | None, + producers, +) -> None: + if reshape is None: + return + shape = [ + int(dim) for dim in numpy_helper.to_array(_constant_tensor(producers, reshape.input[1])) + ] + values = numpy_helper.to_array(weight).reshape(shape) + _replace_initializer(graph, numpy_helper.from_array(values, weight.name)) + marker.input[0] = weight.name + _remove_nodes(graph, [reshape]) + + +def _split_scale( + graph, producers, consumers, marker: onnx.NodeProto, *, materialize_initializer: bool +) -> None: + if len(marker.input) < 2: + return + scale_name = marker.input[1] + tensor = _constant_tensor(producers, scale_name) + if ( + materialize_initializer + and len(consumers.get(scale_name, [])) == 1 + and isinstance(producers.get(scale_name), onnx.TensorProto) + ): + return + array = np.asarray(numpy_helper.to_array(tensor)) + unique_name = f"{scale_name}_{marker.output[0]}" + marker.input[1] = unique_name + scale_tensor = numpy_helper.from_array(array, unique_name) + if materialize_initializer: + _replace_initializer(graph, scale_tensor) + else: + marker_index = next( + index for index, node in enumerate(graph.node) if node.output[0] == marker.output[0] + ) + graph.node.insert( + marker_index, + onnx.helper.make_node( + "Constant", [], [unique_name], name=f"{unique_name}_constant", value=scale_tensor + ), + ) + + +def _remove_identity(graph, path: list[onnx.NodeProto], marker_output: str) -> list[onnx.NodeProto]: + if not path or path[0].op_type != "Identity": + return path + identity = path.pop(0) + replacement = identity.output[0] + for node in graph.node: + for index, input_name in enumerate(node.input): + if input_name == replacement: + node.input[index] = marker_output + _remove_nodes(graph, [identity]) + return path + + +def _normalize_int4_path( + graph: onnx.GraphProto, + marker: onnx.NodeProto, + weight: onnx.TensorProto, + prefix_reshape: onnx.NodeProto | None, + path: list[onnx.NodeProto], + terminal: onnx.NodeProto, + producers, +) -> None: + original_shape = list(weight.dims) + block_size = int(_attribute(marker, "block_size") or 0) + if block_size <= 0 or not original_shape or original_shape[-1] % block_size: + raise NotImplementedError( + f"Unsupported Dynamo INT4 weight '{weight.name}': invalid block size {block_size}." + ) + + if prefix_reshape is not None: + blocked_shape = [ + int(dim) + for dim in numpy_helper.to_array(_constant_tensor(producers, prefix_reshape.input[1])) + ] + if not blocked_shape or blocked_shape[-1] != block_size: + raise NotImplementedError( + f"Unsupported Dynamo INT4 weight '{weight.name}': pre-Reshape must use " + f"block size {block_size}." + ) + else: + blocked_shape = [-1, block_size] + marker_rank = len(blocked_shape if prefix_reshape is not None else original_shape) + marker_axis = _attribute(marker, "axis") + if (1 if marker_axis is None else marker_axis) not in {-1, marker_rank - 1}: + raise NotImplementedError( + f"Unsupported Dynamo INT4 weight '{weight.name}': only last-axis blocks are supported." + ) + _validate_restoring_reshape(path, weight, producers, "INT4") + blocked_weight = numpy_helper.to_array(weight).reshape(blocked_shape) + _replace_initializer(graph, numpy_helper.from_array(blocked_weight, weight.name)) + marker.input[0] = weight.name + + path = _remove_identity(graph, path, marker.output[0]) + reshape = next((node for node in path if node.op_type == "Reshape"), None) + + block_axis = len(original_shape) - 1 + marker.attribute.append(onnx.helper.make_attribute("_target_shape", original_shape)) + transpose = next((node for node in path if node.op_type == "Transpose"), None) + if transpose is not None: + perm = list(_attribute(transpose, "perm") or []) + if sorted(perm) != list(range(len(original_shape))): + raise NotImplementedError( + f"Unsupported Dynamo INT4 weight '{weight.name}': invalid transpose permutation." + ) + marker.attribute.append(onnx.helper.make_attribute("_transpose_perm", perm)) + block_axis = perm.index(block_axis) + axis = next((attr for attr in marker.attribute if attr.name == "axis"), None) + if axis is None: + marker.attribute.append(onnx.helper.make_attribute("axis", block_axis)) + else: + axis.i = block_axis + + cast = next((node for node in path if node.op_type == "Cast"), None) + if cast is not None: + cast.input[0] = marker.output[0] + weight_output = cast.output[0] + else: + weight_output = marker.output[0] + terminal.input[1] = weight_output + _remove_nodes( + graph, + [node for node in (prefix_reshape, reshape, transpose) if node is not None], + ) + + scale = next(item for item in graph.initializer if item.name == marker.input[1]) + scale_values = numpy_helper.to_array(scale).reshape(-1, 1) + _replace_initializer(graph, numpy_helper.from_array(scale_values, scale.name)) + + +def _normalize_nvfp4_path( + graph: onnx.GraphProto, + marker: onnx.NodeProto, + weight: onnx.TensorProto, + prefix_reshape: onnx.NodeProto | None, + path: list[onnx.NodeProto], + terminal: onnx.NodeProto, + producers, +) -> None: + path = _remove_identity(graph, path, marker.output[0]) + removable = [] + prefix_nodes = [] + reshape = next((node for node in path if node.op_type == "Reshape"), None) + if prefix_reshape is None and reshape is not None: + raise NotImplementedError( + f"Unsupported Dynamo NVFP4 weight '{weight.name}': restoring Reshape requires a " + "pre-Reshape." + ) + if prefix_reshape is not None: + if reshape is None: + raise NotImplementedError( + f"Unsupported Dynamo NVFP4 weight '{weight.name}': missing restoring Reshape." + ) + _validate_restoring_reshape(path, weight, producers, "NVFP4") + marker.input[0] = weight.name + prefix_nodes.append(prefix_reshape) + removable.append(reshape) + path.remove(reshape) + cast = next((node for node in path if node.op_type == "Cast"), None) + if cast is not None: + removable.append(cast) + path.remove(cast) + + for node in removable: + output = node.output[0] + for candidate in graph.node: + for index, input_name in enumerate(candidate.input): + if input_name == output: + candidate.input[index] = marker.output[0] + _remove_nodes(graph, [*prefix_nodes, *removable]) + + if path and path[0].op_type == "Transpose" and terminal.op_type == "Gemm": + transpose = path[0] + if list(_attribute(transpose, "perm") or []) != [1, 0]: + raise NotImplementedError( + f"Unsupported Dynamo NVFP4 weight '{weight.name}': only a rank-2 transpose is supported." + ) + trans_b = next((item for item in terminal.attribute if item.name == "transB"), None) + if trans_b is None: + terminal.attribute.append(onnx.helper.make_attribute("transB", 1)) + else: + trans_b.i = 1 - trans_b.i + terminal.input[1] = marker.output[0] + _remove_nodes(graph, [transpose]) + + +def _validate_marker_domain(marker: onnx.NodeProto, expected_domain: str) -> None: + if marker.domain != expected_domain: + raise NotImplementedError( + f"Unsupported Dynamo quantization marker '{marker.op_type}': unexpected domain " + f"'{marker.domain}'." + ) + + +def normalize_dynamo_weight_paths(model: onnx.ModelProto) -> onnx.ModelProto: + """Normalize supported Dynamo static-weight paths for the legacy format packers.""" + graph = model.graph + producers = get_tensor_producer_nodes(graph, get_initializer_producers=True) + consumers = get_tensor_consumer_nodes(graph) + static_values: dict[str, bool] = {} + + for marker in list(graph.node): + kind = None + marker_output = marker.output[0] + if marker.op_type == "TRT_FP4QDQ": + _validate_marker_domain(marker, "trt") + kind = "nvfp4" + elif marker.op_type == "TRT_MXFP8DequantizeLinear" and _has_static_weight_source( + producers, marker.input[0], static_values + ): + _validate_marker_domain(marker, "trt") + kind = "mxfp8" + elif marker.op_type == "DequantizeLinear" and (_attribute(marker, "block_size") or 0) > 0: + producer = producers.get(marker.input[0]) + if ( + not isinstance(producer, onnx.NodeProto) + or producer.op_type != "TRT_FP4DynamicQuantize" + ): + if marker.domain != "trt": + _validate_marker_domain(marker, "") + kind = "int4" + elif marker.op_type in { + "TRT_FP8QuantizeLinear", + "QuantizeLinear", + } and _has_static_weight_source(producers, marker.input[0], static_values): + expected_domain = "trt" if marker.op_type == "TRT_FP8QuantizeLinear" else "" + _validate_marker_domain(marker, expected_domain) + dq = _single_consumer(consumers, marker.output[0], marker.input[0]) + expected_dq = ( + "TRT_FP8DequantizeLinear" + if marker.op_type == "TRT_FP8QuantizeLinear" + else "DequantizeLinear" + ) + if dq.op_type != expected_dq or dq.domain != marker.domain: + raise NotImplementedError( + f"Unsupported Dynamo quantized weight '{marker.input[0]}': expected a matching " + f"{expected_dq}." + ) + if list(dq.input[1:]) != list(marker.input[1:]): + raise NotImplementedError( + f"Unsupported Dynamo quantized weight '{marker.input[0]}': " + "Q/DQ scale and zero-point inputs must match." + ) + kind = "fp8" if marker.op_type.startswith("TRT_") else "int8" + marker_output = dq.output[0] + if kind is None: + continue + + if kind in {"fp8", "int8"}: + direct_consumers = consumers.get(marker_output, []) + if ( + len(direct_consumers) == 1 + and direct_consumers[0].op_type == "Conv" + and len(direct_consumers[0].input) > 1 + and direct_consumers[0].input[1] == marker_output + ): + continue + + weight, prefix_reshape = _weight_source(graph, producers, consumers, marker) + path, terminal = _path_after_marker(consumers, marker_output, weight.name) + if kind in {"fp8", "int8", "mxfp8"}: + _validate_restoring_reshape(path, weight, producers, kind.upper()) + _fold_prefix_reshape(graph, marker, weight, prefix_reshape, producers) + if kind in {"int4", "mxfp8"}: + _split_scale( + graph, + producers, + consumers, + marker, + materialize_initializer=kind == "int4", + ) + if kind == "int4": + if not marker.name.startswith(_DYNAMO_INT4_NODE_PREFIX): + marker.name = f"{_DYNAMO_INT4_NODE_PREFIX}{marker.name}" + _normalize_int4_path(graph, marker, weight, prefix_reshape, path, terminal, producers) + marker.domain = "" + elif kind == "mxfp8" and not marker.name.startswith(_DYNAMO_MXFP8_NODE_PREFIX): + marker.name = f"{_DYNAMO_MXFP8_NODE_PREFIX}{marker.name}" + elif kind == "nvfp4": + _normalize_nvfp4_path(graph, marker, weight, prefix_reshape, path, terminal, producers) + + return model diff --git a/modelopt/onnx/quantization/autotune/benchmark.py b/modelopt/onnx/quantization/autotune/benchmark.py index ba5cf1142bf..6c62ee4e628 100644 --- a/modelopt/onnx/quantization/autotune/benchmark.py +++ b/modelopt/onnx/quantization/autotune/benchmark.py @@ -358,8 +358,7 @@ def __init__( trt.init_libnvinfer_plugins(self.trt_logger, "") self._plugin_registry = trt.get_plugin_registry() - self.network_flags = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) - self.network_flags |= 1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED) + self.network_flags = 1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED) self._timing_cache = None self._load_timing_cache() self._shape_configs = {} diff --git a/modelopt/onnx/quantization/qdq_graph.py b/modelopt/onnx/quantization/qdq_graph.py index 47ee6ef1464..c712b7e385c 100755 --- a/modelopt/onnx/quantization/qdq_graph.py +++ b/modelopt/onnx/quantization/qdq_graph.py @@ -842,6 +842,10 @@ def union(x, y): # Find quantized inputs to concat quantized_inputs = [] for input_tensor in node.inputs: + # A constant (initializer) input is not an activation: the calibrator collects no range + # for it, so it cannot share a scaling factor with the other tensors of the group. + if isinstance(input_tensor, gs.Constant): + continue input_name = input_tensor.name input_consumers = input_name_to_nodes[input_name] if any(consumer.name in nodes_to_quantize for consumer in input_consumers): diff --git a/modelopt/onnx/trt_utils.py b/modelopt/onnx/trt_utils.py index b407fdc5411..88e199b1788 100644 --- a/modelopt/onnx/trt_utils.py +++ b/modelopt/onnx/trt_utils.py @@ -139,7 +139,7 @@ def get_custom_layers( for plugin in trt_plugins: _load_trt_plugin(plugin, registry) - logger.debug(f"Total registered plugin creators: {len(registry.plugin_creator_list)}") + logger.debug(f"Total registered plugin creators: {len(registry.all_creators)}") builder = trt.Builder(trt_logger) network = ( builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)) diff --git a/modelopt/torch/distill/plugins/megatron.py b/modelopt/torch/distill/plugins/megatron.py index 3aac2666650..26b3c9cc5c8 100644 --- a/modelopt/torch/distill/plugins/megatron.py +++ b/modelopt/torch/distill/plugins/megatron.py @@ -71,6 +71,8 @@ class DistillationConfig: Only the smallest prefix of the (sorted) Top-K whose cumulative teacher probability reaches this value contributes to the loss. Requires ``logit_kl_topk``. Must be in (0, 1]. logit_kl_top_p_min_k: Minimum number of Top-K entries kept per token when top-P is active. + freeze_mtp: Keep the student's MTP head fixed: skip its loss, which otherwise also trains the + shared layers it reads from, and freeze its parameters. Default: ``False``. """ intermediate_layer_pairs: list[tuple[str, ...]] = field(default_factory=list) @@ -80,6 +82,7 @@ class DistillationConfig: logit_kl_topk: int | None = None logit_kl_top_p: float | None = None logit_kl_top_p_min_k: int = 1 + freeze_mtp: bool = False criterion: Criterion | None = None loss_balancer: mtd.DistillationLossBalancer | None = None @@ -701,21 +704,32 @@ def _sharded_state_dict(self, *args, **kwargs) -> "ShardedStateDict": # skip_lm_loss=True — only the final main-head call is zeroed. # An MTP head left out of quantization is exempt from that: there is no quantization # error to recover there, and its CE materialises an fp32 [seq, vocab] tensor. - skip_mtp_loss = _mtp_excluded_from_quantization(model) + # ``freeze_mtp`` skips the MTP loss unconditionally. skip_lm_loss = distill_cfg.skip_lm_loss - if skip_lm_loss and skip_mtp_loss: + skip_mtp_loss = distill_cfg.freeze_mtp or ( + skip_lm_loss and _mtp_excluded_from_quantization(model) + ) + if skip_mtp_loss: # Freeze the untrained MTP head: DDP's overlapped grad reduce asserts on params with no grad. - warn_rank_0("MTP head is outside quantization and its loss is skipped: freezing it.") + num_frozen = 0 with model.hide_teacher_model(): for name, param in model.named_parameters(): if "mtp" in name.split("."): param.requires_grad_(False) + num_frozen += 1 + if num_frozen: + warn_rank_0( + f"MTP head is frozen ({num_frozen} params) and its loss is skipped" + + (" (freeze_mtp)." if distill_cfg.freeze_mtp else ": it is outside quantization.") + ) + elif parallel_state.is_pipeline_last_stage(): # earlier stages never hold the MTP head + warn_rank_0("freeze_mtp is set but the model has no MTP head: nothing to freeze.") def _compute_student_lm_loss(self, labels, logits) -> Tensor: self._lm_loss_call_count += 1 mtp_num_layers = self.config.mtp_num_layers or 0 is_mtp_call = self._lm_loss_call_count <= mtp_num_layers - if skip_lm_loss and self.training and (not is_mtp_call or skip_mtp_loss): + if self.training and (skip_mtp_loss if is_mtp_call else skip_lm_loss): return torch.zeros_like(labels, dtype=logits.dtype) return type(self).compute_language_model_loss(self, labels, logits) diff --git a/modelopt/torch/export/__init__.py b/modelopt/torch/export/__init__.py index 874fa71b2b9..44d1ad2557e 100644 --- a/modelopt/torch/export/__init__.py +++ b/modelopt/torch/export/__init__.py @@ -38,6 +38,7 @@ from .moe_utils import * from .plugins import * from .quant_format import * +from .quantized_weight_export import * from .registry import * from .shard_cast_utils import * from .transformer_engine import * diff --git a/modelopt/torch/export/convert_hf_config.py b/modelopt/torch/export/convert_hf_config.py index d78410c661e..bbcc25a28da 100644 --- a/modelopt/torch/export/convert_hf_config.py +++ b/modelopt/torch/export/convert_hf_config.py @@ -43,9 +43,25 @@ def _quant_algo_to_group_config(quant_algo: str, group_size: int | None = None) } elif quant_algo == "FP8_PER_CHANNEL_PER_TOKEN": return { - "input_activations": {"dynamic": False, "num_bits": 8, "type": "float"}, + "input_activations": { + "dynamic": True, + "num_bits": 8, + "type": "float", + "strategy": "token", + }, "weights": {"dynamic": False, "num_bits": 8, "type": "float", "strategy": "channel"}, } + elif quant_algo == "FP8_PB_WO": + block_size = group_size or 128 + return { + "weights": { + "dynamic": False, + "num_bits": 8, + "type": "float", + "strategy": "block", + "block_structure": [block_size, block_size], + }, + } elif quant_algo == "NVFP4": gs = group_size or 16 return { @@ -115,13 +131,40 @@ def _quant_algo_to_group_config(quant_algo: str, group_size: int | None = None) gs = group_size or 32 return { "input_activations": { + "dynamic": True, + "num_bits": 8, + "type": "float", + "group_size": gs, + "strategy": "group", + }, + "weights": { "dynamic": False, "num_bits": 8, "type": "float", "group_size": gs, + "strategy": "group", }, - "weights": {"dynamic": False, "num_bits": 8, "type": "float", "group_size": gs}, } + elif quant_algo in ("MXFP4", "W4A16_MXFP4"): + gs = group_size or 32 + config = { + "weights": { + "dynamic": False, + "num_bits": 4, + "type": "float", + "group_size": gs, + "strategy": "group", + }, + } + if quant_algo == "MXFP4": + config["input_activations"] = { + "dynamic": True, + "num_bits": 4, + "type": "float", + "group_size": gs, + "strategy": "group", + } + return config elif quant_algo.lower() in GGML_FORMATS: ggml_format = GGML_FORMAT_REGISTRY[quant_algo.lower()] block_size, payload_bytes = ggml_format.block_size, ggml_format.block_bytes @@ -260,6 +303,19 @@ def convert_hf_quant_config_format(input_config: dict[str, Any]) -> dict[str, An if lora_rank is not None: config_group_details["lora_rank"] = lora_rank new_config["config_groups"] = {"group_0": config_group_details} + elif quant_algo_value in { + "FP8_PER_CHANNEL_PER_TOKEN", + "FP8_PB_WO", + "MXFP4", + "W4A16_MXFP4", + "MXFP8", + "W4A8_MXFP4_FP8", + "W4A8_NVFP4_FP8", + }: + group_size = original_quantization_details.get("group_size") + config_group_details = _quant_algo_to_group_config(quant_algo_value, group_size) + config_group_details["targets"] = ["Linear"] + new_config["config_groups"] = {"group_0": config_group_details} elif quant_algo_value == "MIXED_PRECISION": quantized_layers = original_quantization_details.get("quantized_layers", {}) @@ -292,6 +348,11 @@ def convert_hf_quant_config_format(input_config: dict[str, Any]) -> dict[str, An if quant_algo_value: new_config["quant_algo"] = quant_algo_value + group_size = original_quantization_details.get("group_size") + if group_size is not None: + # ModelOpt consumers use this top-level value to size block scales. + new_config["group_size"] = group_size + kv_cache_quant_algo = original_quantization_details.get("kv_cache_quant_algo") if kv_cache_quant_algo: if kv_cache_quant_algo == "FP8": diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index 7c526d8255f..200ec6e5dae 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/modelopt/torch/export/hf_export_handlers.py @@ -108,12 +108,13 @@ def _is_quant_moe_linear(module: nn.Module) -> bool: name gets uniquified — would bypass this handler and export without the input-amax fallback. The name check is kept as a fallback so stand-in modules match too. - The wrapper lives in the optional transformers plugin, hence the lazy import. + The wrapper lives in the Step PTQ modeling module, imported lazily like other export + modules do to avoid an import cycle with ``modelopt.torch.quantization``. """ if not hasattr(module, "experts"): return False try: - from modelopt.torch.quantization.plugins.huggingface import _QuantMoELinear + from modelopt.torch.models.step3p5.modeling_ptq import _QuantMoELinear if isinstance(module, _QuantMoELinear): return True diff --git a/modelopt/torch/export/layerwise_export.py b/modelopt/torch/export/layerwise_export.py index 8bd9e6cc3f8..92f31011ea3 100644 --- a/modelopt/torch/export/layerwise_export.py +++ b/modelopt/torch/export/layerwise_export.py @@ -286,8 +286,8 @@ def export_layer( a fusing format can rediscover which modules share an input; omit them only when nothing fuses. """ - # Local, as in every other export module: the plugin imports transformers. - from modelopt.torch.quantization.plugins.huggingface import _reconstruct_fused_moe_linear + # Local, as in every other export module: see pack_quantized_weights. + from modelopt.torch.models.step3p5.modeling_ptq import _reconstruct_fused_moe_linear assert self._bound, "export_layer() before bind()" assert not self._finalized, "export_layer() called after finalize()" diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index 734302690f1..02660142736 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -40,11 +40,76 @@ def _delete_fused_moe_source_attrs(module: nn.Module) -> None: "down_proj", "down_proj_weight_quantizers", "down_proj_input_quantizer", + f"{first_proj_attr}_scale_inv", + "down_proj_scale_inv", + "gate_up_proj_activation_scale", + "up_proj_activation_scale", + "down_proj_activation_scale", ): if hasattr(module, attr): delattr(module, attr) +def _tensor_quantizers(quantizer) -> list: + """The ``TensorQuantizer`` stages of ``quantizer`` (several for W4A8's SequentialQuantizer). + + A SequentialQuantizer has no ``_amax`` of its own and forwards ``amax`` writes to every + stage, so each stage must be checked and set on its own. + """ + # Deferred: ``export/__init__`` star-imports this module, so keep quantization out of its + # import-time dependencies. + from modelopt.torch.quantization.nn import SequentialQuantizer + + return list(quantizer) if isinstance(quantizer, SequentialQuantizer) else [quantizer] + + +def _is_uncalibrated(quantizer) -> bool: + amax = getattr(quantizer, "_amax", None) + return getattr(quantizer, "is_enabled", False) and (amax is None or bool(torch.all(amax == 0))) + + +def _fallback_amax(weight: torch.Tensor, quantizer) -> torch.Tensor: + """Weight-derived amax for an uncalibrated expert, in the shape export expects. + + The amax setter rejects shape changes, and a block/channel-wise quantizer (e.g. the INT4 + stage of W4A8 AWQ) keeps a zero ``_amax`` of that shape, so a scalar fallback would raise. + A missing ``_amax`` on an INT4 block quantizer still needs per-block ``(rows, blocks)``. + """ + like = getattr(quantizer, "_amax", None) + block_sizes = getattr(quantizer, "block_sizes", None) + block = block_sizes.get(-1) if block_sizes and block_sizes.get("type") == "static" else None + # Integer static-block stages (INT4 of W4A8) need a per-block amax. Float-block formats + # (NVFP4) keep a per-tensor amax plus separate block scales, so they take the scalar path. + if like is None and isinstance(quantizer.num_bits, int) and block: + if weight.shape[-1] % block: + raise ValueError( + f"Weight of shape {tuple(weight.shape)} is not divisible into blocks of {block}." + ) + like = weight.new_empty(weight.shape[0], weight.shape[-1] // block) + if like is None or like.numel() == 1: + amax = torch.maximum(weight.amax(), -weight.amin()).float() + return amax if like is None else amax.view_as(like) + # Block-quant calibration records the flat (rows * blocks, 1); also accept (rows, blocks) and + # a per-row (rows,). Refuse any other layout (e.g. per input channel) rather than guess its + # grouping. + rows = weight.shape[0] + if like.dim() == 1 and like.shape[0] == rows: + blocks = 1 + elif like.dim() == 2 and like.shape[-1] == 1 and like.shape[0] % rows == 0: + blocks = like.shape[0] // rows + elif like.dim() == 2 and like.shape[0] == rows: + blocks = like.shape[1] + else: + blocks = 0 + if not blocks or weight.shape[-1] % blocks: + raise ValueError( + f"Cannot derive amax of shape {tuple(like.shape)} from a weight of shape " + f"{tuple(weight.shape)}." + ) + groups = weight.reshape(rows, blocks, -1) + return torch.maximum(groups.amax(dim=-1), -groups.amin(dim=-1)).float().reshape(like.shape) + + def _export_fused_experts( module: nn.Module, dtype: torch.dtype, @@ -100,6 +165,8 @@ def _export_fused_experts( for idx in range(n): expert = nn.Module() + first_weight = module.get_expert_weight(first_proj_attr, idx, dtype) + down_weight = module.get_expert_weight("down_proj", idx, dtype) # If the gate_up source quantizer was never calibrated (rare expert # that received no calibration tokens), derive its amax once from the @@ -111,17 +178,18 @@ def _export_fused_experts( # mismatched weight_scale_2 and garbled MoE output at inference. # Non-gated experts have no gate/up fusion, so this shared-amax step is # skipped — their single up_proj uses the generic per-projection fallback. - first_proj_q = first_proj_weight_quantizers[idx] - if ( - is_gated - and getattr(first_proj_q, "is_enabled", False) - and ( - not hasattr(first_proj_q, "_amax") - or first_proj_q._amax is None - or torch.all(first_proj_q._amax == 0) - ) - ): - first_proj_q.amax = first_proj[idx].abs().amax().to(torch.float32) + uncalibrated_stages = ( + [ + q + for q in _tensor_quantizers(first_proj_weight_quantizers[idx]) + if _is_uncalibrated(q) + ] + if is_gated + else [] + ) + for q in uncalibrated_stages: + q.amax = _fallback_amax(first_weight, q) + if uncalibrated_stages: warnings.warn( f"Expert {idx} gate_up_proj weight quantizer was not calibrated " f"(amax missing or zero). Using fused-tensor amax as fallback " @@ -132,16 +200,16 @@ def _export_fused_experts( if is_gated: projections = [ - ("gate_proj", first_proj[idx, :expert_dim, :], 0, fused_dim0, True), - ("up_proj", first_proj[idx, expert_dim:, :], expert_dim, fused_dim0, True), - ("down_proj", down[idx], 0, down.shape[1], False), + ("gate_proj", first_weight[:expert_dim, :], 0, fused_dim0, True), + ("up_proj", first_weight[expert_dim:, :], expert_dim, fused_dim0, True), + ("down_proj", down_weight, 0, down.shape[1], False), ] else: # Non-gated: the single up_proj maps 1:1 to its weight quantizer, so it # is exported whole (no dim-0 split, no shared gate/up weight_scale_2). projections = [ - ("up_proj", first_proj[idx], 0, fused_dim0, True), - ("down_proj", down[idx], 0, down.shape[1], False), + ("up_proj", first_weight, 0, fused_dim0, True), + ("down_proj", down_weight, 0, down.shape[1], False), ] for ( @@ -163,53 +231,46 @@ def _export_fused_experts( copy.deepcopy(w_quantizer_src) if uses_first_proj_quantizers else w_quantizer_src ) - # For per-channel amax (dim >= 1), proportionally slice dim-0 - # to match the split weight. - if ( - hasattr(w_quantizer, "_amax") - and w_quantizer._amax is not None - and w_quantizer._amax.dim() >= 1 - ): - amax = w_quantizer._amax - # Per-block _amax (NVFP4 static) collapses the row axis we want - # to slice on; restore it so dim-0 slicing splits gate/up. - if amax.numel() != fused_total and amax.numel() % fused_total == 0: - amax = amax.contiguous().view(fused_total, amax.numel() // fused_total) - amax_dim0 = amax.shape[0] - if fused_total % amax_dim0 == 0: - slice_start = fused_start * amax_dim0 // fused_total - slice_end = (fused_start + weight_slice.shape[0]) * amax_dim0 // fused_total - sliced = amax[slice_start:slice_end].contiguous() - # The amax setter refuses shape changes; drop _amax first. - if hasattr(w_quantizer, "_amax"): - delattr(w_quantizer, "_amax") - w_quantizer.amax = sliced - else: + for stage in _tensor_quantizers(w_quantizer): + # For per-channel amax (dim >= 1), proportionally slice dim-0 + # to match the split weight. + if getattr(stage, "_amax", None) is not None and stage._amax.numel() > 1: + amax = stage._amax + # Per-block _amax (NVFP4 static) collapses the row axis we want + # to slice on; restore it so dim-0 slicing splits gate/up. + if amax.numel() != fused_total and amax.numel() % fused_total == 0: + amax = amax.contiguous().view(fused_total, amax.numel() // fused_total) + amax_dim0 = amax.shape[0] + if fused_total % amax_dim0 == 0: + slice_start = fused_start * amax_dim0 // fused_total + slice_end = (fused_start + weight_slice.shape[0]) * amax_dim0 // fused_total + sliced = amax[slice_start:slice_end].contiguous() + # The amax setter refuses shape changes; drop _amax first. + delattr(stage, "_amax") + stage.amax = sliced + # Block-quant calibration recorded the fused amax's row count here, and + # export_amax() reshapes to it, undoing the slice. + export_shape = getattr(stage, "_amax_shape_for_export", None) + if export_shape is not None and export_shape[0] == amax_dim0: + stage._amax_shape_for_export = (sliced.shape[0], *export_shape[1:]) + else: + warnings.warn( + f"Expert {idx} {proj_name}: fused amax dim0 ({amax_dim0}) does not " + f"evenly divide fused_total ({fused_total}). Skipping amax slicing, " + f"which may produce incorrect quantization scales.", + stacklevel=2, + ) + + # If the weight quantizer was never calibrated, compute amax from weights. + if _is_uncalibrated(stage): + stage.amax = _fallback_amax(weight_slice, stage) warnings.warn( - f"Expert {idx} {proj_name}: fused amax dim0 ({amax_dim0}) does not " - f"evenly divide fused_total ({fused_total}). Skipping amax slicing, " - f"which may produce incorrect quantization scales.", + f"Expert {idx} {proj_name} weight quantizer was not calibrated " + f"(amax missing or zero). Using weight-derived amax as fallback. " + f"Consider using more calibration data to activate all experts.", stacklevel=2, ) - # If the weight quantizer was never calibrated, compute amax from weights. - if ( - hasattr(w_quantizer, "is_enabled") - and w_quantizer.is_enabled - and ( - not hasattr(w_quantizer, "_amax") - or w_quantizer._amax is None - or torch.all(w_quantizer._amax == 0) - ) - ): - w_quantizer.amax = weight_slice.abs().amax().to(torch.float32) - warnings.warn( - f"Expert {idx} {proj_name} weight quantizer was not calibrated " - f"(amax missing or zero). Using weight-derived amax as fallback. " - f"Consider using more calibration data to activate all experts.", - stacklevel=2, - ) - wrapper = nn.Module() wrapper.weight = nn.Parameter(weight_slice.contiguous(), requires_grad=False) wrapper.weight_quantizer = w_quantizer diff --git a/modelopt/torch/export/quant_aware_conversion.py b/modelopt/torch/export/quant_aware_conversion.py index 482bb95e94b..ad186df9300 100644 --- a/modelopt/torch/export/quant_aware_conversion.py +++ b/modelopt/torch/export/quant_aware_conversion.py @@ -62,13 +62,21 @@ (in-memory-name) behavior rather than emit a silently-wrong checkpoint. """ +import logging import re from collections.abc import Iterable from dataclasses import dataclass from fnmatch import fnmatchcase +from pathlib import Path from typing import Any import torch +from huggingface_hub import try_to_load_from_cache +from huggingface_hub.errors import HFValidationError + +from modelopt.torch.utils.plugins.hf_checkpoint_utils import indexed_weight_map + +logger = logging.getLogger(__name__) __all__ = [ "QuantConversionUnsupportedError", @@ -251,9 +259,25 @@ def _apply_merge_rule(state_dict: dict[str, torch.Tensor], rule: _MergeRule) -> state_dict[target_key] = merged +def _idempotent_pattern(pattern: str, repl: str) -> str: + """Stop a rename rule from re-matching its own replacement. + + Tensors copied from the source checkpoint skip the forward conversion and reach the reverse + pass already in hub naming, so a rule like ``input_conditioner`` -> + ``radio_model.input_conditioner`` would add its prefix twice. When the replacement is + `` + ``, skip keys that already carry the prefix. + """ + if len(repl) > len(pattern) and repl.endswith(pattern): + return f"(? str: @@ -305,6 +329,7 @@ def apply_reverse_rules( split_rules: list[SplitRule], rename_rules: list[RenameRule], merge_rules: list[_MergeRule] | None = None, + source_keys: set[str] | None = None, ) -> dict[str, torch.Tensor]: """Apply quant-aware reverse conversion: tensor transforms first, then renames. @@ -313,6 +338,10 @@ def apply_reverse_rules( ``merge_rules`` is an internal hook for model-derived rules: unlike the public rename/split descriptors, each rule retains a live transformers converter for scoped matching. Call :func:`revert_weight_conversion_quant_aware` to derive them. + + ``source_keys`` are the tensor names of the checkpoint the model was loaded from. + They let a rename that would otherwise invent a name resolve to the spelling the + checkpoint actually carries; omitting them applies the rules unchanged. """ out = dict(state_dict) for rule in split_rules: @@ -320,10 +349,10 @@ def apply_reverse_rules( for merge_rule in merge_rules or []: _apply_merge_rule(out, merge_rule) - compiled = _compile_rename_rules(rename_rules) + rename = _SourceAwareRenamer(rename_rules, source_keys) renamed: dict[str, torch.Tensor] = {} for key, value in out.items(): - new_key = _apply_rename_rules(key, compiled) + new_key = rename(key) if new_key in renamed: raise QuantConversionUnsupportedError(f"rename collision on '{new_key}'") renamed[new_key] = value @@ -338,11 +367,13 @@ def revert_weight_conversion_quant_aware(model, state_dict: dict[str, torch.Tens when the mapping uses an op that cannot be reversed quant-aware yet, so the caller can fall back to the legacy behavior. """ - split_rules, merge_rules, rename_rules, expert_fused_leaves = _build_reverse_rules(model) + split_rules, merge_rules, rename_rules, expert_fused_leaves, source_keys = _build_reverse_rules( + model + ) if not split_rules and not merge_rules and not rename_rules: return state_dict _assert_experts_pre_expanded(state_dict, expert_fused_leaves) - return apply_reverse_rules(state_dict, split_rules, rename_rules, merge_rules) + return apply_reverse_rules(state_dict, split_rules, rename_rules, merge_rules, source_keys) def build_reverse_name_mapper(model, *, tensor_keys: bool = False): @@ -366,15 +397,15 @@ def build_reverse_name_mapper(model, *, tensor_keys: bool = False): :class:`QuantConversionUnsupportedError` when the mapping can't be reversed, so the caller can keep the in-memory names for BOTH weights and config (mutually consistent). """ - _, merge_rules, rename_rules, _ = _build_reverse_rules(model) + _, merge_rules, rename_rules, _, source_keys = _build_reverse_rules(model) if not merge_rules and not rename_rules: return None - compiled = _compile_rename_rules(rename_rules) + rename = _SourceAwareRenamer(rename_rules, source_keys) def _apply(text: str) -> str: for rule in merge_rules: text, _ = rule.converter.rename_source_key(text) - return _apply_rename_rules(text, compiled) + return rename(text) def _map(name: str) -> str: if tensor_keys: @@ -488,6 +519,160 @@ def _scope_prefixes(rev) -> tuple[str, ...]: return tuple(seen) +_SAFETENSORS_INDEX = "model.safetensors.index.json" +_SAFETENSORS_SINGLE = "model.safetensors" + + +def _checkpoint_dir(path: str, revision: str | None = None) -> Path | None: + """Local directory holding ``path``'s safetensors checkpoint, or ``None``. + + ``path`` is a :attr:`name_or_path`, so it is either a local directory or a hub id. + A hub id is resolved against the local HF cache only -- this never downloads, and + an uncached hub id simply yields ``None``. + + ``revision`` pins the lookup to the commit the model was actually loaded from. + Without it the cache lookup defaults to ``main``, which for a model loaded from a + pinned revision could read a *different* snapshot's names and filter the wrong + rules. ``None`` keeps the hub default, which is correct when the revision is + unknown because nothing else identifies the snapshot. + + A non-default ``cache_dir`` is *not* covered: ``transformers`` does not record it + on the model, so a model loaded from one resolves here against the default cache + and most likely finds nothing -- which leaves the guard inert and behavior + unchanged, the safe direction. + """ + directory = Path(path) + if directory.is_dir(): + return directory + for filename in (_SAFETENSORS_INDEX, _SAFETENSORS_SINGLE): + try: + cached = try_to_load_from_cache(path, filename, revision=revision) + except HFValidationError: # not a well-formed hub id either + return None + # Returns the file path on a hit, and a sentinel object (not a str) for a + # cached "this file does not exist" result. + if isinstance(cached, str): + return Path(cached).parent + return None + + +def _source_checkpoint_keys(model) -> set[str] | None: + """Tensor names of the checkpoint ``model`` was loaded from, or ``None``. + + ``None`` means "unknown" and leaves the reverse conversion exactly as it was: it is + returned for a model built from a config, a hub id absent from the local HF cache, + or a checkpoint layout this does not recognize. + """ + config = getattr(model, "config", None) + path = getattr(model, "name_or_path", None) or getattr(config, "_name_or_path", None) + if not path: + return None + try: + # ``transformers`` records the resolved commit on the config; a hub-id lookup + # must use it rather than falling back to ``main``. + directory = _checkpoint_dir(path, getattr(config, "_commit_hash", None)) + if directory is None: + return None + # ``indexed_weight_map`` is the repo's one index reader: ``weight_map`` for a + # sharded checkpoint, ``safe_open`` keys for a single file, ``{}`` for neither. + # An empty result means "nothing recorded", which is exactly our unknown. + return set(indexed_weight_map(directory)) or None + except Exception as exc: # unreadable or malformed checkpoint -> treat as unknown + # Deliberately no path and no exception text: SECURITY.md warns that logs may + # carry sensitive paths, and the type alone identifies the failure mode. + logger.debug("Cannot read source checkpoint keys: %s", type(exc).__name__) + return None + + +def _non_idempotent_renames(rules: list[RenameRule], source_keys: set[str] | None) -> set[int]: + """Indices of reverse renames that only ever invent names the source lacks. + + A reverse rename maps an in-memory name back to its source-checkpoint name, so it + must be a *no-op* on names that are already in source form -- the same invariant + :func:`_idempotent_pattern` enforces for prefix rules. A rule that breaks it may be + inventing a name rather than restoring one. + + That happens when a model declares a backwards-compatibility *load* rename as a + reversible ``WeightRenaming``. Reversing it rewrites the canonical name to a legacy + spelling the published checkpoint never used, and a deployment loader keyed on the + canonical name then cannot find the tensor. + + A rule is flagged only on positive evidence: it must rewrite at least one source key, + and every source key it rewrites must land on a name absent from the checkpoint. + + Flagging is not a verdict. A rule can be *needed* for one namespace and wrong for + another -- a mixed-namespace checkpoint whose sparse layers are stored under one + container name and whose dense and vision layers are already stored under the + in-memory one. :class:`_SourceAwareRenamer` therefore decides per key, and a flagged + rule still applies wherever it produces a name the checkpoint actually carries. + """ + if not source_keys: + return set() + suspect: set[int] = set() + for index, rule in enumerate(rules): + pattern = re.compile(_idempotent_pattern(rule.pattern, rule.repl)) + rewrites_a_source_key = False + every_result_absent = True + for key in source_keys: + mapped = _sub_scoped(pattern, rule.repl, key, rule.scope_prefixes) + if mapped == key: + continue + rewrites_a_source_key = True + if mapped in source_keys: + every_result_absent = False + break + if rewrites_a_source_key and every_result_absent: + suspect.add(index) + return suspect + + +class _SourceAwareRenamer: + """Apply reverse renames, preferring names the source checkpoint actually carries. + + The reverse of a rename is ambiguous: given a canonical in-memory name, both the + canonical and a legacy spelling are valid pre-images, and nothing inside the model + breaks the tie. The checkpoint the model was loaded from does, so a key whose + rewritten name is absent from it is resolved against the rules that do not invent + names (see :func:`_non_idempotent_renames`). + + Deciding per key rather than per rule matters: no rule is ever discarded, so a rule + that a mixed-namespace checkpoint needs for one set of layers keeps applying there + while the same rule stops rewriting layers already stored in their in-memory form. + + Without source evidence, or when neither candidate names a source tensor, the full + rule chain applies unchanged -- the behavior this class replaced. + """ + + def __init__(self, rename_rules: list[RenameRule], source_keys: set[str] | None): + self._full = _compile_rename_rules(rename_rules) + self._source_keys = source_keys or set() + suspect = _non_idempotent_renames(rename_rules, source_keys) + self._reduced = ( + _compile_rename_rules( + [rule for index, rule in enumerate(rename_rules) if index not in suspect] + ) + if suspect + else None + ) + + def _names_a_source_tensor(self, name: str) -> bool: + if name in self._source_keys: + return True + # A tensor the export adds -- a quantizer scale, or a bias the source omits -- + # has no counterpart in the checkpoint, so it is judged by the weight of the + # module it belongs to. That keeps every tensor of one module on the same + # decision instead of splitting the module across two namespaces. + parent, _, leaf = name.rpartition(".") + return bool(leaf) and f"{parent}.weight" in self._source_keys + + def __call__(self, key: str) -> str: + renamed = _apply_rename_rules(key, self._full) + if self._reduced is None or self._names_a_source_tensor(renamed): + return renamed + reduced = _apply_rename_rules(key, self._reduced) + return reduced if self._names_a_source_tensor(reduced) else renamed + + def _drop_shadowed_prefix_renames(model, rules: list[RenameRule]) -> list[RenameRule]: """Drop child-model reverse renames when the child namespace already exists. @@ -533,13 +718,16 @@ def _drop_shadowed_prefix_renames(model, rules: list[RenameRule]) -> list[Rename def _build_reverse_rules( model, -) -> tuple[list[SplitRule], list[_MergeRule], list[RenameRule], list[str]]: +) -> tuple[list[SplitRule], list[_MergeRule], list[RenameRule], list[str], set[str] | None]: """Derive reverse rules from the model's transformers conversion mapping. - Returns tensor-transform rules, rename rules, and fused expert leaves. The last is - used to guard against experts that were not pre-expanded. Returns empty lists when - no mapping applies (export unchanged). Uses transformers' own - ``reverse_transform()`` to get correctly-reversed name patterns, then translates them: + Returns tensor-transform rules, rename rules, fused expert leaves, and the key set + of the checkpoint the model was loaded from (``None`` when it cannot be read). The + expert leaves guard against experts that were not pre-expanded; the key set lets + :class:`_SourceAwareRenamer` resolve a rename against the names the checkpoint has. + Returns empty lists when no mapping applies (export unchanged). Uses transformers' + own ``reverse_transform()`` to get correctly-reversed name patterns, then translates + them: * ``WeightRenaming`` -> :class:`RenameRule` (carries scale siblings for free). * Expert ``WeightConverter`` (reverse contains ``SplitModulelist``): ModelOpt's @@ -564,7 +752,7 @@ def _build_reverse_rules( raise QuantConversionUnsupportedError(f"could not read conversion mapping: {exc}") from exc if not conversions: - return [], [], [], [] + return [], [], [], [], None try: from transformers.core_model_loading import ( @@ -644,7 +832,13 @@ def _build_reverse_rules( # substrings and are applied first. weight_renamings = _drop_shadowed_prefix_renames(model, weight_renamings) rename_rules = leaf_renamings + list(reversed(weight_renamings)) - return split_rules, merge_rules, rename_rules, expert_fused_leaves + return ( + split_rules, + merge_rules, + rename_rules, + expert_fused_leaves, + _source_checkpoint_keys(model), + ) # ModelOpt's export splits a fused ``gate_up_proj`` into these per-expert linears, diff --git a/modelopt/torch/export/quant_format.py b/modelopt/torch/export/quant_format.py index 8a520f972b4..8c503341760 100644 --- a/modelopt/torch/export/quant_format.py +++ b/modelopt/torch/export/quant_format.py @@ -31,6 +31,7 @@ QUANTIZATION_NVFP4_SVDQUANT = "nvfp4_svdquant" QUANTIZATION_W4A8_NVFP4_FP8 = "w4a8_nvfp4_fp8" QUANTIZATION_MXFP4 = "mxfp4" +QUANTIZATION_W4A16_MXFP4 = "w4a16_mxfp4" QUANTIZATION_MXFP8 = "mxfp8" QUANTIZATION_W4A8_MXFP4_FP8 = "w4a8_mxfp4_fp8" QUANTIZATION_W4A16_NVFP4 = "w4a16_nvfp4" diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index 6450c1afc74..88092f689dc 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -42,7 +42,6 @@ QTensorWrapper, ) from modelopt.torch.quantization.utils import ( - QuantizerAttrNames, quantizer_attr_names, representative_weight_quantizer, weight_attr_names, @@ -74,6 +73,7 @@ QUANTIZATION_W4A8_AWQ, QUANTIZATION_W4A8_MXFP4_FP8, QUANTIZATION_W4A8_NVFP4_FP8, + QUANTIZATION_W4A16_MXFP4, QUANTIZATION_W4A16_NVFP4, ) @@ -225,7 +225,11 @@ def get_weight_scaling_factor(module: nn.Module, weight_name: str = "weight") -> weight_scaling_factor_2.to(weight.device), )[0] - if quantization_format in [QUANTIZATION_W4A8_MXFP4_FP8, QUANTIZATION_MXFP4]: + if quantization_format in [ + QUANTIZATION_W4A8_MXFP4_FP8, + QUANTIZATION_MXFP4, + QUANTIZATION_W4A16_MXFP4, + ]: return MXFP4QTensor.quantize(weight, block_size=weight_quantizer.block_sizes[-1])[ 1 ].reshape(*weight.shape[:-1], -1) @@ -466,134 +470,120 @@ def uses_iq_quantization(module) -> bool: return any(uses_iq_quantization(child) for _, child in module.named_children()) -def get_quantization_format(module) -> str | None: - """Gets the quantization string. - - Gets the quantization string by iterating through the module and its children. - The first non-None quantization string is returned. - """ - - def _get_quantization_from_layer(layer, quantizer_attr_names: QuantizerAttrNames): - # Singular form first, plural ModuleList fallback (fused-experts). - # Strip the "_weight_quantizer" suffix to recover the weight attr name. - weight_attr = quantizer_attr_names.weight_quantizer - weight_name = weight_attr[: -len("_weight_quantizer")].rstrip("_") or "weight" - weight_quantizer = representative_weight_quantizer(layer, weight_name) - input_quantizer = getattr(layer, quantizer_attr_names.input_quantizer, None) - - if weight_quantizer is None or not weight_quantizer.is_enabled: - return QUANTIZATION_NONE - - # Handle SequentialQuantizer - if isinstance(weight_quantizer, SequentialQuantizer): - assert ( - len(weight_quantizer) == 2 - and weight_quantizer[0].num_bits == 4 - and weight_quantizer[1].num_bits == (4, 3) - ), "Unsupported SequentialQuantizer configuration" - assert ( - weight_quantizer[0].block_sizes - and len(weight_quantizer[0].block_sizes) > 0 - and weight_quantizer[0].block_sizes[-1] > 0 - ), "Invalid block_sizes for SequentialQuantizer" - - return QUANTIZATION_W4A8_AWQ - - # Handle individual num_bits cases - if weight_quantizer.num_bits in GGML_FORMATS: - if weight_quantizer.backend != "ggml": - raise ValueError("GGML formats require the built-in 'ggml' quantization backend") - # Both exporters return before collecting input_scale and before the pre_quant_scale - # handling below, so an enabled activation quantizer would be dropped without a trace - # and the checkpoint would load as weight-only. Refuse instead. - if input_quantizer is not None and input_quantizer.is_enabled: - raise NotImplementedError( - "GGML export is weight-only, but this layer has an enabled input " - "quantizer. The GGML block payload carries no activation scale, so the " - "activation quantization would be silently lost." - ) - if input_quantizer is not None and hasattr(input_quantizer, "_pre_quant_scale"): - raise NotImplementedError( - "GGML export does not support an AWQ-style pre_quant_scale." - ) - return weight_quantizer.num_bits +def _get_quantization_from_quantizers( + layer: nn.Module, + weight_quantizer: TensorQuantizer | SequentialQuantizer | None, + input_quantizer: TensorQuantizer | SequentialQuantizer | None, +) -> str | None: + if weight_quantizer is None or not weight_quantizer.is_enabled: + return QUANTIZATION_NONE - if weight_quantizer.num_bits == 4: - assert len(weight_quantizer.block_sizes) > 0 and weight_quantizer.block_sizes[-1] > 0, ( - "Invalid block_sizes for INT4 quantizer" + if isinstance(weight_quantizer, SequentialQuantizer): + assert ( + len(weight_quantizer) == 2 + and weight_quantizer[0].num_bits == 4 + and weight_quantizer[1].num_bits == (4, 3) + ), "Unsupported SequentialQuantizer configuration" + assert ( + weight_quantizer[0].block_sizes + and len(weight_quantizer[0].block_sizes) > 0 + and weight_quantizer[0].block_sizes[-1] > 0 + ), "Invalid block_sizes for SequentialQuantizer" + return QUANTIZATION_W4A8_AWQ + + if weight_quantizer.num_bits in GGML_FORMATS: + if weight_quantizer.backend != "ggml": + raise ValueError("GGML formats require the built-in 'ggml' quantization backend") + # Both exporters return before collecting input_scale and before the pre_quant_scale + # handling below, so an enabled activation quantizer would be dropped without a trace + # and the checkpoint would load as weight-only. Refuse instead. + if input_quantizer is not None and input_quantizer.is_enabled: + raise NotImplementedError( + "GGML export is weight-only, but this layer has an enabled input " + "quantizer. The GGML block payload carries no activation scale, so the " + "activation quantization would be silently lost." ) - return QUANTIZATION_INT4_AWQ + if input_quantizer is not None and hasattr(input_quantizer, "_pre_quant_scale"): + raise NotImplementedError("GGML export does not support an AWQ-style pre_quant_scale.") + return weight_quantizer.num_bits - if weight_quantizer.num_bits == 8: - if input_quantizer is not None and input_quantizer.is_enabled: - return QUANTIZATION_INT8_SQ - else: - return QUANTIZATION_INT8_WO + if weight_quantizer.num_bits == 4: + assert len(weight_quantizer.block_sizes) > 0 and weight_quantizer.block_sizes[-1] > 0, ( + "Invalid block_sizes for INT4 quantizer" + ) + return QUANTIZATION_INT4_AWQ - if weight_quantizer.num_bits == (4, 3): - if weight_quantizer.block_sizes: - assert weight_quantizer.block_sizes[-1] > 0, "Invalid block_sizes for FP8 quantizer" - # Check if this is MXFP8 (dynamic block quantization with scale_bits (8, 0)) - block_sizes = getattr(weight_quantizer, "block_sizes") - if ( - isinstance(block_sizes, dict) - and block_sizes.get("type", "static") == "dynamic" - and block_sizes.get("scale_bits") == (8, 0) - ): - return QUANTIZATION_MXFP8 - if weight_quantizer.fake_quant: - return QUANTIZATION_FP8_PB_WO - else: - return QUANTIZATION_FP8_PB_REAL - if weight_quantizer.axis == 0: - return QUANTIZATION_FP8_PC_PT - return QUANTIZATION_FP8 + if weight_quantizer.num_bits == 8: + if input_quantizer is not None and input_quantizer.is_enabled: + return QUANTIZATION_INT8_SQ + return QUANTIZATION_INT8_WO - if weight_quantizer.num_bits == (2, 1): - # FP4 formats are all block quantization + if weight_quantizer.num_bits == (4, 3): + if weight_quantizer.block_sizes: + assert weight_quantizer.block_sizes[-1] > 0, "Invalid block_sizes for FP8 quantizer" block_sizes = getattr(weight_quantizer, "block_sizes") - scale_bits = block_sizes.get("scale_bits") - - if input_quantizer is not None and hasattr(weight_quantizer, "svdquant_lora_a"): - return QUANTIZATION_NVFP4_SVDQUANT - if input_quantizer is not None and hasattr(input_quantizer, "_pre_quant_scale"): - return QUANTIZATION_NVFP4_AWQ - if getattr(layer, "fused_with_prequant", False): - return QUANTIZATION_NVFP4_AWQ - if input_quantizer is None or not input_quantizer.is_enabled: - if scale_bits == (4, 3): - return QUANTIZATION_W4A16_NVFP4 - assert input_quantizer is not None, ( - f"input_quantizer is None for {quantizer_attr_names}" - ) - if ( - block_sizes.get("type", "static") == "dynamic" - and scale_bits == (8, 0) - and input_quantizer.is_enabled - and input_quantizer.num_bits == (4, 3) - and input_quantizer.block_sizes is None - ): - return QUANTIZATION_W4A8_MXFP4_FP8 if ( - block_sizes.get("type", "static") == "dynamic" - and scale_bits == (4, 3) - and input_quantizer.is_enabled - and input_quantizer.num_bits == (4, 3) - and input_quantizer.block_sizes is None + isinstance(block_sizes, dict) + and block_sizes.get("type", "static") == "dynamic" + and block_sizes.get("scale_bits") == (8, 0) ): - return QUANTIZATION_W4A8_NVFP4_FP8 + return QUANTIZATION_MXFP8 + if weight_quantizer.fake_quant: + return QUANTIZATION_FP8_PB_WO + return QUANTIZATION_FP8_PB_REAL + if weight_quantizer.axis == 0: + return QUANTIZATION_FP8_PC_PT + return QUANTIZATION_FP8 + + if weight_quantizer.num_bits == (2, 1): + block_sizes = getattr(weight_quantizer, "block_sizes") + scale_bits = block_sizes.get("scale_bits") + + if input_quantizer is not None and hasattr(weight_quantizer, "svdquant_lora_a"): + return QUANTIZATION_NVFP4_SVDQUANT + if input_quantizer is not None and hasattr(input_quantizer, "_pre_quant_scale"): + return QUANTIZATION_NVFP4_AWQ + if getattr(layer, "fused_with_prequant", False): + return QUANTIZATION_NVFP4_AWQ + if input_quantizer is None or not input_quantizer.is_enabled: if scale_bits == (4, 3): - return QUANTIZATION_NVFP4 - elif scale_bits == (8, 0): - return QUANTIZATION_MXFP4 + return QUANTIZATION_W4A16_NVFP4 + if scale_bits == (8, 0): + return QUANTIZATION_W4A16_MXFP4 + assert input_quantizer is not None, "input_quantizer is required for weight-activation FP4" + if ( + block_sizes.get("type", "static") == "dynamic" + and scale_bits == (8, 0) + and input_quantizer.is_enabled + and input_quantizer.num_bits == (4, 3) + and input_quantizer.block_sizes is None + ): + return QUANTIZATION_W4A8_MXFP4_FP8 + if ( + block_sizes.get("type", "static") == "dynamic" + and scale_bits == (4, 3) + and input_quantizer.is_enabled + and input_quantizer.num_bits == (4, 3) + and input_quantizer.block_sizes is None + ): + return QUANTIZATION_W4A8_NVFP4_FP8 + if scale_bits == (4, 3): + return QUANTIZATION_NVFP4 + if scale_bits == (8, 0): + return QUANTIZATION_MXFP4 + + raise NotImplementedError(f"Unsupported quantizer with num_bits: {weight_quantizer.num_bits}") - # Raise error for unsupported num_bits - raise NotImplementedError( - f"Unsupported quantizer with num_bits: {weight_quantizer.num_bits}" - ) +def get_quantization_format(module) -> str | None: + """Return the first enabled weight quantization format in a module tree.""" for weight_name in weight_attr_names(module): - quantization = _get_quantization_from_layer(module, quantizer_attr_names(weight_name)) + attr_names = quantizer_attr_names(weight_name) + quantization = _get_quantization_from_quantizers( + module, + representative_weight_quantizer(module, weight_name), + getattr(module, attr_names.input_quantizer, None), + ) if quantization != QUANTIZATION_NONE: return quantization @@ -707,6 +697,11 @@ def process_layer_quant_config(layer_config_dict): layer_config = {"quant_algo": "FP8"} elif v == "fp8_pc_pt": layer_config = {"quant_algo": "FP8_PER_CHANNEL_PER_TOKEN"} + elif v == "fp8_pb_wo": + layer_config = { + "quant_algo": "FP8_PB_WO", + "group_size": block_size_value, + } elif v == "int4_awq": layer_config = { "quant_algo": "W4A16_AWQ", @@ -766,6 +761,11 @@ def process_layer_quant_config(layer_config_dict): "quant_algo": "MXFP8", "group_size": block_size_value, } + elif v in ["mxfp4", "w4a16_mxfp4"]: + layer_config = { + "quant_algo": v.upper(), + "group_size": block_size_value, + } elif v in GGML_FORMATS: ggml_format = GGML_FORMAT_REGISTRY[v] block_size, payload_bytes = ggml_format.block_size, ggml_format.block_bytes @@ -794,6 +794,11 @@ def process_layer_quant_config(layer_config_dict): # If we have more than one quantization format, infer MIXED_PRECISION if len(quantization_formats) > 1: per_layer_config["quant_algo"] = "MIXED_PRECISION" + per_layer_config["exclude_modules"] = sorted( + _prefix_wildcard_summarize_exclude_modules( + exclude_modules, per_layer_config["quantized_layers"].keys() + ) + ) elif len(quantization_formats) == 1 and quantization_config is not None: per_layer_config.update(quantization_config) per_layer_config["exclude_modules"] = sorted( @@ -961,7 +966,11 @@ def to_quantized_weight( else weights_scaling_factor2, )[0]._quantized_data - if quantization in [QUANTIZATION_W4A8_MXFP4_FP8, QUANTIZATION_MXFP4]: + if quantization in [ + QUANTIZATION_W4A8_MXFP4_FP8, + QUANTIZATION_MXFP4, + QUANTIZATION_W4A16_MXFP4, + ]: return MXFP4QTensor.quantize(weight, block_size=block_size)[0]._quantized_data raise NotImplementedError(f"quantization format {quantization} not supported") @@ -1817,6 +1826,20 @@ def get_quant_config( for weight_name in weight_names: weight_block_size = get_weight_block_size(module, weight_name) + if quantization_format == QUANTIZATION_FP8_PB_WO: + # One group size can describe only square two-dimensional blocks. + weight_quantizer = representative_weight_quantizer(module, weight_name) + if weight_quantizer is not None and weight_quantizer.is_enabled: + block_axes = { + axis: size + for axis, size in weight_quantizer.block_sizes.items() + if isinstance(axis, int) + } + if block_axes != {-2: weight_block_size, -1: weight_block_size}: + raise NotImplementedError( + "Offline FP8 block export requires square blocks over " + f"the last two dimensions, got {block_axes} for {name}" + ) if weight_block_size > 0: block_size = weight_block_size break diff --git a/modelopt/torch/export/quantized_weight_export.py b/modelopt/torch/export/quantized_weight_export.py new file mode 100644 index 00000000000..1b1b3e5a5f2 --- /dev/null +++ b/modelopt/torch/export/quantized_weight_export.py @@ -0,0 +1,444 @@ +# 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. + +"""Capture state for functional quantized-weight export.""" + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Any + +import torch +import torch.nn as nn + +from modelopt.torch.quantization.qtensor import NVFP4QTensor, QTensorWrapper +from modelopt.torch.quantization.utils import quantizer_attr_names, representative_weight_quantizer + +from ..quantization.nn import SequentialQuantizer, TensorQuantizer +from .quant_format import ( + QUANTIZATION_FP8, + QUANTIZATION_FP8_PB_WO, + QUANTIZATION_FP8_PC_PT, + QUANTIZATION_MXFP4, + QUANTIZATION_MXFP8, + QUANTIZATION_NVFP4, + QUANTIZATION_W4A8_MXFP4_FP8, + QUANTIZATION_W4A8_NVFP4_FP8, + QUANTIZATION_W4A16_MXFP4, + QUANTIZATION_W4A16_NVFP4, +) +from .quant_utils import _get_quantization_from_quantizers, get_scaling_factor + +__all__ = [ + "capture_quantized_weight_export_state", + "get_quantized_weight_export_spec", +] + +_NVFP4_EXPORT_FORMATS = { + QUANTIZATION_NVFP4, + QUANTIZATION_W4A8_NVFP4_FP8, + QUANTIZATION_W4A16_NVFP4, +} +_MXFP4_EXPORT_FORMATS = {QUANTIZATION_MXFP4, QUANTIZATION_W4A8_MXFP4_FP8} +_FUNCTIONAL_WEIGHT_EXPORT_FORMATS = ( + _NVFP4_EXPORT_FORMATS + | _MXFP4_EXPORT_FORMATS + | { + QUANTIZATION_FP8, + QUANTIZATION_FP8_PC_PT, + QUANTIZATION_FP8_PB_WO, + QUANTIZATION_MXFP8, + } +) + + +class _UnsupportedQuantizedWeightExportFormatError(NotImplementedError): + pass + + +@dataclass(frozen=True) +class _ExportStateTensor: + name: str + value: torch.Tensor + axes: tuple[int, ...] = () + block_sizes: tuple[int, ...] = () + + +@dataclass(frozen=True) +class _QuantizedWeightExportState: + """Opaque state required to export one logical quantized weight.""" + + quantization_format: str + block_size: int + weight_shape: tuple[int, ...] + tensors: tuple[_ExportStateTensor, ...] + packing_permutation: tuple[int, ...] + input_quantizer_enabled: bool + static_nvfp4: bool = False + four_over_six: bool = False + + +@dataclass(frozen=True) +class _QuantizedWeightExportSpec: + """Stable deployment metadata for one logical quantized weight.""" + + quantization_format: str + block_size: int + input_quantizer_enabled: bool + + +def _same_storage(left: object, right: object) -> bool: + if left is right: + return True + if not isinstance(left, torch.Tensor) or not isinstance(right, torch.Tensor): + return False + if left.device.type == "meta" or right.device.type == "meta": + return False + return ( + left.device == right.device + and left.numel() == right.numel() + and left.untyped_storage().data_ptr() == right.untyped_storage().data_ptr() + and left.storage_offset() == right.storage_offset() + ) + + +def _resolve_weight_quantizer( + module: nn.Module, weight_name: str +) -> tuple[torch.Tensor, TensorQuantizer | SequentialQuantizer] | None: + weight = getattr(module, weight_name) + iter_weights = getattr(module, "iter_weights_for_calibration", None) + if iter_weights is not None: + for candidate, quantizer in iter_weights(): + if _same_storage(candidate, weight): + return candidate, quantizer + + quantizer = representative_weight_quantizer(module, weight_name) + if quantizer is None and weight_name.startswith("weight"): + quantizer = representative_weight_quantizer(module) + if quantizer is not None and weight_name.removeprefix("weight").isdigit(): + raise NotImplementedError( + f"Unable to resolve quantizer for numbered weight {weight_name!r}" + ) + if quantizer is None: + return None + # A representative quantizer cannot supply another expert's calibrated state. + if iter_weights is not None: + raise NotImplementedError( + f"Calibration iterator does not expose a matching full weight view for {weight_name!r}" + ) + return weight, quantizer + + +def _packing_permutation(weight: torch.Tensor, quantized_view: torch.Tensor) -> tuple[int, ...]: + ndim = weight.ndim + identity = tuple(range(ndim)) + if ndim >= 2: + transposed = weight.transpose(-1, -2) + if ( + tuple(quantized_view.shape) == tuple(transposed.shape) + and quantized_view.stride() == transposed.stride() + ): + return (*range(ndim - 2), ndim - 1, ndim - 2) + if ( + tuple(quantized_view.shape) == tuple(weight.shape) + and quantized_view.stride() == weight.stride() + ): + return identity + raise NotImplementedError( + "Unsupported quantized weight view " + f"shape/stride={tuple(quantized_view.shape)}/{quantized_view.stride()} for " + f"shape/stride={tuple(weight.shape)}/{weight.stride()}" + ) + + +def _state_tensor( + name: str, + value: torch.Tensor, + *, + axes: tuple[int, ...] = (), + block_sizes: tuple[int, ...] = (), + cpu: bool = True, +) -> _ExportStateTensor: + value = value.detach().cpu().clone() if cpu else value.detach().clone() + if value.numel() == 1 and not axes and not block_sizes: + return _ExportStateTensor(name, value.reshape(())) + if len(axes) != value.ndim or len(block_sizes) != value.ndim: + raise ValueError( + f"Explicit axis and block metadata is required for non-scalar {name}: " + f"shape={tuple(value.shape)}, axes={axes}, block_sizes={block_sizes}" + ) + return _ExportStateTensor(name, value, axes, block_sizes) + + +def _axis_scale_state( + name: str, + value: torch.Tensor, + packed_shape: tuple[int, ...], + axis: int | Sequence[int], + *, + cpu: bool = True, +) -> _ExportStateTensor: + axes = (axis,) if isinstance(axis, int) else tuple(axis) + axes = tuple(_normalize_weight_dim(dim, len(packed_shape)) for dim in axes) + expected_shape = tuple(packed_shape[dim] for dim in axes) + if value.numel() != torch.Size(expected_shape).numel(): + raise RuntimeError( + f"Scale {name!r} shape {tuple(value.shape)} does not match axes {axes} " + f"of packed weight shape {packed_shape}" + ) + return _state_tensor( + name, + value.reshape(expected_shape), + axes=axes, + block_sizes=(1,) * len(axes), + cpu=cpu, + ) + + +def _block_scale_state( + name: str, + value: torch.Tensor, + packed_shape: tuple[int, ...], + block_config: Mapping[int | str, Any], + *, + cpu: bool = True, +) -> _ExportStateTensor: + ndim = len(packed_shape) + block_sizes = [1] * ndim + for dim, block_size in block_config.items(): + if isinstance(dim, int) and block_size is not None: + block_sizes[_normalize_weight_dim(dim, ndim)] = int(block_size) + expected_shape = tuple( + (size + block_size - 1) // block_size for size, block_size in zip(packed_shape, block_sizes) + ) + if value.numel() != torch.Size(expected_shape).numel(): + raise RuntimeError( + f"Scale {name!r} shape {tuple(value.shape)} does not match block sizes " + f"{tuple(block_sizes)} of packed weight shape {packed_shape}" + ) + + expanded_shape = [] + expanded_axes = [] + expanded_block_sizes = [] + for axis, (size, block_size) in enumerate(zip(expected_shape, block_sizes)): + expanded_shape.append(size) + expanded_axes.append(axis) + expanded_block_sizes.append(block_size) + if block_size != 1: + expanded_shape.append(1) + expanded_axes.append(axis) + expanded_block_sizes.append(1) + if tuple(value.shape) == tuple(expanded_shape): + return _state_tensor( + name, + value, + axes=tuple(expanded_axes), + block_sizes=tuple(expanded_block_sizes), + cpu=cpu, + ) + return _state_tensor( + name, + value.reshape(expected_shape), + axes=tuple(range(ndim)), + block_sizes=tuple(block_sizes), + cpu=cpu, + ) + + +def _input_quantizer(module: nn.Module, weight_name: str): + quantizer = getattr(module, quantizer_attr_names(weight_name).input_quantizer, None) + if quantizer is None: + quantizer = getattr(module, "input_quantizer", None) + return quantizer + + +def _resolve_quantized_weight_export_inputs( + module: nn.Module, + weight_name: str, +) -> ( + tuple[ + torch.Tensor, + torch.Tensor, + TensorQuantizer, + TensorQuantizer | None, + str, + ] + | None +): + weight = getattr(module, weight_name) + if isinstance(weight, QTensorWrapper): + raise NotImplementedError("Functional export requires an uncompressed source weight") + + resolved = _resolve_weight_quantizer(module, weight_name) + if resolved is None: + return None + quantized_view, weight_quantizer = resolved + input_quantizer = _input_quantizer(module, weight_name) + if not weight_quantizer.is_enabled: + return None + quantization_format = _get_quantization_from_quantizers( + module, weight_quantizer, input_quantizer + ) + if quantization_format is None: + raise RuntimeError(f"Unable to resolve quantization format for {weight_name!r}") + if isinstance(weight_quantizer, SequentialQuantizer): + raise _UnsupportedQuantizedWeightExportFormatError( + f"Functional export does not support {quantization_format!r}" + ) + if quantization_format == QUANTIZATION_W4A16_MXFP4: + raise _UnsupportedQuantizedWeightExportFormatError( + "Functional export does not support weight-only MXFP4" + ) + if quantization_format == QUANTIZATION_FP8_PB_WO and weight_quantizer.block_sizes != { + -2: 128, + -1: 128, + }: + raise _UnsupportedQuantizedWeightExportFormatError( + "Functional export supports FP8 2D blockwise weight quantization only with " + "128x128 blocks" + ) + if quantization_format not in _FUNCTIONAL_WEIGHT_EXPORT_FORMATS: + raise _UnsupportedQuantizedWeightExportFormatError( + f"Functional export does not support {quantization_format!r}" + ) + return weight, quantized_view, weight_quantizer, input_quantizer, quantization_format + + +def capture_quantized_weight_export_state( + module: nn.Module, + weight_name: str = "weight", + *, + cpu: bool = True, +) -> _QuantizedWeightExportState | None: + """Capture detached state for one weight, or ``None`` when it is not quantized.""" + resolved = _resolve_quantized_weight_export_inputs(module, weight_name) + if resolved is None: + return None + weight, quantized_view, weight_quantizer, input_quantizer, quantization_format = resolved + + permutation = _packing_permutation(weight, quantized_view) + packed_shape = tuple(weight.shape[axis] for axis in permutation) + block_config = getattr(weight_quantizer, "block_sizes", None) or {} + block_size = int(block_config.get(-1, 0)) if isinstance(block_config, dict) else 0 + tensors = [] + static_nvfp4 = ( + quantization_format in _NVFP4_EXPORT_FORMATS + and NVFP4QTensor._is_static_quantizer(weight_quantizer) + ) + + if static_nvfp4: + if not block_size or packed_shape[-1] % block_size: + raise RuntimeError(f"Invalid static NVFP4 block size for {weight_name!r}") + per_block_amax = weight_quantizer.amax + global_amax = NVFP4QTensor._get_static_global_amax(weight_quantizer) + if per_block_amax is None or global_amax is None: + raise RuntimeError(f"Missing calibrated static NVFP4 state for {weight_name!r}") + block_shape = (*packed_shape[:-1], packed_shape[-1] // block_size) + tensors.append( + _state_tensor( + "weight_block_amax", + per_block_amax.reshape(block_shape), + axes=tuple(range(len(packed_shape))), + block_sizes=(1,) * (len(packed_shape) - 1) + (block_size,), + cpu=cpu, + ) + ) + tensors.append(_state_tensor("weight_global_amax", global_amax, cpu=cpu)) + elif quantization_format in _NVFP4_EXPORT_FORMATS: + weight_scale_2 = ( + weight_quantizer._amax.float() / 448.0 + if quantization_format == QUANTIZATION_W4A8_NVFP4_FP8 + else NVFP4QTensor.get_weights_scaling_factor_2_from_quantizer(weight_quantizer) + ) + tensors.append(_state_tensor("weight_scale_2", weight_scale_2, cpu=cpu)) + elif quantization_format == QUANTIZATION_FP8: + weight_scale = get_scaling_factor(weight_quantizer) + if weight_scale is None: + raise RuntimeError(f"Missing calibrated weight scale for {weight_name!r}") + tensors.append(_state_tensor("weight_scale", weight_scale, cpu=cpu)) + elif quantization_format == QUANTIZATION_FP8_PC_PT: + weight_scale = get_scaling_factor(weight_quantizer) + if weight_scale is None or weight_quantizer.axis is None: + raise RuntimeError(f"Missing calibrated per-channel scale for {weight_name!r}") + tensors.append( + _axis_scale_state( + "weight_scale", + weight_scale, + packed_shape, + weight_quantizer.axis, + cpu=cpu, + ) + ) + elif quantization_format == QUANTIZATION_FP8_PB_WO: + weight_scale = get_scaling_factor(weight_quantizer) + if weight_scale is None: + raise RuntimeError(f"Missing calibrated block scale for {weight_name!r}") + tensors.append( + _block_scale_state("weight_scale", weight_scale, packed_shape, block_config, cpu=cpu) + ) + elif quantization_format == QUANTIZATION_MXFP8: + cached_scale = getattr(weight_quantizer, "_scale", None) + if cached_scale is not None: + tensors.append( + _block_scale_state( + "weight_scale", cached_scale, packed_shape, block_config, cpu=cpu + ) + ) + elif quantization_format not in _MXFP4_EXPORT_FORMATS: + raise AssertionError(f"Unexpected functional export format {quantization_format!r}") + + if input_quantizer is not None and input_quantizer.is_enabled: + input_amax = input_quantizer.export_amax() + if input_amax is not None: + input_scale = ( + NVFP4QTensor.get_activation_scaling_factor(input_quantizer) + if quantization_format == QUANTIZATION_NVFP4 + else get_scaling_factor(input_quantizer) + ) + tensors.append(_state_tensor("input_scale", input_scale, cpu=cpu)) + + return _QuantizedWeightExportState( + quantization_format=quantization_format, + block_size=block_size, + weight_shape=tuple(weight.shape), + tensors=tuple(tensors), + packing_permutation=permutation, + input_quantizer_enabled=input_quantizer is not None and input_quantizer.is_enabled, + static_nvfp4=static_nvfp4, + four_over_six=bool(block_config.get("four_over_six", False)), + ) + + +def get_quantized_weight_export_spec( + module: nn.Module, + weight_name: str = "weight", +) -> _QuantizedWeightExportSpec | None: + """Return stable deployment metadata without retaining quantizer tensors.""" + resolved = _resolve_quantized_weight_export_inputs(module, weight_name) + if resolved is None: + return None + _, _, weight_quantizer, input_quantizer, quantization_format = resolved + block_config = getattr(weight_quantizer, "block_sizes", None) or {} + block_size = int(block_config.get(-1, 0)) if isinstance(block_config, dict) else 0 + return _QuantizedWeightExportSpec( + quantization_format, + block_size, + input_quantizer is not None and input_quantizer.is_enabled, + ) + + +def _normalize_weight_dim(weight_dim: int, ndim: int) -> int: + if not -ndim <= weight_dim < ndim: + raise IndexError(f"Weight dimension {weight_dim} is invalid for a rank-{ndim} weight") + return weight_dim % ndim diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 4c39674f585..f605ddaf636 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -1027,10 +1027,9 @@ def _prepare_model_for_export(model, dtype, is_modelopt_qlora): def pack_quantized_weights(model, dtype, is_modelopt_qlora: bool = False) -> None: """Quantize every module's weight in place, then rebuild the fused MoE linears.""" - # Deferred: modelopt.torch.quantization.plugins.huggingface imports transformers at module - # scope, and transformers is an optional extra -- importing it here keeps - # ``import modelopt.torch.export`` working without it. - from modelopt.torch.quantization.plugins.huggingface import _reconstruct_fused_moe_linear + # Deferred: modelopt.torch.quantization imports modelopt.torch.export, so importing the + # per-model PTQ module at module scope here would risk an import cycle. + from modelopt.torch.models.step3p5.modeling_ptq import _reconstruct_fused_moe_linear _process_quantized_modules(model, dtype, is_modelopt_qlora) _reconstruct_fused_moe_linear(model) @@ -1448,12 +1447,20 @@ def _export_diffusers_checkpoint( if not is_partial_export and not model_index_path.exists() and hasattr(pipe, "save_config"): pipe.save_config(export_dir) - # Last resort: synthesize a minimal model_index.json from exported components. + # Preserve pipeline settings while listing only the exported components. if not model_index_path.exists() and hasattr(pipe, "config") and pipe.config is not None: + pipeline_components = pipe.components model_index = { - "_class_name": type(pipe).__name__, - "_diffusers_version": diffusers.__version__, + key: value + for key, value in json.loads(pipe.to_json_string()).items() + if key not in pipeline_components and not key.startswith("_") } + model_index.update( + { + "_class_name": type(pipe).__name__, + "_diffusers_version": diffusers.__version__, + } + ) for name, comp in all_components.items(): module = type(comp).__module__ library = module.split(".")[0] @@ -1654,8 +1661,12 @@ def _source_checkpoint(model: nn.Module) -> str | None: tensors never reach ``exclude_modules`` because the two halves disagreed about where the checkpoint was. """ - return getattr(model, "_modelopt_source_checkpoint", None) or getattr( - getattr(model, "config", None), "_name_or_path", None + # ``or None``: a model built with ``from_config`` has ``_name_or_path == ""``, which would + # otherwise name the current directory. + return ( + getattr(model, "_modelopt_source_checkpoint", None) + or getattr(getattr(model, "config", None), "_name_or_path", None) + or None ) @@ -1881,7 +1892,7 @@ def export_hf_checkpoint( export_dir: The target export path. save_modelopt_state: Whether to save the modelopt state_dict. components: Only used for diffusers pipelines. Optional list of component names - to export. If None, all quantized components are exported. + to export. If None, all components are exported. extra_state_dict: Extra state dictionary to add to the exported model. max_shard_size: Maximum size of each safetensors shard file. Defaults to "10GB". **kwargs: Runtime-specific post-processing options forwarded to diff --git a/modelopt/torch/export/unified_export_hf_streaming.py b/modelopt/torch/export/unified_export_hf_streaming.py index 7f692194fa6..69cdc53117f 100644 --- a/modelopt/torch/export/unified_export_hf_streaming.py +++ b/modelopt/torch/export/unified_export_hf_streaming.py @@ -211,7 +211,7 @@ def _assert_no_split_rules(model: nn.Module) -> None: reversal cannot do. """ try: - split_rules, merge_rules, _, _ = _build_reverse_rules(model) + split_rules, merge_rules, _, _, _ = _build_reverse_rules(model) except Exception: return # build_reverse_name_mapper reports the failure with a warning if split_rules or merge_rules: @@ -346,9 +346,9 @@ def _export_transformers_checkpoint_streaming( NotImplementedError: if the model's conversion mapping regroups tensors. RuntimeError: if decoder layers cannot be discovered for layer-wise materialization. """ - # Deferred: the huggingface plugin imports transformers at module scope, and transformers - # is an optional extra -- keep ``import modelopt.torch.export`` working without it. - from modelopt.torch.quantization.plugins.huggingface import _reconstruct_fused_moe_linear + # Deferred: modelopt.torch.quantization imports modelopt.torch.export, so importing the + # per-model PTQ module at module scope here would risk an import cycle. + from modelopt.torch.models.step3p5.modeling_ptq import _reconstruct_fused_moe_linear export_dir = Path(export_dir) # Materialization dispatch walks the module tree from the root; without these maps each @@ -545,11 +545,20 @@ def collect_export_tensors( All the gathers finish before this returns, so the caller can write or postprocess without stalling anyone. Returning a list rather than a generator is what guarantees that. """ - # Deferred: the huggingface plugin imports transformers at module scope, and transformers - # is an optional extra -- keep ``import modelopt.torch.export`` working without it. - from modelopt.torch.quantization.plugins.huggingface import _reconstruct_fused_moe_linear + # Deferred: modelopt.torch.quantization imports modelopt.torch.export, so importing the + # per-model PTQ module at module scope here would risk an import cycle. + from modelopt.torch.models.step3p5.modeling_ptq import _reconstruct_fused_moe_linear _assert_fsdp2_owns_every_mesh_dim(model) + # A forward leaves the root's parameters unsharded (FSDP2 skips the root's post-forward + # reshard), and FSDP2's state_dict() pre-hook reshards a module's whole parameter group. Only + # the owner calls state_dict() below, so it would reshard the root on its rank alone, and the + # next gather window would then gather parameters on that rank that the others see as plain + # tensors -- a collective they never join. Resharding locally on every rank first makes the + # pre-hook a no-op. + for module in model.modules(): + if isinstance(module, FSDPModule): + module.reshard() my_rank, world = _dist.rank(), _dist.size() names = module_name_maps(model) ctx = ExportContext( diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py index 1d6a0fc99d7..502c8f863de 100644 --- a/modelopt/torch/export/unified_export_megatron.py +++ b/modelopt/torch/export/unified_export_megatron.py @@ -18,7 +18,6 @@ """Code that export quantized Megatron Core models for deployment.""" -import io import json import os import re @@ -38,6 +37,7 @@ from modelopt.torch.quantization.ggml import GGML_FORMAT_REGISTRY from modelopt.torch.quantization.nn.modules.tensor_quantizer import GroupedQuantizer from modelopt.torch.utils import import_plugin, warn_rank_0 +from modelopt.torch.utils.distributed import gather_state_dicts from modelopt.torch.utils.plugins.hf_checkpoint_utils import ( copy_hf_ckpt_remote_code, copy_non_safetensor_files_from_ckpt, @@ -1709,29 +1709,14 @@ def _grouped_mlp_slicing( self._record_layer_quant_config(expert_prefix, seen_qformat, seen_block_size) if ep_size > 1: - # all_gather_object pickles trip on quantized uint8 tensors whose - # UntypedStorage has no dtype attr -- round-trip through torch.save bytes. - _buf = io.BytesIO() - torch.save(local_expert_state, _buf) - local_bytes = _buf.getvalue() - del _buf # Gather to EP rank 0 only, which writes the shards: holding every expert on every EP # rank multiplies host memory by EP (a full GLM-5.3-Flash export OOMs at EP4). + # Quantized uint8 tensors trip all_gather_object pickles, so go through torch.save bytes. ep_group = get_expert_model_parallel_group() - gathered_bytes: list | None = [None] * ep_size if ep_rank == 0 else None - torch.distributed.gather_object( - local_bytes, - gathered_bytes, - dst=torch.distributed.get_global_rank(ep_group, 0), - group=ep_group, - ) - del local_bytes - for b in gathered_bytes or (): - # weights_only=False: our own torch.save output from a sibling EP rank - # in this job's collective, not user-supplied. - s_loaded = torch.load(io.BytesIO(b), map_location="cpu", weights_only=False) + for s_loaded in gather_state_dicts( + local_expert_state, torch.distributed.get_global_rank(ep_group, 0), ep_group + ): self._state_dict.update(s_loaded) - del gathered_bytes else: self._state_dict.update(local_expert_state) diff --git a/modelopt/torch/kernels/quantization/gemm/tensor_quant_gpu_fp8.cu b/modelopt/torch/kernels/quantization/gemm/tensor_quant_gpu_fp8.cu index 9db141f0cb4..b4493ac8939 100644 --- a/modelopt/torch/kernels/quantization/gemm/tensor_quant_gpu_fp8.cu +++ b/modelopt/torch/kernels/quantization/gemm/tensor_quant_gpu_fp8.cu @@ -17,12 +17,17 @@ #include #include +#include #include #include #include #define BLOCK_SIZE 128 +// Only an amax whose FP8 scale (448 / amax) overflows is treated as zero. Any other positive amax +// is valid, e.g. the tiny global_amax / 6 used to quantize NVFP4 block scales. +#define ZERO_AMAX_THRESHOLD (448.f / FLT_MAX) + #define AT_DISPATCH_CASE_FLOATING_TYPES(...) \ AT_DISPATCH_CASE(at::ScalarType::Double, __VA_ARGS__) \ AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ @@ -71,8 +76,7 @@ at::Tensor fake_e4m3fy_with_axis(at::Tensor inputs, at::Tensor amax, int axis) { int axis_size = inputs.size(axis); int outer_size = inputs.stride(axis); - const float epsilon = 1.0f / (1 << 24); - auto zero_amax_mask = amax <= epsilon; + auto zero_amax_mask = amax <= ZERO_AMAX_THRESHOLD; auto safe_amax = at::where(zero_amax_mask, torch::ones_like(amax), amax); auto scale = 448.f / safe_amax; auto inv_scale = 1.f / scale; @@ -91,8 +95,7 @@ at::Tensor fake_e4m3fy(at::Tensor inputs, at::Tensor amax) { inputs = inputs.contiguous(); amax = amax.view(-1).to(at::kFloat); size_t numel = inputs.numel(); - const float epsilon = 1.0f / (1 << 24); - auto zero_amax_mask = amax <= epsilon; + auto zero_amax_mask = amax <= ZERO_AMAX_THRESHOLD; auto safe_amax = at::where(zero_amax_mask, torch::ones_like(amax), amax); at::Tensor scale = 448.f / safe_amax; auto inv_scale = 1.f / scale; diff --git a/modelopt/torch/models/README.md b/modelopt/torch/models/README.md index 287fd7f7da1..c0caf7f6e01 100644 --- a/modelopt/torch/models/README.md +++ b/modelopt/torch/models/README.md @@ -53,18 +53,19 @@ an intra-family detail, not a package API. ## PTQ modeling Quantized-module wrappers and registrations that only one model needs (e.g. Llama4's fused -BMM experts) live in `/modeling_ptq.py`; generic +BMM experts, Step's expert-indexed `MoELinear`) live in `/modeling_ptq.py`; generic ones (fused/sequential MoE auto-detection, attention, `FP8Linear`) stay in `modelopt/torch/quantization/plugins/huggingface.py`. The HF plugin imports every `modeling_ptq` from an explicit list at its end, so add the new model type there. Leave the model's `__init__.py` importing only `specs`: `modeling_ptq` pulls in quantization and transformers, which `import modelopt.torch.models` must not. -A `modeling_ptq.py` registers on import — static classes via `QuantModuleRegistry.register` -behind a `try`/`except ImportError`, remote-code classes via a callback added to +A `modeling_ptq.py` registers on import — static classes via `QuantModuleRegistry.register`, +imported from transformers at module scope; remote-code classes via a callback added to `CUSTOM_MODEL_PLUGINS` — and exports nothing (`__all__ = []`). Every package still has a `specs.py`, even when its spec only records `model_type` and where the modeling code comes -from. +from. A family revision that shares another type's PTQ code (`step3p7` uses `step3p5`'s) still +gets a `modeling_ptq.py` that imports the shared one, so loading it by model type works. ## Sections diff --git a/modelopt/torch/models/gpt_oss/modeling_ptq.py b/modelopt/torch/models/gpt_oss/modeling_ptq.py new file mode 100644 index 00000000000..e92780ef34e --- /dev/null +++ b/modelopt/torch/models/gpt_oss/modeling_ptq.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. + +"""GPT-OSS PTQ modeling (HF model type ``gpt_oss``).""" + +from contextlib import contextmanager +from functools import partial + +import torch +from transformers.models.gpt_oss.modeling_gpt_oss import GptOssExperts + +from modelopt.torch.quantization.nn import QuantModuleRegistry, TensorQuantizer +from modelopt.torch.quantization.plugins.custom import _QuantFunctionalMixin +from modelopt.torch.quantization.plugins.huggingface import ( + _transposed_quantize, + _TransposedExpertsCalibMixin, +) + +__all__: list[str] = [] + + +class _QuantGptOssExperts(_TransposedExpertsCalibMixin, _QuantFunctionalMixin): + """Quantized wrapper for `transformers.GptOssExperts`. + + Quantizes `gate_up_proj` and `down_proj` weights via dynamic attributes inside `quantize_weight()`. + Activations into `gate_up_proj` are quantized by `gate_up_proj_input_quantizer`. For `down_proj` + activation quantization, we intercept `torch.Tensor.__matmul__`/`torch.bmm` and quantize inputs + on every second call (since the first call computes `gate_up_proj` outputs and second call + computes `down_proj` outputs). + """ + + @staticmethod + def _get_quantized_weight(quantizer, module, weight): + # MoE weight is accessed for each expert in one forward pass. so lets cache it + if module._enable_weight_quantization: + if hasattr(quantizer, "_cached_quant_val"): + return getattr(quantizer, "_cached_quant_val") + quantizer._cached_quant_val = _transposed_quantize(weight, quantizer) + return quantizer._cached_quant_val + return weight + + def _setup_for_weight_quantization(self): + self._register_dynamic_attribute( + "gate_up_proj", partial(self._get_quantized_weight, self.gate_up_proj_weight_quantizer) + ) + self._register_dynamic_attribute( + "down_proj", partial(self._get_quantized_weight, self.down_proj_weight_quantizer) + ) + + def _setup(self): + assert not hasattr(self, "kernel_layer_name"), ( + "ModelOpt quantization does not support patched forward for kernel_hub" + ) + self.gate_up_proj_input_quantizer = TensorQuantizer() + self.gate_up_proj_weight_quantizer = TensorQuantizer() + self.down_proj_input_quantizer = TensorQuantizer() + self.down_proj_weight_quantizer = TensorQuantizer() + + self._register_temp_attribute("_enable_weight_quantization", False) + self._register_temp_attribute("_down_proj_mul", False) + self._setup_for_weight_quantization() + + @property + def functionals_to_replace(self): + # Use torch.ops.aten to bypass Python dispatch and avoid RecursionError + # (torch.matmul / __matmul__ can dispatch to each other) + _aten_bmm = torch.ops.aten.bmm + _aten_matmul = torch.ops.aten.matmul + + def _quantized_bmm(batch1, batch2, *, out=None): + batch1 = self.down_proj_input_quantizer(batch1) if self._down_proj_mul else batch1 + self._down_proj_mul = not self._down_proj_mul # toggle the flag + if out is not None: + return torch.ops.aten.bmm.out(batch1, batch2, out=out) + return _aten_bmm(batch1, batch2) + + def _tensor_matmul(self_t, other): + self_t = self.down_proj_input_quantizer(self_t) if self._down_proj_mul else self_t + self._down_proj_mul = not self._down_proj_mul + return _aten_matmul(self_t, other) + + return [ + (torch, "bmm", _quantized_bmm), + (torch.Tensor, "__matmul__", _tensor_matmul), + ] + + @contextmanager + def quantize_weight(self): + """Context in which MoE weight is quantized.""" + self._enable_weight_quantization = True + try: + yield + finally: + for module in self.modules(): + if isinstance(module, TensorQuantizer) and hasattr(module, "_cached_quant_val"): + delattr(module, "_cached_quant_val") + self._enable_weight_quantization = False + + def forward( + self, hidden_states: torch.Tensor, router_indices=None, routing_weights=None + ) -> torch.Tensor: + """Forward method to add quantization.""" + hidden_states = self.gate_up_proj_input_quantizer(hidden_states) + with self.quantize_weight(): + return super().forward(hidden_states, router_indices, routing_weights) + + +if GptOssExperts not in QuantModuleRegistry: + QuantModuleRegistry.register({GptOssExperts: "hf.GptOssExperts"})(_QuantGptOssExperts) diff --git a/modelopt/torch/models/step3p5/modeling_ptq.py b/modelopt/torch/models/step3p5/modeling_ptq.py new file mode 100644 index 00000000000..9e3d40d49b8 --- /dev/null +++ b/modelopt/torch/models/step3p5/modeling_ptq.py @@ -0,0 +1,264 @@ +# 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. + +"""Step-family PTQ modeling (HF model types ``step3p5``, ``step3p7``). + +The package is named for the first Step release; later revisions share its expert-indexed +``MoELinear`` layout and are matched by :func:`_is_step_family_model`. +""" + +import inspect +import re + +import torch +import torch.nn as nn +from torch import Tensor + +from modelopt.torch.models import hf_model_type +from modelopt.torch.quantization.nn import QuantModule, QuantModuleRegistry +from modelopt.torch.quantization.plugins.custom import CUSTOM_MODEL_PLUGINS + +__all__: list[str] = [] + + +class _QuantMoELinear(QuantModule): + """Quantization wrapper for expert-indexed MoELinear modules (fused expert weights). + + MoELinear has weight shape [num_experts, out_features, in_features] with + forward(x, expert_id). We expand it into per-expert nn.Linear modules so + each expert gets its own weight_quantizer and input_quantizer, calibrated + only on tokens actually routed to that expert. + + On export, _reconstruct_fused_moe_linear() stacks the per-expert quantized + weights and scales back into the original 3D format. + + Note: we use expansion-then-reconstruction rather than the add_module() approach + because vLLM requires stacked 3D scaling factors; per-expert expanded keys are + not accepted by the downstream serving engine. + """ + + def _setup(self): + # Deferred: accelerate is a separate optional dependency, and the HF plugin (which + # imports this module) must stay importable without it. + from accelerate import init_empty_weights + + # Accelerate's CPU/disk offload (`device_map="auto"`, `--offload_folder`) leaves + # `weight` as a meta tensor and keeps the real value in the module's offload hook, + # keyed on the original `weight` name. Expanding that would copy meta storage into + # every expert and then delete the key the hook restores into, silently producing a + # checkpoint of zeros. Refuse instead of corrupting. + if self.weight.is_meta or getattr(getattr(self, "_hf_hook", None), "offload", False): + raise NotImplementedError( + f"{type(self).__name__}: expert-indexed MoELinear weights cannot be quantized " + "while offloaded by Accelerate (the weight is a meta tensor whose value lives " + "in the offload hook). Load the model without CPU/disk offload — more GPUs, or " + "a device_map that keeps the MoE layers resident — and re-run." + ) + + dtype, device = self.weight.dtype, self.weight.device + + with init_empty_weights(): + experts = nn.ModuleList( + [ + nn.Linear(self.in_features, self.out_features, bias=False) + for _ in range(self.num_experts) + ] + ) + + for i in range(self.num_experts): + experts[i].to_empty(device=device) + with torch.no_grad(): + experts[i].weight.data = self.weight[i].detach().to(dtype=dtype, device=device) + + delattr(self, "weight") + self.experts = experts + + def forward(self, x, expert_id): + # experts[expert_id] is a _QuantLinear after quantization wrapping, providing + # per-expert input_quantizer and weight_quantizer. + # + # MoELinear.forward always promotes to fp32 for the matmul regardless of storage + # dtype (`F.linear(x.float(), self.weight[expert_id].float())`), so leaving the + # expert's weight at its native storage dtype (e.g. bf16) and downcasting x to + # match before the matmul would compute in bf16 and change the model's output even + # with every quantizer disabled -- the bf16 rounding this class exists to quantize + # *past*, not to reintroduce as a side effect of conversion. + # + # A prior version of this fix instead expanded every expert's weight in fp32 + # permanently in `_setup`. That reproduces Step's fp32 compute but turns a per-call + # transient promotion into persistent model state: on Step-3.7's full routed-expert + # set (42 layers x 3 projections x 288 experts x 4096 x 1280), doubling from bf16 to + # fp32 adds roughly 354 GiB held throughout calibration, on top of device placement + # already sized for bf16 -- a model that loaded successfully can then OOM. It also + # left disabled/unquantized experts reconstructed at fp32 in the exported checkpoint. + # + # Instead, only the one expert actually being called is promoted, transiently, for + # the duration of this one call -- matching Step's own per-call `.float()` memory + # profile instead of Step-3.7's full expert set. `expert.weight` is read here + # outside any `quantize_weight()` context, so `_get_quantized_weight` passes it + # through unchanged and this is the real underlying nn.Parameter (the same pattern + # `_setup` above uses), not a value computed by the quantizer -- so reassigning its + # `.data` genuinely mutates the persisted storage, not a transient wrapper. + # + # This must keep calling `expert(x)` (`__call__`, not `.forward()`) rather than + # reimplementing the input/weight-quantize/output-quantize sequence inline: some + # calibration algorithms (e.g. `local_hessian_calibrate`) register a + # `forward_pre_hook` directly on the quantized Linear module, which only fires + # through standard `nn.Module.__call__` dispatch. + expert = self.experts[expert_id] + original_weight = expert.weight.data + with torch.no_grad(): + expert.weight.data = original_weight.float() + try: + out = expert(x.float()) + finally: + with torch.no_grad(): + expert.weight.data = original_weight + return out.float() + + +def _is_expert_indexed_moe_linear(module: nn.Module) -> bool: + """Whether ``module`` packs one projection's experts into an expert-indexed 3-D weight. + + The Step family (``stepfun-ai/Step-3.5-Flash``, ``stepfun-ai/Step-3.7-Flash``) ships a + custom ``MoELinear`` via ``trust_remote_code``: a plain ``nn.Module`` holding a single + ``weight`` of shape ``[num_experts, out_features, in_features]``, whose + ``forward(x, expert_id)`` runs ``F.linear`` against the selected expert's slice. The + weights therefore live on the projection submodule rather than on the expert container, + which is what :func:`_fused_experts_wrapper_class` looks for, and the module is not an + ``nn.Linear``, so neither the fused-experts path nor the plain linear path claims it. + + Detection is structural rather than keyed on class names so new Step revisions are picked + up without another hardcoded name, but the shape alone is not a sufficient contract: the + replacement forward indexes ``self.experts[expert_id]``, so it only works for callers that + pass a **scalar expert index**. Grouped-GEMM MoE layers share the exact same 3-D weight and + attribute set while passing a per-expert token-count *tensor* instead (e.g. Moondream3's + ``MoeFusedLinear.forward(input, m_sizes)``, which would raise ``TypeError: only integer + tensors of a single element can be converted to an index`` on the first calibration + forward). The second parameter must therefore be named ``expert_id``, which is the + scalar-index contract both Step revisions declare. + """ + weight = getattr(module, "weight", None) + if not isinstance(weight, (nn.Parameter, Tensor)) or weight.dim() != 3: + return False + if not all(hasattr(module, attr) for attr in ("num_experts", "in_features", "out_features")): + return False + # The wrapper rebuilds the weight as `num_experts` Linears of (out_features, in_features), + # so a 3-D weight laid out any other way would silently copy the wrong slices. + if tuple(weight.shape) != (module.num_experts, module.out_features, module.in_features): + return False + try: + params = list(inspect.signature(type(module).forward).parameters.values())[1:] + except (TypeError, ValueError): + return False + # The replacement forward is exactly `(x, expert_id)`, so anything the caller could pass + # beyond those two — a keyword-only `router_state`, *args, **kwargs — would raise once + # converted. Require the signature to match what the wrapper can honour. + return ( + len(params) == 2 + and all(p.kind is p.POSITIONAL_OR_KEYWORD and p.default is p.empty for p in params) + and params[1].name == "expert_id" + ) + + +_STEP_FAMILY_RE = re.compile(r"(?i)^step\d") + + +def _is_step_family_model(model: nn.Module) -> bool: + """Whether ``model`` is a Step-family root model (Step-3.5, Step-3.7, or a future revision). + + Matched against the ``step`` convention shared by ``model_type`` (``"step3p5"``, + ``"step3p7"``) and the remote-code class name (``Step3p5ForCausalLM``, + ``Step3p7ForConditionalGeneration``), not an exact revision, so a new Step release is + still picked up without another hardcoded name. This is deliberately narrower than the + shape/signature check in :func:`_is_expert_indexed_moe_linear` alone: that check accepts + any module with a matching 3-D weight and an ``(x, expert_id)`` forward, which is a + coincidence risk on its own -- an unrelated architecture happening to reuse the parameter + name ``expert_id`` with different semantics (a per-expert bias or post-scale, say) would + be claimed and have that behavior silently dropped by the replacement wrapper. Gating on + the model family keeps the shape check doing what it is actually good at: telling + Step revisions apart without a class-name allowlist, rather than distinguishing Step + from arbitrary third-party MoE code. + """ + model_type = hf_model_type(model) or "" + return bool(_STEP_FAMILY_RE.match(model_type) or _STEP_FAMILY_RE.match(type(model).__name__)) + + +def register_moe_linear_on_the_fly(model): + """Register expert-indexed ``MoELinear`` modules (Step-3.5 / Step-3.7) for quantization. + + Without this the routed experts carry no quantizer at all: an experts-only recipe matches + nothing and the export writes a checkpoint with ``quant_algo: null``. + """ + if not _is_step_family_model(model): + return + visited_types = set() + for name, module in model.named_modules(): + mod_type = type(module) + if mod_type in visited_types or QuantModuleRegistry.get(mod_type) is not None: + continue + visited_types.add(mod_type) + + if _is_expert_indexed_moe_linear(module): + print( + f"\033[1mDetected expert-indexed MoE linear '{name}' of type " + f"{mod_type.__name__}, registering with _QuantMoELinear.\033[0m" + ) + QuantModuleRegistry.register({mod_type: f"hf.{mod_type.__name__}"})(_QuantMoELinear) + + +def _reconstruct_fused_moe_linear(model: nn.Module) -> None: + """Reconstruct :class:`_QuantMoELinear` per-expert weights back to the 3-D MoELinear format. + + After _process_quantized_modules, each expert's nn.Linear inside the wrapper has: + - weight: fp4-quantized tensor [out_features, in_features] + - weight_scale, weight_scale_2: per-block / global scales + - input_scale: activation scale (if calibrated) + + This stacks them back into the original MoELinear layout so the exported state_dict + uses the original key names (e.g. moe.up_proj.weight with shape [N, out, in]). + + Matched by wrapper type rather than by the dynamically generated class name (``Quant`` + + the model's own class name): a model whose class is not spelled ``MoELinear`` would + otherwise quantize normally but export unusable per-expert keys. + """ + for _name, module in model.named_modules(): + if not isinstance(module, _QuantMoELinear): + continue + + n = module.num_experts + experts = module.experts + + # Reconstruct 3D weight: [num_experts, out_features, in_features] + module.weight = nn.Parameter( + torch.stack([experts[i].weight.data for i in range(n)]), + requires_grad=False, + ) + + # Stack per-expert scales back under the original attribute names. + # Check all experts: some may lack input_scale if they were never routed + # during calibration, so only stack when every expert has the attribute. + for attr in ("weight_scale", "weight_scale_2", "input_scale"): + if all(hasattr(experts[i], attr) for i in range(n)): + module.register_buffer( + attr, + torch.stack([getattr(experts[i], attr) for i in range(n)]), + ) + + # Remove expanded experts — the reconstructed 3D tensors replace them + del module.experts + + +CUSTOM_MODEL_PLUGINS.add(register_moe_linear_on_the_fly) diff --git a/modelopt/torch/models/step3p7/modeling_ptq.py b/modelopt/torch/models/step3p7/modeling_ptq.py new file mode 100644 index 00000000000..e43c5237278 --- /dev/null +++ b/modelopt/torch/models/step3p7/modeling_ptq.py @@ -0,0 +1,26 @@ +# 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. + +"""Step-3.7 PTQ modeling (HF model type ``step3p7``). + +Step-3.7 shares Step-3.5's expert-indexed ``MoELinear`` layout, so its PTQ support is the +Step-family code in ``step3p5/modeling_ptq.py``, which matches every Step revision. This +module imports it, so that loading the PTQ modeling for model type ``step3p7`` by name +registers that support too. +""" + +from ..step3p5 import modeling_ptq as _step3p5_modeling_ptq # noqa: F401 + +__all__: list[str] = [] diff --git a/modelopt/torch/models/step3p7/specs.py b/modelopt/torch/models/step3p7/specs.py index f38c0339989..74db68aed4b 100644 --- a/modelopt/torch/models/step3p7/specs.py +++ b/modelopt/torch/models/step3p7/specs.py @@ -15,8 +15,8 @@ """Step-3.7 specs (HF model type ``step3p7``). -Shares Step-3.5's architecture; its PTQ modeling (in the HF quantization plugin) -matches every Step revision. +Shares Step-3.5's architecture; its PTQ modeling lives in ``step3p5/modeling_ptq.py``, +which matches every Step revision. """ from ..specs import ModelSpec, register diff --git a/modelopt/torch/nas/hparams/concat.py b/modelopt/torch/nas/hparams/concat.py index 31274ab052e..9234939e1ed 100644 --- a/modelopt/torch/nas/hparams/concat.py +++ b/modelopt/torch/nas/hparams/concat.py @@ -317,3 +317,6 @@ def build_concat_hp(inputs: list[TracedHp]): concat_hp._importance_estimators = None return concat_hp + + +torch.serialization.add_safe_globals([ConcatTracedHp]) diff --git a/modelopt/torch/nas/hparams/container.py b/modelopt/torch/nas/hparams/container.py index dc7ee17f2f2..61548b29852 100644 --- a/modelopt/torch/nas/hparams/container.py +++ b/modelopt/torch/nas/hparams/container.py @@ -17,6 +17,8 @@ from collections.abc import Callable +import torch + from modelopt.torch.trace import Symbol, SymDepth from ..traced_hp import TracedHp, TracedHpRegistry @@ -46,3 +48,6 @@ def _resolve_dependencies( self.choices = list(choices & set(self.choices)) return super()._resolve_dependencies(sym, get_hp) + + +torch.serialization.add_safe_globals([DepthHparam]) diff --git a/modelopt/torch/nas/plugins/megatron.py b/modelopt/torch/nas/plugins/megatron.py index c2d31b6f3de..9daa81f0cf4 100644 --- a/modelopt/torch/nas/plugins/megatron.py +++ b/modelopt/torch/nas/plugins/megatron.py @@ -1100,6 +1100,9 @@ def _resolve_dependencies( raise NotImplementedError("MambaDInnerHp does not support `_resolve_dependencies`!") +torch.serialization.add_safe_globals([NumAttentionHeadsHp, MambaNumHeadsHp, MambaDInnerHp]) + + class _DynamicExtendedRMSNorm(DynamicModule): """An ExtendedRMSNorm (GroupNorm) layer with dynamic hyperparams. diff --git a/modelopt/torch/nas/traced_hp.py b/modelopt/torch/nas/traced_hp.py index 2f424bef0ce..0a29842f80e 100644 --- a/modelopt/torch/nas/traced_hp.py +++ b/modelopt/torch/nas/traced_hp.py @@ -17,6 +17,8 @@ from collections.abc import Callable +import torch + from modelopt.torch.opt.hparam import Hparam from modelopt.torch.trace import Symbol @@ -147,3 +149,6 @@ def _resolve_dependencies( self.choices = [self.original] return mapping + + +torch.serialization.add_safe_globals([TracedHp]) diff --git a/modelopt/torch/opt/hparam.py b/modelopt/torch/opt/hparam.py index c36e2ec84e4..df537c68352 100644 --- a/modelopt/torch/opt/hparam.py +++ b/modelopt/torch/opt/hparam.py @@ -82,6 +82,19 @@ def __init__(self, choices: Sequence[HPType], original: HPType | None = None) -> # an optional order to enforce for active_slice self._slice_order: torch.LongTensor | None = None + def __getstate__(self): + # A bound method pickles via ``getattr``, which ``weights_only`` loading rightly rejects. + state = self.__dict__.copy() + # Custom estimators stay as-is (deepcopy relies on this) and are rejected by weights_only loads. + if state.get("_importance_estimators") == [self._default_get_importance]: + state["_importance_estimators"] = "default" + return state + + def __setstate__(self, state): + self.__dict__.update(state) + if state.get("_importance_estimators") == "default": + self._importance_estimators = [self._default_get_importance] + def __iter__(self) -> Generator[HPType, None, None]: """Iterate over choices.""" yield from self.choices @@ -284,3 +297,6 @@ def __iand__(self, hp: "Hparam"): self._importance_estimators = None return self + + +torch.serialization.add_safe_globals([Hparam]) diff --git a/modelopt/torch/quantization/_dynamo_onnx.py b/modelopt/torch/quantization/_dynamo_onnx.py new file mode 100644 index 00000000000..55439488184 --- /dev/null +++ b/modelopt/torch/quantization/_dynamo_onnx.py @@ -0,0 +1,266 @@ +# 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. + +"""Private ONNXScript translations for quantized Dynamo export.""" + +import onnx +import onnxscript +import torch +from onnxscript.function_libs.torch_lib.tensor_typing import TFloat +from onnxscript.onnx_types import FLOAT4E2M1, FLOAT8E4M3FN + +from .export_onnx import onnx_dtype_map + +__all__ = [] + +_OPSET = onnxscript.opset23 +_TRT_OPSET = onnxscript.values.Opset(domain="trt", version=1) + + +@onnxscript.script(_TRT_OPSET) +def _fp8_qdq(inputs: TFloat, scale: TFloat) -> TFloat: + quantized = _TRT_OPSET.TRT_FP8QuantizeLinear(inputs, scale) + return _TRT_OPSET.TRT_FP8DequantizeLinear(quantized, scale) + + +@onnxscript.script(_TRT_OPSET) +def _int4_dq(inputs: TFloat, scale: TFloat, axis: int, block_size: int) -> TFloat: + return _TRT_OPSET.DequantizeLinear(inputs, scale, axis=axis, block_size=block_size) + + +@onnxscript.script(_TRT_OPSET) +def _fp4_qdq(inputs: TFloat, block_size: int) -> TFloat: + return _TRT_OPSET.TRT_FP4QDQ(inputs, block_size=block_size) + + +@onnxscript.script(_TRT_OPSET) +def _fp4_dynamic_quantize( + inputs: TFloat, scale: TFloat, block_size: int +) -> tuple[FLOAT4E2M1[...], FLOAT8E4M3FN[...]]: + quantized, dynamic_scale = _TRT_OPSET.TRT_FP4DynamicQuantize( + inputs, scale, axis=-1, block_size=block_size, scale_type=17 + ) + return quantized, dynamic_scale + + +@onnxscript.script(_TRT_OPSET) +def _mxfp8_dynamic_qdq(inputs: TFloat, block_size: int, output_dtype: int) -> TFloat: + quantized, scale = _TRT_OPSET.TRT_MXFP8DynamicQuantize( + inputs, axis=-1, block_size=block_size, output_dtype=17 + ) + return _TRT_OPSET.TRT_MXFP8DequantizeLinear( + quantized, + scale, + axis=-1, + block_size=block_size, + output_dtype=output_dtype, + ) + + +@onnxscript.script(_TRT_OPSET) +def _mxfp8_static_dq(inputs: TFloat, scale: TFloat, block_size: int, output_dtype: int) -> TFloat: + return _TRT_OPSET.TRT_MXFP8DequantizeLinear( + inputs, + scale, + axis=-1, + block_size=block_size, + output_dtype=output_dtype, + ) + + +def _cast(inputs, dtype: int): + return inputs if int(inputs.dtype) == dtype else _OPSET.Cast(inputs, to=dtype) + + +def _resolve_dtype(inputs, high_precision_dtype: str | None) -> int: + return ( + int(inputs.dtype) if high_precision_dtype is None else onnx_dtype_map[high_precision_dtype] + ) + + +def _static_shape(value) -> list[int]: + try: + return [int(dim) for dim in value.shape] + except (AttributeError, TypeError, ValueError): + raise NotImplementedError("Dynamo ONNX export does not support dynamic shapes.") from None + + +def _block_activation_shape( + inputs, block_size: int, quantizer_type: str | None +) -> list[int] | None: + if quantizer_type != "dynamic": + return None + shape = _static_shape(inputs) + if len(shape) not in (2, 3): + raise NotImplementedError("Dynamo ONNX block activation export supports rank 2 or 3 only.") + if shape[-1] % block_size: + raise NotImplementedError( + f"Dynamo ONNX block size {block_size} must divide the last dimension {shape[-1]}." + ) + return shape + + +def _translate_quantize_op( + inputs, + amax, + num_bits: int, + exponent_bits: int, + unsigned: bool, + narrow_range: bool, + high_precision_dtype: str | None = None, + block_size: int | None = None, + axis: int | None = None, +): + if num_bits == 8 and exponent_bits == 4: + if amax is None: + scale = _OPSET.CastLike(1.0, inputs) + elif any(dim != 1 for dim in _static_shape(amax)): + raise AssertionError( + "E4M3 supports ONNX export only for per-tensor quantization with scalar amax." + ) + else: + scale = _OPSET.Div(_OPSET.CastLike(amax, inputs), _OPSET.CastLike(448.0, inputs)) + return _fp8_qdq(inputs, scale) + + output_dtype = _resolve_dtype(inputs, high_precision_dtype) + if num_bits == 8 and exponent_bits == 0: + source_dtype = int(inputs.dtype) + input_shape = _static_shape(inputs) + if not unsigned: + assert not narrow_range, "ONNX does not support signed narrow-range INT8." + assert output_dtype in ( + source_dtype, + onnx.TensorProto.FLOAT, + onnx.TensorProto.BFLOAT16, + ), "TensorRT strongly typed mode requires Q/DQ in the input dtype, FP32, or BF16." + inputs = _cast(inputs, output_dtype) + input_rank = len(input_shape) + amax_shape = _static_shape(amax) + quantized_axes = [index for index, dim in enumerate(amax_shape) if dim != 1] + if len(quantized_axes) > 1: + raise AssertionError("ONNX does not support multi-axis quantization.") + if axis is not None: + if not -input_rank <= axis < input_rank: + raise ValueError(f"INT8 axis {axis} is out of bounds for input rank {input_rank}.") + axis %= input_rank + if not quantized_axes: + axis = None + elif axis is None: + if len(amax_shape) != input_rank: + raise ValueError("INT8 per-channel amax requires an explicit input axis.") + axis = quantized_axes[0] + elif len(amax_shape) == input_rank and quantized_axes[0] != axis: + inferred_axis = quantized_axes[0] + raise ValueError(f"INT8 axis {axis} does not match amax axis {inferred_axis}.") + amax = _OPSET.Cast(amax, to=output_dtype) + amax = _OPSET.Squeeze(amax) if axis is None else _OPSET.Reshape(amax, [-1]) + scale = _OPSET.Div(amax, float((1 << (7 + int(unsigned))) - 1)) + scale = _OPSET.Where(_OPSET.Equal(scale, 0.0), _OPSET.CastLike(1.0, scale), scale) + zero_point_dtype = onnx.TensorProto.UINT8 if unsigned else onnx.TensorProto.INT8 + zero_point = _OPSET.Cast(_OPSET.Mul(amax, 0.0), to=zero_point_dtype) + if axis is None: + quantized = _OPSET.QuantizeLinear(inputs, scale, zero_point) + output = _OPSET.DequantizeLinear(quantized, scale, zero_point) + else: + quantized = _OPSET.QuantizeLinear(inputs, scale, zero_point, axis=axis) + output = _OPSET.DequantizeLinear(quantized, scale, zero_point, axis=axis) + return output if output_dtype == source_dtype else _OPSET.Cast(output, to=source_dtype) + + if num_bits == 4 and exponent_bits == 0: + if unsigned: + raise NotImplementedError("Dynamo ONNX export supports signed INT4 only.") + if block_size is None or axis is None: + raise ValueError("INT4 ONNX export requires block_size and axis.") + if amax is None: + raise ValueError("INT4 ONNX export requires amax.") + source_dtype = int(inputs.dtype) + amax = _OPSET.Cast(amax, to=output_dtype) + scale = _OPSET.Div(amax, _OPSET.CastLike(7.0, amax)) + output = _int4_dq(inputs, scale, axis, block_size) + # The TRT carrier has no schema to infer metadata when input and scale dtypes differ. + output.dtype = onnxscript.ir.DataType(output_dtype) + output.shape = inputs.shape + return output if output_dtype == source_dtype else _OPSET.Cast(output, to=source_dtype) + + raise NotImplementedError( + f"Unsupported num_bits={num_bits}, exponent_bits={exponent_bits} for ONNX export." + ) + + +def _translate_dynamic_block_quantize_op( + inputs, + block_size: int, + amax, + num_bits: int, + exponent_bits: int, + scale_num_bits: int, + scale_exponent_bits: int, + high_precision_dtype: str | None = None, + quantizer_type: str | None = None, +): + activation_shape = _block_activation_shape(inputs, block_size, quantizer_type) + format_bits = (num_bits, exponent_bits, scale_num_bits, scale_exponent_bits) + + if format_bits == (4, 2, 8, 4): + if quantizer_type != "dynamic": + # Keep a standard-domain op so older ONNXScript retains the default opset import. + return _OPSET.Identity(_fp4_qdq(inputs, block_size)) + assert activation_shape is not None + source_dtype = int(inputs.dtype) + output_dtype = _resolve_dtype(inputs, high_precision_dtype) + inputs = _cast(inputs, output_dtype) + if amax is None: + scale = _OPSET.Constant(value_float=1.0) + else: + scale = _OPSET.Div(_OPSET.Cast(amax, to=onnx.TensorProto.FLOAT), 2688.0) + scale = _OPSET.Where(_OPSET.Equal(scale, 0.0), _OPSET.CastLike(1.0, scale), scale) + quantized, dynamic_scale = _fp4_dynamic_quantize(inputs, scale, block_size) + # ONNXScript before 0.6 may drop custom multi-output metadata during inlining. + quantized.dtype = onnxscript.ir.DataType.FLOAT4E2M1 + dynamic_scale.dtype = onnxscript.ir.DataType.FLOAT8E4M3FN + quantized.shape = onnxscript.ir.Shape(activation_shape) + scale_shape = [*activation_shape[:-1], activation_shape[-1] // block_size] + dynamic_scale.shape = onnxscript.ir.Shape(scale_shape) + dequantized_scale = _OPSET.DequantizeLinear(dynamic_scale, scale) + output = _OPSET.DequantizeLinear( + quantized, dequantized_scale, axis=-1, block_size=block_size + ) + return _OPSET.Cast(output, to=source_dtype) + + if format_bits == (8, 4, 9, 8): + output_dtype = int(inputs.dtype) + if quantizer_type == "dynamic": + output = _mxfp8_dynamic_qdq(inputs, block_size, output_dtype) + else: + output = _mxfp8_static_dq( + inputs, _OPSET.CastLike(1.0, inputs), block_size, output_dtype + ) + # Keep a standard-domain op so older ONNXScript retains the default opset import. + return _OPSET.Identity(output) + + raise NotImplementedError( + "Unsupported block format " + f"({exponent_bits}, {num_bits - exponent_bits - 1}) with scale format " + f"({scale_exponent_bits}, {scale_num_bits - scale_exponent_bits - 1})." + ) + + +def _get_dynamo_onnx_translation_table() -> dict: + dynamic_op = torch.ops.tensorrt.dynamic_block_quantize_op + return { + torch.ops.tensorrt.quantize_op.default: _translate_quantize_op, + dynamic_op.default: _translate_dynamic_block_quantize_op, + dynamic_op.overload: _translate_dynamic_block_quantize_op, + } diff --git a/modelopt/torch/quantization/compress.py b/modelopt/torch/quantization/compress.py index 2a5cbbee9fc..ee7e642a0e8 100644 --- a/modelopt/torch/quantization/compress.py +++ b/modelopt/torch/quantization/compress.py @@ -27,6 +27,7 @@ from modelopt.torch.opt.conversion import ModelLikeModule, ModeloptStateManager from modelopt.torch.opt.dynamic import _DMRegistryCls from modelopt.torch.opt.mode import ConvertReturnType, MetadataDict +from modelopt.torch.utils.import_utils import import_plugin from .backends.gemm_registry import disable_real_quant_gemm, enable_real_quant_gemm from .config import CompressCfgType, CompressConfig @@ -35,7 +36,8 @@ from .qtensor import QTensorWrapper, pack_real_quantize_weight from .utils import is_quantized_linear -try: +mcore_available = False +with import_plugin("megatron"): from .plugins.megatron import ( _MegatronColumnParallelLinear, _MegatronRowParallelLinear, @@ -44,8 +46,6 @@ ) mcore_available = True -except ImportError: - mcore_available = False RealQuantModuleRegistry = _DMRegistryCls("RealQuant") diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 77eeb9371b5..3741ed27883 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -2339,6 +2339,7 @@ def _make_gptq_handle(name, m): for handle in gptq_handles.values(): handle.cleanup() + handle.sync_hessian_across_dp() print_rank_0("Updating weights using GPTQ algorithm...") names = module_name_maps(model) diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 5b53526aba6..f40df709a04 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -1078,6 +1078,9 @@ def fold_weight(model: nn.Module, keep_attrs: bool = False): Any weight-quantizer rotation is folded into the weights and disabled so subsequent forwards do not re-rotate the already-folded weights. """ + for module in model.modules(): + if isinstance(module, QuantModule): + module._validate_fold_weight() for name, module in model.named_modules(): if isinstance(module, QuantModule): module.fold_weight(keep_attrs) diff --git a/modelopt/torch/quantization/nn/modules/quant_module.py b/modelopt/torch/quantization/nn/modules/quant_module.py index 522f2d0db09..8b1d8a4d349 100644 --- a/modelopt/torch/quantization/nn/modules/quant_module.py +++ b/modelopt/torch/quantization/nn/modules/quant_module.py @@ -168,6 +168,9 @@ def _fold_weight_quantizer( if hasattr(quantizer, attr_name): delattr(quantizer, attr_name) + def _validate_fold_weight(self): + """Validate that weight folding is supported before any module is modified.""" + def fold_weight(self, keep_attrs: bool = False): """Bake each fake-quant weight quantizer into its weight for faster eval. diff --git a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py index a698514d53a..454d18bdaf6 100644 --- a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py +++ b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py @@ -63,7 +63,7 @@ static_blockwise_fp4_fake_quant, ) from ...utils import is_torch_export_mode -from ...utils.numeric_utils import fp8_max_for_normalization +from ...utils.numeric_utils import E4M3_MAX, fp8_max_for_normalization from ..functional import normalized_hadamard_transform # torch.finfo(...).tiny gives the smallest normal E4M3 value; scale clamping needs @@ -749,8 +749,9 @@ def _get_amax(self, inputs): reduce_axis = quant_utils.convert_quantization_axis_to_reduce_axis(inputs, self._axis) amax = quant_utils.reduce_amax(inputs, axis=reduce_axis, keepdims=True).detach() - amax = amax.detach() if is_torch_export_mode() else amax.data - return amax + if is_torch_export_mode() or torch.compiler.is_exporting(): + return amax.detach() + return amax.data def validate_attr( self, attr_value=None, attr_name="amax", raise_error=False, warn_error=False, name="" @@ -1729,7 +1730,9 @@ def _block_scale_from_amax(self, amax: torch.Tensor, quantize: bool) -> torch.Te """Compute the per-block scale from a per-block amax, optionally FP8-quantizing it.""" if quantize: per_tensor_scale = _amax_to_scale(self.global_amax, self._quant_max_bound) - min_value = _FP8_E4M3_MIN_POSITIVE * per_tensor_scale.view(-1) + # scaled_e4m3 maps per_tensor_scale to E4M3_MAX (448); apply the inverse + # mapping to express the minimum encoded FP8 scale in physical scale units. + min_value = (_FP8_E4M3_MIN_POSITIVE / E4M3_MAX) * per_tensor_scale.view(-1) scale = _amax_to_scale(amax, self._quant_max_bound, min_value=min_value) return scaled_e4m3(scale, per_tensor_scale, None, 4, 3) return _amax_to_scale(amax, self._quant_max_bound, min_value=1e-8) diff --git a/modelopt/torch/quantization/plugins/custom.py b/modelopt/torch/quantization/plugins/custom.py index d16596908c9..54dad627457 100644 --- a/modelopt/torch/quantization/plugins/custom.py +++ b/modelopt/torch/quantization/plugins/custom.py @@ -173,10 +173,15 @@ def _has_complete_static_nvfp4_weight_state(quantizer, weight): _check_unsupported_states( quantizer if isinstance(quantizer, TensorQuantizer) else quantizer[0] ) + # Grouped linears validate each expert against its corresponding weight in their + # override. Resetting the container here would clear every expert amax before this + # single-weight compatibility path recalibrates only the representative quantizer. # Skip max_calibrate when saved static NVFP4 state is intact; else MSE scales get overwritten. - if _has_state( - self.weight_quantizer, "_amax" - ) and not _has_complete_static_nvfp4_weight_state(self.weight_quantizer, self.weight): + if ( + not isinstance(self.weight_quantizer, GroupedQuantizer) + and _has_state(self.weight_quantizer, "_amax") + and not _has_complete_static_nvfp4_weight_state(self.weight_quantizer, self.weight) + ): self.weight_quantizer.reset_amax() max_calibrate(self.weight_quantizer, lambda wq: wq(self.weight), distributed_sync=False) if _has_state(self.input_quantizer, "_pre_quant_scale"): diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 5eae0ae937a..6a6c5d18bfc 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# 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"); @@ -18,7 +18,6 @@ import importlib import inspect import logging -import re import warnings from contextlib import contextmanager from functools import partial @@ -29,8 +28,7 @@ import transformers from torch import Tensor from torch.nn.functional import linear -from transformers.integrations.finegrained_fp8 import FP8Linear -from transformers.models.gpt_oss.modeling_gpt_oss import GptOssExperts +from transformers.integrations.finegrained_fp8 import FP8Experts, FP8Linear from transformers.models.t5.modeling_t5 import T5Attention from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_FA_AVAILABLE @@ -39,7 +37,6 @@ validate_triton_attention_envelope, ) from modelopt.torch.kernels.quantization.gemm import IS_AVAILABLE as IS_TRITON_AVAILABLE -from modelopt.torch.models import hf_model_type from modelopt.torch.opt.dynamic import DynamicModule from modelopt.torch.utils.distributed import ParallelState @@ -871,12 +868,16 @@ def _quantized_linear(input, weight, bias=None): if self._down_proj_linear: idx = self._current_expert_idx input = self.down_proj_input_quantizer(input) - weight = self.down_proj_weight_quantizers[idx](weight) + weight = self.down_proj_weight_quantizers[idx]( + self.get_expert_weight("down_proj", idx, input.dtype) + ) else: idx = self._get_expert_idx_from_first_proj(weight) self._current_expert_idx = idx input = getattr(self, self._first_proj_input_quantizer_attr)(input) - weight = getattr(self, self._first_proj_weight_quantizers_attr)[idx](weight) + weight = getattr(self, self._first_proj_weight_quantizers_attr)[idx]( + self.get_expert_weight(self._first_proj_attr, idx, input.dtype) + ) self._down_proj_linear = not self._down_proj_linear return _orig_linear(input, weight, bias) @@ -888,6 +889,10 @@ def forward(self, *args, **kwargs): self._down_proj_linear = False return super().forward(*args, **kwargs) + def get_expert_weight(self, name, idx, dtype=None): + """Return one expert's projection in the calibration/export weight domain.""" + return getattr(self, name)[idx] + def iter_weights_for_calibration(self): """Yield ``(weight_slice, quantizer)`` per-expert pairs. @@ -904,7 +909,7 @@ def iter_weights_for_calibration(self): if weight is None or quantizers is None: continue for idx, q in enumerate(quantizers): - yield weight[idx], q + yield self.get_expert_weight(weight_name, idx), q def iter_weight_quantizers_for_calibration(self): """Yield the per-expert weight quantizers without slicing the fused 3-D weight. @@ -950,6 +955,58 @@ class _QuantNonGatedFusedExperts(_QuantFusedExperts): _is_gated = False +class _QuantFP8Experts(_QuantFusedExperts): + """Scale-aware PTQ for native Transformers FP8 experts, dequantized per expert.""" + + @property + def _first_proj_attr(self): + return "gate_up_proj" if self.has_gate else "up_proj" + + @property + def _is_gated(self): + return self.has_gate + + def _setup(self): + if getattr(self.config, "expert_dtype", "fp8") != "fp8": + raise NotImplementedError("Only FP8 E4M3 native expert weights are supported") + super()._setup() + + def get_expert_weight(self, name, idx, dtype=None): + weight = getattr(self, name)[idx] + if weight.element_size() > 1: + return weight if dtype is None else weight.to(dtype) + dtype = dtype or torch.get_default_dtype() + scales = getattr(self, f"{name}_scale_inv")[idx].float() + rows, cols = self.block_size or (weight.shape[0] // scales.shape[0], weight.shape[1]) + scales = scales.repeat_interleave(rows, 0).repeat_interleave(cols, 1) + return (weight.float() * scales[: weight.shape[0], : weight.shape[1]]).to(dtype) + + def linear(self, input, weight, weight_scale_inv, activation_scale=None): + # Preserve the FP8 view's storage offset for the per-expert quantizer lookup. + return torch.nn.functional.linear(input, weight) + + def _validate_fold_weight(self): + if any( + getattr(self, name).element_size() == 1 for name in (self._first_proj_attr, "down_proj") + ): + raise ValueError("Call unpack_weight() before folding native FP8 expert weights.") + super()._validate_fold_weight() + + def fold_weight(self, keep_attrs: bool = False): + self._validate_fold_weight() + super().fold_weight(keep_attrs) + + def unpack_weight(self): + """Materialize floating weights before folding; retain scales for native forward.""" + for name in (self._first_proj_attr, "down_proj"): + weight = getattr(self, name) + if weight.element_size() == 1: + unpacked = torch.empty_like(weight, dtype=torch.get_default_dtype()) + for idx in range(self.num_experts): + unpacked[idx].copy_(self.get_expert_weight(name, idx)) + setattr(self, name, nn.Parameter(unpacked, requires_grad=False)) + + def _get_fused_experts_quantizer_attr_names(module): """Return quantizer attribute names for a converted fused-experts module.""" first_proj_attr = getattr(module, "_first_proj_attr", "gate_up_proj") @@ -1202,100 +1259,13 @@ def unpack_weight(self): except ImportError: pass +if FP8Experts not in QuantModuleRegistry: + QuantModuleRegistry.register({FP8Experts: "hf.FP8Experts"})(_QuantFP8Experts) + if FP8Linear not in QuantModuleRegistry: QuantModuleRegistry.register({FP8Linear: "hf.FP8Linear"})(_QuantFP8Linear) -class _QuantGptOssExperts(_TransposedExpertsCalibMixin, _QuantFunctionalMixin): - """Quantized wrapper for `transformers.GptOssExperts`. - - Quantizes `gate_up_proj` and `down_proj` weights via dynamic attributes inside `quantize_weight()`. - Activations into `gate_up_proj` are quantized by `gate_up_proj_input_quantizer`. For `down_proj` - activation quantization, we intercept `torch.Tensor.__matmul__`/`torch.bmm` and quantize inputs - on every second call (since the first call computes `gate_up_proj` outputs and second call - computes `down_proj` outputs). - """ - - @staticmethod - def _get_quantized_weight(quantizer, module, weight): - # MoE weight is accessed for each expert in one forward pass. so lets cache it - if module._enable_weight_quantization: - if hasattr(quantizer, "_cached_quant_val"): - return getattr(quantizer, "_cached_quant_val") - quantizer._cached_quant_val = _transposed_quantize(weight, quantizer) - return quantizer._cached_quant_val - return weight - - def _setup_for_weight_quantization(self): - self._register_dynamic_attribute( - "gate_up_proj", partial(self._get_quantized_weight, self.gate_up_proj_weight_quantizer) - ) - self._register_dynamic_attribute( - "down_proj", partial(self._get_quantized_weight, self.down_proj_weight_quantizer) - ) - - def _setup(self): - assert not hasattr(self, "kernel_layer_name"), ( - "ModelOpt quantization does not support patched forward for kernel_hub" - ) - self.gate_up_proj_input_quantizer = TensorQuantizer() - self.gate_up_proj_weight_quantizer = TensorQuantizer() - self.down_proj_input_quantizer = TensorQuantizer() - self.down_proj_weight_quantizer = TensorQuantizer() - - self._register_temp_attribute("_enable_weight_quantization", False) - self._register_temp_attribute("_down_proj_mul", False) - self._setup_for_weight_quantization() - - @property - def functionals_to_replace(self): - # Use torch.ops.aten to bypass Python dispatch and avoid RecursionError - # (torch.matmul / __matmul__ can dispatch to each other) - _aten_bmm = torch.ops.aten.bmm - _aten_matmul = torch.ops.aten.matmul - - def _quantized_bmm(batch1, batch2, *, out=None): - batch1 = self.down_proj_input_quantizer(batch1) if self._down_proj_mul else batch1 - self._down_proj_mul = not self._down_proj_mul # toggle the flag - if out is not None: - return torch.ops.aten.bmm.out(batch1, batch2, out=out) - return _aten_bmm(batch1, batch2) - - def _tensor_matmul(self_t, other): - self_t = self.down_proj_input_quantizer(self_t) if self._down_proj_mul else self_t - self._down_proj_mul = not self._down_proj_mul - return _aten_matmul(self_t, other) - - return [ - (torch, "bmm", _quantized_bmm), - (torch.Tensor, "__matmul__", _tensor_matmul), - ] - - @contextmanager - def quantize_weight(self): - """Context in which MoE weight is quantized.""" - self._enable_weight_quantization = True - try: - yield - finally: - for module in self.modules(): - if isinstance(module, TensorQuantizer) and hasattr(module, "_cached_quant_val"): - delattr(module, "_cached_quant_val") - self._enable_weight_quantization = False - - def forward( - self, hidden_states: torch.Tensor, router_indices=None, routing_weights=None - ) -> torch.Tensor: - """Forward method to add quantization.""" - hidden_states = self.gate_up_proj_input_quantizer(hidden_states) - with self.quantize_weight(): - return super().forward(hidden_states, router_indices, routing_weights) - - -if GptOssExperts not in QuantModuleRegistry: - QuantModuleRegistry.register({GptOssExperts: "hf.GptOssExperts"})(_QuantGptOssExperts) - - def _has_num_experts(obj): # n_routed_experts: NemotronH-style MoE return hasattr(obj, "num_experts") or hasattr(obj, "n_routed_experts") @@ -1572,8 +1542,11 @@ def _is_param_grad_enabled_for_auto_quantize(pname, model): # Nemotron-H's is the more specific one. for _model_type in ( "falcon", + "gpt_oss", "llama4", "nemotron_h", + "step3p5", + "step3p7", ): importlib.import_module(f"modelopt.torch.models.{_model_type}.modeling_ptq") @@ -1582,235 +1555,8 @@ def _is_param_grad_enabled_for_auto_quantize(pname, model): ) -class _QuantMoELinear(QuantModule): - """Quantization wrapper for expert-indexed MoELinear modules (fused expert weights). - - MoELinear has weight shape [num_experts, out_features, in_features] with - forward(x, expert_id). We expand it into per-expert nn.Linear modules so - each expert gets its own weight_quantizer and input_quantizer, calibrated - only on tokens actually routed to that expert. - - On export, _reconstruct_fused_moe_linear() stacks the per-expert quantized - weights and scales back into the original 3D format. - - Note: we use expansion-then-reconstruction rather than the add_module() approach - because vLLM requires stacked 3D scaling factors; per-expert expanded keys are - not accepted by the downstream serving engine. - """ - - def _setup(self): - from accelerate import init_empty_weights - - # Accelerate's CPU/disk offload (`device_map="auto"`, `--offload_folder`) leaves - # `weight` as a meta tensor and keeps the real value in the module's offload hook, - # keyed on the original `weight` name. Expanding that would copy meta storage into - # every expert and then delete the key the hook restores into, silently producing a - # checkpoint of zeros. Refuse instead of corrupting. - if self.weight.is_meta or getattr(getattr(self, "_hf_hook", None), "offload", False): - raise NotImplementedError( - f"{type(self).__name__}: expert-indexed MoELinear weights cannot be quantized " - "while offloaded by Accelerate (the weight is a meta tensor whose value lives " - "in the offload hook). Load the model without CPU/disk offload — more GPUs, or " - "a device_map that keeps the MoE layers resident — and re-run." - ) - - dtype, device = self.weight.dtype, self.weight.device - - with init_empty_weights(): - experts = nn.ModuleList( - [ - nn.Linear(self.in_features, self.out_features, bias=False) - for _ in range(self.num_experts) - ] - ) - - for i in range(self.num_experts): - experts[i].to_empty(device=device) - with torch.no_grad(): - experts[i].weight.data = self.weight[i].detach().to(dtype=dtype, device=device) - - delattr(self, "weight") - self.experts = experts - - def forward(self, x, expert_id): - # experts[expert_id] is a _QuantLinear after quantization wrapping, providing - # per-expert input_quantizer and weight_quantizer. - # - # MoELinear.forward always promotes to fp32 for the matmul regardless of storage - # dtype (`F.linear(x.float(), self.weight[expert_id].float())`), so leaving the - # expert's weight at its native storage dtype (e.g. bf16) and downcasting x to - # match before the matmul would compute in bf16 and change the model's output even - # with every quantizer disabled -- the bf16 rounding this class exists to quantize - # *past*, not to reintroduce as a side effect of conversion. - # - # A prior version of this fix instead expanded every expert's weight in fp32 - # permanently in `_setup`. That reproduces Step's fp32 compute but turns a per-call - # transient promotion into persistent model state: on Step-3.7's full routed-expert - # set (42 layers x 3 projections x 288 experts x 4096 x 1280), doubling from bf16 to - # fp32 adds roughly 354 GiB held throughout calibration, on top of device placement - # already sized for bf16 -- a model that loaded successfully can then OOM. It also - # left disabled/unquantized experts reconstructed at fp32 in the exported checkpoint. - # - # Instead, only the one expert actually being called is promoted, transiently, for - # the duration of this one call -- matching Step's own per-call `.float()` memory - # profile instead of Step-3.7's full expert set. `expert.weight` is read here - # outside any `quantize_weight()` context, so `_get_quantized_weight` passes it - # through unchanged and this is the real underlying nn.Parameter (the same pattern - # `_setup` above uses), not a value computed by the quantizer -- so reassigning its - # `.data` genuinely mutates the persisted storage, not a transient wrapper. - # - # This must keep calling `expert(x)` (`__call__`, not `.forward()`) rather than - # reimplementing the input/weight-quantize/output-quantize sequence inline: some - # calibration algorithms (e.g. `local_hessian_calibrate`) register a - # `forward_pre_hook` directly on the quantized Linear module, which only fires - # through standard `nn.Module.__call__` dispatch. - expert = self.experts[expert_id] - original_weight = expert.weight.data - with torch.no_grad(): - expert.weight.data = original_weight.float() - try: - out = expert(x.float()) - finally: - with torch.no_grad(): - expert.weight.data = original_weight - return out.float() - - -def _is_expert_indexed_moe_linear(module: nn.Module) -> bool: - """Whether ``module`` packs one projection's experts into an expert-indexed 3-D weight. - - The Step family (``stepfun-ai/Step-3.5-Flash``, ``stepfun-ai/Step-3.7-Flash``) ships a - custom ``MoELinear`` via ``trust_remote_code``: a plain ``nn.Module`` holding a single - ``weight`` of shape ``[num_experts, out_features, in_features]``, whose - ``forward(x, expert_id)`` runs ``F.linear`` against the selected expert's slice. The - weights therefore live on the projection submodule rather than on the expert container, - which is what :func:`_fused_experts_wrapper_class` looks for, and the module is not an - ``nn.Linear``, so neither the fused-experts path nor the plain linear path claims it. - - Detection is structural rather than keyed on class names so new Step revisions are picked - up without another hardcoded name, but the shape alone is not a sufficient contract: the - replacement forward indexes ``self.experts[expert_id]``, so it only works for callers that - pass a **scalar expert index**. Grouped-GEMM MoE layers share the exact same 3-D weight and - attribute set while passing a per-expert token-count *tensor* instead (e.g. Moondream3's - ``MoeFusedLinear.forward(input, m_sizes)``, which would raise ``TypeError: only integer - tensors of a single element can be converted to an index`` on the first calibration - forward). The second parameter must therefore be named ``expert_id``, which is the - scalar-index contract both Step revisions declare. - """ - weight = getattr(module, "weight", None) - if not isinstance(weight, (nn.Parameter, Tensor)) or weight.dim() != 3: - return False - if not all(hasattr(module, attr) for attr in ("num_experts", "in_features", "out_features")): - return False - # The wrapper rebuilds the weight as `num_experts` Linears of (out_features, in_features), - # so a 3-D weight laid out any other way would silently copy the wrong slices. - if tuple(weight.shape) != (module.num_experts, module.out_features, module.in_features): - return False - try: - params = list(inspect.signature(type(module).forward).parameters.values())[1:] - except (TypeError, ValueError): - return False - # The replacement forward is exactly `(x, expert_id)`, so anything the caller could pass - # beyond those two — a keyword-only `router_state`, *args, **kwargs — would raise once - # converted. Require the signature to match what the wrapper can honour. - return ( - len(params) == 2 - and all(p.kind is p.POSITIONAL_OR_KEYWORD and p.default is p.empty for p in params) - and params[1].name == "expert_id" - ) - - -_STEP_FAMILY_RE = re.compile(r"(?i)^step\d") - - -def _is_step_family_model(model: nn.Module) -> bool: - """Whether ``model`` is a Step-family root model (Step-3.5, Step-3.7, or a future revision). - - Matched against the ``step`` convention shared by ``model_type`` (``"step3p5"``, - ``"step3p7"``) and the remote-code class name (``Step3p5ForCausalLM``, - ``Step3p7ForConditionalGeneration``), not an exact revision, so a new Step release is - still picked up without another hardcoded name. This is deliberately narrower than the - shape/signature check in :func:`_is_expert_indexed_moe_linear` alone: that check accepts - any module with a matching 3-D weight and an ``(x, expert_id)`` forward, which is a - coincidence risk on its own -- an unrelated architecture happening to reuse the parameter - name ``expert_id`` with different semantics (a per-expert bias or post-scale, say) would - be claimed and have that behavior silently dropped by the replacement wrapper. Gating on - the model family keeps the shape check doing what it is actually good at: telling - Step revisions apart without a class-name allowlist, rather than distinguishing Step - from arbitrary third-party MoE code. - """ - model_type = hf_model_type(model) or "" - return bool(_STEP_FAMILY_RE.match(model_type) or _STEP_FAMILY_RE.match(type(model).__name__)) - - -def register_moe_linear_on_the_fly(model): - """Register expert-indexed ``MoELinear`` modules (Step-3.5 / Step-3.7) for quantization. - - Without this the routed experts carry no quantizer at all: an experts-only recipe matches - nothing and the export writes a checkpoint with ``quant_algo: null``. - """ - if not _is_step_family_model(model): - return - visited_types = set() - for name, module in model.named_modules(): - mod_type = type(module) - if mod_type in visited_types or QuantModuleRegistry.get(mod_type) is not None: - continue - visited_types.add(mod_type) - - if _is_expert_indexed_moe_linear(module): - print( - f"\033[1mDetected expert-indexed MoE linear '{name}' of type " - f"{mod_type.__name__}, registering with _QuantMoELinear.\033[0m" - ) - QuantModuleRegistry.register({mod_type: f"hf.{mod_type.__name__}"})(_QuantMoELinear) - - -def _reconstruct_fused_moe_linear(model: nn.Module) -> None: - """Reconstruct :class:`_QuantMoELinear` per-expert weights back to the 3-D MoELinear format. - - After _process_quantized_modules, each expert's nn.Linear inside the wrapper has: - - weight: fp4-quantized tensor [out_features, in_features] - - weight_scale, weight_scale_2: per-block / global scales - - input_scale: activation scale (if calibrated) - - This stacks them back into the original MoELinear layout so the exported state_dict - uses the original key names (e.g. moe.up_proj.weight with shape [N, out, in]). - - Matched by wrapper type rather than by the dynamically generated class name (``Quant`` + - the model's own class name): a model whose class is not spelled ``MoELinear`` would - otherwise quantize normally but export unusable per-expert keys. - """ - for _name, module in model.named_modules(): - if not isinstance(module, _QuantMoELinear): - continue - - n = module.num_experts - experts = module.experts - - # Reconstruct 3D weight: [num_experts, out_features, in_features] - module.weight = nn.Parameter( - torch.stack([experts[i].weight.data for i in range(n)]), - requires_grad=False, - ) - - # Stack per-expert scales back under the original attribute names. - # Check all experts: some may lack input_scale if they were never routed - # during calibration, so only stack when every expert has the attribute. - for attr in ("weight_scale", "weight_scale_2", "input_scale"): - if all(hasattr(experts[i], attr) for i in range(n)): - module.register_buffer( - attr, - torch.stack([getattr(experts[i], attr) for i in range(n)]), - ) - - # Remove expanded experts — the reconstructed 3D tensors replace them - del module.experts - - CUSTOM_MODEL_PLUGINS.update( [ - register_moe_linear_on_the_fly, register_fused_experts_on_the_fly, force_eager_experts_impl_on_the_fly, register_sparse_moe_on_the_fly, diff --git a/modelopt/torch/quantization/qtensor/int8_tensor.py b/modelopt/torch/quantization/qtensor/int8_tensor.py index 890b88e9ae9..fe3126279d8 100644 --- a/modelopt/torch/quantization/qtensor/int8_tensor.py +++ b/modelopt/torch/quantization/qtensor/int8_tensor.py @@ -60,11 +60,15 @@ def quantize( else: reduce_axis = convert_quantization_axis_to_reduce_axis(input, axis) amax = reduce_amax(input, axis=reduce_axis) - scales = amax / 127.0 + # Keep the scales in FP32: amax / 127 of a small FP16 channel lands in the FP16 + # subnormal range, where it loses precision or rounds to zero. + scales = amax.float() / 127.0 # Calculate the scale shape and make sure it aligns with input and block_sizes expected_shape = list(input.shape) - expanded_scales = scales.clone() + # A zero scale means an all-zero channel or block; divide it by one so it stays zero + # instead of becoming NaN. + expanded_scales = torch.where(scales == 0, 1.0, scales) if block_sizes: for dim, block_size in block_sizes.items(): dim = dim if dim >= 0 else len(input.shape) + dim # Convert negative index @@ -84,8 +88,10 @@ def quantize( for dim, block_size in block_sizes.items(): expanded_scales = expanded_scales.repeat_interleave(block_size, dim=dim) - # Quantization - quantized_data = (input / expanded_scales).round().clamp(-128, 127).to(torch.int8) + # Quantize in FP32. `input / expanded_scales` alone would run in the input dtype when the + # scale is 0-dim (per-tensor), rounding the quotient before round(). + quantized_data = input.to(torch.float32, copy=True).div_(expanded_scales) + quantized_data = quantized_data.round_().clamp_(-128, 127).to(torch.int8) return cls(original_input.shape, original_input.dtype, quantized_data), scales @@ -121,4 +127,8 @@ def dequantize(self, dtype: torch.dtype = None, **kwarg): slices = tuple(slice(0, dim) for dim in self.metadata["shape"]) scales = scales.to(self._quantized_data.device) - return (self._quantized_data.view(torch.int8).to(dtype) * scales.to(dtype))[slices] + # Multiply in FP32 and cast once: casting a small scale to FP16 first underflows it. + # Writing into a `dtype` output casts each product as it is stored, with no FP32 copy. + dequantized = torch.empty(shape, dtype=dtype, device=self._quantized_data.device) + torch.mul(self._quantized_data.view(torch.int8), scales.float(), out=dequantized) + return dequantized[slices] diff --git a/modelopt/torch/quantization/tensor_quant.py b/modelopt/torch/quantization/tensor_quant.py index 20e083491aa..7ba13264576 100644 --- a/modelopt/torch/quantization/tensor_quant.py +++ b/modelopt/torch/quantization/tensor_quant.py @@ -47,8 +47,8 @@ def _fp8_eager(x, amax=None): dtype = x.dtype if amax is not None: amax = amax.to(torch.float32) - epsilon = 1.0 / (1 << 24) - zero_amax_mask = amax <= epsilon + # Like the CUDA extension, only an amax whose FP8 scale overflows is treated as zero. + zero_amax_mask = amax <= 448.0 / torch.finfo(torch.float32).max safe_amax = torch.where(zero_amax_mask, torch.ones_like(amax), amax) scale = 448.0 / safe_amax scale_inv = 1 / scale @@ -119,11 +119,16 @@ def _quantize_impl( exponent_bits: int = 0, unsigned: bool = False, narrow_range: bool = True, + trt_high_precision_dtype: str | None = None, + block_size: int | None = None, + axis: int | None = None, ): if num_bits == 8 and exponent_bits == 4: return scaled_e4m3_impl(inputs=inputs, amax=amax) elif isinstance(num_bits, int): - return fake_quant_impl( + # CPU export executes this carrier directly; use the same math as the eager CPU fallback. + quantize_impl = fake_quant_impl if inputs.is_cuda else _tensor_quant + return quantize_impl( inputs=inputs, amax=amax, num_bits=num_bits, @@ -143,14 +148,15 @@ def _quantize_impl_abstract( exponent_bits: int = 0, unsigned: bool = False, narrow_range: bool = True, + trt_high_precision_dtype: str | None = None, + block_size: int | None = None, + axis: int | None = None, ) -> torch.Tensor: """Register an abstract implementation for quantizing tensor. This abstract function returns an empty tensor with the same shape and dtype. """ - output = torch.empty_like(input) - - return output + return torch.empty_like(input) # Argument types: Tensor, int, NoneType, int, int, int, int, @@ -162,6 +168,8 @@ def _dynamic_block_quantize_impl( exponent_bits: int, scale_num_bits: int, scale_exponent_bits: int, + trt_high_precision_dtype: str | None = None, + onnx_quantizer_type: str | None = None, ): scale_bits = (scale_exponent_bits, scale_num_bits - scale_exponent_bits - 1) if exponent_bits != 0: @@ -203,14 +211,14 @@ def _dynamic_block_quantize_impl_abstract( exponent_bits: int, scale_num_bits: int, scale_exponent_bits: int, + trt_high_precision_dtype: str | None = None, + onnx_quantizer_type: str | None = None, ): """Register an abstract implementation for dynamic block quantization. This abstract function returns an empty tensor with the same shape and dtype. """ - output = torch.empty_like(inputs) - - return output + return torch.empty_like(inputs) quantize_op = _quantize_impl @@ -223,17 +231,20 @@ def _dynamic_block_quantize_impl_abstract( torch.library.define( "tensorrt::quantize_op", "(Tensor input, Tensor amax, int num_bits, int exponent_bits, " - "bool unsigned, bool narrow_range) -> Tensor", + "bool unsigned, bool narrow_range, str? trt_high_precision_dtype=None, " + "int? block_size=None, int? axis=None) -> Tensor", ) torch.library.define( "tensorrt::dynamic_block_quantize_op", "(Tensor input, int block_size, Tensor amax, int num_bits, int exponent_bits, " - "int scale_num_bits, int scale_exponent_bits) -> Tensor", + "int scale_num_bits, int scale_exponent_bits, str? trt_high_precision_dtype=None, " + "str? onnx_quantizer_type=None) -> Tensor", ) torch.library.define( "tensorrt::dynamic_block_quantize_op.overload", "(Tensor input, int block_size, None amax, int num_bits, int exponent_bits, " - "int scale_num_bits, int scale_exponent_bits) -> Tensor", + "int scale_num_bits, int scale_exponent_bits, str? trt_high_precision_dtype=None, " + "str? onnx_quantizer_type=None) -> Tensor", ) # Implement the None amax case @@ -245,6 +256,8 @@ def _dynamic_block_quantize_impl_none_amax( exponent_bits: int, scale_num_bits: int, scale_exponent_bits: int, + trt_high_precision_dtype: str | None = None, + onnx_quantizer_type: str | None = None, ): return torch.empty_like(inputs) @@ -368,8 +381,7 @@ def forward( def legacy_quant_func(): # The LegacyFakeTensorQuantFunction support cpu and amax with any shape that can be broadcasted to inputs. - outputs = _tensor_quant(inputs, amax, num_bits, unsigned, narrow_range) - return outputs + return _tensor_quant(inputs, amax, num_bits, unsigned, narrow_range) if not inputs.is_cuda: outputs = legacy_quant_func() @@ -382,6 +394,9 @@ def legacy_quant_func(): exponent_bits=0, unsigned=unsigned, narrow_range=narrow_range, + trt_high_precision_dtype=trt_high_precision_dtype, + block_size=block_size, + axis=axis, ) except (AttributeError, ValueError): # AttributeError: cuda_ext is not imported, possibly due to CPU only installation @@ -447,6 +462,7 @@ def forward( exponent_bits=4, unsigned=False, narrow_range=False, + trt_high_precision_dtype=trt_high_precision_dtype, ) if bias is not None: @@ -461,7 +477,6 @@ def backward(ctx, grad_outputs): def _dynamic_block_quantize_forward( - ctx, inputs, block_size, amax, @@ -469,7 +484,6 @@ def _dynamic_block_quantize_forward( scale_bits, trt_high_precision_dtype=None, onnx_quantizer_type="dynamic", - pass_through_bwd=True, ): """Forward method.""" if isinstance(num_bits, int): @@ -482,7 +496,7 @@ def _dynamic_block_quantize_forward( assert isinstance(scale_bits, tuple) and len(scale_bits) == 2 scale_exponent_bits = scale_bits[0] scale_num_bits = scale_bits[0] + scale_bits[1] + 1 - outputs = dynamic_block_quantize_op( + return dynamic_block_quantize_op( inputs, block_size, amax, @@ -490,8 +504,9 @@ def _dynamic_block_quantize_forward( exponent_bits, scale_num_bits, scale_exponent_bits, + trt_high_precision_dtype, + onnx_quantizer_type, ) - return outputs class DynamicBlockQuantizationFunction(Function): @@ -551,7 +566,6 @@ def forward( """Forward method.""" _save_for_backward_if_needed(ctx, pass_through_bwd, inputs, amax) return _dynamic_block_quantize_forward( - ctx, inputs, block_size, amax, @@ -559,7 +573,6 @@ def forward( scale_bits, trt_high_precision_dtype, onnx_quantizer_type, - pass_through_bwd, ) @staticmethod diff --git a/modelopt/torch/quantization/utils/calib_utils.py b/modelopt/torch/quantization/utils/calib_utils.py index b1ade8e7809..f489bdd5d3e 100644 --- a/modelopt/torch/quantization/utils/calib_utils.py +++ b/modelopt/torch/quantization/utils/calib_utils.py @@ -41,6 +41,7 @@ import math import torch +import torch.distributed as dist from modelopt.torch.utils import print_rank_0 from modelopt.torch.utils.network import bind_forward_method, unpatch_forward_method @@ -170,6 +171,25 @@ def cleanup(self): """Unpatch the module's forward method.""" unpatch_forward_method(self.module, self.CACHE_NAME) + def sync_hessian_across_dp(self): + """Merge the Hessians of this module's data-parallel replicas, which saw different data.""" + parallel_state = getattr(self.module, "parallel_state", None) + group = parallel_state.data_parallel_group if parallel_state is not None else None + if group is None or not group.is_initialized() or group.world_size() == 1: + return + on_cuda = "nccl" in dist.get_backend(group.group) + device = ( + torch.device("cuda", torch.cuda.current_device()) if on_cuda else torch.device("cpu") + ) + n_samples = torch.tensor(self.n_samples, dtype=torch.float64, device=device) + hessian = self.hessian.to(device).mul_(self.n_samples) + dist.all_reduce(n_samples, group=group.group) + dist.all_reduce(hessian, group=group.group) + self.n_samples = int(n_samples.item()) + if self.n_samples: + hessian /= self.n_samples + self.hessian = hessian.to(self.hessian.device) + def free(self): """Release Hessian and working tensors to reclaim memory.""" self.hessian = None diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py index 3c327d48bd0..c4f4fd79afa 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py @@ -844,6 +844,19 @@ def build(*args, **kwargs): return True +def _flashinfer_kv_cache_views(kv_cache: torch.Tensor, head_size: int): + """Return logical [blocks, page, heads, dim] K/V views for FlashInfer.""" + if kv_cache.ndim == 5 and kv_cache.shape[1] == 2: + return kv_cache[:, 0], kv_cache[:, 1] + if kv_cache.ndim == 4 and kv_cache.shape[-1] == 2 * head_size: + return kv_cache.transpose(1, 2).split(head_size, dim=-1) + raise ValueError( + "FlashInfer KV cache must have logical shape [blocks, 2, page, heads, dim] " + "or [blocks, heads, page, 2 * dim], " + f"got {tuple(kv_cache.shape)}" + ) + + def _flashinfer_cache_write(layer, key, value, kv_cache, attn_metadata, impl) -> None: """Issue FlashInfer's native paged K/V cache write.""" torch.ops._C_cache_ops.reshape_and_cache_flash( @@ -924,18 +937,15 @@ def prepare_modelopt(): "FlashInfer metadata is missing the ModelOpt calibration " f"fields: {', '.join(missing)}" ) - if kv_cache.ndim != 5 or kv_cache.shape[1] != 2: - raise ValueError( - "FlashInfer KV cache must have logical shape [blocks, 2, page, heads, dim]" - ) + key_cache, value_cache = _flashinfer_kv_cache_views(kv_cache, impl.head_size) # Order matters: releases that update the KV cache inside forward must # write the current K/V before the calibrate kernel reads the cache. prepare_modelopt() return _forward_calibrate( impl, query=query, - key_cache=kv_cache[:, 0], - value_cache=kv_cache[:, 1], + key_cache=key_cache, + value_cache=value_cache, block_table=attn_metadata._modelopt_block_table, seq_lens=attn_metadata._modelopt_seq_lens, cu_seqlens_q=attn_metadata._modelopt_query_start_loc, @@ -954,13 +964,7 @@ def prepare_modelopt(): if resolved is None: return dense_fallback() - if kv_cache.ndim != 5 or kv_cache.shape[1] != 2: - raise ValueError( - "FlashInfer KV cache must have logical shape [blocks, 2, page, heads, dim]" - ) - - key_cache = kv_cache[:, 0] - value_cache = kv_cache[:, 1] + key_cache, value_cache = _flashinfer_kv_cache_views(kv_cache, impl.head_size) max_query_len = attn_metadata._modelopt_max_query_len is_decode_only = max_query_len <= 1 common_kw = { diff --git a/modelopt/torch/utils/dataset_utils.py b/modelopt/torch/utils/dataset_utils.py index 6202fd96076..f882a6e7364 100644 --- a/modelopt/torch/utils/dataset_utils.py +++ b/modelopt/torch/utils/dataset_utils.py @@ -888,6 +888,8 @@ def get_dataset_dataloader( all_samples, return_tensors="pt", padding=True, + # Left padding creates fully masked causal rows, which can yield NaNs in cuDNN backward. + padding_side="right" if include_labels else None, truncation=True, max_length=max_sample_length, ) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index c3cea6acaf4..181bcff632d 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -21,7 +21,7 @@ import sys import time import traceback -from collections.abc import Callable +from collections.abc import Callable, Iterator from contextlib import suppress from datetime import timedelta from typing import Any @@ -38,6 +38,8 @@ "backend", "barrier", "fsdp2_wrap", + "gather_bytes", + "gather_state_dicts", "is_available", "is_fsdp2_model", "is_initialized", @@ -108,9 +110,9 @@ def _deserialize(tensor: torch.Tensor, size: int | None = None) -> Any: buffer = tensor.numpy().tobytes() if size is not None: buffer = buffer[:size] - # Security NOTE: weights_only=False is used here on internally-generated buffer, not on untrusted user input - obj = torch.load(io.BytesIO(buffer), weights_only=False) - return obj + # Explicit so TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD cannot re-enable pickle for peer payloads. Custom + # classes sent through broadcast/allgather must be registered with add_safe_globals. + return torch.load(io.BytesIO(buffer), weights_only=True) def _broadcast(tensor: torch.Tensor, src: int = 0, group=None) -> None: @@ -144,6 +146,60 @@ def broadcast(obj: Any, src: int = 0, group=None) -> Any: return obj +def gather_bytes( + data: bytes, dst: int, group=None, chunk_bytes: int = 1 << 28 +) -> list[bytearray] | None: + """Gathers byte strings to global rank ``dst`` only, without unpickling anything on the way. + + Unlike ``torch.distributed.gather_object`` it never runs pickle, and it moves ``chunk_bytes`` per + rank at a time to bound device memory. Returns the per-rank bytes on ``dst`` and ``None`` elsewhere. + """ + n = size(group) + device = torch.device("cuda" if torch.distributed.get_backend(group) == "nccl" else "cpu") + local = ( + torch.frombuffer(bytearray(data), dtype=torch.uint8) + if data + else torch.empty(0, dtype=torch.uint8) + ) + sizes_t = [torch.zeros(1, dtype=torch.long, device=device) for _ in range(n)] + torch.distributed.all_gather(sizes_t, torch.tensor([len(data)], device=device), group=group) + sizes = [int(t.item()) for t in sizes_t] + + is_dst = torch.distributed.get_rank() == dst + out = [bytearray(s) for s in sizes] if is_dst else None + for off in range(0, max(sizes), chunk_bytes): + length = min(chunk_bytes, max(sizes) - off) + buf = torch.zeros(length, dtype=torch.uint8, device=device) + piece = local[off : off + length] + buf[: piece.numel()] = piece.to(device) + bufs = [torch.empty_like(buf) for _ in range(n)] if is_dst else None + torch.distributed.gather(buf, bufs, dst=dst, group=group) + if bufs is not None and out is not None: + for i, b in enumerate(bufs): + k = min(length, max(0, sizes[i] - off)) + out[i][off : off + k] = b[:k].cpu().numpy().tobytes() + return out + + +def gather_state_dicts(state: dict, dst: int, group=None) -> Iterator[dict]: + """Gathers ``state`` from every rank of ``group`` to global rank ``dst``, loaded with ``weights_only``. + + Yields one rank's dict at a time on ``dst`` (to bound host memory) and nothing elsewhere. The collective + runs eagerly, so every rank must call this before iterating. + """ + buf = io.BytesIO() + torch.save(state, buf) + gathered = gather_bytes(buf.getvalue(), dst, group) + del buf + + def _load() -> Iterator[dict]: + while gathered: + # Explicit weights_only so TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD cannot re-enable pickle for peers. + yield torch.load(io.BytesIO(gathered.pop(0)), map_location="cpu", weights_only=True) + + return _load() + + def _allgather(tensors: list[torch.Tensor], tensor: torch.Tensor, group=None) -> None: if backend() == "torch": torch.distributed.all_gather(tensors, tensor, group) diff --git a/modelopt/torch/utils/plugins/hf_checkpoint_utils.py b/modelopt/torch/utils/plugins/hf_checkpoint_utils.py index de8532e5297..dc8138d4286 100644 --- a/modelopt/torch/utils/plugins/hf_checkpoint_utils.py +++ b/modelopt/torch/utils/plugins/hf_checkpoint_utils.py @@ -187,10 +187,22 @@ def sanitize_hf_config_for_deployment(config_data: dict[str, Any], model: Any) - Fix conservative deployment-only config incompatibilities: + * remove legacy config/model auto mappings from native NemotronH exports; * add missing llama3 ``rope_theta`` metadata when available; * trim trailing MTP/next-token-prediction ``layer_types`` entries only when the mismatch is exactly explained by next-token-prediction metadata. """ + model_config = getattr(model, "config", None) + auto_map = config_data.get("auto_map") + if isinstance(auto_map, dict) and type(model_config).__module__.startswith( + "transformers.models.nemotron_h." + ): + # The native schema cannot reload through the source's legacy read-only config. + for auto_class in ("AutoConfig", "AutoModel", "AutoModelForCausalLM"): + auto_map.pop(auto_class, None) + if not auto_map: + config_data.pop("auto_map") + _sanitize_llama3_rope_config(config_data, model) num_hidden_layers = _as_nonnegative_int(config_data.get("num_hidden_layers")) diff --git a/modelopt/torch/utils/plugins/megatron_preprocess_data.py b/modelopt/torch/utils/plugins/megatron_preprocess_data.py index 50768abb0d8..a3601b02e67 100644 --- a/modelopt/torch/utils/plugins/megatron_preprocess_data.py +++ b/modelopt/torch/utils/plugins/megatron_preprocess_data.py @@ -90,6 +90,7 @@ import multiprocessing import time import warnings +from multiprocessing.context import BaseContext from pathlib import Path from datasets import get_dataset_config_names, get_dataset_split_names, load_dataset @@ -117,6 +118,35 @@ def _is_main_or_first_worker() -> bool: return not identity or identity[0] == 1 +_PRELOAD_MODULE = "modelopt.torch.utils.plugins.megatron_preprocess_data" + + +def _worker_context() -> BaseContext: + """Return the multiprocessing context for the tokenization worker pool. + + ``fork`` can segfault the children once CUDA is initialized in this process, so workers are + forked from a separate ``forkserver`` process instead. It preloads this module, so the heavy + ``torch`` / ``megatron`` import is paid once rather than per worker as with ``spawn``. + Windows has no ``forkserver`` and uses ``spawn``. + """ + if "forkserver" not in multiprocessing.get_all_start_methods(): + return multiprocessing.get_context("spawn") + ctx = multiprocessing.get_context("forkserver") + ctx.set_forkserver_preload([_PRELOAD_MODULE]) + return ctx + + +def _prestart_forkserver() -> None: + """Start the forkserver without waiting for it, so its preload overlaps with other work. + + The later ``Pool()`` then blocks only for what remains of it. No-op where ``spawn`` is used. + """ + if _worker_context().get_start_method() == "forkserver": + from multiprocessing import forkserver # POSIX-only module + + forkserver.ensure_running() + + class _Encoder: tokenizer: AutoTokenizer = None _chat_template_logged: set[str] = set() @@ -280,12 +310,10 @@ def _encode_in_batches(pool, encoder: "_Encoder", lines, batch_size: int): yield from encoded_batch def _encode_docs(self, encoder: "_Encoder", lines, may_stop_early: bool = False): - """Tokenize ``lines``, forking worker processes only when ``workers > 1``. + """Tokenize ``lines``, using worker processes only when ``workers > 1``. - ``multiprocessing.Pool`` always ``fork()``s, even for a single worker. Forking a - process that has already initialized a CUDA context / many native threads (e.g. when - this is called in-process after GPU work) is unsafe and can segfault the children. - The single-worker path avoids the fork entirely by tokenizing inline in this process. + Workers start through ``_worker_context``; ``_Encoder`` holds only picklable fields, as each + worker loads its tokenizer in ``initializer``. The single-worker path tokenizes inline. When ``may_stop_early`` is true, wait for finite batches so no worker results are pending if the caller stops consuming documents after reaching its token limit. @@ -295,7 +323,7 @@ def _encode_docs(self, encoder: "_Encoder", lines, may_stop_early: bool = False) if self.workers == 1: encoder.initializer() return None, map(encoder.encode, lines) - pool = multiprocessing.Pool(self.workers, initializer=encoder.initializer) + pool = _worker_context().Pool(self.workers, initializer=encoder.initializer) if may_stop_early: batch_size = self.workers * 4 # Balance throughput against unused final-batch work. encoded_docs = self._encode_in_batches(pool, encoder, lines, batch_size) @@ -556,7 +584,8 @@ def megatron_preprocess_data( json_keys: Key or list of keys to extract from json. Defaults to ["text"]. append_eod: Append an token to the end of a document. Defaults to False. max_sequence_length: Maximum tokenized sequence length. Defaults to None. - workers: Number of worker processes to launch. Defaults to 1. + workers: Number of worker processes to launch. Defaults to 1. With more than one, call from + under an ``if __name__ == "__main__":`` guard, since workers re-import the main script. log_interval: Interval between progress updates. Defaults to 100000. reasoning_content: How to handle the ``reasoning_content`` field present in many Nemotron Post-Training v3 datasets. One of: @@ -591,6 +620,10 @@ def megatron_preprocess_data( stacklevel=2, ) + if workers > 1: + # Start the worker server now so its import overlaps the tokenizer and dataset loading below. + _prestart_forkserver() + Path(output_dir).mkdir(parents=True, exist_ok=True) vocab_size = AutoTokenizer.from_pretrained(tokenizer_name_or_path).vocab_size diff --git a/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py b/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py index fa5875a5e2c..2bc7aacb8c9 100644 --- a/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py +++ b/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py @@ -16,17 +16,42 @@ """Prepare a weighted Megatron data blend from a YAML configuration.""" import argparse +import gzip +import hashlib +import json +import math import os +import random +import re import shutil +from collections.abc import Iterator from pathlib import Path from typing import Any, cast import huggingface_hub import yaml +from datasets import load_dataset + +from modelopt.torch.utils import ( + REASONING_CONTENT_MODES, + ReasoningContentMode, + prepare_messages_for_chat_template, +) from .megatron_preprocess_data import megatron_preprocess_data -__all__ = ["prepare_megatron_data_blend"] +__all__ = ["prepare_chat_data_blend", "prepare_megatron_data_blend"] + +# A turn of a ``text`` row in the NeMo chat format: ``User\n...``. Any marker ends the +# turn, so a trailing delimiter is not kept as content. The role must be followed by a newline, a +# marker or the end, so ``Userguide`` is no role at all. +_CHAT_TURN = re.compile( + r"(User|Assistant|System)(?:\n|(?=|\Z))(.*?)(?=|\Z)", + re.DOTALL, +) +_CHAT_ROLES = {"User": "user", "Assistant": "assistant", "System": "system"} +_CHAT_MARKER = re.compile(r"") +_CHAT_TRAILING_MARKER = re.compile(r"\s*\Z") def load_config(path: Path) -> dict[str, Any]: @@ -53,6 +78,10 @@ def load_config(path: Path) -> dict[str, Any]: optional ``target_tokens``. Each source has ``hf_dataset``, ``content_field``, and ``weight``; it uses ``split`` with optional ``config`` and ``max_samples``, or selects repository ``files``. + + With ``output_format: chat_jsonl``, the blend is written as chat JSONL instead (see + :func:`prepare_chat_data_blend`): ``tokenizer`` and ``target_tokens`` are not used, and the + optional ``target_samples``, ``shuffle``, ``seed`` and ``reasoning_content`` keys apply. """ with path.open(encoding="utf-8") as stream: return cast("dict[str, Any]", yaml.safe_load(stream)) @@ -76,6 +105,18 @@ def _copy_config(source: Path, destination: Path) -> None: shutil.copyfile(source, destination) +def _download_files(source: dict[str, Any], output_dir: Path) -> list[str]: + """Download the repository ``files`` of a source and return their local paths.""" + dataset = source["hf_dataset"] + raw_dir = output_dir.parent / "raw" / dataset.replace("/", "--") + return [ + huggingface_hub.hf_hub_download( + repo_id=dataset, filename=file, repo_type="dataset", local_dir=raw_dir + ) + for file in source["files"] + ] + + def _prepare_sources( sources: list[dict[str, Any]], output_dir: Path, @@ -104,17 +145,7 @@ def _prepare_sources( content_field = source["content_field"] input_args: dict[str, Any] if "files" in source: - raw_dir = output_dir.parent / "raw" / dataset.replace("/", "--") - paths = [ - huggingface_hub.hf_hub_download( - repo_id=dataset, - filename=file, - repo_type="dataset", - local_dir=raw_dir, - ) - for file in source["files"] - ] - input_args = {"jsonl_paths": paths} + input_args = {"jsonl_paths": _download_files(source, output_dir)} else: input_args = { "hf_dataset": dataset, @@ -146,9 +177,252 @@ def _prepare_sources( return blend +def _to_chat( + row: dict[str, Any], content_field: str, reasoning_content: ReasoningContentMode = "native" +) -> dict[str, Any] | None: + """Convert a source row to ``{"messages": [...]}``, or ``None`` if it has nothing to train on. + + ``messages`` go through :func:`prepare_messages_for_chat_template` with ``reasoning_content``, + and non-empty ``tools`` and ``system`` fields are kept for the chat path. A ``text`` row in the NeMo chat + format (``User`` / ``Assistant`` markers) is split into turns; any other ``text`` + row becomes one assistant turn, so all of it is trained. + """ + if content_field == "messages": + messages = row.get("messages") + if not messages: + return None + messages = prepare_messages_for_chat_template(messages, reasoning_content=reasoning_content) + chat: dict[str, Any] = {"messages": messages} + # As in the tokenizing mode, the chat template renders the tool definitions into the prompt; + # Megatron-Bridge turns a top-level system prompt into the leading system message. + for key in ("tools", "system"): + if row.get(key): + chat[key] = row[key] + else: + text = row.get("text") or "" + if not text: + return None + if " delimiters as assistant content. Text before the first + # marker, a marker without a known role (a Tool turn, or a SteerLM label line) or a known turn + # left empty would silently cut it out of the conversation, so drop the row. Only a bare + # delimiter at the very end needs no role, and only the System block, which real rows leave + # empty, may be empty. + roleless = len(_CHAT_MARKER.findall(text)) - len(matches) + if ( + not _CHAT_MARKER.match(text.lstrip()) + or roleless > (1 if _CHAT_TRAILING_MARKER.search(text) else 0) + or any(not content.strip() and role != "System" for role, content in matches) + ): + return None + turns = [ + {"role": _CHAT_ROLES[role], "content": content.strip()} + for role, content in matches + if content.strip() + ] + chat = {"messages": turns} + # Nothing in a row without an assistant turn is trained under assistant-only loss. + if not any(turn.get("role") == "assistant" for turn in chat["messages"]): + return None + return chat + + +def _iter_rows( + source: dict[str, Any], output_dir: Path, shuffle: bool, seed: int +) -> Iterator[dict[str, Any]]: + """Yield the rows of a source: its ``files`` in order, or its Hugging Face split, streamed.""" + if "files" in source: + for path in _download_files(source, output_dir): + opener = gzip.open if path.endswith(".gz") else open + with opener(path, "rt", encoding="utf-8") as f: + for line in f: + if not line.strip(): + continue + try: + yield json.loads(line) + except ValueError as e: + # As in the tokenizing reader: some Nemotron files carry a malformed record. + print(f"[WARN] Skipping malformed JSONL record in {path}: {e}", flush=True) + return + dataset = load_dataset( + source["hf_dataset"], source.get("config"), split=source["split"], streaming=True + ) + yield from dataset.shuffle(seed=seed) if shuffle else dataset + + +def _take_chat_rows( + source: dict[str, Any], + output_dir: Path, + want: int, + shuffle: bool, + seed: int, + reasoning_content: ReasoningContentMode, +) -> list[dict[str, Any]]: + """Return the first ``want`` usable chat rows of a source. + + Fails if the source cannot be read or runs out of usable rows, so the blend cannot change + silently; gives up after ``10 * want + 100`` rejected rows, a floor that keeps a small budget from + misjudging a source by its first few rows. + """ + # Megatron-Bridge's chat SFT path consumes this output; keep exactly the rows it accepts. Imported + # here, not at module scope, because the tokenizing mode does not need Megatron-Bridge installed. + from megatron.bridge.data.conversation_processing import is_text_only_chat_example + + name = f"{source['hf_dataset']} {source.get('split') or source.get('files')}" + rows: list[dict[str, Any]] = [] + rejected = 0 + max_rejected = 10 * want + 100 + try: + for row in _iter_rows(source, output_dir, shuffle, seed): + chat = _to_chat(row, source["content_field"], reasoning_content) + # Judge the source row's top-level media keys too, which the conversion does not keep. + if chat is None or not is_text_only_chat_example({**row, **chat}): + rejected += 1 + if rejected > max_rejected: + break + continue + rows.append(chat) + if len(rows) >= want: + break + except Exception as e: # e.g. a gated split, or a schema the streaming reader cannot cast + raise RuntimeError(f"{name}: failed to read or convert") from e + if len(rows) < want: + gave_up = rejected > max_rejected + reason = f"gave up after {rejected} rejected rows" if gave_up else "source exhausted" + raise RuntimeError(f"{name}: only {len(rows)}/{want} usable rows ({reason})") + print(f"{name}: {len(rows)} rows ({rejected} rejected)", flush=True) + return rows + + +def _write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None: + """Write ``rows`` atomically, so an interrupted write never leaves a partial file.""" + tmp = path.with_name(path.name + ".tmp") + with tmp.open("w", encoding="utf-8") as f: + f.writelines(json.dumps(row) + "\n" for row in rows) + os.replace(tmp, path) + + +def _read_jsonl(path: Path) -> list[dict[str, Any]] | None: + """Read a cached part, or return ``None`` if it is missing or unreadable.""" + try: + with path.open(encoding="utf-8") as f: + return [json.loads(line) for line in f] + except (OSError, ValueError): + return None + + +def _part_path(parts_dir: Path, source: dict[str, Any], settings: dict[str, Any]) -> Path: + """Return the cache file of a source's selected rows, keyed by every input that selects them. + + A part holds the first usable rows of the source, so it serves any budget it covers. + """ + selection = {key: source.get(key) for key in ("hf_dataset", "config", "split", "files")} + selection.update(settings, content_field=source["content_field"]) + digest = hashlib.sha256(json.dumps(selection, sort_keys=True).encode()).hexdigest()[:16] + label = source.get("split") or Path(source["files"][0]).name.split(".")[0] + return parts_dir / f"{source['hf_dataset'].replace('/', '--')}_{label}_{digest}.jsonl" + + +def _chat_budgets(config: dict[str, Any]) -> list[int]: + """Return each source's row budget; validate the inputs.""" + weights = [float(source["weight"]) for source in config["sources"]] + if not all(math.isfinite(w) and w >= 0 for w in weights) or sum(weights) <= 0: + raise ValueError(f"Source weights must be finite, nonnegative and not all zero: {weights}.") + # Required: the selected rows are held in memory to be mixed, which a whole source may not fit. + target = config.get("target_samples") + if isinstance(target, bool) or not isinstance(target, int) or target <= 0: + raise ValueError( + f"target_samples is required and must be a positive integer, got {target!r}." + ) + return [round(target * w / sum(weights)) for w in weights] + + +def prepare_chat_data_blend(config_path: Path) -> dict[str, int]: + """Write a weighted chat blend as ``train.jsonl`` / ``validation.jsonl``; return their sizes. + + For Megatron-Bridge's chat SFT path (``distill.py --sft --sft_hf_dataset``), which renders each + conversation with the model's chat template and can restrict the loss to assistant tokens. + Each source contributes ``round(target_samples * weight / sum(weights))`` rows (a source rounded + to zero rows is skipped) that the chat path accepts as text-only chat. ``content_field`` must be + ``messages`` or ``text``; ``messages`` keep ``reasoning_content`` for the chat template by + default (``reasoning_content: native``; set ``inline`` for a template that ignores the field, or + ``strip``). Hugging Face splits are buffer-shuffled with ``seed`` unless ``shuffle: false``; + ``files`` (``.jsonl`` or ``.jsonl.gz``) are read from the beginning, skipping malformed records. + The rows are mixed with ``seed`` and ``min(1000, n // 20)`` of them are held out for validation. + Each source is cached under ``/parts``, keyed by the settings that select its rows, + so an interrupted run resumes. + """ + config = load_config(config_path) + if config.get("output_format") != "chat_jsonl": + raise ValueError( + "prepare_chat_data_blend writes output_format: chat_jsonl, " + f"got {config.get('output_format')!r}." + ) + ignored = {"tokenizer", "target_tokens"} & config.keys() + if ignored: + raise ValueError(f"output_format: chat_jsonl does not use {sorted(ignored)}.") + if any("max_samples" in source for source in config["sources"]): + raise ValueError("output_format: chat_jsonl does not use max_samples; set target_samples.") + fields = {source["content_field"] for source in config["sources"]} - {"messages", "text"} + if fields: + raise ValueError( + f"output_format: chat_jsonl reads content_field messages or text: {fields}." + ) + reasoning_content = config.get("reasoning_content", "native") + if reasoning_content not in REASONING_CONTENT_MODES: + raise ValueError( + f"reasoning_content must be one of {sorted(REASONING_CONTENT_MODES)}, " + f"got {reasoning_content!r}." + ) + budgets = _chat_budgets(config) + output_dir = Path(config["output_dir"]) + parts_dir = output_dir / "parts" + parts_dir.mkdir(parents=True, exist_ok=True) + shuffle = bool(config.get("shuffle", True)) + seed = int(config.get("seed", 1234)) + # Everything besides the source itself that decides which rows a part holds. + settings = {"shuffle": shuffle, "reasoning_content": reasoning_content} + if shuffle: + settings["seed"] = seed + + rows: list[dict[str, Any]] = [] + for source, want in zip(config["sources"], budgets): + if want == 0: + print(f"{source['hf_dataset']}: weight rounds to 0 rows, skipped", flush=True) + continue + part = _part_path(parts_dir, source, settings) + got = _read_jsonl(part) + if got is None or len(got) < want: + got = _take_chat_rows(source, output_dir, want, shuffle, seed, reasoning_content) + _write_jsonl(part, got) + rows.extend(got[:want]) + + random.Random(seed).shuffle(rows) + num_validation = min(1000, len(rows) // 20) + splits = {"validation": rows[:num_validation], "train": rows[num_validation:]} + if not splits["train"] or not splits["validation"]: + raise RuntimeError(f"Too few rows ({len(rows)}) for a train/validation split.") + for name, split_rows in splits.items(): + _write_jsonl(output_dir / f"{name}.jsonl", split_rows) + _copy_config(config_path, output_dir / "config.yaml") + return {name: len(split_rows) for name, split_rows in splits.items()} + + def prepare_megatron_data_blend(config_path: Path) -> list[tuple[float, str]]: """Download and tokenize the configured weighted data sources.""" config = load_config(config_path) + output_format = config.get("output_format", "megatron") + if output_format != "megatron": + raise ValueError( + f"Unknown output_format {output_format!r} for a tokenized blend; use megatron, or " + "chat_jsonl with prepare_chat_data_blend." + ) + chat_keys = {"target_samples", "shuffle", "seed", "reasoning_content"} & config.keys() + if chat_keys: + raise ValueError(f"Tokenized blends do not use {sorted(chat_keys)}.") output_dir = Path(config["output_dir"]) output_dir.mkdir(parents=True, exist_ok=True) target_tokens = config.get("target_tokens") @@ -165,6 +439,11 @@ def main() -> None: """Prepare a data blend from the supplied configuration.""" parser = _build_parser() args = parser.parse_args() + # Each entry point rejects a config written for the other. + if load_config(args.config).get("output_format") == "chat_jsonl": + sizes = prepare_chat_data_blend(args.config) + print(f"Prepared chat JSONL: {sizes}. See train.jsonl, validation.jsonl and config.yaml.") + return blend = prepare_megatron_data_blend(args.config) print(f"Prepared {len(blend)} data paths. See data_blend.txt and config.yaml in the output.") diff --git a/modelopt_recipes/model_type/qwen3_5/ptq/fp8_vision-kv_none.yaml b/modelopt_recipes/model_type/qwen3_5/ptq/fp8_vision-kv_none.yaml index d6ed27b32d9..8854b376af9 100644 --- a/modelopt_recipes/model_type/qwen3_5/ptq/fp8_vision-kv_none.yaml +++ b/modelopt_recipes/model_type/qwen3_5/ptq/fp8_vision-kv_none.yaml @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -# Dense Qwen3.5 Vision Encoder FP8 with a high-precision language model and KV cache. +# Dense Qwen3.5 vision-merger FP8 with a high-precision language model and KV cache. # modelopt-schema: modelopt.recipe.config.ModelOptPTQRecipe imports: @@ -22,9 +22,8 @@ imports: metadata: description: >- - FP8 quantization of dense Qwen3.5 Vision Encoder Linear layers, including any merger Linears; - the language model, KV cache, patch embedding, and vision-attention operands remain in high - precision. + FP8 quantization of dense Qwen3.5 vision-merger Linears; the vision blocks, language model, + KV cache, and patch embedding remain in high precision for vLLM compatibility. quantize: algorithm: max diff --git a/modelopt_recipes/model_type/qwen3_5/ptq/fp8_vision_lm-kv_fp8_cast.yaml b/modelopt_recipes/model_type/qwen3_5/ptq/fp8_vision_lm-kv_fp8_cast.yaml index 173d32b286a..7ccf42855a0 100644 --- a/modelopt_recipes/model_type/qwen3_5/ptq/fp8_vision_lm-kv_fp8_cast.yaml +++ b/modelopt_recipes/model_type/qwen3_5/ptq/fp8_vision_lm-kv_fp8_cast.yaml @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -# Dense Qwen3.5 Vision Encoder and language model FP8 with FP8 KV-cache cast. +# Dense Qwen3.5 vision-merger and language-model FP8 with FP8 KV-cache cast. # modelopt-schema: modelopt.recipe.config.ModelOptPTQRecipe imports: @@ -25,9 +25,9 @@ imports: metadata: description: >- - W8A8 FP8 quantization of dense Qwen3.5 Vision Encoder and language-model Linear layers, - including any merger Linears, with FP8 KV-cache cast; patch embedding and vision-attention - operands remain in high precision. + W8A8 FP8 quantization of dense Qwen3.5 vision-merger and language-model Linears, with FP8 + KV-cache cast; vision blocks and patch embedding remain in high precision for vLLM + compatibility. quantize: algorithm: max @@ -36,6 +36,6 @@ quantize: - $import: w8a8_fp8_fp8 - $import: kv_fp8_cast - $import: default_disabled_quantizers - # Re-enable visual Linears after the standard vision exclusion. The imported snippet also + # Re-enable merger Linears after the standard vision exclusion. The imported snippet also # keeps vision-attention BMM quantizers disabled as a final defense-in-depth rule. - $import: vision_fp8 diff --git a/modelopt_recipes/model_type/qwen3_vl/ptq/fp8_vision-kv_none.yaml b/modelopt_recipes/model_type/qwen3_vl/ptq/fp8_vision-kv_none.yaml index 9fc2e4ee590..f80aad5962c 100644 --- a/modelopt_recipes/model_type/qwen3_vl/ptq/fp8_vision-kv_none.yaml +++ b/modelopt_recipes/model_type/qwen3_vl/ptq/fp8_vision-kv_none.yaml @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -# Qwen3-VL Vision Encoder FP8 with a high-precision language model and KV cache. +# Qwen3-VL vision-merger FP8 with a high-precision language model and KV cache. # modelopt-schema: modelopt.recipe.config.ModelOptPTQRecipe imports: @@ -22,9 +22,8 @@ imports: metadata: description: >- - FP8 quantization of Qwen3-VL Vision Encoder Linear layers, including primary and deepstack - merger Linears; the language model, KV cache, patch embedding, and vision-attention operands - remain in high precision. + FP8 quantization of Qwen3-VL primary and deepstack vision-merger Linears; the vision blocks, + language model, KV cache, and patch embedding remain in high precision for vLLM compatibility. quantize: algorithm: max diff --git a/modelopt_recipes/model_type/qwen3_vl/ptq/fp8_vision_lm-kv_fp8_cast.yaml b/modelopt_recipes/model_type/qwen3_vl/ptq/fp8_vision_lm-kv_fp8_cast.yaml index fbd449d2d42..131859ab2c6 100644 --- a/modelopt_recipes/model_type/qwen3_vl/ptq/fp8_vision_lm-kv_fp8_cast.yaml +++ b/modelopt_recipes/model_type/qwen3_vl/ptq/fp8_vision_lm-kv_fp8_cast.yaml @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -# Qwen3-VL Vision Encoder and language model FP8 with FP8 KV-cache cast. +# Qwen3-VL vision-merger and language-model FP8 with FP8 KV-cache cast. # modelopt-schema: modelopt.recipe.config.ModelOptPTQRecipe imports: @@ -25,9 +25,9 @@ imports: metadata: description: >- - W8A8 FP8 quantization of Qwen3-VL Vision Encoder and language-model Linear layers, including - primary and deepstack merger Linears, with FP8 KV-cache cast; patch embedding and - vision-attention operands remain in high precision. + W8A8 FP8 quantization of Qwen3-VL primary and deepstack vision-merger and language-model + Linears, with FP8 KV-cache cast; vision blocks and patch embedding remain in high precision + for vLLM compatibility. quantize: algorithm: max @@ -36,6 +36,6 @@ quantize: - $import: w8a8_fp8_fp8 - $import: kv_fp8_cast - $import: default_disabled_quantizers - # Re-enable visual Linears after the standard vision exclusion. The imported snippet also + # Re-enable merger Linears after the standard vision exclusion. The imported snippet also # keeps vision-attention BMM quantizers disabled as a final defense-in-depth rule. - $import: vision_fp8 diff --git a/modelopt_recipes/model_type/qwen3_vl/ptq/vision_fp8.quant_cfg.yaml b/modelopt_recipes/model_type/qwen3_vl/ptq/vision_fp8.quant_cfg.yaml index 9486d5a7fa2..a650aa895a8 100644 --- a/modelopt_recipes/model_type/qwen3_vl/ptq/vision_fp8.quant_cfg.yaml +++ b/modelopt_recipes/model_type/qwen3_vl/ptq/vision_fp8.quant_cfg.yaml @@ -13,23 +13,21 @@ # See the License for the specific language governing permissions and # limitations under the License. -# FP8 Linear weights and inputs in the Qwen visual branch. Patch embedding -# and vision-attention operands remain in high precision. +# Keep vision blocks in high precision: vLLM's FP8 GEMM flattens the 3D +# attention inputs expected by the vision encoder. Merger Linears use 2D inputs. # modelopt-schema: modelopt.torch.quantization.config.QuantizerCfgListConfig imports: fp8: configs/numerics/fp8 --- - - quantizer_name: '*visual.*weight_quantizer' + - quantizer_name: '*visual.*merger*.weight_quantizer' parent_class: 'nn.Linear' cfg: $import: fp8 - - quantizer_name: '*visual.*input_quantizer' + - quantizer_name: '*visual.*merger*.input_quantizer' parent_class: 'nn.Linear' cfg: $import: fp8 - - quantizer_name: '*visual.*patch_embed*' - enable: false # Keep this defense-in-depth rule last when the snippet follows a KV-cache unit. - quantizer_name: '*visual.*_bmm_quantizer' enable: false diff --git a/modelopt_recipes/models/Qwen/Qwen3.8-27B/ptq/nvfp4_w4a4_mlp_fp8_attn_local_hessian.yaml b/modelopt_recipes/models/Qwen/Qwen3.8-27B/ptq/nvfp4_w4a4_mlp_fp8_attn_local_hessian.yaml index 90decca0902..08894fec91d 100644 --- a/modelopt_recipes/models/Qwen/Qwen3.8-27B/ptq/nvfp4_w4a4_mlp_fp8_attn_local_hessian.yaml +++ b/modelopt_recipes/models/Qwen/Qwen3.8-27B/ptq/nvfp4_w4a4_mlp_fp8_attn_local_hessian.yaml @@ -20,7 +20,7 @@ imports: metadata: description: >- Qwen3.8-27B local-Hessian PTQ using the NVFP4 W4A4 and FP8 W8A8 assignment exported by - the 5.5-bit NVFP4-max AutoQuantize sweep. + the 5.5-bit NVFP4-max AutoQuantize sweep, as published in nvidia/Qwen3.8-27B-NVFP4. quantize: $import: base diff --git a/modelopt_recipes/models/google/diffusiongemma-26B-A4B-it/ptq/nvfp4_experts_only-kv_fp8_cast.yaml b/modelopt_recipes/models/google/diffusiongemma-26B-A4B-it/ptq/nvfp4_experts_only-kv_fp8_cast.yaml new file mode 100644 index 00000000000..d41599700b3 --- /dev/null +++ b/modelopt_recipes/models/google/diffusiongemma-26B-A4B-it/ptq/nvfp4_experts_only-kv_fp8_cast.yaml @@ -0,0 +1,30 @@ +# 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. + +# Alias recipe for nvidia/diffusiongemma-26B-A4B-it-NVFP4 (source google/diffusiongemma-26B-A4B-it). Same recipe as `general/ptq/nvfp4_experts_only-kv_fp8_cast`: NVFP4 on the routed experts +# only, max calibration, with an FP8 KV cache in cast mode -- the same scheme as the +# gemma-4-26B-A4B-it release. This model nests its decoder under `model.decoder.layers` +# rather than `model.language_model.layers`, which the recipe's `*.experts.*` patterns +# match either way. + +imports: + base: general/ptq/nvfp4_experts_only-kv_fp8_cast + +$import: base +metadata: + description: >- + google/diffusiongemma-26B-A4B-it quantized with the general expert-only NVFP4 scheme + (max calibration) and an FP8 KV cache in cast mode, as published in + nvidia/diffusiongemma-26B-A4B-it-NVFP4. diff --git a/modelopt_recipes/models/zai-org/GLM-5.1/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast.yaml b/modelopt_recipes/models/google/gemma-4-26B-A4B-it/ptq/nvfp4_experts_only-kv_fp8_cast.yaml similarity index 57% rename from modelopt_recipes/models/zai-org/GLM-5.1/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast.yaml rename to modelopt_recipes/models/google/gemma-4-26B-A4B-it/ptq/nvfp4_experts_only-kv_fp8_cast.yaml index 876a47584e0..4f38d435a43 100644 --- a/modelopt_recipes/models/zai-org/GLM-5.1/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast.yaml +++ b/modelopt_recipes/models/google/gemma-4-26B-A4B-it/ptq/nvfp4_experts_only-kv_fp8_cast.yaml @@ -13,16 +13,16 @@ # See the License for the specific language governing permissions and # limitations under the License. -# Alias recipe for nvidia/GLM-5.1-NVFP4 (source zai-org/GLM-5.1). Same recipe as `general/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast`: NVFP4 on the -# routed experts with the expert `input_scale` pinned to 1.0, and an FP8 KV cache in cast -# mode. +# Alias recipe for nvidia/Gemma-4-26B-A4B-NVFP4 (source google/gemma-4-26B-A4B-it). Same recipe as `general/ptq/nvfp4_experts_only-kv_fp8_cast`: NVFP4 on the routed experts +# only, max calibration, with an FP8 KV cache in cast mode. The dense MLPs, router, +# attention, vision tower and `lm_head` stay BF16. imports: - base: general/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast + base: general/ptq/nvfp4_experts_only-kv_fp8_cast $import: base metadata: description: >- - zai-org/GLM-5.1 quantized with the general expert-only NVFP4 scheme, expert - `input_scale` pinned to 1.0, and an FP8 KV cache in cast mode, as published in - nvidia/GLM-5.1-NVFP4. + google/gemma-4-26B-A4B-it quantized with the general expert-only NVFP4 scheme (max + calibration) and an FP8 KV cache in cast mode, as published in + nvidia/Gemma-4-26B-A4B-NVFP4. diff --git a/modelopt_recipes/models/moonshotai/Kimi-K2.7-Code/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast.yaml b/modelopt_recipes/models/moonshotai/Kimi-K2.7-Code/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast.yaml new file mode 100644 index 00000000000..119cf553713 --- /dev/null +++ b/modelopt_recipes/models/moonshotai/Kimi-K2.7-Code/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast.yaml @@ -0,0 +1,30 @@ +# 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. + +# Alias recipe for nvidia/Kimi-K2.7-Code-NVFP4 (source moonshotai/Kimi-K2.7-Code). Same recipe as +# `general/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast`: NVFP4 on the routed experts +# with the expert `input_scale` pinned to 1.0 -- no activation calibration -- and an FP8 KV +# cache in cast mode. The release's exported expert `input_scale` tensors are exactly 1.0, +# and its shared experts, attention, `lm_head` and vision tower are all left unquantized. + +imports: + base: general/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast + +$import: base +metadata: + description: >- + moonshotai/Kimi-K2.7-Code quantized with the general expert-only NVFP4 scheme, expert + `input_scale` pinned to 1.0 (calibration-free), and an FP8 KV cache in cast mode, as + published in nvidia/Kimi-K2.7-Code-NVFP4. diff --git a/modelopt_recipes/models/zai-org/GLM-5.1/ptq/nvfp4_experts_only-kv_fp8_cast.yaml b/modelopt_recipes/models/zai-org/GLM-5.1/ptq/nvfp4_experts_only-kv_fp8_cast.yaml new file mode 100644 index 00000000000..dbb98cac5f5 --- /dev/null +++ b/modelopt_recipes/models/zai-org/GLM-5.1/ptq/nvfp4_experts_only-kv_fp8_cast.yaml @@ -0,0 +1,33 @@ +# 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. + +# Alias recipe for nvidia/GLM-5.1-NVFP4 (source zai-org/GLM-5.1). Same recipe as +# `general/ptq/nvfp4_experts_only-kv_fp8_cast`: NVFP4 on the routed experts with calibrated +# activation scales, and an FP8 KV cache in cast mode. +# +# Not the `input_scale1` variant, despite the sibling GLM-5.2 entry using it. That recipe +# pins the expert activation amax to a constant 2688.0 (E2M1_MAX * E4M3_MAX) so every +# exported expert `input_scale` is exactly 1.0. GLM-5.2's are; GLM-5.1's are calibrated +# values (sampled 0 of 8 equal to 1.0, spanning 3.4e-4 to 1.3e-2), so this release used +# ordinary activation calibration. + +imports: + base: general/ptq/nvfp4_experts_only-kv_fp8_cast + +$import: base +metadata: + description: >- + zai-org/GLM-5.1 quantized with the general expert-only NVFP4 scheme (max calibration) + and an FP8 KV cache in cast mode, as published in nvidia/GLM-5.1-NVFP4. diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index 33c43e42359..af3c10a117d 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -334,8 +334,9 @@ linear layers and MSE-calibrated NVFP4 to routed experts, with expert remain unquantized. **`qwen3_vl/ptq`** and **`qwen3_5/ptq`** provide FP8 recipes for the `visual` branch of -validated Qwen3-VL and dense Qwen3.5 checkpoints. `fp8_vision-kv_none` enables only Vision -Encoder `nn.Linear` weights and inputs, including the primary merger and any deepstack mergers. +validated Qwen3-VL and dense Qwen3.5 checkpoints. `fp8_vision-kv_none` enables FP8 only on the +primary and any deepstack vision-merger `nn.Linear` weights and inputs. Vision blocks remain in +high precision because vLLM's FP8 GEMM flattens their 3D attention inputs. `fp8_vision_lm-kv_fp8_cast` combines that visual configuration with the standard W8A8 FP8 model and FP8 KV-cache-cast units. Both keep patch embedding and vision-attention BMM operands in high precision. The shared visual snippet lives under `qwen3_vl`; thin wrappers remain discoverable @@ -563,11 +564,24 @@ entry is a thin **alias** that imports that recipe wholesale and overrides only `model_type/qwen3_5_moe/ptq/nvfp4_experts_mse-fp8_rest-kv_fp8` — NVFP4 (MSE static weights) on the routed experts, ModelOpt-default FP8 elsewhere, and an FP8 KV cache — as published in `nvidia/Qwen3.5-397B-A17B-NVFP4-V2`. -- **`models/zai-org/GLM-5.1/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast`** and - **`models/zai-org/GLM-5.2/ptq/...`** alias +- **`models/moonshotai/Kimi-K2.7-Code/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast`** + aliases `general/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast` — expert-only NVFP4 with + the expert `input_scale` pinned to 1.0 (no activation calibration) and an FP8 KV cache in + cast mode — as published in `nvidia/Kimi-K2.7-Code-NVFP4`. +- **`models/google/gemma-4-26B-A4B-it/ptq/nvfp4_experts_only-kv_fp8_cast`** and + **`models/google/diffusiongemma-26B-A4B-it/ptq/...`** alias + `general/ptq/nvfp4_experts_only-kv_fp8_cast` — expert-only NVFP4 with max calibration and an + FP8 KV cache in cast mode — as published in `nvidia/Gemma-4-26B-A4B-NVFP4` and + `nvidia/diffusiongemma-26B-A4B-it-NVFP4`. +- **`models/zai-org/GLM-5.2/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast`** aliases `general/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast` — expert-only NVFP4 with the expert `input_scale` pinned to 1.0 and an FP8 KV cache in cast mode — as published in - `nvidia/GLM-5.1-NVFP4` and `nvidia/GLM-5.2-NVFP4`. + `nvidia/GLM-5.2-NVFP4`. +- **`models/zai-org/GLM-5.1/ptq/nvfp4_experts_only-kv_fp8_cast`** aliases + `general/ptq/nvfp4_experts_only-kv_fp8_cast` — expert-only NVFP4 with calibrated + activation scales and an FP8 KV cache in cast mode — as published in + `nvidia/GLM-5.1-NVFP4`. Unlike GLM-5.2, this release's exported expert `input_scale` + tensors are calibrated rather than pinned to 1.0. - **`models/MiniMaxAI/MiniMax-M2.5/ptq/nvfp4_experts_only-kv_fp8_cast`** aliases `general/ptq/nvfp4_experts_only-kv_fp8_cast` — expert-only NVFP4 with an FP8 KV cache in cast mode — as published in `nvidia/MiniMax-M2.5-NVFP4`. diff --git a/noxfile.py b/noxfile.py index afa96010ac4..edd0544e1cd 100644 --- a/noxfile.py +++ b/noxfile.py @@ -69,7 +69,20 @@ def unit(session, torch_ver, tf_ver): tf_pins = TRANSFORMERS_VERSIONS[tf_ver] if tf_pins: session.install(*tf_pins) - session.run("python", "-m", "pytest", "tests/unit", *_cov_args(), env=_CPU_ONLY_ENV) + run_manual = "--run-manual" in session.posargs + pytest_args = [arg for arg in session.posargs if arg != "--run-manual"] + session.run( + "python", "-m", "pytest", "tests/unit", *_cov_args(), *pytest_args, env=_CPU_ONLY_ENV + ) + if run_manual: + # A separate process: some manual tests (e.g. multi-threaded from_pretrained) leak global + # state that breaks later tests. No coverage, so it cannot overwrite the main run's data. + session.run( + "python", + # Caller args first: the last ``-m`` wins, and it must be ours. + *("-m", "pytest", "tests/unit", *pytest_args, "--run-manual", "-m", "manual"), + env={**_CPU_ONLY_ENV, "COVERAGE_PROCESS_START": ""}, + ) @nox.session(python="3.12") diff --git a/plugins/modelopt/skills/ptq/references/checkpoint-validation.md b/plugins/modelopt/skills/ptq/references/checkpoint-validation.md index 3d755933134..c836681147f 100644 --- a/plugins/modelopt/skills/ptq/references/checkpoint-validation.md +++ b/plugins/modelopt/skills/ptq/references/checkpoint-validation.md @@ -201,6 +201,6 @@ Layers silently skipped because the quantization config patterns don't match the ## What to do when warnings appear -- **Layers should have been quantized** (e.g., MoE experts with `nvfp4_mlp_only`): the quantization config patterns missed them. Fix by adding the missing pattern to the config and re-running PTQ. Check if ModelOpt already has a plugin for the model in `modelopt/torch/quantization/plugins/huggingface.py`. +- **Layers should have been quantized** (e.g., MoE experts with `nvfp4_mlp_only`): the quantization config patterns missed them. Fix by adding the missing pattern to the config and re-running PTQ. Check if ModelOpt already has a plugin for the model in `modelopt/torch/models//modeling_ptq.py`. - **Layers are intentionally unquantized** (e.g., attention layers with `nvfp4_mlp_only`): they should be in the `exclude_modules` list but the export didn't add them. Add them manually to both `hf_quant_config.json` and `config.json` `quantization_config.ignore` in the checkpoint to prevent deployment failures. diff --git a/plugins/modelopt/skills/ptq/references/unsupported-models.md b/plugins/modelopt/skills/ptq/references/unsupported-models.md index 4ed7959c71b..09514ebf4be 100644 --- a/plugins/modelopt/skills/ptq/references/unsupported-models.md +++ b/plugins/modelopt/skills/ptq/references/unsupported-models.md @@ -66,7 +66,7 @@ Manual dequantization is only needed for **non-standard parameter names** (e.g., Read the model source to identify how weights are stored. **If all linear layers are plain `nn.Linear`, no custom code is needed** — ModelOpt quantizes them automatically. -**For HuggingFace models**, check `modelopt/torch/quantization/plugins/huggingface.py` first — it already registers patches for common non-standard modules (`Llama4TextExperts`, `FP8Linear`, `FalconLinear`, `Conv1D`, `Qwen3_5MoeExperts`, etc.). If your model's non-standard class is already registered there, no extra code is needed. +**For HuggingFace models**, check `modelopt/torch/models//modeling_ptq.py` (model-specific patches, e.g. `Llama4TextExperts`, `FalconLinear`) and `modelopt/torch/quantization/plugins/huggingface.py` (generic ones: `FP8Linear`, `Conv1D`, auto-detected MoE such as `Qwen3_5MoeExperts`) first. If your model's non-standard class is already registered, no extra code is needed. Custom patches are required when: @@ -94,7 +94,7 @@ Compare against the `enable`/`disable` patterns in the config. Add custom overri After Steps A-D: - **No patches needed** (all standard modules) → run `hf_ptq.py` with a smoke test (`--calib_size 4`). If it succeeds, proceed with full calibration. If it fails, read the error and revisit Steps C/D. -- **Patches needed** → patch ModelOpt directly using the patterns below (add `QuantModule` in `modelopt/torch/quantization/plugins/huggingface.py`, update `modelopt/torch/export/` if needed), then run `hf_ptq.py` with a smoke test. This is preferred over writing a standalone script because it reuses all existing `hf_ptq.py` logic. Debug failures iteratively — quantization errors often reveal additional modules that need patching. +- **Patches needed** → patch ModelOpt directly using the patterns below (add `QuantModule` in `modelopt/torch/models//modeling_ptq.py`, update `modelopt/torch/export/` if needed), then run `hf_ptq.py` with a smoke test. This is preferred over writing a standalone script because it reuses all existing `hf_ptq.py` logic. Debug failures iteratively — quantization errors often reveal additional modules that need patching. --- @@ -138,7 +138,7 @@ class QuantCustomModule(OriginalModule): - **Fused experts** (`gate_up_proj` + `down_proj` 3D tensors, all native transformers MoEs) → auto-detected by `register_fused_experts_on_the_fly`, handled by `_QuantFusedExperts`. Covers Mixtral, Qwen, DeepSeek, Jamba, OlMoE, etc. - **Sequential per-expert `nn.Linear`** with `gate` + `experts` (e.g. remote-code models) → auto-detected by `register_sparse_moe_on_the_fly`. -**Custom MoE** (non-standard layout not matching auto-detection) requires patching. Find the closest pattern in the plugin (`modelopt/torch/quantization/plugins/huggingface.py`): +**Custom MoE** (non-standard layout not matching auto-detection) requires patching. Find the closest pattern in `modelopt/torch/models/*/modeling_ptq.py`: | MoE design | Strategy | Plugin example | | --- | --- | --- | @@ -163,7 +163,7 @@ for name, module in model.named_modules(): ## Pattern 3: Registering with ModelOpt -**When patching the plugin directly** (preferred): Use `QuantModuleRegistry.register` in `modelopt/torch/quantization/plugins/huggingface.py`, following existing examples: +**When patching ModelOpt directly** (preferred): Use `QuantModuleRegistry.register` in `modelopt/torch/models//modeling_ptq.py` and add `` to the `modeling_ptq` import loop at the end of `modelopt/torch/quantization/plugins/huggingface.py`, following existing examples: ```python from modelopt.torch.quantization.nn import QuantModuleRegistry @@ -179,6 +179,8 @@ def register_my_model_on_the_fly(model): if QuantModuleRegistry.get(mod_type) is None: QuantModuleRegistry.register({mod_type: f"hf.{mod_type.__name__}"})(QuantCustomModule) break + +CUSTOM_MODEL_PLUGINS.add(register_my_model_on_the_fly) # from ...quantization.plugins.custom ``` **When writing a standalone script** (fallback): Use `mtq.register()`: diff --git a/tests/_test_utils/onnx/lib_test_models.py b/tests/_test_utils/onnx/lib_test_models.py index 1b669317e8a..57a0621c5da 100644 --- a/tests/_test_utils/onnx/lib_test_models.py +++ b/tests/_test_utils/onnx/lib_test_models.py @@ -1344,3 +1344,35 @@ def build_matmul_1xn_model(): onnx.checker.check_model(model_inferred) return model_inferred + + +# x ----- relu -------------+ +# | +# cls (initializer, or --- +--- concat --- matmul --- out +# relu(y) if not constant_input) +# +# ``cls`` mimics a ViT class token: a constant Concat input that has no producer node. +def build_concat_with_constant_input_model(constant_input=True): + x = helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT, [1, 4, 8]) + y = helper.make_tensor_value_info("y", onnx.TensorProto.FLOAT, [1, 1, 8]) + out = helper.make_tensor_value_info("out", onnx.TensorProto.FLOAT, [1, 5, 8]) + nodes = [helper.make_node("Relu", ["x"], ["relu_out"], name="relu")] + inputs = [x] + initializers = [ + helper.make_tensor("weight", onnx.TensorProto.FLOAT, [8, 8], np.ones(64, dtype=np.float32)) + ] + if constant_input: + cls = helper.make_tensor( + "cls", onnx.TensorProto.FLOAT, [1, 1, 8], np.ones(8, dtype=np.float32) + ) + initializers.append(cls) + else: + inputs.append(y) + nodes.append(helper.make_node("Relu", ["y"], ["cls"], name="relu_cls")) + nodes += [ + helper.make_node("Concat", ["cls", "relu_out"], ["concat_out"], name="concat", axis=1), + helper.make_node("MatMul", ["concat_out", "weight"], ["out"], name="matmul"), + ] + graph = helper.make_graph(nodes, "concat_with_constant_input", inputs, [out], initializers) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + return onnx.shape_inference.infer_shapes(model) diff --git a/tests/_test_utils/torch/distributed/utils.py b/tests/_test_utils/torch/distributed/utils.py index eae7c2ad38c..ea954efbcf5 100644 --- a/tests/_test_utils/torch/distributed/utils.py +++ b/tests/_test_utils/torch/distributed/utils.py @@ -13,8 +13,13 @@ # See the License for the specific language governing permissions and # limitations under the License. +import contextlib +import faulthandler import os +import queue +import signal import socket +import time import traceback import torch @@ -79,6 +84,18 @@ def spawn_multiprocess_job(size, job, backend="gloo"): assert not p.exitcode +def get_live_pool(pools, key, make_pool): + """Return ``pools[key]``, building it with ``make_pool`` if it is missing or was killed. + + Fixtures hand pools out through this, so a hang kills the pool for one test instead of + failing every later test that uses it. + """ + pool = pools.get(key) + if pool is None or pool.broken: + pools[key] = pool = make_pool() + return pool + + def default_worker_teardown(rank, world_size): """Minimal cleanup between tests in persistent workers.""" try: @@ -117,6 +134,10 @@ def _test_fn(rank, size): ... def __init__(self, world_size, backend="nccl", teardown_fn=default_worker_teardown): assert world_size > 0, "World size must be greater than 0" self.world_size = world_size + # Set once a job was interrupted: the workers are then killed, see ``run``. + self.broken = False + self.job_timeout = 600 + self.failure_grace = 30 # seconds the other ranks get to finish once one has failed ctx = mp.get_context("spawn") self._cmd_queues = [ctx.Queue() for _ in range(world_size)] self._result_queue = ctx.Queue() @@ -145,6 +166,9 @@ def __init__(self, world_size, backend="nccl", teardown_fn=default_worker_teardo @staticmethod def _worker_loop(rank, world_size, backend, port, cmd_queue, result_queue, teardown_fn): + if hasattr(signal, "SIGUSR1"): + # The parent signals this to get the stack of a hung job into the CI log. + faulthandler.register(signal.SIGUSR1, all_threads=True) os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = str(port) os.environ["LOCAL_RANK"] = str(rank) @@ -186,19 +210,50 @@ def run(self, fn, *args, **kwargs): ``fn`` is called as ``fn(rank, world_size, *args, **kwargs)`` and must be picklable (top-level function or ``functools.partial`` of one). - """ - for q in self._cmd_queues: - q.put((fn, args, kwargs)) + If the wait is interrupted (e.g. by ``pytest-timeout``), or a rank fails and its peers do + not finish within ``failure_grace`` (they are typically blocked in a collective waiting for + it), the workers are still inside the job and every later ``run`` would hang until its own + timeout. The workers are killed instead and the pool is marked ``broken``; callers that + want to continue need a new pool. + """ + if self.broken: + raise RuntimeError("Worker pool was killed after an earlier job hung; create a new one") errors = [] - for _ in range(self.world_size): - status, rank, tb = self._result_queue.get(timeout=600) - if status == "error": - errors.append(f"--- Rank {rank} ---\n{tb}") + deadline = None # set when the first rank fails: its peers get ``failure_grace`` in total + try: + for q in self._cmd_queues: + q.put((fn, args, kwargs)) + for _ in range(self.world_size): + timeout = self.job_timeout if deadline is None else deadline - time.monotonic() + status, rank, tb = self._result_queue.get(timeout=max(timeout, 0)) + if status == "error": + errors.append(f"--- Rank {rank} ---\n{tb}") + deadline = deadline or time.monotonic() + self.failure_grace + except queue.Empty: + self._abort() + if not errors: + raise + errors.append("--- The other ranks did not finish after a rank failed; killed ---") + except BaseException: + self._abort() + raise if errors: raise RuntimeError("Worker(s) failed:\n" + "\n".join(errors)) + def _abort(self): + """Print every worker's stack (to stderr, for the CI log), then kill the workers.""" + self.broken = True + if hasattr(signal, "SIGUSR1"): + for p in self._processes: + with contextlib.suppress(OSError): + os.kill(p.pid, signal.SIGUSR1) + time.sleep(3) # let faulthandler write before the workers go away + for p in self._processes: + p.kill() + p.join(timeout=10) + def shutdown(self): """Signal all workers to exit and wait for them to finish.""" for q in self._cmd_queues: diff --git a/tests/examples/hf_ptq/test_hf_ptq_model_type.py b/tests/examples/hf_ptq/test_hf_ptq_model_type.py index 0ab585b19d0..328dcd5e564 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_model_type.py +++ b/tests/examples/hf_ptq/test_hf_ptq_model_type.py @@ -53,6 +53,20 @@ def test_pre_quantize_generation_skip(hf_ptq, model_type: str, skips_generation: assert len(generate_calls) == (0 if skips_generation else 1) +@pytest.mark.parametrize( + ("ids", "mask"), [([0, 0, 1, 2], [0, 0, 1, 1]), ([1, 2, 0, 0], [1, 1, 0, 0])] +) +def test_pre_quantize_preview_removes_padding(hf_ptq, ids, mask): + calls = [] + model = SimpleNamespace(generate=lambda ids, **kwargs: calls.append((ids, kwargs)) or ids) + args = SimpleNamespace(specdec_offline_dataset=None, skip_generate=False) + batch = {"input_ids": torch.tensor([ids]), "attention_mask": torch.tensor([mask])} + prompt, attention_mask, _ = hf_ptq.pre_quantize(args, model, "llama", None, [batch], False) + assert torch.equal(prompt, torch.tensor([[1, 2]])) + assert torch.equal(attention_mask, torch.ones_like(prompt)) + assert torch.equal(calls[0][0], prompt) + + @pytest.mark.parametrize( ("model_type", "expected"), [ diff --git a/tests/examples/torch_trt/test_custom_op_plugin.py b/tests/examples/torch_trt/test_custom_op_plugin.py new file mode 100644 index 00000000000..7af15527b57 --- /dev/null +++ b/tests/examples/torch_trt/test_custom_op_plugin.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import shutil +import subprocess +from pathlib import Path + +import pytest + + +@pytest.mark.skipif( + shutil.which("cmake") is None or shutil.which("nvcc") is None, + reason="Requires CMake, CUDA Toolkit, and TensorRT development files", +) +def test_identity_conv_plugin_build(tmp_path): + """Build the custom-op plugin against the installed TensorRT C++ interfaces.""" + source = Path(__file__).resolve().parents[3] / "examples/onnx_ptq/custom_op_plugin/plugin" + subprocess.run(["cmake", "-S", str(source), "-B", str(tmp_path)], check=True, timeout=60) + subprocess.run(["cmake", "--build", str(tmp_path), "--parallel", "2"], check=True, timeout=120) + assert (tmp_path / "libidentity_conv_plugin.so").is_file() diff --git a/tests/gpu/conftest.py b/tests/gpu/conftest.py index b2b9344945b..40a272e8838 100644 --- a/tests/gpu/conftest.py +++ b/tests/gpu/conftest.py @@ -15,15 +15,29 @@ import pytest import torch -from _test_utils.torch.distributed.utils import DistributedWorkerPool, default_worker_teardown +from _test_utils.torch.distributed.utils import ( + DistributedWorkerPool, + default_worker_teardown, + get_live_pool, +) @pytest.fixture(scope="session") -def dist_workers(): - pool = DistributedWorkerPool( - world_size=torch.cuda.device_count(), - backend="nccl", - teardown_fn=default_worker_teardown, +def _dist_pools(): + pools = {} + yield pools + for pool in pools.values(): + pool.shutdown() + + +@pytest.fixture +def dist_workers(_dist_pools): + """The session-wide pool; replaced if an earlier test hung and killed it.""" + world_size = torch.cuda.device_count() + return get_live_pool( + _dist_pools, + world_size, + lambda: DistributedWorkerPool( + world_size=world_size, backend="nccl", teardown_fn=default_worker_teardown + ), ) - yield pool - pool.shutdown() diff --git a/tests/gpu/onnx/quantization/autotune/test_benchmark.py b/tests/gpu/onnx/quantization/autotune/test_benchmark.py index 925d45fffe6..012bf495ab6 100644 --- a/tests/gpu/onnx/quantization/autotune/test_benchmark.py +++ b/tests/gpu/onnx/quantization/autotune/test_benchmark.py @@ -31,9 +31,11 @@ from unittest.mock import MagicMock, patch import numpy as np +import onnx import pytest import torch from _test_utils.onnx.quantization.autotune.models import _create_simple_conv_onnx_model +from onnx import TensorProto, helper import modelopt.onnx.quantization.autotune.benchmark as bm from modelopt.onnx.quantization.autotune import TensorRTPyBenchmark, TrtExecBenchmark @@ -345,10 +347,50 @@ class TestTensorRTPyBenchmark: @pytest.fixture(autouse=True) def _require_tensorrt_and_torch_cuda(self): + """Skip benchmark integration tests when TensorRT or CUDA is unavailable.""" pytest.importorskip("tensorrt") if not bm.TORCH_CUDA_AVAILABLE: pytest.skip("CUDA not available") + @pytest.mark.parametrize("quant_type", [TensorProto.INT8, TensorProto.FLOAT8E4M3FN]) + def test_run_quantized_model(self, quant_type): + """Quantized candidates build and run with strongly typed networks.""" + if quant_type == TensorProto.FLOAT8E4M3FN and torch.cuda.get_device_capability() < (8, 9): + pytest.skip("FP8 requires SM89 or newer") + + nodes = [] + for name in ("input", "weight"): + nodes.extend( + [ + helper.make_node( + "QuantizeLinear", [name, "scale", "zero_point"], [f"{name}_q"] + ), + helper.make_node( + "DequantizeLinear", [f"{name}_q", "scale", "zero_point"], [f"{name}_dq"] + ), + ] + ) + nodes.append(helper.make_node("MatMul", ["input_dq", "weight_dq"], ["output"])) + graph = helper.make_graph( + nodes, + "quantized_matmul", + [helper.make_tensor_value_info("input", TensorProto.FLOAT, [16, 16])], + [helper.make_tensor_value_info("output", TensorProto.FLOAT, [16, 16])], + initializer=[ + helper.make_tensor("weight", TensorProto.FLOAT, [16, 16], [0.1] * 256), + helper.make_tensor("scale", TensorProto.FLOAT, [], [0.1]), + helper.make_tensor("zero_point", quant_type, [], [0]), + ], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 19)], ir_version=9) + onnx.checker.check_model(model) + + benchmark = TensorRTPyBenchmark(warmup_runs=1, timing_runs=2) + network = benchmark.builder.create_network(benchmark.network_flags) + assert network.get_flag(bm.trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED) + latency_ms = benchmark.run(model.SerializeToString()) + assert np.isfinite(latency_ms) and latency_ms > 0 + def test_run_with_bytes(self, simple_conv_model_bytes): """TensorRTPyBenchmark accepts model bytes and returns finite latency.""" benchmark = TensorRTPyBenchmark(warmup_runs=1, timing_runs=2) diff --git a/tests/gpu/onnx/test_trt_utils.py b/tests/gpu/onnx/test_trt_utils.py new file mode 100644 index 00000000000..a1dae0466ff --- /dev/null +++ b/tests/gpu/onnx/test_trt_utils.py @@ -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. + +import ctypes.util +import logging + +import onnx +import pytest + +from modelopt.onnx import trt_utils + +pytest.importorskip("tensorrt") + + +@pytest.mark.parametrize("log_level", [logging.WARNING, logging.DEBUG]) +def test_get_custom_layers_with_plugin_library(caplog, log_level): + """Loading a plugin library must preserve ONNX parsing at either log level.""" + plugin_library = ctypes.util.find_library("nvinfer_plugin") + assert plugin_library is not None + model = onnx.helper.make_model( + onnx.helper.make_graph( + [onnx.helper.make_node("Identity", ["x"], ["y"])], + "identity", + [onnx.helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT, [1])], + [onnx.helper.make_tensor_value_info("y", onnx.TensorProto.FLOAT, [1])], + ), + opset_imports=[onnx.helper.make_opsetid("", 13)], + ir_version=10, + ) + + with caplog.at_level(log_level, logger=trt_utils.logger.name): + custom_layers, tensor_info = trt_utils.get_custom_layers(model, [plugin_library]) + + assert custom_layers == [] + assert tensor_info == { + name: {"shape": [1], "dtype": trt_utils.trt.float32} for name in ["x", "y"] + } diff --git a/tests/gpu/torch/export/test_export.py b/tests/gpu/torch/export/test_export.py index 6f90c9104c7..c25c135ee11 100644 --- a/tests/gpu/torch/export/test_export.py +++ b/tests/gpu/torch/export/test_export.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import copy import json from fnmatch import fnmatch @@ -32,7 +33,8 @@ partial_nvfp4_config, partial_w4a8_config, ) -from _test_utils.torch.transformers_models import get_tiny_qwen3_moe +from _test_utils.torch.transformers_models import get_tiny_llama, get_tiny_qwen3_moe +from safetensors.torch import load_file import modelopt.torch.quantization as mtq from modelopt.torch.export.quant_format import ( @@ -107,6 +109,7 @@ def test_get_quantization_format(config, expected): { "quant_algo": "MIXED_PRECISION", "kv_cache_quant_algo": None, + "exclude_modules": ["layer8"], "quantized_layers": { "layer1": {"quant_algo": "NVFP4", "group_size": 16}, "layer3": { @@ -509,6 +512,123 @@ def test_get_quant_config(config, expected): assert quant_config["quantization"] == expected +@pytest.mark.parametrize("activation", ["disabled", "mxfp4", "fp8", "mixed"]) +def test_mxfp4_offline_export_activation_schema(tmp_path, activation): + model = get_tiny_llama().to("cuda") + model.config.architectures = ["LlamaForCausalLM"] + config = copy.deepcopy(mtq.MXFP4_MLP_WEIGHT_ONLY_CFG) + if activation in ("mxfp4", "fp8", "mixed"): + config["quant_cfg"].append( + { + "quantizer_name": ( + "*layers.0.mlp*input_quantizer" + if activation == "mixed" + else "*mlp*input_quantizer" + ), + "cfg": ( + {"num_bits": (4, 3)} + if activation == "fp8" + else { + "num_bits": (2, 1), + "block_sizes": {-1: 32, "type": "dynamic", "scale_bits": (8, 0)}, + } + ), + } + ) + config["algorithm"] = "max" + inputs = torch.randint(model.config.vocab_size, (1, 8), device="cuda") + mtq.quantize(model, config, lambda m: m(inputs)) + export_hf_checkpoint(model, export_dir=tmp_path) + with open(tmp_path / "config.json") as file: + exported = json.load(file)["quantization_config"] + tensors = load_file(next(tmp_path.glob("*.safetensors"))) + assert tensors["model.layers.0.mlp.gate_proj.weight"].dtype == torch.uint8 + assert tensors["model.layers.0.mlp.gate_proj.weight"].shape == (32, 16) + assert tensors["model.layers.0.mlp.gate_proj.weight_scale"].shape == (32, 1) + + expected_algo = { + "disabled": "W4A16_MXFP4", + "mxfp4": "MXFP4", + "fp8": "W4A8_MXFP4_FP8", + "mixed": "MIXED_PRECISION", + }[activation] + assert exported["quant_algo"] == expected_algo + if activation == "mixed": + assert {cfg["quant_algo"] for cfg in exported["quantized_layers"].values()} == { + "MXFP4", + "W4A16_MXFP4", + } + assert len(exported["config_groups"]) == 2 + for group in exported["config_groups"].values(): + assert group["weights"]["num_bits"] == 4 + assert group["weights"]["group_size"] == 32 + has_input = activation in ("mxfp4", "fp8") or ( + activation == "mixed" and "layers.0.mlp." in group["targets"][0] + ) + assert ("input_activations" in group) == has_input + if has_input: + assert group["input_activations"]["num_bits"] == (8 if activation == "fp8" else 4) + + +@pytest.mark.parametrize("block_size", [64, 128, "mixed"]) +def test_fp8_block_offline_export_preserves_block_structure(tmp_path, block_size): + model = get_tiny_llama(hidden_size=128, intermediate_size=128).to("cuda") + model.config.architectures = ["LlamaForCausalLM"] + layer_blocks = [64, 128] if block_size == "mixed" else [block_size, block_size] + config = { + "algorithm": "max", + "quant_cfg": [{"quantizer_name": "*", "enable": False}], + } + for layer, size in enumerate(layer_blocks): + config["quant_cfg"].append( + { + "quantizer_name": f"*layers.{layer}.mlp*weight_quantizer", + "enable": True, + "cfg": {"num_bits": (4, 3), "block_sizes": {-2: size, -1: size}}, + } + ) + inputs = torch.randint(model.config.vocab_size, (1, 8), device="cuda") + mtq.quantize(model, config, lambda m: m(inputs)) + export_hf_checkpoint(model, export_dir=tmp_path) + with open(tmp_path / "config.json") as file: + exported = json.load(file)["quantization_config"] + tensors = load_file(next(tmp_path.glob("*.safetensors"))) + + assert exported["quant_algo"] == ("MIXED_PRECISION" if block_size == "mixed" else "FP8_PB_WO") + groups = exported["config_groups"].values() + assert {tuple(group["weights"]["block_structure"]) for group in groups} == { + (size, size) for size in layer_blocks + } + if block_size == "mixed": + for group in groups: + for target in group["targets"]: + layer = 0 if "layers.0." in target else 1 + assert group["weights"]["block_structure"] == [layer_blocks[layer]] * 2 + for layer, size in enumerate(layer_blocks): + name = f"model.layers.{layer}.mlp.gate_proj" + assert tensors[f"{name}.weight"].dtype == torch.float8_e4m3fn + assert tensors[f"{name}.weight_scale"].numel() == (128 // size) ** 2 + + +@pytest.mark.parametrize("block_sizes", [{-2: 64, -1: 128}, {-2: 128, -1: 64}, {-1: 64}]) +def test_offline_fp8_block_export_rejects_unsupported_layout(block_sizes): + model = torch.nn.Linear(128, 128, bias=False, device="cuda") + config = { + "algorithm": "max", + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + { + "quantizer_name": "*weight_quantizer", + "enable": True, + "cfg": {"num_bits": (4, 3), "block_sizes": block_sizes}, + }, + ], + } + mtq.quantize(model, config, lambda m: m(torch.randn(2, 128, device="cuda"))) + with pytest.raises(NotImplementedError, match="square blocks over the last two dimensions"): + get_quant_config(model) + + def test_qwen3_moe_nvfp4_experts_only_export_exclude_modules(tmp_path): """Test that NVFP4_EXPERTS_ONLY_CFG correctly excludes non-expert modules in HF export. diff --git a/tests/gpu/torch/export/test_qwen_vision_recipe_export.py b/tests/gpu/torch/export/test_qwen_vision_recipe_export.py index 9c1a804ea1c..389375428b1 100644 --- a/tests/gpu/torch/export/test_qwen_vision_recipe_export.py +++ b/tests/gpu/torch/export/test_qwen_vision_recipe_export.py @@ -92,7 +92,7 @@ def calibration_forward(_model): scale_names = {name for name in tensor_names if name.endswith(("input_scale", "weight_scale"))} assert scale_names assert any(name.startswith("model.visual.") for name in scale_names) - assert any(".attn.qkv.weight_scale" in name for name in scale_names) + assert not any(".blocks." in name for name in scale_names) assert any(".merger.linear_fc1.weight_scale" in name for name in scale_names) if model_type == "qwen3_vl": assert any( diff --git a/tests/unit/torch/models/test_modeling_ptq_registration.py b/tests/gpu/torch/models/test_modeling_ptq_registration.py similarity index 66% rename from tests/unit/torch/models/test_modeling_ptq_registration.py rename to tests/gpu/torch/models/test_modeling_ptq_registration.py index b97a2294a67..6782b64c344 100644 --- a/tests/unit/torch/models/test_modeling_ptq_registration.py +++ b/tests/gpu/torch/models/test_modeling_ptq_registration.py @@ -21,8 +21,10 @@ so an in-process test would only ever see the import order of whichever test ran first. """ +import os import subprocess import sys +from concurrent.futures import ThreadPoolExecutor import pytest @@ -30,8 +32,11 @@ FIRST_IMPORTS = [ "modelopt.torch.models.falcon.modeling_ptq", + "modelopt.torch.models.gpt_oss.modeling_ptq", "modelopt.torch.models.llama4.modeling_ptq", "modelopt.torch.models.nemotron_h.modeling_ptq", + "modelopt.torch.models.step3p5.modeling_ptq", + "modelopt.torch.models.step3p7.modeling_ptq", "modelopt.torch.quantization", "modelopt.torch.export", ] @@ -45,6 +50,7 @@ from modelopt.torch.models.falcon.modeling_ptq import register_falcon_linears_on_the_fly from modelopt.torch.models.nemotron_h.modeling_ptq import is_nemotron_h_model +from modelopt.torch.models.step3p5.modeling_ptq import register_moe_linear_on_the_fly from modelopt.torch.quantization.nn import QuantModuleRegistry from modelopt.torch.quantization.plugins.custom import CUSTOM_MODEL_PLUGINS from modelopt.torch.quantization.plugins.huggingface import is_homogeneous_hf_model @@ -63,7 +69,15 @@ else: assert Llama4TextExperts in QuantModuleRegistry, "Llama4TextExperts not registered" +try: + from transformers.models.gpt_oss.modeling_gpt_oss import GptOssExperts +except ImportError: + pass +else: + assert GptOssExperts in QuantModuleRegistry, "GptOssExperts not registered" + assert register_falcon_linears_on_the_fly in CUSTOM_MODEL_PLUGINS, "Falcon callback missing" +assert register_moe_linear_on_the_fly in CUSTOM_MODEL_PLUGINS, "Step callback missing" # The first matching decoder discoverer wins, so Nemotron-H's must precede the generic one. predicates = [is_supported for is_supported, _ in LayerActivationCollector._decoder_layer_support] @@ -73,23 +87,29 @@ """ +def _run_check(first_import): + # The child only asserts; tracing its heavy imports for coverage (inherited through + # COVERAGE_PROCESS_START) would just slow it down. + env = {k: v for k, v in os.environ.items() if k != "COVERAGE_PROCESS_START"} + result = subprocess.run( + [sys.executable, "-c", CHECK, first_import], + capture_output=True, + text=True, + env=env, + check=False, + ) + return first_import, result + + +# Several fresh interpreters each import torch and transformers, which outlasts the default +# 60 s per-test cap on a small CI runner. +@pytest.mark.timeout(300) def test_registration_is_independent_of_import_order(): - # Each first import needs its own interpreter, but starting them one after another would - # add over a minute to the unit-test job, so launch them all and then collect the results. - procs = { - first_import: subprocess.Popen( - [sys.executable, "-c", CHECK, first_import], - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - text=True, - ) - for first_import in FIRST_IMPORTS - } - failures = {} - for first_import, proc in procs.items(): - _, stderr = proc.communicate(timeout=600) - if proc.returncode != 0: - failures[first_import] = stderr + # Each first import needs its own interpreter; run them in parallel, at most one per CPU, + # rather than one after another. + with ThreadPoolExecutor(max_workers=os.cpu_count() or 1) as pool: + results = list(pool.map(_run_check, FIRST_IMPORTS)) + failures = {name: result.stderr for name, result in results if result.returncode != 0} assert not failures, "\n\n".join( f"importing {name} first:\n{err}" for name, err in failures.items() ) diff --git a/tests/gpu/torch/quantization/test_fp8_experts.py b/tests/gpu/torch/quantization/test_fp8_experts.py new file mode 100644 index 00000000000..f33e80e7e12 --- /dev/null +++ b/tests/gpu/torch/quantization/test_fp8_experts.py @@ -0,0 +1,81 @@ +# 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. + +"""Native FP8 expert PTQ and packed NVFP4 export.""" + +import pytest +import torch +from transformers import PretrainedConfig +from transformers.integrations.finegrained_fp8 import FP8Experts + +import modelopt.torch.quantization as mtq +from modelopt.torch.export.moe_utils import _export_fused_experts + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("has_gate", [True, False]) +def test_fp8_experts_nvfp4_export(has_gate, dtype): + torch.manual_seed(42) + config = PretrainedConfig( + hidden_size=256, intermediate_size=256, num_experts=2, hidden_act="silu" + ) + experts = FP8Experts(config, block_size=(128, 128), has_gate=has_gate).cuda() + first = "gate_up_proj" if has_gate else "up_proj" + expected = {} + for name in (first, "down_proj"): + weight = getattr(experts, name) + scales = getattr(experts, f"{name}_scale_inv") + with torch.no_grad(): + weight.copy_(torch.randn(weight.shape, device="cuda").to(weight.dtype)) + scales.copy_( + torch.linspace(0.01, 0.08, scales.numel(), device="cuda").reshape(scales.shape) + ) + expected[name] = weight.float() * scales.repeat_interleave(128, 1).repeat_interleave(128, 2) + inputs = torch.randn(2, 256, device="cuda", dtype=dtype) + indices = torch.tensor([[0], [1]], device="cuda") + routing = torch.ones(2, 1, device="cuda") + mtq.quantize(experts, mtq.NVFP4_DEFAULT_CFG, lambda m: m(inputs, indices, routing)) + assert torch.isfinite(experts(inputs, indices, routing)).all() + _export_fused_experts(experts, torch.float32) + levels = torch.tensor( + [0, 0.5, 1, 1.5, 2, 3, 4, 6, 0, -0.5, -1, -1.5, -2, -3, -4, -6], device="cuda" + ) + for idx in range(2): + exported = getattr(experts, str(idx)) + projections = [("down_proj", expected["down_proj"][idx])] + if has_gate: + projections += [ + ("gate_proj", expected[first][idx, :256]), + ("up_proj", expected[first][idx, 256:]), + ] + torch.testing.assert_close( + exported.gate_proj.weight_scale_2, exported.up_proj.weight_scale_2 + ) + else: + projections += [("up_proj", expected[first][idx])] + for name, reference in projections: + proj = getattr(exported, name) + packed = proj.weight + codes = torch.stack([packed & 15, packed >> 4], -1).reshape(reference.shape).long() + decoded = ( + levels[codes] + * proj.weight_scale.float().repeat_interleave(16, -1) + * proj.weight_scale_2 + ) + relative_l2 = torch.linalg.vector_norm(decoded - reference) / torch.linalg.vector_norm( + reference + ) + assert relative_l2 < 0.15 + assert not any("scale_inv" in key for key in experts.state_dict()) diff --git a/tests/gpu/torch/quantization/test_lsq_cuda.py b/tests/gpu/torch/quantization/test_lsq_cuda.py index ffb6fce95d4..465be40467e 100644 --- a/tests/gpu/torch/quantization/test_lsq_cuda.py +++ b/tests/gpu/torch/quantization/test_lsq_cuda.py @@ -152,9 +152,19 @@ def test_lsq_quantize_e2e(config): assert out.shape == (2, 64) -def test_lsq_fp4_fake_quantize_differentiable(): - """Test that _fake_quantize in FP4 LSQ mode is differentiable.""" +@pytest.mark.parametrize("tied_amax", [True, False], ids=["tied", "default_dual"]) +def test_lsq_fp4_fake_quantize_differentiable(tied_amax): + """Representable small scales retain static outputs and analytical LSQ gradients.""" device = torch.device("cuda") + scales = torch.tensor([2.0**-5, 1.0, 448.0], device=device) + codes = torch.tensor( + [0.625, -0.625, 1, -1, 1.5, -1.5, 2, -2, 3, -3, 4, -4, 6, -6, 0, 0], device=device + ) + x = scales[:, None] * codes + rounded = codes.clone() + rounded[:2] = torch.tensor([0.5, -0.5], device=device) + expected = scales[:, None] * rounded + tq = TensorQuantizer() tq._num_bits = (2, 1) tq._unsigned = False @@ -162,24 +172,49 @@ def test_lsq_fp4_fake_quantize_differentiable(): tq._disabled = False tq._block_sizes = {-1: 16, "type": "static", "scale_bits": (4, 3)} tq._pass_through_bwd = True - tq.register_buffer("_amax", torch.ones(4, device=device)) + tq.register_buffer("_amax", 6.0 * scales) tq.to(device) sbsq = StaticBlockScaleQuantizer.from_tensor_quantizer( - tq, global_amax=torch.tensor(1.0, device=device) + tq, global_amax=torch.tensor(6.0 * 448.0, device=device) ) - # global_amax=1.0 with NVFP4 _quant_max_bound=6.0 yields per_tensor_scale = 1/6. - sbsq.amax = torch.ones(4, device=device) * 3.0 - sbsq.enable_lsq( - quantize_scales=True, - learnable_amax=["post"], + static_output = sbsq(x) + torch.testing.assert_close(static_output, expected, rtol=0, atol=0) + if tied_amax: + sbsq.enable_lsq(learnable_amax=["pre", "post"], tied_amax=True) + else: + sbsq.enable_lsq() + + out = sbsq(x) + torch.testing.assert_close(out, static_output, rtol=0, atol=0) + (out[0, 0] + out[1, 0]).backward() + # At x/scale=0.625, Q(x/scale)=0.5 and d(scale)/d(amax)=1/6. + # Tied: (0.5 - 0.625)/6; default dual: 0.5/6 with fixed pre scale. + gradient = -1.0 / 48.0 if tied_amax else 1.0 / 12.0 + torch.testing.assert_close( + sbsq.amax_post.grad, + torch.tensor([gradient, gradient, 0.0], device=device), + rtol=1e-6, + atol=1e-7, ) - x = torch.randn(4, 16, device=device) - out = sbsq._fake_quantize(x) - assert out.shape == x.shape - out.sum().backward() - assert sbsq._amax_post.grad is not None + +@pytest.mark.parametrize("tied_amax", [False, True]) +def test_lsq_fp4_small_global_amax(tied_amax): + """A block representable in NVFP4 survives LSQ fake quant even when global_amax < 6 * 2**-24.""" + device = torch.device("cuda") + x = torch.full((1, 16), 2.0**-22, device=device) + tq = TensorQuantizer() + tq._num_bits = (2, 1) + tq._unsigned = False + tq._narrow_range = True + tq._disabled = False + tq._block_sizes = {-1: 16, "type": "static", "scale_bits": (4, 3)} + tq.register_buffer("_amax", x.abs().amax(-1)) + sbsq = StaticBlockScaleQuantizer.from_tensor_quantizer(tq, global_amax=x.abs().amax()) + sbsq.enable_lsq(quantize_scales=True, tied_amax=tied_amax) + + assert torch.equal(sbsq._fake_quantize(x), x) def test_lsq_fp4_cast_ste(): diff --git a/tests/gpu/torch/quantization/test_qtensor_cuda.py b/tests/gpu/torch/quantization/test_qtensor_cuda.py index cfdf38864fb..c5d89b17d44 100644 --- a/tests/gpu/torch/quantization/test_qtensor_cuda.py +++ b/tests/gpu/torch/quantization/test_qtensor_cuda.py @@ -175,7 +175,7 @@ def test_amax_from_tensor_quantizer( 0, torch.tensor([[0, 1, 2, 3, 4, 5, 6, 7]], dtype=torch.bfloat16), torch.tensor( - [[0.0000, 0.9922, 1.9844, 2.9844, 3.9688, 4.9688, 5.9688, 7.0000]], + [[0.0000, 0.9922, 1.9844, 2.9688, 4.0312, 5.0312, 6.0000, 7.0000]], dtype=torch.bfloat16, ), ), @@ -186,7 +186,7 @@ def test_amax_from_tensor_quantizer( None, torch.tensor([[0, 1, 2, 3], [4, 5, 6, 7]], dtype=torch.bfloat16), torch.tensor( - [[0.0000, 1.0234, 1.9844, 2.9844], [4.0000, 5.0000, 5.9688, 7.0000]], + [[0.0000, 0.9844, 1.9844, 2.9688], [4.0312, 5.0000, 6.0000, 7.0000]], dtype=torch.bfloat16, ), ), diff --git a/tests/gpu/torch/quantization/test_tensor_quant_cuda.py b/tests/gpu/torch/quantization/test_tensor_quant_cuda.py index cc3262e1917..c1cbd6a09fa 100644 --- a/tests/gpu/torch/quantization/test_tensor_quant_cuda.py +++ b/tests/gpu/torch/quantization/test_tensor_quant_cuda.py @@ -178,6 +178,23 @@ def test_zero_amax_per_channel_is_finite(self): xq = tensor_quant.scaled_e4m3(x, amax, None, 4, 3) assert torch.isfinite(xq).all() + @pytest.mark.parametrize("device", ["cuda", "cpu"]) + @pytest.mark.parametrize( + "amax_dims", [None, (1, 2), (2,)], ids=["per_tensor", "per_channel", "multi_axis"] + ) + def test_small_amax_scales_output(self, device, amax_dims): + """Scaling the input and amax by a power of two scales the output by the same factor. + + A positive amax below 2**-24 is not zero and must not flush the tensor to zero. + """ + torch.manual_seed(0) + x = torch.randn(4, 3, 16, device=device) * 10 + amax = x.abs().amax() if amax_dims is None else x.abs().amax(amax_dims, keepdim=True) + reference = tensor_quant.scaled_e4m3(x, amax, None, 4, 3) + for factor in (2.0**-30, 2.0**-60): + scaled = tensor_quant.scaled_e4m3(x * factor, amax * factor, None, 4, 3) + assert torch.equal(scaled, reference * factor) + class Testfp4: @pytest.mark.skipif(not NATIVE_E4M3_AVAILABLE, reason="Native E4M3 requires compute >= 8.9") @@ -373,6 +390,31 @@ def test_fp4_kernel_scales_with_small_inputs(self, dtype): scaled = triton_kernel.fp4_fake_quant_block(x * factor, x.float().abs().amax() * factor) assert torch.equal(scaled, reference * factor) + @pytest.mark.skipif(not triton_kernel.IS_AVAILABLE, reason="triton kernel is not available") + @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) + def test_static_fp4_kernel_scales_with_small_inputs(self, dtype): + """Static counterpart of ``test_fp4_kernel_scales_with_small_inputs``. + + A global amax below 6 * 2**-24 must not zero the FP8-quantized block scales, while an + all-zero tensor still quantizes to zeros. + """ + torch.manual_seed(0) + x = (torch.randn(8, 64, device="cuda") * 10).to(dtype) + amax = x.float().abs().view(-1, 16).amax(-1) + reference = triton_kernel.static_blockwise_fp4_fake_quant(x, amax, amax.max()) + assert reference.abs().amax() > 0 + for factor in (2.0**-10, 2.0**-20, 2.0**-30): + scaled = triton_kernel.static_blockwise_fp4_fake_quant( + x * factor, amax * factor, amax.max() * factor + ) + assert torch.equal(scaled, reference * factor) + + zeros = torch.zeros_like(x) + output = triton_kernel.static_blockwise_fp4_fake_quant( + zeros, torch.zeros_like(amax), torch.zeros_like(amax.max()) + ) + assert torch.equal(output, zeros) + @pytest.mark.skipif( not hasattr(triton_kernel, "fp4_fake_quant_block"), reason="fp4_fake_quant_block requires compute >= 8.9", diff --git a/tests/gpu/torch/speculative/plugins/test_hf_dflash.py b/tests/gpu/torch/speculative/plugins/test_hf_dflash.py index 0cb740c0e9f..69f0a117b63 100644 --- a/tests/gpu/torch/speculative/plugins/test_hf_dflash.py +++ b/tests/gpu/torch/speculative/plugins/test_hf_dflash.py @@ -287,20 +287,20 @@ def sharded_forward(*args, input_ids=None, **kwargs): model._base_model.forward = sharded_forward def test_generate_gathers_sharded_hidden_states(self): - model = get_tiny_llama(num_hidden_layers=4) + # 8 layers so the default target_layer_ids ([1, 5]) straddle the split; 4 gives [1, 1]. + model = get_tiny_llama(num_hidden_layers=8) mtsp.convert(model, [("dflash", get_dflash_config(block_size=4))]) model = model.to("cuda:1").eval() - # The fix is only exercised if the split actually separates two target layers. - assert ( - min(model.target_layer_ids) - < len(model._base_model.layers) // 2 - <= max(model.target_layer_ids) - ), model.target_layer_ids - self._shard_hidden_states(model, spare_device=torch.device("cuda:0")) input_ids = torch.tensor([[1, 2, 3, 4]], device="cuda:0") + # The fix is only exercised if the selected hidden states really sit on different devices. + with torch.no_grad(): + hidden = model._base_model(input_ids=input_ids, output_hidden_states=True).hidden_states + devices = {hidden[lid + 1].device for lid in model.target_layer_ids} + assert len(devices) > 1, (model.target_layer_ids, devices) + with torch.no_grad(): base_token, draft_tokens = model.pseudo_speculative_generate(input_ids, steps=2) diff --git a/tests/gpu/torch/utils/test_distributed.py b/tests/gpu/torch/utils/test_distributed.py index 4091c7465dc..d1581e1ca03 100644 --- a/tests/gpu/torch/utils/test_distributed.py +++ b/tests/gpu/torch/utils/test_distributed.py @@ -15,6 +15,8 @@ """GPU/distributed tests for ``modelopt.torch.utils.distributed``.""" +import os +import pickle from functools import partial import pytest @@ -25,7 +27,14 @@ 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 +from modelopt.torch.nas.traced_hp import TracedHp +from modelopt.torch.utils.distributed import ( + allgather, + broadcast, + fsdp2_wrap, + gather_bytes, + gather_state_dicts, +) VOCAB_SIZE = 32 N_EXPERTS = 4 @@ -129,3 +138,81 @@ 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) ) + + +def _test_broadcast_allgather_objects(rank, size): + hp = TracedHp([1, 2, 4]) if rank == 0 else None + out = broadcast([True, {"latency": 1.5}, hp], src=0) + assert out[0] is True + assert out[1] == {"latency": 1.5} + assert out[2].choices == [1, 2, 4] + assert allgather({"rank": rank}) == [{"rank": r} for r in range(size)] + + +def test_broadcast_allgather_objects(dist_workers): + dist_workers.run(_test_broadcast_allgather_objects) + + +def _test_gather_bytes(rank, size): + data = bytes([rank + 1]) * (5 + 3 * rank) if rank else b"" + out = gather_bytes(data, dst=1, group=None, chunk_bytes=4) + if rank == 1: + assert out == [ + bytearray(b"" if r == 0 else bytes([r + 1]) * (5 + 3 * r)) for r in range(size) + ] + else: + assert out is None + + +def test_gather_bytes(need_2_gpus, dist_workers): + dist_workers.run(_test_gather_bytes) + + +class _Gadget: + def __init__(self, marker): + self.marker = marker + + def __reduce__(self): + return (os.mkdir, (self.marker,)) + + +def _test_gather_state_dicts(rank, size): + out = list(gather_state_dicts({f"w{rank}": torch.full((2,), float(rank))}, dst=0)) + if rank == 0: + assert [list(d) for d in out] == [[f"w{r}"] for r in range(size)] + assert torch.equal(out[1]["w1"], torch.ones(2)) + else: + assert out == [] + + +def test_gather_state_dicts(need_2_gpus, dist_workers): + dist_workers.run(_test_gather_state_dicts) + + +def _test_gather_state_dicts_rejects_gadget(rank, size, marker, force_no_weights_only): + if force_no_weights_only: + os.environ["TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD"] = "1" + # A non-root peer sends a pickle gadget; the root must refuse it without running it. + state = {"w": _Gadget(marker)} if rank == 1 else {"w": torch.zeros(1)} + gathered = gather_state_dicts(state, dst=0) + if rank == 0: + with pytest.raises(pickle.UnpicklingError): + list(gathered) + assert not os.path.exists(marker) + else: + assert list(gathered) == [] + + +@pytest.mark.parametrize("force_no_weights_only", [False, True]) +def test_gather_state_dicts_rejects_gadget( + need_2_gpus, dist_workers, tmp_path, force_no_weights_only +): + marker = str(tmp_path / "gadget_ran") + dist_workers.run( + partial( + _test_gather_state_dicts_rejects_gadget, + marker=marker, + force_no_weights_only=force_no_weights_only, + ) + ) + assert not os.path.exists(marker) diff --git a/tests/gpu/torch/utils/test_distributed_worker_pool.py b/tests/gpu/torch/utils/test_distributed_worker_pool.py new file mode 100644 index 00000000000..7ae3a966278 --- /dev/null +++ b/tests/gpu/torch/utils/test_distributed_worker_pool.py @@ -0,0 +1,81 @@ +# 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. + +"""A hung job must not leave a pool that stalls every later job.""" + +import queue +import time +from types import SimpleNamespace + +import pytest +import torch.distributed as dist +from _test_utils.torch.distributed.utils import DistributedWorkerPool, get_live_pool + + +def _hang(rank, size): + time.sleep(600) + + +def _rank_0_fails_and_rank_1_waits_for_it(rank, size): + if rank == 0: + raise ValueError("rank 0 failed") + dist.barrier() # never completes: its partner left + + +def test_a_hung_job_kills_the_pool_instead_of_stalling_later_runs(capfd): + pool = DistributedWorkerPool(world_size=1, backend="gloo", teardown_fn=None) + try: + pool.job_timeout = 1 + with pytest.raises(queue.Empty): + pool.run(_hang) + + assert pool.broken + assert not any(p.is_alive() for p in pool._processes) + # Fails at once rather than waiting out another timeout. + with pytest.raises(RuntimeError, match="killed after an earlier job hung"): + pool.run(_hang) + # The stuck worker's stack reached stderr, which is what makes the next CI hang diagnosable. + assert "_hang" in capfd.readouterr().err + finally: + pool.shutdown() + + +def test_a_failed_rank_does_not_make_its_peers_wait_out_the_job_timeout(): + pool = DistributedWorkerPool(world_size=2, backend="gloo", teardown_fn=None) + try: + pool.failure_grace = 2 + start = time.monotonic() + with pytest.raises(RuntimeError, match="rank 0 failed"): + pool.run(_rank_0_fails_and_rank_1_waits_for_it) + + assert time.monotonic() - start < 60 # not the 600 s job timeout + assert pool.broken + finally: + pool.shutdown() + + +def test_a_killed_pool_is_replaced_and_a_live_one_is_reused(): + pools = {} + made = [] + + def make(): + made.append(SimpleNamespace(broken=False)) + return made[-1] + + first = get_live_pool(pools, 2, make) + assert get_live_pool(pools, 2, make) is first + first.broken = True + second = get_live_pool(pools, 2, make) + assert second is not first and len(made) == 2 diff --git a/tests/gpu/torch/utils/test_hf_checkpoint_utils_dist.py b/tests/gpu/torch/utils/test_hf_checkpoint_utils_dist.py new file mode 100644 index 00000000000..72322ad50b4 --- /dev/null +++ b/tests/gpu/torch/utils/test_hf_checkpoint_utils_dist.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Multi-rank tests for ``ensure_local_checkpoint`` (each case spawns a 2-process gloo job).""" + +import os +from functools import partial + +import pytest +import torch.distributed as dist +from _test_utils.torch.distributed.utils import spawn_multiprocess_job + +from modelopt.torch.utils.plugins import hf_checkpoint_utils + +pytest.importorskip("huggingface_hub") + + +class _FakeHfApi: + """``HfApi`` whose ``repo_info`` succeeds, so the Hub counts as reachable.""" + + def repo_info(self, repo_id): + pass + + +def _ensure_local_checkpoint_job(tmp_path, fail, group_ranks, rank, size): + utils = hf_checkpoint_utils + snapshot = tmp_path / "snapshot" + + def fake_snapshot_download(repo_id, **kwargs): + with (tmp_path / f"downloaded-by-rank{rank}").open("a") as marker: + marker.write(f"{repo_id}\n") # one line per call, so a repeat download shows + if fail: + raise OSError("hub unreachable") + snapshot.mkdir(exist_ok=True) + return str(snapshot) + + utils.snapshot_download = fake_snapshot_download + utils.HfApi = _FakeHfApi + os.environ.pop("HF_HUB_OFFLINE", None) # inherited from tests/unit/conftest.py + # new_group is collective over every rank, members or not. + group = None if group_ranks is None else dist.new_group(group_ranks, backend="gloo") + if group_ranks is not None and rank not in group_ranks: + return + if fail: + with pytest.raises(RuntimeError, match="hub unreachable"): + utils.ensure_local_checkpoint("org/model", group=group) + else: + result = utils.ensure_local_checkpoint("org/model", group=group) + assert result == ("org/model", str(snapshot)) + + +def _ensure_local_checkpoint_local_path_job(tmp_path, hidden_from_rank1, rank, size): + utils = hf_checkpoint_utils + + def fail(*args, **kwargs): + raise AssertionError("a local path must not be downloaded") + + utils.snapshot_download = fail + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir(exist_ok=True) + if hidden_from_rank1 and rank == 1: + # As if the directory sat on rank 0's node-local disk only. + isdir = os.path.isdir + utils.os.path.isdir = lambda path: path != str(checkpoint) and isdir(path) + with pytest.raises(RuntimeError, match="this rank cannot see"): + utils.ensure_local_checkpoint(str(checkpoint)) + else: + assert utils.ensure_local_checkpoint(str(checkpoint)) == (None, str(checkpoint)) + + +@pytest.mark.parametrize("hidden_from_rank1", [False, True]) +def test_ensure_local_checkpoint_lets_rank_0_resolve_a_local_path(tmp_path, hidden_from_rank1): + """Ranks that each checked a local path could disagree about it and strand one another in the + collective; rank 0 decides, and a rank that cannot see its answer raises instead.""" + job = partial(_ensure_local_checkpoint_local_path_job, tmp_path, hidden_from_rank1) + spawn_multiprocess_job(2, job, backend="gloo") + + +@pytest.mark.parametrize( + ("fail", "group_ranks", "downloader"), + [ + (False, None, "rank0"), + (True, None, "rank0"), + (False, [0, 1], "rank0"), + # Rank 0 of the group, not of the world: global rank 1. + (False, [1], "rank1"), + ], +) +def test_ensure_local_checkpoint_downloads_once_per_group(tmp_path, fail, group_ranks, downloader): + """Rank 0 of the group downloads once and every member gets its result -- including its + failure, so the waiting ranks raise instead of hanging.""" + job = partial(_ensure_local_checkpoint_job, tmp_path, fail, group_ranks) + spawn_multiprocess_job(2, job, backend="gloo") + assert [p.name for p in tmp_path.glob("downloaded-by-*")] == [f"downloaded-by-{downloader}"] + assert (tmp_path / f"downloaded-by-{downloader}").read_text().splitlines() == ["org/model"] diff --git a/tests/gpu_megatron/conftest.py b/tests/gpu_megatron/conftest.py index df061553001..5760b9fe22d 100644 --- a/tests/gpu_megatron/conftest.py +++ b/tests/gpu_megatron/conftest.py @@ -17,12 +17,13 @@ import pytest import torch from _test_utils.fs_utils import assert_unmodified_tree -from _test_utils.torch.distributed.utils import DistributedWorkerPool +from _test_utils.torch.distributed.utils import DistributedWorkerPool, get_live_pool from _test_utils.torch.transformers_models import get_tiny_tokenizer from megatron.core.parallel_state import destroy_model_parallel import modelopt.torch.quantization.extensions as ext import modelopt.torch.utils.distributed as dist +from modelopt.torch.utils.plugins.megatron_preprocess_data import _prestart_forkserver @pytest.fixture(scope="session") @@ -54,6 +55,17 @@ def _prebuild_quant_cuda_extensions(): ext.precompile() +@pytest.fixture(scope="session", autouse=True) +def _prestart_preprocess_forkserver(): + """Start the forkserver that ``megatron_preprocess_data`` workers fork from, at session start. + + Its first start cold-imports torch/megatron, which would otherwise land inside the first + multi-worker preprocess test and its per-test timeout. Starting it here, without waiting, + overlaps that import with the earlier tests. The server's environment is fixed at start. + """ + _prestart_forkserver() + + def megatron_worker_teardown(rank, world_size): """Clean up model-parallel state between tests in persistent workers.""" if dist.is_initialized(): @@ -96,31 +108,31 @@ def _pool_cache(): def _get_pool(cache, world_size): - if world_size not in cache: - cache[world_size] = _make_pool(world_size) - return cache[world_size] + # A hung test kills its pool; a fresh one keeps a single hang from timing out every later + # test in the module (that cascade cost the nightly 20+ min). + return get_live_pool(cache, world_size, lambda: _make_pool(world_size)) -@pytest.fixture(scope="module") +@pytest.fixture def dist_workers(_pool_cache): - """Module-scoped pool with world_size=torch.cuda.device_count().""" + """Pool (cached per module) with world_size=torch.cuda.device_count().""" return _get_pool(_pool_cache, torch.cuda.device_count()) -@pytest.fixture(scope="module") +@pytest.fixture def dist_workers_size_1(_pool_cache): - """Module-scoped pool with world_size=1 for tests that require a single process.""" + """Pool (cached per module) with world_size=1 for tests that require a single process.""" return _get_pool(_pool_cache, 1) -@pytest.fixture(scope="module") +@pytest.fixture def dist_workers_size_2(_pool_cache): if torch.cuda.device_count() < 2: pytest.skip("Need at least 2 GPUs") return _get_pool(_pool_cache, 2) -@pytest.fixture(scope="module") +@pytest.fixture def dist_workers_size_4(_pool_cache): if torch.cuda.device_count() < 4: pytest.skip("Need at least 4 GPUs") diff --git a/tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py b/tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py index 9985f9778b0..2bbcc108d69 100644 --- a/tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py +++ b/tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py @@ -14,6 +14,7 @@ # limitations under the License. import dataclasses +import warnings from functools import partial from types import SimpleNamespace @@ -233,8 +234,11 @@ def _test_topk_logits_kl_loss(kd_kwargs, rank, size): loss["kd_loss"].backward() -def _test_skip_lm_loss_with_mtp(quantized, skip_lm_loss, rank, size): - """Test that skip_lm_loss only zeroes the main LM head, unless MTP is left out of quantization.""" +def _test_skip_lm_loss_with_mtp(quantized, skip_lm_loss, freeze_mtp, rank, size): + """Test that skip_lm_loss only zeroes the main LM head, unless MTP is left out of quantization. + + ``freeze_mtp`` zeroes the MTP loss and freezes the MTP head regardless. + """ set_seed(SEED) num_layers = 2 @@ -282,7 +286,9 @@ def _test_skip_lm_loss_with_mtp(quantized, skip_lm_loss, rank, size): distill_cfg = setup_distillation_config( # kd_loss_alpha == 1.0 skips the LM loss; any alpha < 1 keeps it. - config_or_path=DistillationConfig(kd_loss_alpha=1.0 if skip_lm_loss else 0.5), + config_or_path=DistillationConfig( + kd_loss_alpha=1.0 if skip_lm_loss else 0.5, freeze_mtp=freeze_mtp + ), student_cfg=student_model.config, teacher_cfg=teacher_model.config, ) @@ -295,7 +301,7 @@ def _test_skip_lm_loss_with_mtp(quantized, skip_lm_loss, rank, size): adjust_distillation_model_for_mcore(distillation_model, distill_cfg) # An untrained MTP head is frozen so DDP doesn't wait for its grads. - mtp_skipped = quantized and skip_lm_loss + mtp_skipped = freeze_mtp or (quantized and skip_lm_loss) with distillation_model.hide_teacher_model(): mtp_params = [p for n, p in distillation_model.named_parameters() if "mtp" in n.split(".")] assert mtp_params @@ -337,6 +343,93 @@ def _recording_loss(labels, logits): assert recorded_losses[-1].any() != skip_lm_loss, "Main LM head loss is zero iff skip_lm_loss" +def _test_freeze_mtp_with_quantized_mtp(freeze_mtp, kd_loss_alpha, rank, size): + """With the MTP head quantized too, only ``freeze_mtp`` keeps it out of a training step.""" + set_seed(SEED) + + vocab_size = 32 + batch_size = 2 + max_sequence_length = 8 + model_kwargs = { + "tensor_model_parallel_size": size, + "pipeline_model_parallel_size": 1, + "num_layers": 2, + "hidden_size": 8, + "num_attention_heads": 4, + "num_query_groups": 2, + "ffn_hidden_size": 8, + "max_sequence_length": max_sequence_length, + "vocab_size": vocab_size, + "activation_func": "squared_relu", + "mtp_num_layers": 1, + } + teacher_model = get_mcore_gpt_model(initialize_megatron=True, **model_kwargs).cuda() + student_model = get_mcore_gpt_model(initialize_megatron=False, **model_kwargs).cuda() + # A recipe that quantizes the MTP head as well as the body. + student_model.decoder.weight_quantizer = TensorQuantizer() + student_model.mtp.weight_quantizer = TensorQuantizer() + assert not _mtp_excluded_from_quantization(student_model) + + distill_cfg = setup_distillation_config( + config_or_path=DistillationConfig(kd_loss_alpha=kd_loss_alpha, freeze_mtp=freeze_mtp), + student_cfg=student_model.config, + teacher_cfg=teacher_model.config, + ) + kd_config = { + "teacher_model": teacher_model, + "criterion": distill_cfg.criterion, + "loss_balancer": distill_cfg.loss_balancer, + } + distillation_model = mtd.convert(student_model, mode=[("kd_loss", kd_config)]) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + adjust_distillation_model_for_mcore(distillation_model, distill_cfg) + if rank == 0: + frozen_msg = [w for w in caught if "MTP head is frozen" in str(w.message)] + assert bool(frozen_msg) == freeze_mtp + + with distillation_model.hide_teacher_model(): + named_params = dict(distillation_model.named_parameters()) + mtp_params = {n: p for n, p in named_params.items() if "mtp" in n.split(".")} + body_params = {n: p for n, p in named_params.items() if n.startswith("decoder.")} + assert mtp_params and body_params + mtp_before = {n: p.detach().clone() for n, p in mtp_params.items()} + optimizer = torch.optim.AdamW( + [p for p in named_params.values() if p.requires_grad], lr=1e-2, weight_decay=0.1 + ) + + distillation_model.train() + prompt_tokens = torch.randint(0, vocab_size, (batch_size, max_sequence_length)).cuda() + labels = torch.randint(0, vocab_size, (batch_size, max_sequence_length)).cuda() + position_ids = ( + torch.arange(max_sequence_length, dtype=torch.long) + .unsqueeze(0) + .repeat(batch_size, 1) + .cuda() + ) + attention_mask = torch.tril( + torch.ones((batch_size, 1, max_sequence_length, max_sequence_length), dtype=torch.bool) + ).cuda() + student_loss = distillation_model(prompt_tokens, position_ids, attention_mask, labels=labels) + loss = distillation_model.compute_kd_loss( + student_loss=student_loss.mean(), loss_reduction_fn=lambda x: x[0].mean() + ) + loss["kd_loss"].backward() + optimizer.step() + + # The main LM loss follows kd_loss_alpha either way. + assert student_loss.any() == (kd_loss_alpha < 1) + expected = (1 - kd_loss_alpha) * student_loss.mean() + kd_loss_alpha * loss["logits_loss"] + assert torch.allclose(loss["kd_loss"], expected) + # The rest of the student keeps training. + assert any(p.grad is not None and p.grad.any() for p in body_params.values()) + # The MTP head trains unless frozen; frozen, it gets no gradient and the step leaves it as is. + mtp_trained = any(p.grad is not None and p.grad.any() for p in mtp_params.values()) + assert mtp_trained != freeze_mtp + if freeze_mtp: + assert all(torch.equal(p, mtp_before[n]) for n, p in mtp_params.items()) + + def test_logits_kl_loss(dist_workers): """Test LogitsKLLoss with TP parallelism.""" dist_workers.run(_test_logits_kl_loss) @@ -460,11 +553,24 @@ def test_distillation_config_top_p_validation(): @pytest.mark.parametrize( - ("quantized", "skip_lm_loss"), [(False, True), (True, True), (True, False)] + ("quantized", "skip_lm_loss", "freeze_mtp"), + [ + (False, True, False), + (True, True, False), + (True, False, False), + (False, True, True), + (False, False, True), + ], ) -def test_skip_lm_loss_with_mtp(dist_workers, quantized, skip_lm_loss): +def test_skip_lm_loss_with_mtp(dist_workers, quantized, skip_lm_loss, freeze_mtp): """Test that skip_lm_loss only zeroes the main LM head, unless MTP is left out of quantization.""" - dist_workers.run(partial(_test_skip_lm_loss_with_mtp, quantized, skip_lm_loss)) + dist_workers.run(partial(_test_skip_lm_loss_with_mtp, quantized, skip_lm_loss, freeze_mtp)) + + +@pytest.mark.parametrize(("freeze_mtp", "kd_loss_alpha"), [(True, 1.0), (True, 0.5), (False, 1.0)]) +def test_freeze_mtp_with_quantized_mtp(dist_workers, freeze_mtp, kd_loss_alpha): + """With the MTP head quantized too, only ``freeze_mtp`` keeps it out of a training step.""" + dist_workers.run(partial(_test_freeze_mtp_with_quantized_mtp, freeze_mtp, kd_loss_alpha)) def test_mtp_excluded_from_quantization(): diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py index e2702f7cf30..177ce61fc4b 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py @@ -97,6 +97,7 @@ ) from modelopt.torch.quantization.qtensor import QTensorWrapper from modelopt.torch.quantization.utils import is_quantized_linear +from modelopt.torch.quantization.utils.calib_utils import GPTQHelper from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector try: @@ -1524,6 +1525,53 @@ def _test_gptq_mamba_hybrid(rank, size): assert torch.isfinite(forward(model)).all() +def test_gptq_hessian_sync_across_dp(dist_workers_size_2): + """GPTQ at DP2 must use the Hessian of both ranks' calibration data and agree across ranks.""" + dist_workers_size_2.run(_test_gptq_hessian_sync_across_dp) + + +def _test_gptq_hessian_sync_across_dp(rank, size): + initialize_for_megatron(tensor_model_parallel_size=1, seed=SEED) + model = get_mcore_gpt_model(tensor_model_parallel_size=1, num_layers=1).cuda() + ref_model = get_mcore_gpt_model(tensor_model_parallel_size=1, num_layers=1).cuda() + ref_model.load_state_dict(model.state_dict()) + + torch.manual_seed(SEED + rank) + batch = get_batch(model, batch_size=2) + all_input_ids = [torch.empty_like(batch[0]) for _ in range(size)] + torch.distributed.all_gather(all_input_ids, batch[0]) + + def forward(model, input_ids_list): + for input_ids in input_ids_list: + model(input_ids, *batch[2:4], labels=batch[1], loss_mask=batch[4]) + + hessians = {} + update_weights = GPTQHelper.update_weights + + def capture_hessian(self, *args, **kwargs): + hessians[self.name] = self.hessian.clone() + return update_weights(self, *args, **kwargs) + + quant_cfg = copy.deepcopy(mtq.NVFP4_DEFAULT_CFG) + quant_cfg["algorithm"] = {"method": "gptq"} + with patch.object(GPTQHelper, "update_weights", capture_hessian): + mtq.quantize(model, quant_cfg, partial(forward, input_ids_list=[batch[0]])) + dp_hessians = dict(hessians) + # Every rank calibrates on the union of the data, so its Hessian is the full-data one. + mtq.quantize(ref_model, quant_cfg, partial(forward, input_ids_list=all_input_ids)) + + assert dp_hessians + for name, hessian in dp_hessians.items(): + # TF32 matmuls in the Hessian update leave ~1e-4 relative noise between the two runs. + assert (hessian - hessians[name]).norm() < 1e-3 * hessians[name].norm(), name + + for name, m in model.named_modules(): + if is_quantized_linear(m) and m.weight_quantizer.is_enabled: + weights = [torch.empty_like(m.weight) for _ in range(size)] + torch.distributed.all_gather(weights, m.weight.contiguous()) + assert all(torch.equal(w, weights[0]) for w in weights), f"{name} differs across ranks" + + def _auto_quantize_mamba_hybrid_cost_helper(rank, size, expert_model_parallel_size, result_path): initialize_for_megatron( tensor_model_parallel_size=1, diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_transformer_engine.py b/tests/gpu_megatron/torch/quantization/plugins/test_transformer_engine.py index 13ecc941c76..e98293c9f56 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_transformer_engine.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_transformer_engine.py @@ -118,6 +118,26 @@ def test_quantize(model_cls, config): quantize_model_and_forward(model, config, calib_data) +def test_grouped_post_restore_preserves_weight_amax(): + """Grouped restore must not clear calibrated state for non-representative experts.""" + model = TEGroupedLinear().cuda() + with torch.no_grad(): + for index in range(model.num_gemms): + getattr(model.net, f"weight{index}").fill_(0.25 + index) + + data = model.get_input().cuda() + mtq.quantize(model, copy.deepcopy(mtq.FP8_DEFAULT_CFG), lambda module: module(data)) + + quantizers = model.net.weight_quantizer + assert isinstance(quantizers, GroupedQuantizer) + before_restore = [quantizer.amax.clone() for quantizer in quantizers] + + model.net.modelopt_post_restore() + + for quantizer, expected in zip(quantizers, before_restore): + torch.testing.assert_close(quantizer.amax, expected) + + @pytest.mark.parametrize("share_weight_quantizer", [False, True]) def test_fold_weight_grouped_linear(share_weight_quantizer): model = TEGroupedLinear().cuda() diff --git a/tests/gpu_megatron/torch/utils/plugins/test_megatron_preprocess_data.py b/tests/gpu_megatron/torch/utils/plugins/test_megatron_preprocess_data.py index 009c638f8bd..5dd0815080c 100644 --- a/tests/gpu_megatron/torch/utils/plugins/test_megatron_preprocess_data.py +++ b/tests/gpu_megatron/torch/utils/plugins/test_megatron_preprocess_data.py @@ -15,11 +15,29 @@ import gzip import json +import multiprocessing +from multiprocessing import forkserver from pathlib import Path import pytest -from modelopt.torch.utils.plugins.megatron_preprocess_data import megatron_preprocess_data +from modelopt.torch.utils.plugins.megatron_preprocess_data import ( + _prestart_forkserver, + _worker_context, + megatron_preprocess_data, +) + + +def test_worker_context_uses_forkserver(): + assert _worker_context().get_start_method() == "forkserver" + + +def test_worker_context_falls_back_to_spawn_without_forkserver(monkeypatch): + monkeypatch.setattr(multiprocessing, "get_all_start_methods", lambda: ["spawn"]) + monkeypatch.setattr(forkserver, "ensure_running", lambda: pytest.fail("started a forkserver")) + + assert _worker_context().get_start_method() == "spawn" + _prestart_forkserver() # Depends on the external HF datasets-server, which intermittently returns 5xx diff --git a/tests/gpu_megatron/torch/utils/plugins/test_prepare_megatron_data_blend.py b/tests/gpu_megatron/torch/utils/plugins/test_prepare_megatron_data_blend.py index f531be27ad4..6306ecef292 100644 --- a/tests/gpu_megatron/torch/utils/plugins/test_prepare_megatron_data_blend.py +++ b/tests/gpu_megatron/torch/utils/plugins/test_prepare_megatron_data_blend.py @@ -13,8 +13,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +import gzip +import importlib import json import sys +from collections import Counter from pathlib import Path from unittest.mock import Mock @@ -24,6 +27,9 @@ from modelopt.torch.utils.plugins.prepare_megatron_data_blend import main +# The plugins package re-exports the function under the module's name, so import the module itself. +blend_module = importlib.import_module("modelopt.torch.utils.plugins.prepare_megatron_data_blend") + def _setup_test( tiny_tokenizer_path: str, tmp_path: Path, monkeypatch, target_tokens: int | None @@ -103,3 +109,273 @@ def test_prepare_megatron_data_blend_with_split_and_files_sources( f"competitive_programming_python_00_messages{token_suffixes[1]}", ] assert (output_dir / "config.yaml").read_bytes() == config_path.read_bytes() + + +def _chat(i: int) -> dict: + return { + "messages": [ + {"role": "user", "content": f"q{i}"}, + {"role": "assistant", "content": f"a{i}"}, + ] + } + + +def _image_chat(i: int) -> dict: + return { + "messages": [ + {"role": "user", "content": [{"type": "image", "image": "x.png"}]}, + {"role": "assistant", "content": f"a{i}"}, + ] + } + + +def _source(name: str, **fields) -> dict: + source = { + "hf_dataset": f"org/{name}", + "split": "train", + "content_field": "messages", + "weight": 1, + } + source.update(fields) + return {k: v for k, v in source.items() if v is not None} + + +def _read_jsonl(path: Path) -> list[dict]: + return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines()] + + +@pytest.fixture +def chat_blend(tmp_path: Path, monkeypatch): + """Sources of 100 ``messages`` rows each; returns (config writer, stream counts per source).""" + calls = Counter() + + def iter_rows(source, output_dir, shuffle, seed): + calls[source["hf_dataset"]] += 1 + yield from (_chat(i) for i in range(100)) + + monkeypatch.setattr(blend_module, "_iter_rows", iter_rows) + + def write_config(**overrides) -> Path: + config = { + "output_dir": str(tmp_path / "blend"), + "output_format": "chat_jsonl", + "target_samples": 100, + "shuffle": False, + "seed": 0, + "sources": [_source("a"), _source("b")], + } + config.update(overrides) + path = tmp_path / "config.yaml" + path.write_text(yaml.safe_dump(config), encoding="utf-8") + return path + + return write_config, calls + + +def test_to_chat_handles_all_row_shapes(): + msgs = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}] + assert blend_module._to_chat({"messages": msgs}, "messages") == {"messages": msgs} + assert blend_module._to_chat({"messages": []}, "messages") is None + # Nothing would be trained under assistant-only loss. + assert blend_module._to_chat({"messages": msgs[:1]}, "messages") is None + # Tool definitions and a top-level system prompt reach the chat path. + tools = [{"type": "function", "function": {"name": "add"}}] + row = {"messages": msgs, "tools": tools, "system": "Be brief.", "id": 7} + chat = blend_module._to_chat(row, "messages") + assert chat == {"messages": msgs, "tools": tools, "system": "Be brief."} + # Unlike the tokenizing mode, reasoning_content stays on the message by default (native). + reasoning = [{"role": "assistant", "content": "4", "reasoning_content": "2+2"}] + assert blend_module._to_chat({"messages": reasoning}, "messages")["messages"] == reasoning + inline = blend_module._to_chat({"messages": reasoning}, "messages", "inline") + assert inline["messages"] == [{"role": "assistant", "content": "\n2+2\n\n4"}] + + expected = { + "messages": [ + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": "world"}, + ] + } + text = "System\nUser\nhello\nAssistant\nworld\n" + assert blend_module._to_chat({"text": text}, "text") == expected + # Any marker ends a turn, so a trailing marker is not trained as content. + assert blend_module._to_chat({"text": text + ""}, "text") == expected + assert blend_module._to_chat({"text": "def f(): pass"}, "text") == { + "messages": [{"role": "assistant", "content": "def f(): pass"}] + } + assert blend_module._to_chat({"text": ""}, "text") is None + # Markers that parse to no turns: dropped, never trained as assistant text. + assert blend_module._to_chat({"text": "\nraw body\n"}, "text") is None + # A role without a mapping would cut its turn out of the middle: the whole row is dropped. + text = "User\nq\nTool\nresult\nAssistant\na\n" + assert blend_module._to_chat({"text": text}, "text") is None + for text in ( + # Text before the first marker would be cut from the conversation. + "preamble\nUser\nq\nAssistant\na\n", + # A SteerLM label line is no role either. + "User\nq\nAssistant\nquality:4\na\n", + # A role must end at a newline: this is no User turn. + "Userguide\nq\nAssistant\na\n", + # An empty known turn would be cut out of the middle of the conversation. + "User\nq1\nAssistant\nUser\nq2\nAssistant\na2\n", + # No assistant turn to train on. + "User\nq\nAssistant\n", + ): + assert blend_module._to_chat({"text": text}, "text") is None, text + + +def test_take_chat_rows_skips_non_text_rows_and_fails_without_enough(tmp_path: Path, monkeypatch): + source = _source("a") + # Media is rejected inside message content and as a top-level key the conversion drops. + rows = [_chat(0), _image_chat(1), {**_chat(9), "images": ["a.png"]}, _chat(2), _chat(3)] + monkeypatch.setattr(blend_module, "_iter_rows", lambda *a: iter(rows)) + got = blend_module._take_chat_rows(source, tmp_path, 3, False, 0, "native") + assert [c["messages"][0]["content"] for c in got] == ["q0", "q2", "q3"] + + images = [_image_chat(i) for i in range(200)] + monkeypatch.setattr(blend_module, "_iter_rows", lambda *a: iter(images)) + with pytest.raises(RuntimeError, match=r"only 0/2 usable rows \(gave up after 121 rejected"): + blend_module._take_chat_rows(source, tmp_path, 2, False, 0, "native") + + +def test_chat_blend_splits_weighted_rows(tmp_path: Path, chat_blend, monkeypatch): + write_config, calls = chat_blend + sources = [_source("a", weight=3), _source("b"), _source("c", weight=0)] + config_path = write_config(sources=sources) + monkeypatch.setattr( + sys, "argv", ["prepare_megatron_data_blend.py", "--config", str(config_path)] + ) + main() + out = tmp_path / "blend" + assert len(_read_jsonl(out / "validation.jsonl")) == 5 + rows = _read_jsonl(out / "train.jsonl") + _read_jsonl(out / "validation.jsonl") + # 3:1 weights: the first 75 rows of org/a and 25 of org/b; org/c rounds to zero and is skipped. + expected = Counter({f"q{i}": 1 for i in range(75)}) + Counter(f"q{i}" for i in range(25)) + assert Counter(r["messages"][0]["content"] for r in rows) == expected + assert calls["org/c"] == 0 + assert (out / "config.yaml").read_bytes() == config_path.read_bytes() + + +def test_chat_blend_reuses_and_rebuilds_cached_parts(tmp_path: Path, chat_blend): + write_config, calls = chat_blend + config_path = write_config() + blend_module.prepare_chat_data_blend(config_path) + first = (tmp_path / "blend" / "train.jsonl").read_text(encoding="utf-8") + blend_module.prepare_chat_data_blend(config_path) + assert calls == {"org/a": 1, "org/b": 1} # the second run streams nothing + + (part,) = (tmp_path / "blend" / "parts").glob("org--a_*.jsonl") + part.write_text(part.read_text(encoding="utf-8")[:-20], encoding="utf-8") # cut mid-record + blend_module.prepare_chat_data_blend(config_path) + assert calls == {"org/a": 2, "org/b": 1} + assert (tmp_path / "blend" / "train.jsonl").read_text(encoding="utf-8") == first + + +@pytest.mark.parametrize( + ("first", "second"), + [ + ({}, {"sources": [_source("a", split="test")]}), + ( + {"sources": [_source("a", split=None, files=["a.jsonl"])]}, + {"sources": [_source("a", split=None, files=["b.jsonl"])]}, + ), + ({}, {"shuffle": True}), + ({"shuffle": True, "seed": 0}, {"shuffle": True, "seed": 1}), + ({}, {"reasoning_content": "inline"}), + ], + ids=["split", "files", "shuffle", "seed", "reasoning_content"], +) +def test_chat_blend_reselects_after_a_selection_change(chat_blend, first, second): + write_config, calls = chat_blend + blend_module.prepare_chat_data_blend(write_config(**first)) + blend_module.prepare_chat_data_blend(write_config(**second)) + assert calls["org/a"] == 2 + + +@pytest.mark.parametrize( + ("overrides", "match"), + [ + ({"target_samples": 0}, "positive integer"), + ({"target_samples": 1.5}, "positive integer"), + # A whole source would be held in memory. + ({"target_samples": None}, "target_samples is required"), + ({"output_format": "megatron"}, "writes output_format: chat_jsonl"), + ({"sources": [_source("a", weight=-1)]}, "nonnegative"), + ({"sources": [_source("a", weight=0)]}, "not all zero"), + ({"sources": [_source("a", weight=float("inf"))]}, "finite"), + ({"tokenizer": "x"}, "does not use"), + ({"target_tokens": 1000}, "does not use"), + ({"reasoning_content": "drop"}, "reasoning_content must be one of"), + # A row with both columns must not silently be read from ``text``. + ({"sources": [_source("a", content_field="conversations")]}, "messages or text"), + ], +) +def test_chat_blend_rejects_invalid_configs(chat_blend, overrides, match): + write_config, calls = chat_blend + with pytest.raises(ValueError, match=match): + blend_module.prepare_chat_data_blend(write_config(**overrides)) + assert not calls + + +@pytest.mark.parametrize( + ("overrides", "match"), + [ + ({}, "Unknown output_format 'chat_jsonl'"), + ({"output_format": "chat-jsonl"}, "Unknown output_format 'chat-jsonl'"), + ({"output_format": "megatron"}, r"Tokenized blends do not use \['seed', 'shuffle'"), + ], +) +def test_tokenized_blend_rejects_chat_options(chat_blend, overrides, match): + write_config, _ = chat_blend + with pytest.raises(ValueError, match=match): + blend_module.prepare_megatron_data_blend(write_config(**overrides)) + + +def test_chat_blend_fails_on_a_bad_short_or_tiny_source(tmp_path: Path, chat_blend, monkeypatch): + write_config, _ = chat_blend + with pytest.raises(RuntimeError, match="org/a train: only 100/250 usable rows"): + blend_module.prepare_chat_data_blend(write_config(target_samples=500)) + with pytest.raises(RuntimeError, match="Too few rows"): + blend_module.prepare_chat_data_blend(write_config(target_samples=10)) + + def iter_rows(source, output_dir, shuffle, seed): + yield from (_chat(i) for i in range(10)) # yields some rows, then the reader fails + raise TypeError("cannot cast") + + monkeypatch.setattr(blend_module, "_iter_rows", iter_rows) + with pytest.raises(RuntimeError, match="org/a train: failed to read or convert"): + blend_module.prepare_chat_data_blend(write_config()) + assert not (tmp_path / "blend" / "train.jsonl").exists() + + +def test_write_jsonl_is_atomic(tmp_path: Path): + path = tmp_path / "part.jsonl" + with pytest.raises(TypeError): + blend_module._write_jsonl(path, [_chat(0), {"bad": object()}]) + assert not path.exists() + + +def test_chat_blend_reads_files_sources(tmp_path: Path, monkeypatch): + plain, packed = tmp_path / "a.jsonl", tmp_path / "b.jsonl.gz" + lines = [json.dumps(_chat(i)) for i in range(20)] + # A malformed record is skipped with a warning, as in the tokenizing reader. + plain.write_text("\n".join([*lines[:10], '{"messages": [', *lines[10:]]) + "\n", "utf-8") + with gzip.open(packed, "wt", encoding="utf-8") as f: + f.writelines(json.dumps(_chat(i)) + "\n" for i in range(20, 40)) + paths = {"data/a.jsonl": str(plain), "data/b.jsonl.gz": str(packed)} + monkeypatch.setattr( + huggingface_hub, "hf_hub_download", lambda repo_id, filename, **kwargs: paths[filename] + ) + config_path = tmp_path / "config.yaml" + config = { + "output_dir": str(tmp_path / "blend"), + "output_format": "chat_jsonl", + "target_samples": 40, + "sources": [_source("a", split=None, files=list(paths))], + } + config_path.write_text(yaml.safe_dump(config), encoding="utf-8") + assert blend_module.prepare_chat_data_blend(config_path) == {"validation": 2, "train": 38} + out = tmp_path / "blend" + rows = _read_jsonl(out / "train.jsonl") + _read_jsonl(out / "validation.jsonl") + assert sorted(r["messages"][0]["content"] for r in rows) == sorted(f"q{i}" for i in range(40)) + # The part is named after its first file. + assert len(list((out / "parts").glob("org--a_a_*.jsonl"))) == 1 diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py index e45c4dcf54f..10af727381f 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py @@ -395,12 +395,32 @@ def fake_decode(query, key_cache, value_cache, block_table, seq_lens, **kwargs): assert finalize_kw["page_size"] == page_size -def test_flashinfer_sparse_prefill_uses_shared_triton_kernel(monkeypatch): +@pytest.mark.parametrize("layout", ["HND", "packed"]) +def test_flashinfer_sparse_prefill_uses_shared_triton_kernel(monkeypatch, layout): """FlashInfer must route 2:4 prefill through the current ModelOpt kernel.""" impl = _make_flashinfer_impl(sparse=True) page_size, num_heads, head_dim = 16, 2, 64 - physical = torch.zeros(4, 2, num_heads, page_size, head_dim, dtype=torch.float16) - kv_cache = physical.permute(0, 1, 3, 2, 4) + expected_key = ( + torch.arange(4 * page_size * num_heads * head_dim) + .remainder(127) + .reshape(4, page_size, num_heads, head_dim) + .to(torch.float16) + ) + expected_value = -expected_key - 1 + if layout == "packed": + kv_cache = torch.cat((expected_key, expected_value), dim=-1).transpose(1, 2).contiguous() + expected_stride = ( + 2 * num_heads * page_size * head_dim, + 2 * head_dim, + 2 * page_size * head_dim, + 1, + ) + else: + physical = ( + torch.stack((expected_key, expected_value), dim=1).permute(0, 1, 3, 2, 4).contiguous() + ) + kv_cache = physical.permute(0, 1, 3, 2, 4) + expected_stride = kv_cache[:, 0].stride() metadata = _flashinfer_metadata(query_lens=(2, 2), max_query_len=4) query = torch.zeros(4, num_heads, head_dim, dtype=torch.float16) captured = {} @@ -419,8 +439,11 @@ def fake_attention(query, **kwargs): assert captured["sparsity_n"] == 2 assert captured["sparsity_m"] == 4 assert captured["page_size"] == page_size - assert captured["k_cache"].stride() == kv_cache[:, 0].stride() - assert captured["v_cache"].stride() == kv_cache[:, 1].stride() + for name, expected in (("k_cache", expected_key), ("v_cache", expected_value)): + cache = captured[name] + torch.testing.assert_close(cache, expected) + assert cache.stride() == expected_stride + assert cache.untyped_storage().data_ptr() == kv_cache.untyped_storage().data_ptr() assert captured["block_table"] is metadata._modelopt_block_table assert captured["b_seq_len_k"] is metadata._modelopt_seq_lens diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py index 4ed2637ef5d..5661ffa43fd 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py @@ -487,7 +487,7 @@ def test_rejects_non_16bit_cache(self, monkeypatch): # FlashInfer adapter: cache write must precede the calibrate-kernel read # --------------------------------------------------------------------------- class TestFlashInferCalibrationOrdering: - @pytest.mark.parametrize("layout", ["NHD", "HND"]) + @pytest.mark.parametrize("layout", ["NHD", "HND", "packed"]) def test_cache_write_happens_before_calibrate_read(self, monkeypatch, layout): calls = [] monkeypatch.setattr( @@ -499,6 +499,11 @@ def test_cache_write_happens_before_calibrate_read(self, monkeypatch, layout): def fake_calibrate(q, *args, **kwargs): calls.append("calibrate") calls.append(kwargs["k_cache"].stride()) + for name, expected in (("k_cache", expected_key), ("v_cache", expected_value)): + cache = kwargs[name] + torch.testing.assert_close(cache, expected) + assert cache.stride() == expected_stride + assert cache.untyped_storage().data_ptr() == kv_cache.untyped_storage().data_ptr() counters = torch.zeros(len(TRIALS), 2, dtype=torch.int64) return torch.zeros_like(q), counters @@ -513,12 +518,28 @@ def fake_calibrate(q, *args, **kwargs): _calib_threshold_trials=list(TRIALS), _calib_records=[], ) - shape = (3, 2, page, num_kv_heads, head_dim) - kv_cache = torch.zeros(shape, dtype=torch.bfloat16) - if layout == "HND": - kv_cache = torch.zeros( - shape[0], shape[1], shape[3], shape[2], shape[4], dtype=torch.bfloat16 - ).permute(0, 1, 3, 2, 4) + expected_key = ( + torch.arange(3 * page * num_kv_heads * head_dim) + .remainder(127) + .reshape(3, page, num_kv_heads, head_dim) + .to(torch.bfloat16) + ) + expected_value = -expected_key - 1 + if layout == "packed": + kv_cache = ( + torch.cat((expected_key, expected_value), dim=-1).transpose(1, 2).contiguous() + ) + expected_stride = ( + 2 * num_kv_heads * page * head_dim, + 2 * head_dim, + 2 * page * head_dim, + 1, + ) + else: + kv_cache = torch.stack((expected_key, expected_value), dim=1) + if layout == "HND": + kv_cache = kv_cache.permute(0, 1, 3, 2, 4).contiguous().permute(0, 1, 3, 2, 4) + expected_stride = kv_cache[:, 0].stride() attn_metadata = SimpleNamespace( _modelopt_block_table=torch.zeros(1, 1, dtype=torch.int32), _modelopt_seq_lens=torch.tensor([8], dtype=torch.int32), @@ -543,7 +564,7 @@ def fake_calibrate(q, *args, **kwargs): output=torch.empty_like(q), ) - assert calls == ["cache_write", "calibrate", kv_cache[:, 0].stride()] + assert calls == ["cache_write", "calibrate", expected_stride] assert torch.isfinite(out).all() assert len(impl._calib_records) == 1 assert impl._calib_records[0]["phase"] == "decode" diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index f397205f022..c8ee193c624 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -15,6 +15,9 @@ import contextlib import os +import platform + +import pytest # Enforce no HuggingFace Hub network access for unit tests os.environ["HF_HUB_OFFLINE"] = "1" @@ -25,3 +28,30 @@ import huggingface_hub.constants as _hf_constants _hf_constants.HF_HUB_OFFLINE = True + + +@pytest.fixture(scope="session") +def _cpu_pools(): + pools = {} + yield pools + for pool in pools.values(): + pool.shutdown() + + +@pytest.fixture +def cpu_dist_workers(_cpu_pools): + """Two persistent gloo ranks shared by every multi-rank unit test. + + A fresh spawn per test re-imports torch and transformers in each child (~10-20s on CI), so the + workers start once (a killed pool is rebuilt). Jobs must leave the process group intact and + undo any global state they set. + """ + if platform.system() == "Windows": + pytest.skip("Multi-rank tests are not run on Windows") + # Lazy: the onnx-only partial-install job has no torch. + from _test_utils.torch.distributed.utils import DistributedWorkerPool, get_live_pool + + # Replaced if an earlier test hung and killed it, so one hang fails one test. + return get_live_pool( + _cpu_pools, 2, lambda: DistributedWorkerPool(world_size=2, backend="gloo", teardown_fn=None) + ) diff --git a/tests/unit/onnx/quantization/test_calibrated_quantization.py b/tests/unit/onnx/quantization/test_calibrated_quantization.py index e726f667d39..485099f9555 100644 --- a/tests/unit/onnx/quantization/test_calibrated_quantization.py +++ b/tests/unit/onnx/quantization/test_calibrated_quantization.py @@ -78,6 +78,8 @@ def _initializer_arrays(model: onnx.ModelProto) -> dict[str, np.ndarray]: 0.007997047156095505, [0.019685039296746254, 0.023622047156095505, 0.027559055015444756], id="int8-entropy", + # ~13s each; entropy calibration is also covered in tests/gpu/onnx + marks=pytest.mark.manual, ), pytest.param( "int8", @@ -94,6 +96,8 @@ def _initializer_arrays(model: onnx.ModelProto) -> dict[str, np.ndarray]: 0.028210056945681572, [0.06944013386964798, 0.08332816511392593, 0.09721619635820389], id="fp8-entropy", + # ~13s each; entropy calibration is also covered in tests/gpu/onnx + marks=pytest.mark.manual, ), pytest.param( "fp8", diff --git a/tests/unit/onnx/quantization/test_concat_elimination.py b/tests/unit/onnx/quantization/test_concat_elimination.py new file mode 100644 index 00000000000..9edf9712a79 --- /dev/null +++ b/tests/unit/onnx/quantization/test_concat_elimination.py @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from _test_utils.onnx.lib_test_models import build_concat_with_constant_input_model + +from modelopt.onnx.quantization.qdq_graph import get_concat_eliminated_tensors + + +def test_concat_elimination_groups_activation_inputs(): + model = build_concat_with_constant_input_model(constant_input=False) + groups = get_concat_eliminated_tensors(model, ["concat", "matmul", "relu", "relu_cls"]) + assert groups["concat_out"] == {"cls", "relu_out", "concat_out"} + assert groups["cls"] == groups["concat_out"] + + +def test_concat_elimination_skips_constant_inputs(): + """A constant Concat input has no calibration range, so it must not join the scale group.""" + model = build_concat_with_constant_input_model(constant_input=True) + groups = get_concat_eliminated_tensors(model, ["concat", "matmul", "relu"]) + assert "cls" not in groups + assert groups["concat_out"] == {"relu_out", "concat_out"} + assert all("cls" not in group for group in groups.values()) diff --git a/tests/unit/onnx/test_dynamo_adapter.py b/tests/unit/onnx/test_dynamo_adapter.py new file mode 100644 index 00000000000..98bf9df004c --- /dev/null +++ b/tests/unit/onnx/test_dynamo_adapter.py @@ -0,0 +1,733 @@ +# 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. + +"""Behavior-level tests for the private Dynamo ONNX graph adapter.""" + +import numpy as np +import pytest +from onnx import TensorProto, helper, numpy_helper + +from modelopt.onnx.export._dynamo_adapter import normalize_dynamo_weight_paths + + +def _tensor(name, values, dtype=np.float32): + return numpy_helper.from_array(np.asarray(values, dtype=dtype), name) + + +def _model( + nodes, + *, + initializers=(), + inputs=(), + outputs=(), + value_info=(), + opset=23, + trt_opset=None, +): + graph = helper.make_graph( + nodes, + "dynamo_adapter", + list(inputs), + list(outputs), + list(initializers), + value_info=list(value_info), + ) + imports = [helper.make_opsetid("", opset)] + if trt_opset is not None: + imports.append(helper.make_opsetid("trt", trt_opset)) + return helper.make_model(graph, opset_imports=imports) + + +def _value(name, dtype=TensorProto.FLOAT, shape=(2, 4)): + return helper.make_tensor_value_info(name, dtype, list(shape)) + + +def _nvfp4_marker(weight="weight", output="weight_dq"): + return helper.make_node("TRT_FP4QDQ", [weight], [output], domain="trt", block_size=16) + + +@pytest.mark.parametrize("static_source", [False, True]) +def test_normalize_handles_deep_shared_provenance(static_source): + nodes = [] + source = "source" + for index in range(2000): + output = f"value_{index}" + nodes.append(helper.make_node("Add" if index % 2 else "Mul", [source, source], [output])) + source = output + for index in range(2): + nodes.extend( + [ + helper.make_node("QuantizeLinear", [source, "scale", "zero"], [f"q_{index}"]), + helper.make_node( + "DequantizeLinear", [f"q_{index}", "scale", "zero"], [f"dq_{index}"] + ), + helper.make_node("MatMul", ["other", f"dq_{index}"], [f"output_{index}"]), + ] + ) + model = _model( + nodes, + initializers=[ + _tensor("scale", 0.25), + _tensor("zero", 0, np.uint8), + *([_tensor("source", np.ones((4, 4)))] if static_source else []), + ], + inputs=[_value("other"), *([] if static_source else [_value("source", shape=(4, 4))])], + outputs=[_value(f"output_{index}") for index in range(2)], + ) + original = model.SerializeToString() + + if static_source: + with pytest.raises(NotImplementedError, match="before the quantization marker"): + normalize_dynamo_weight_paths(model) + else: + assert normalize_dynamo_weight_paths(model) is model + assert model.SerializeToString() == original + + +@pytest.mark.parametrize( + ("quantize_op", "dequantize_op", "domain"), + [ + ("TRT_FP8QuantizeLinear", "TRT_FP8DequantizeLinear", "trt"), + ("QuantizeLinear", "DequantizeLinear", ""), + ], +) +def test_normalize_accepts_canonical_qdq_weight_path(quantize_op, dequantize_op, domain): + weight = _tensor("weight", np.ones((4, 4))) + blocked_shape = _tensor("blocked_shape", [2, 2, 4], np.int64) + restored_shape = _tensor("restored_shape", [4, 4], np.int64) + scale = _tensor("scale", 0.25) + zero = _tensor("zero", 0, np.uint8) + model = _model( + [ + helper.make_node("Reshape", ["weight", "blocked_shape"], ["blocked"]), + helper.make_node( + quantize_op, ["blocked", "scale", "zero"], ["weight_q"], domain=domain + ), + helper.make_node( + dequantize_op, ["weight_q", "scale", "zero"], ["weight_dq"], domain=domain + ), + helper.make_node("Reshape", ["weight_dq", "restored_shape"], ["restored"]), + helper.make_node("MatMul", ["input", "restored"], ["output"]), + ], + initializers=[weight, blocked_shape, restored_shape, scale, zero], + inputs=[_value("input", shape=(2, 4))], + outputs=[_value("output", shape=(2, 4))], + trt_opset=1 if domain else None, + ) + + assert normalize_dynamo_weight_paths(model) is model + assert [node.op_type for node in model.graph.node] == [ + quantize_op, + dequantize_op, + "Reshape", + "MatMul", + ] + quantizer = model.graph.node[0] + assert quantizer.input[0] == "weight" + assert list(next(item for item in model.graph.initializer if item.name == "weight").dims) == [ + 2, + 2, + 4, + ] + + +def test_normalize_rejects_mismatched_qdq_parameters(): + model = _model( + [ + helper.make_node("QuantizeLinear", ["weight", "scale", "zero"], ["weight_q"]), + helper.make_node( + "DequantizeLinear", ["weight_q", "other_scale", "zero"], ["weight_dq"] + ), + helper.make_node("MatMul", ["input", "weight_dq"], ["output"]), + ], + initializers=[ + _tensor("weight", np.ones((4, 4))), + _tensor("scale", 0.25), + _tensor("other_scale", 0.5), + _tensor("zero", 0, np.uint8), + ], + inputs=[_value("input")], + outputs=[_value("output")], + ) + + with pytest.raises(NotImplementedError, match="scale and zero-point inputs must match"): + normalize_dynamo_weight_paths(model) + + +@pytest.mark.parametrize( + ("view_order", "trans_b"), + [("reshape_cast", None), ("cast_reshape", 0), ("cast_reshape", 1)], +) +def test_normalize_canonical_nvfp4_path_removes_export_only_views(view_order, trans_b): + weight = _tensor("weight", np.ones((2, 4))) + blocked_shape = _tensor("blocked_shape", [1, 2, 4], np.int64) + restored_shape = _tensor("restored_shape", [2, 4], np.int64) + views = { + "reshape_cast": [ + helper.make_node("Reshape", ["identity", "restored_shape"], ["restored"]), + helper.make_node("Cast", ["restored"], ["cast"], to=TensorProto.FLOAT16), + ], + "cast_reshape": [ + helper.make_node("Cast", ["identity"], ["cast"], to=TensorProto.FLOAT16), + helper.make_node("Reshape", ["cast", "restored_shape"], ["restored"]), + ], + }[view_order] + view_output = views[-1].output[0] + model = _model( + [ + helper.make_node("Reshape", ["weight", "blocked_shape"], ["blocked"]), + _nvfp4_marker("blocked", "marked"), + helper.make_node("Identity", ["marked"], ["identity"]), + *views, + helper.make_node("Transpose", [view_output], ["transposed"], perm=[1, 0]), + helper.make_node( + "Gemm", + ["input", "transposed"], + ["output"], + **({} if trans_b is None else {"transB": trans_b}), + ), + ], + initializers=[weight, blocked_shape, restored_shape], + inputs=[_value("input", TensorProto.FLOAT16, (1, 2 if trans_b else 4))], + outputs=[_value("output", TensorProto.FLOAT16, (1, 4 if trans_b else 2))], + trt_opset=1, + ) + + normalize_dynamo_weight_paths(model) + + assert [node.op_type for node in model.graph.node] == ["TRT_FP4QDQ", "Gemm"] + marker, gemm = model.graph.node + assert marker.input[0] == "weight" + assert gemm.input[1] == marker.output[0] + assert helper.get_attribute_value( + next(attr for attr in gemm.attribute if attr.name == "transB") + ) == 1 - (trans_b or 0) + + +@pytest.mark.parametrize( + ("view_op", "weight_source"), + [ + ("Slice", "initializer"), + ("Transpose", "initializer"), + ("Expand", "initializer"), + ("Gather", "initializer"), + ("GatherND", "initializer"), + ("GatherND", "constant"), + ], +) +def test_normalize_rejects_static_view_before_quantizer(view_op, weight_source): + if view_op == "Slice": + view = helper.make_node("Slice", ["weight"], ["view"]) + elif view_op == "Transpose": + view = helper.make_node("Transpose", ["weight"], ["view"], perm=[1, 0]) + elif view_op == "Expand": + view = helper.make_node("Expand", ["weight", "expanded_shape"], ["view"]) + else: + view = helper.make_node(view_op, ["weight", "runtime_indices"], ["view"]) + weight = _tensor("weight", np.ones((4, 4))) + weight_nodes = ( + [helper.make_node("Constant", [], ["weight"], value=weight)] + if weight_source == "constant" + else [] + ) + inputs = [_value("input")] + if view_op in {"Gather", "GatherND"}: + shape = (4, 1) if view_op == "GatherND" else (4,) + inputs.append(_value("runtime_indices", TensorProto.INT64, shape)) + model = _model( + [ + *weight_nodes, + view, + helper.make_node("QuantizeLinear", ["view", "scale", "zero"], ["weight_q"]), + helper.make_node("DequantizeLinear", ["weight_q", "scale", "zero"], ["weight_dq"]), + helper.make_node("MatMul", ["input", "weight_dq"], ["output"]), + ], + initializers=[ + *([weight] if weight_source == "initializer" else []), + _tensor("scale", 0.25), + _tensor("zero", 0, np.uint8), + _tensor("expanded_shape", [4, 4], np.int64), + ], + inputs=inputs, + outputs=[_value("output")], + ) + + with pytest.raises(NotImplementedError, match="before the quantization marker"): + normalize_dynamo_weight_paths(model) + + +def test_normalize_rejects_static_quantizer_without_matching_dequantizer(): + model = _model( + [ + helper.make_node("QuantizeLinear", ["weight", "scale", "zero"], ["weight_q"]), + helper.make_node("Identity", ["weight_q"], ["weight_dq"]), + helper.make_node("MatMul", ["input", "weight_dq"], ["output"]), + ], + initializers=[ + _tensor("weight", np.ones((4, 4))), + _tensor("scale", 0.25), + _tensor("zero", 0, np.uint8), + ], + inputs=[_value("input")], + outputs=[_value("output")], + ) + + with pytest.raises(NotImplementedError, match="matching DequantizeLinear"): + normalize_dynamo_weight_paths(model) + + +def test_normalize_rejects_wrong_marker_domain(): + model = _model( + [ + helper.make_node( + "QuantizeLinear", ["weight", "scale", "zero"], ["weight_q"], domain="custom" + ), + helper.make_node( + "DequantizeLinear", + ["weight_q", "scale", "zero"], + ["weight_dq"], + domain="custom", + ), + helper.make_node("MatMul", ["input", "weight_dq"], ["output"]), + ], + initializers=[ + _tensor("weight", np.ones((4, 4))), + _tensor("scale", 0.25), + _tensor("zero", 0, np.uint8), + ], + inputs=[_value("input")], + outputs=[_value("output")], + ) + + with pytest.raises(NotImplementedError, match="unexpected domain"): + normalize_dynamo_weight_paths(model) + + +def test_normalize_rejects_mismatched_post_marker_reshape(): + model = _model( + [ + helper.make_node("QuantizeLinear", ["weight", "scale", "zero"], ["weight_q"]), + helper.make_node("DequantizeLinear", ["weight_q", "scale", "zero"], ["weight_dq"]), + helper.make_node("Reshape", ["weight_dq", "wrong_shape"], ["restored"]), + helper.make_node("MatMul", ["input", "restored"], ["output"]), + ], + initializers=[ + _tensor("weight", np.ones((4, 4))), + _tensor("scale", 0.25), + _tensor("zero", 0, np.uint8), + _tensor("wrong_shape", [2, 8], np.int64), + ], + inputs=[_value("input")], + outputs=[_value("output")], + ) + + with pytest.raises(NotImplementedError, match="restoring Reshape does not match"): + normalize_dynamo_weight_paths(model) + + +def test_normalize_materializes_mxfp8_prefix_reshape(): + model = _model( + [ + helper.make_node("Reshape", ["linear.weight", "blocked_shape"], ["blocked"]), + helper.make_node( + "TRT_MXFP8DequantizeLinear", + ["blocked", "scale"], + ["weight_dq"], + domain="trt", + axis=-1, + block_size=32, + output_dtype=TensorProto.FLOAT, + ), + helper.make_node("Reshape", ["weight_dq", "restored_shape"], ["restored"]), + helper.make_node("MatMul", ["input", "restored"], ["output"]), + ], + initializers=[ + _tensor("linear.weight", np.ones((2, 64))), + _tensor("blocked_shape", [4, 32], np.int64), + _tensor("restored_shape", [2, 64], np.int64), + _tensor("scale", 1.0), + ], + inputs=[_value("input", shape=(1, 2))], + outputs=[_value("output", shape=(1, 64))], + trt_opset=1, + ) + + normalize_dynamo_weight_paths(model) + + marker = next(node for node in model.graph.node if node.op_type.startswith("TRT_MXFP8")) + assert marker.input[0] == "linear.weight" + weight = next(item for item in model.graph.initializer if item.name == "linear.weight") + assert list(weight.dims) == [4, 32] + + +def _unsupported_nvfp4_model(kind): + weight = _tensor("weight", np.ones((2, 4))) + marker = _nvfp4_marker() + common = { + "initializers": [weight], + "inputs": [_value("input")], + "outputs": [_value("output")], + "trt_opset": 1, + } + if kind == "shared weight": + nodes = [ + marker, + helper.make_node("Identity", ["weight"], ["other_weight_use"]), + helper.make_node("MatMul", ["input", "weight_dq"], ["output"]), + ] + elif kind == "fanout": + nodes = [ + marker, + helper.make_node("MatMul", ["input", "weight_dq"], ["output"]), + helper.make_node("Identity", ["weight_dq"], ["other_output"]), + ] + elif kind == "Slice": + nodes = [ + marker, + helper.make_node("Slice", ["weight_dq"], ["view"]), + helper.make_node("MatMul", ["input", "view"], ["output"]), + ] + else: + nodes = [marker, helper.make_node("Relu", ["weight_dq"], ["output"])] + return _model(nodes, **common) + + +@pytest.mark.parametrize("kind", ["shared weight", "fanout", "Slice", "non-compute"]) +def test_normalize_rejects_unsupported_weight_topologies(kind): + with pytest.raises(NotImplementedError, match="Unsupported Dynamo quantized weight topology"): + normalize_dynamo_weight_paths(_unsupported_nvfp4_model(kind)) + + +@pytest.mark.parametrize("domain", ["", "trt"]) +@pytest.mark.parametrize( + ("prefix_shape", "transpose", "axis"), + [ + (None, False, -1), + (None, False, 1), + (None, False, None), + ((2, 2, 2), False, 2), + ((2, 2, 2), True, -1), + ], +) +def test_normalize_int4_path_preserves_values_and_layout(prefix_shape, transpose, axis, domain): + values = np.arange(-4, 4, dtype=np.float32).reshape(2, 4) + scale_shape = (*prefix_shape[:-1], 1) if prefix_shape else (2, 2) + initializers = [_tensor("weight", values), _tensor("scale", np.ones(scale_shape))] + nodes = [] + if prefix_shape: + initializers.extend( + [ + _tensor("blocked_shape", prefix_shape, np.int64), + _tensor("restored_shape", values.shape, np.int64), + ] + ) + nodes.append(helper.make_node("Reshape", ["weight", "blocked_shape"], ["blocked"])) + nodes.append( + helper.make_node( + "DequantizeLinear", + ["blocked" if prefix_shape else "weight", "scale"], + ["weight_dq"], + name="weight_dequantizer", + domain=domain, + block_size=2, + **({} if axis is None else {"axis": axis}), + ) + ) + output = "weight_dq" + if prefix_shape: + nodes.extend( + [ + helper.make_node("Reshape", [output, "restored_shape"], ["restored"]), + helper.make_node("Cast", ["restored"], ["cast"], to=TensorProto.FLOAT16), + ] + ) + output = "cast" + if transpose: + nodes.append(helper.make_node("Transpose", [output], ["transposed"], perm=[1, 0])) + output = "transposed" + nodes.append(helper.make_node("MatMul", ["input", output], ["output"])) + dtype = TensorProto.FLOAT16 if prefix_shape else TensorProto.FLOAT + model = _model( + nodes, + initializers=initializers, + inputs=[_value("input", dtype, (1, 4 if transpose else 2))], + outputs=[_value("output", dtype, (1, 2 if transpose else 4))], + trt_opset=1 if domain else None, + ) + + normalize_dynamo_weight_paths(model) + + assert [node.op_type for node in model.graph.node] == [ + "DequantizeLinear", + *(["Cast"] if prefix_shape else []), + "MatMul", + ] + marker = model.graph.node[0] + assert marker.name == "__modelopt_dynamo_int4__weight_dequantizer" + assert marker.domain == "" + assert marker.input[0] == "weight" + attributes = {attr.name: helper.get_attribute_value(attr) for attr in marker.attribute} + assert attributes["_target_shape"] == [2, 4] + assert attributes["axis"] == (0 if transpose else 1) + if transpose: + assert attributes["_transpose_perm"] == [1, 0] + else: + assert "_transpose_perm" not in attributes + tensors = {tensor.name: tensor for tensor in model.graph.initializer} + np.testing.assert_array_equal( + numpy_helper.to_array(tensors["weight"]), values.reshape(prefix_shape or (4, 2)) + ) + np.testing.assert_array_equal(numpy_helper.to_array(tensors[marker.input[1]]), np.ones((4, 1))) + if prefix_shape: + assert model.graph.node[1].input[0] == marker.output[0] + assert model.graph.node[-1].input[1] == ("cast" if prefix_shape else marker.output[0]) + + +@pytest.mark.parametrize( + ("weight_shape", "prefix_shape", "axis", "message"), + [ + ((2, 4), None, 0, "axis"), + ((2, 4), (2, 2, 2), 0, "axis"), + ((2, 4), (2, 2, 2), 1, "axis"), + ((2, 4), (2, 2, 2), None, "axis"), + ((2, 3), None, -1, "block size"), + ((2, 3), (3, 2), -1, "block size"), + ], +) +def test_normalize_rejects_unsupported_int4_blocks(weight_shape, prefix_shape, axis, message): + initializers = [_tensor("weight", np.ones(weight_shape)), _tensor("scale", 1.0)] + nodes = [] + if prefix_shape: + initializers.append(_tensor("blocked_shape", prefix_shape, np.int64)) + nodes.append(helper.make_node("Reshape", ["weight", "blocked_shape"], ["blocked"])) + nodes.extend( + [ + helper.make_node( + "DequantizeLinear", + ["blocked" if prefix_shape else "weight", "scale"], + ["weight_dq"], + block_size=2, + **({} if axis is None else {"axis": axis}), + ), + helper.make_node("MatMul", ["input", "weight_dq"], ["output"]), + ] + ) + model = _model( + nodes, + initializers=initializers, + inputs=[_value("input", shape=(1, weight_shape[0]))], + outputs=[_value("output", shape=(1, weight_shape[1]))], + ) + + with pytest.raises(NotImplementedError, match=message): + normalize_dynamo_weight_paths(model) + + +@pytest.mark.parametrize("format_name", ["int4", "nvfp4"]) +def test_normalize_rejects_unmatched_restoring_reshape(format_name): + weight = _tensor("weight", np.ones((2, 4))) + scale = _tensor("scale", np.ones((1, 2))) + blocked_shape = _tensor("blocked_shape", [2, 2, 2], np.int64) + wrong_shape = _tensor("wrong_shape", [4, 2], np.int64) + model = _model( + [ + helper.make_node("Reshape", ["weight", "blocked_shape"], ["blocked"]), + ( + helper.make_node( + "DequantizeLinear", + ["blocked", "scale"], + ["weight_dq"], + axis=-1, + block_size=2, + ) + if format_name == "int4" + else _nvfp4_marker("blocked") + ), + helper.make_node("Reshape", ["weight_dq", "wrong_shape"], ["restored"]), + helper.make_node("MatMul", ["input", "restored"], ["output"]), + ], + initializers=[weight, scale, blocked_shape, wrong_shape], + inputs=[_value("input")], + outputs=[_value("output")], + trt_opset=1 if format_name == "nvfp4" else None, + ) + + with pytest.raises(NotImplementedError, match="restoring Reshape does not match"): + normalize_dynamo_weight_paths(model) + + +def test_normalize_rejects_unpaired_nvfp4_restoring_reshape(): + model = _model( + [ + _nvfp4_marker(), + helper.make_node("Reshape", ["weight_dq", "view_shape"], ["view"]), + helper.make_node("MatMul", ["input", "view"], ["output"]), + ], + initializers=[ + _tensor("weight", np.ones((2, 4))), + _tensor("view_shape", [4, 2], np.int64), + ], + inputs=[_value("input")], + outputs=[_value("output")], + trt_opset=1, + ) + + with pytest.raises(NotImplementedError, match="restoring Reshape requires a pre-Reshape"): + normalize_dynamo_weight_paths(model) + + +@pytest.mark.parametrize("format_name", ["int4", "mxfp8"]) +def test_normalize_splits_shared_scale_without_mutating_source(format_name): + marker_op = "DequantizeLinear" if format_name == "int4" else "TRT_MXFP8DequantizeLinear" + source_scale = np.asarray([[0.25, 0.5], [0.75, 1.0]], dtype=np.float32) + initializers = [ + _tensor("weight_a", np.ones((2, 64))), + _tensor("weight_b", np.ones((2, 64))), + _tensor("shared_scale", source_scale), + ] + nodes = [] + for suffix in ("a", "b"): + nodes.extend( + [ + helper.make_node( + marker_op, + [f"weight_{suffix}", "shared_scale"], + [f"weight_{suffix}_dq"], + name=f"dequantize_{suffix}", + domain="" if format_name == "int4" else "trt", + axis=-1, + block_size=32, + ), + helper.make_node( + "MatMul", [f"input_{suffix}", f"weight_{suffix}_dq"], [f"output_{suffix}"] + ), + ] + ) + model = _model( + nodes, + initializers=initializers, + inputs=[_value("input_a", shape=(1, 2)), _value("input_b", shape=(1, 2))], + outputs=[_value("output_a", shape=(1, 64)), _value("output_b", shape=(1, 64))], + trt_opset=1 if format_name == "mxfp8" else None, + ) + + normalize_dynamo_weight_paths(model) + + dequantizers = [node for node in model.graph.node if node.op_type == marker_op] + assert [node.name for node in dequantizers] == [ + f"__modelopt_dynamo_{format_name}__dequantize_{suffix}" for suffix in ("a", "b") + ] + split_names = [node.input[1] for node in dequantizers] + assert len(set(split_names)) == 2 + assert "shared_scale" not in split_names + tensors = {tensor.name: tensor for tensor in model.graph.initializer} + np.testing.assert_array_equal(numpy_helper.to_array(tensors["shared_scale"]), source_scale) + constants = { + node.output[0]: helper.get_attribute_value(node.attribute[0]) + for node in model.graph.node + if node.op_type == "Constant" + } + for name in split_names: + if format_name == "int4": + np.testing.assert_array_equal( + numpy_helper.to_array(tensors[name]), source_scale.reshape(-1, 1) + ) + else: + assert name not in tensors + np.testing.assert_array_equal(numpy_helper.to_array(constants[name]), source_scale) + + +@pytest.mark.parametrize("domain", ["", "trt"]) +@pytest.mark.parametrize("operand", ["query", "key", "value", "projected_key"]) +def test_normalize_leaves_fp8_attention_activation_qdq_intact(domain, operand): + transposed = operand in {"key", "projected_key"} + activation_shape = (2, 8, 4) if operand == "value" else (2, 4, 8) + other_shape = (2, 8, 4) if operand == "query" else (2, 4, 8) + initializers = [ + _tensor("scale", 0.25), + helper.make_tensor("zero", TensorProto.FLOAT8E4M3FN, [], [0]), + ] + nodes = [] + quantizer_input = "activation" + if operand == "projected_key": + initializers.append(_tensor("projection_weight", np.ones((8, 8)))) + nodes.extend( + [ + helper.make_node("MatMul", ["activation", "projection_weight"], ["projected"]), + helper.make_node("Identity", ["projected"], ["prepared"]), + ] + ) + quantizer_input = "prepared" + prefix = "TRT_FP8" if domain else "" + nodes.extend( + [ + helper.make_node( + f"{prefix}QuantizeLinear", + [quantizer_input, "scale", "zero"], + ["activation_q"], + domain=domain, + ), + helper.make_node( + f"{prefix}DequantizeLinear", + ["activation_q", "scale", "zero"], + ["activation_dq"], + domain=domain, + ), + ] + ) + activation_output = "activation_dq" + if transposed: + nodes.append( + helper.make_node("Transpose", [activation_output], ["transposed"], perm=[0, 2, 1]) + ) + activation_output = "transposed" + operands = [activation_output, "other"] if operand == "query" else ["other", activation_output] + nodes.append(helper.make_node("MatMul", operands, ["output"])) + model = _model( + nodes, + initializers=initializers, + inputs=[_value("activation", shape=activation_shape), _value("other", shape=other_shape)], + outputs=[_value("output", shape=(2, 4, 4))], + trt_opset=1 if domain else None, + ) + original = model.SerializeToString() + + assert normalize_dynamo_weight_paths(model) is model + assert model.SerializeToString() == original + + +def test_normalize_leaves_standard_int8_conv_weight_qdq_intact(): + model = _model( + [ + helper.make_node("QuantizeLinear", ["weight", "scale", "zero"], ["weight_q"]), + helper.make_node("DequantizeLinear", ["weight_q", "scale", "zero"], ["weight_dq"]), + helper.make_node("Conv", ["input", "weight_dq"], ["output"]), + ], + initializers=[ + _tensor("weight", np.ones((4, 3, 1, 1))), + _tensor("scale", np.ones(4)), + _tensor("zero", np.zeros(4), np.int8), + ], + inputs=[_value("input", shape=(1, 3, 2, 2))], + outputs=[_value("output", shape=(1, 4, 2, 2))], + ) + + normalize_dynamo_weight_paths(model) + + assert [node.op_type for node in model.graph.node] == [ + "QuantizeLinear", + "DequantizeLinear", + "Conv", + ] diff --git a/tests/unit/recipe/test_presets.py b/tests/unit/recipe/test_presets.py index f97b2d59bd4..3474cd51a43 100644 --- a/tests/unit/recipe/test_presets.py +++ b/tests/unit/recipe/test_presets.py @@ -22,9 +22,6 @@ """ import argparse -import subprocess -import sys -import textwrap import pytest @@ -202,31 +199,3 @@ def test_store_true_style_flag_stores_const_not_an_empty_list(): with pytest.warns(FutureWarning, match="--weight_only is deprecated"): args = parser.parse_args(["--weight_only"]) assert args.weight_only is True - - -def test_deprecation_reaches_stderr_under_the_real_default_filters(): - """The warning has to reach an actual CLI user, not just a test run. - - pytest enables every warning, so a category CPython suppresses looks healthy here and says - nothing in production. argparse invokes the action from its own module, so a - ``DeprecationWarning`` would be dropped by the default ``ignore::DeprecationWarning`` filter -- - the flag would keep working with nothing said. A subprocess is the only honest check: it uses - the interpreter's real filters rather than a reconstruction of them. - """ - script = textwrap.dedent( - """ - import argparse - from modelopt.recipe.presets import RecipeSupersededAction - - parser = argparse.ArgumentParser() - parser.add_argument("--qformat", action=RecipeSupersededAction) - parser.parse_args(["--qformat", "nvfp4"]) - """ - ) - proc = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True) - - assert proc.returncode == 0, proc.stderr - assert "--qformat is deprecated" in proc.stderr, ( - "the deprecation is filtered out under Python's default filters, so a CLI user would " - f"never see it; stderr was: {proc.stderr!r}" - ) diff --git a/tests/unit/recipe/test_qwen_vision_recipe.py b/tests/unit/recipe/test_qwen_vision_recipe.py index aeae38749bf..264e3b73239 100644 --- a/tests/unit/recipe/test_qwen_vision_recipe.py +++ b/tests/unit/recipe/test_qwen_vision_recipe.py @@ -52,14 +52,6 @@ def test_qwen_vision_recipes_select_expected_quantizers( enabled = {name for name, module in modules.items() if getattr(module, "is_enabled", False)} expected_visual_suffixes = { - "blocks.0.attn.qkv.weight_quantizer", - "blocks.0.attn.qkv.input_quantizer", - "blocks.0.attn.proj.weight_quantizer", - "blocks.0.attn.proj.input_quantizer", - "blocks.0.mlp.linear_fc1.weight_quantizer", - "blocks.0.mlp.linear_fc1.input_quantizer", - "blocks.0.mlp.linear_fc2.weight_quantizer", - "blocks.0.mlp.linear_fc2.input_quantizer", "merger.linear_fc1.weight_quantizer", "merger.linear_fc1.input_quantizer", "merger.linear_fc2.weight_quantizer", @@ -74,11 +66,11 @@ def test_qwen_vision_recipes_select_expected_quantizers( "deepstack_merger_list.0.linear_fc2.input_quantizer", } ) - assert {f"model.visual.{suffix}" for suffix in expected_visual_suffixes} <= enabled - assert not any("patch_embed" in name for name in enabled) + visual_quantizers = {name for name in enabled if name.startswith("model.visual.")} + assert visual_quantizers == {f"model.visual.{suffix}" for suffix in expected_visual_suffixes} - vision_weight = modules["model.visual.blocks.0.attn.qkv.weight_quantizer"] - vision_input = modules["model.visual.blocks.0.attn.qkv.input_quantizer"] + vision_weight = modules["model.visual.merger.linear_fc1.weight_quantizer"] + vision_input = modules["model.visual.merger.linear_fc1.input_quantizer"] assert vision_weight.num_bits == (4, 3) assert vision_input.num_bits == (4, 3) diff --git a/tests/unit/torch/export/test_export_diffusers.py b/tests/unit/torch/export/test_export_diffusers.py index 753c81a4b0e..126550a978c 100644 --- a/tests/unit/torch/export/test_export_diffusers.py +++ b/tests/unit/torch/export/test_export_diffusers.py @@ -14,8 +14,11 @@ # limitations under the License. import copy +import inspect import json +from pathlib import Path +import numpy as np import pytest import torch import torch.nn as nn @@ -24,12 +27,23 @@ get_tiny_flux, get_tiny_flux2, get_tiny_unet, + get_tiny_wan22_transformer, + get_tiny_wan22_vae, ) pytest.importorskip("diffusers") +pytest.importorskip("transformers") +import diffusers from safetensors import safe_open from safetensors.torch import save_file +from transformers import ( + CLIPTextConfig, + CLIPTextModel, + CLIPTextModelWithProjection, + UMT5Config, + UMT5EncoderModel, +) import modelopt.torch.export.unified_export_hf as unified_export_hf import modelopt.torch.quantization as mtq @@ -88,6 +102,119 @@ def test_export_diffusers_models_non_quantized(tmp_path, model_factory): assert "quantization_config" not in config_data +@pytest.fixture +def tiny_pipeline(request, tmp_path, tiny_tokenizer): + """Build and reload a tiny pipeline without accessing the Hugging Face Hub.""" + if request.param == "wan22": + wan_pipeline = getattr(diffusers, "WanPipeline", None) + if ( + wan_pipeline is None + or "boundary_ratio" not in inspect.signature(wan_pipeline.__init__).parameters + ): + pytest.skip("Requires a Diffusers version with Wan 2.2 support") + text_encoder = UMT5EncoderModel( + UMT5Config( + vocab_size=len(tiny_tokenizer), + d_model=32, + d_kv=8, + d_ff=32, + num_layers=1, + num_heads=4, + ) + ) + pipe = wan_pipeline( + tokenizer=tiny_tokenizer, + text_encoder=text_encoder, + vae=get_tiny_wan22_vae(), + scheduler=diffusers.UniPCMultistepScheduler( + prediction_type="flow_prediction", use_flow_sigmas=True, flow_shift=3.0 + ), + transformer=get_tiny_wan22_transformer(), + transformer_2=get_tiny_wan22_transformer(), + boundary_ratio=0.875, + ) + settings = ("boundary_ratio", "expand_timesteps") + else: + text_config = CLIPTextConfig( + vocab_size=len(tiny_tokenizer), + hidden_size=32, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=4, + projection_dim=32, + ) + pipe = diffusers.StableDiffusionXLPipeline( + vae=diffusers.AutoencoderKL(block_out_channels=(32,), norm_num_groups=8), + text_encoder=CLIPTextModel(text_config), + text_encoder_2=CLIPTextModelWithProjection(text_config), + tokenizer=tiny_tokenizer, + tokenizer_2=tiny_tokenizer, + unet=get_tiny_unet(in_channels=4, out_channels=4, cross_attention_dim=64), + scheduler=diffusers.DDIMScheduler(), + # Use a non-default value so reload cannot hide a missing setting. + force_zeros_for_empty_prompt=False, + add_watermarker=False, + ) + settings = ("force_zeros_for_empty_prompt",) + source_dir = tmp_path / "source_pipeline" + pipe.save_pretrained(source_dir) + return type(pipe).from_pretrained(source_dir), settings + + +@pytest.mark.parametrize( + ("tiny_pipeline", "selection"), + [("wan22", "default"), ("wan22", "all"), ("wan22", "transformer_2"), ("sdxl", "all")], + indirect=["tiny_pipeline"], +) +def test_export_diffusers_pipeline_config(tmp_path, tiny_pipeline, selection): + pipe, settings = tiny_pipeline + component_names = {name for name, component in pipe.components.items() if component is not None} + components = None if selection == "default" else sorted(component_names) + if selection == "transformer_2": + components = [selection] + expected_components = component_names if components is None else set(components) + export_dir = tmp_path / "export_pipeline" + + export_hf_checkpoint(pipe, export_dir=export_dir, components=components) + + model_index = _load_config(export_dir / "model_index.json") + for setting in settings: + assert model_index[setting] == pipe.config[setting] + assert set(model_index) & set(pipe.components) == expected_components + assert {path.name for path in export_dir.iterdir() if path.is_dir()} == expected_components + + # Loading the export must restore the settings; pass in the components it omits. + omitted = {name: pipe.components[name] for name in component_names - expected_components} + reloaded = type(pipe).from_pretrained(export_dir, **omitted) + for setting in settings: + assert reloaded.config[setting] == pipe.config[setting] + + +@pytest.mark.parametrize( + ("setting", "expected"), + [ + pytest.param(Path("settings") / "config.json", "settings/config.json", id="path"), + pytest.param(np.array([1, 2, 3]), [1, 2, 3], id="ndarray"), + ], +) +def test_export_diffusers_pipeline_config_serialization(tmp_path, setting, expected): + class ConfigPipeline(diffusers.DiffusionPipeline): + def __init__(self, unet, scheduler, setting=None): + super().__init__() + self.register_modules(unet=unet, scheduler=scheduler) + self.register_to_config(setting=setting) + + pipe = ConfigPipeline(get_tiny_unet(), diffusers.DDIMScheduler(), setting=setting) + export_dir = tmp_path / "export_pipeline" + + export_hf_checkpoint(pipe, export_dir=export_dir, components=["unet"]) + + model_index = _load_config(export_dir / "model_index.json") + assert model_index["setting"] == expected + assert set(model_index) & set(pipe.components) == {"unet"} + assert {path.name for path in export_dir.iterdir() if path.is_dir()} == {"unet"} + + def test_export_diffusers_unet_quantized_matches_llm_config(tmp_path, monkeypatch): model = get_tiny_unet() export_dir = tmp_path / "export_unet_quant" diff --git a/tests/unit/torch/export/test_fsdp2_parallel_export.py b/tests/unit/torch/export/test_fsdp2_parallel_export.py index 18846c9844d..c80f704bb6d 100644 --- a/tests/unit/torch/export/test_fsdp2_parallel_export.py +++ b/tests/unit/torch/export/test_fsdp2_parallel_export.py @@ -18,16 +18,20 @@ The pipeline is world-agnostic: with ``torch.distributed`` uninitialized, ``rank/size`` are ``0/1``, the unshard window is a no-op, and rank 0 owns every unit. So these cover unit enumeration, per-tensor export, the shard writer, index naming and tied-weight dropping without a -GPU. The multi-rank unshard collective needs real FSDP2 and is covered in ``tests/gpu``. +GPU. ``test_streaming_export_two_cpu_ranks`` adds one two-rank gloo run; the multi-rank paths +on GPU are covered in ``tests/gpu``. """ import copy +import faulthandler import json from pathlib import Path import pytest import torch from safetensors.torch import load_file +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.fsdp import fully_shard import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq @@ -38,7 +42,7 @@ from transformers import AutoModelForCausalLM, LlamaConfig -def _tiny_quantized_llama(quant_cfg=None, tie=True): +def _tiny_quantized_llama(quant_cfg=None, tie=True, mesh=None): mto.enable_huggingface_checkpointing() torch.manual_seed(0) cfg = LlamaConfig( @@ -52,6 +56,10 @@ def _tiny_quantized_llama(quant_cfg=None, tie=True): architectures=["LlamaForCausalLM"], ) model = AutoModelForCausalLM.from_config(cfg).eval() + if mesh is not None: + for layer in model.model.layers: + fully_shard(layer, mesh=mesh) + fully_shard(model, mesh=mesh) calib = [torch.randint(0, 128, (1, 8)) for _ in range(4)] return mtq.quantize(model, quant_cfg or mtq.FP8_DEFAULT_CFG, lambda m: [m(x) for x in calib]) @@ -143,3 +151,35 @@ def test_streaming_export_extra_state_dict_mtp(tmp_path): loaded = _load_all(d) assert "model.layers.99.mtp.weight" in loaded assert torch.equal(loaded["model.layers.99.mtp.weight"], extra["model.layers.99.mtp.weight"]) + + +def _two_rank_streaming_export(rank, size, export_dir): + # A deadlocked rank exits with its stack instead of hanging the run on the join. Cancelled in + # ``finally`` because the worker is reused: a leftover timer would kill it in a later test. + faulthandler.dump_traceback_later(45, exit=True) + try: + # The unsharded reference is built and exported on every rank, so any collective its + # calibration or export issues has a partner. + expected, _ = _export_transformers_checkpoint( + _tiny_quantized_llama(tie=False), torch.bfloat16 + ) + # Calibrating after sharding leaves the root's parameters unsharded: FSDP2 skips the + # root's post-forward reshard. That is the state in which the owner's state_dict() used + # to reshard them on its rank alone, so the next gather window ran a collective the + # other rank skipped. + model = _tiny_quantized_llama(tie=False, mesh=init_device_mesh("cpu", (size,))) + _export_fsdp2_checkpoint_streaming(model, torch.bfloat16, export_dir=export_dir) + torch.distributed.barrier() + finally: + faulthandler.cancel_dump_traceback_later() + if rank != 0: + return + exported = _load_all(Path(export_dir)) + assert set(exported) == set(expected) + for key, tensor in expected.items(): + assert torch.equal(exported[key].float(), tensor.float()), key + + +def test_streaming_export_two_cpu_ranks(tmp_path, cpu_dist_workers): + """Two gloo ranks export what a single process does, without a rank waiting on the other.""" + cpu_dist_workers.run(_two_rank_streaming_export, export_dir=str(tmp_path)) diff --git a/tests/unit/torch/export/test_get_quantization.py b/tests/unit/torch/export/test_get_quantization.py index edf55bf4bea..f726f01f4e1 100644 --- a/tests/unit/torch/export/test_get_quantization.py +++ b/tests/unit/torch/export/test_get_quantization.py @@ -38,6 +38,7 @@ QUANTIZATION_NVFP4, QUANTIZATION_Q8_0, QUANTIZATION_W4A8_AWQ, + QUANTIZATION_W4A16_MXFP4, ) from modelopt.torch.export.quant_utils import ( _get_carried_over_module_names, @@ -50,6 +51,7 @@ seed_carried_over_exclusions, uses_iq_quantization, ) +from modelopt.torch.quantization.config import QuantizerAttributeConfig from modelopt.torch.quantization.nn import ( GroupedQuantizer, NVFP4StaticQuantizer, @@ -339,6 +341,17 @@ def test_get_quantization_format(config, expected): assert get_quantization_format(model) == expected +def test_mxfp4_weight_only_without_input_quantizer(): + module = torch.nn.Linear(32, 32, bias=False) + module.weight_quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits=(2, 1), + block_sizes={-1: 32, "type": "dynamic", "scale_bits": (8, 0)}, + ) + ) + assert get_quantization_format(module) == QUANTIZATION_W4A16_MXFP4 + + def test_nvfp4_static_quantizer_export(): """NVFP4StaticQuantizer: get_quantization_format returns NVFP4 and get_quant_config returns export config.""" model = ToyModel() diff --git a/tests/unit/torch/export/test_quant_aware_conversion.py b/tests/unit/torch/export/test_quant_aware_conversion.py index 209dafb64b3..1a798b2c1fd 100644 --- a/tests/unit/torch/export/test_quant_aware_conversion.py +++ b/tests/unit/torch/export/test_quant_aware_conversion.py @@ -29,8 +29,9 @@ import pytest import torch -from safetensors.torch import load_file +from safetensors.torch import load_file, save_file +import modelopt.torch.export.quant_aware_conversion as quant_aware_conversion import modelopt.torch.quantization as mtq from modelopt.torch.export.layerwise_export import LayerwiseExporter from modelopt.torch.export.quant_aware_conversion import ( @@ -38,6 +39,9 @@ RenameRule, SplitRule, _assert_experts_pre_expanded, + _build_reverse_rules, + _source_checkpoint_keys, + _SourceAwareRenamer, apply_reverse_rules, build_reverse_name_mapper, revert_quant_config_names, @@ -158,6 +162,326 @@ def _fp8_llama(*quantizer_names): ) +def test_reverse_rename_is_idempotent_when_repl_contains_pattern(): + """A rule whose replacement keeps the pattern must not re-apply to its own output.""" + rules = [RenameRule("input_conditioner", "radio_model.input_conditioner", ("vision_model.",))] + key = "vision_model.input_conditioner.norm_mean" + expected = "vision_model.radio_model.input_conditioner.norm_mean" + + once = apply_reverse_rules({key: torch.zeros(3)}, [], rules) + assert list(once) == [expected] + + twice = apply_reverse_rules(once, [], rules) + assert list(twice) == [expected], "reverse rename must be idempotent" + + +def test_reverse_rename_idempotent_for_already_hub_named_key(): + """An input already carrying the reverted prefix is left untouched.""" + rules = [RenameRule("summary_idxs", "radio_model.summary_idxs", ("vision_model.",))] + already = "vision_model.radio_model.summary_idxs" + out = apply_reverse_rules({already: torch.zeros(1)}, [], rules) + assert list(out) == [already] + + +def test_idempotency_guard_leaves_disjoint_rules_alone(): + """A replacement that cannot re-match keeps its original behaviour.""" + rules = [RenameRule("encoder.layer", "radio_model.model.blocks", ("vision_model.",))] + key = "vision_model.encoder.layer.0.attn.weight" + out = apply_reverse_rules({key: torch.zeros(2)}, [], rules) + assert list(out) == ["vision_model.radio_model.model.blocks.0.attn.weight"] + + +# A NemotronH-shaped source checkpoint: the canonical embedding name is plural, and the +# projector weights sit at the top level under their fused ``mlp1.`` names. +_SOURCE_KEYS = { + "backbone.embeddings.weight", + "backbone.layers.0.mixer.in_proj.weight", + "mlp1.0.weight", + "mlp1.1.weight", + "vision_projector.vision_final_layernorm.weight", +} + + +# A MiniMax-M3-shaped source checkpoint: the sparse text layers are stored under +# ``block_sparse_moe``, while the dense text layers and the vision tower are already +# stored under the in-memory ``mlp`` name. One rename rule serves both namespaces. +_MIXED_SOURCE_KEYS = { + "language_model.model.layers.0.mlp.down_proj.weight", + "language_model.model.layers.1.block_sparse_moe.gate.weight", + "language_model.model.layers.1.block_sparse_moe.experts.0.w1.weight", + "vision_tower.encoder.layers.0.mlp.fc1.weight", +} + +# The NemotronH reverse chain, in the order ``_build_reverse_rules`` emits it: the +# backwards-compatibility leaf rename first, then the ``model.`` -> ``backbone.`` prefix. +_NEMOTRON_H_REVERSE = [ + RenameRule(r"embeddings\.weight", "embedding.weight"), + RenameRule(r"^model\.", "backbone."), +] + + +def test_rename_inventing_a_name_absent_from_source_is_suppressed(): + """A BC *load* rename declared reversible would rewrite the canonical name. + + Reversing ``embedding.weight -> embeddings.weight`` turns the published + ``backbone.embeddings.weight`` into a legacy spelling the checkpoint never used, + which a deployment loader keyed on the canonical name cannot find. The prefix rule + in the same chain is still needed, so resolution has to span the whole chain. + """ + rename = _SourceAwareRenamer(_NEMOTRON_H_REVERSE, _SOURCE_KEYS) + assert rename("model.embeddings.weight") == "backbone.embeddings.weight" + + +def test_rename_restoring_a_source_name_is_applied(): + """A rename whose output is a real source name is left alone.""" + rule = RenameRule(r"^vision_projector\.mlp1\.norm\.", "mlp1.0.") + rename = _SourceAwareRenamer([rule], _SOURCE_KEYS) + assert rename("vision_projector.mlp1.norm.weight") == "mlp1.0.weight" + + +def test_rename_matching_no_source_key_is_applied(): + """No evidence means no suppression: a rule that never fires on source names runs.""" + rule = RenameRule(r"^sound_projector\.sound_encoder\.", "sound_encoder.") + rename = _SourceAwareRenamer([rule], _SOURCE_KEYS) + assert rename("sound_projector.sound_encoder.0.weight") == "sound_encoder.0.weight" + + +@pytest.mark.parametrize("source_keys", [None, set()]) +def test_guard_is_inert_when_source_is_unknown(source_keys): + """An unreadable or absent source checkpoint keeps the previous behaviour.""" + rename = _SourceAwareRenamer(_NEMOTRON_H_REVERSE, source_keys) + assert rename("model.embeddings.weight") == "backbone.embedding.weight" + + +# The same container rename shows up in two shapes depending on how the model declares it: +# one bare container rule, or a rule per anchored leaf. The per-key resolution must give +# the same answer either way, so every mixed-namespace assertion runs against both. +_MIXED_RULE_SHAPES = { + "bare-container": [RenameRule(r"\.mlp\.", ".block_sparse_moe.")], + "anchored-leaves": [ + RenameRule(r"\.mlp\.experts\.", ".block_sparse_moe.experts."), + RenameRule(r"\.mlp\.gate\.", ".block_sparse_moe.gate."), + ], +} + + +@pytest.mark.parametrize("shape", list(_MIXED_RULE_SHAPES)) +def test_mixed_namespace_rename_applies_only_where_the_source_needs_it(shape): + """One rule, two namespaces: suppression is per key, never per rule. + + A rule that rewrites only already-canonical source keys would be discarded by a + rule-level guard, taking the sparse-layer renames it is actually needed for with it. + """ + rename = _SourceAwareRenamer(_MIXED_RULE_SHAPES[shape], _MIXED_SOURCE_KEYS) + + # Sparse text layers: the rename is what restores their source names. + assert ( + rename("language_model.model.layers.1.mlp.experts.0.w1.weight") + == "language_model.model.layers.1.block_sparse_moe.experts.0.w1.weight" + ) + assert ( + rename("language_model.model.layers.1.mlp.gate.weight") + == "language_model.model.layers.1.block_sparse_moe.gate.weight" + ) + # Dense text layer and vision tower are already canonical: left untouched. + assert ( + rename("language_model.model.layers.0.mlp.down_proj.weight") + == "language_model.model.layers.0.mlp.down_proj.weight" + ) + assert ( + rename("vision_tower.encoder.layers.0.mlp.fc1.weight") + == "vision_tower.encoder.layers.0.mlp.fc1.weight" + ) + + +@pytest.mark.parametrize("shape", list(_MIXED_RULE_SHAPES)) +@pytest.mark.parametrize("suffix", [".weight", ".weight_scale", ".input_scale"]) +def test_mixed_namespace_scales_follow_their_weight(suffix, shape): + """A module's scale tensors must not be split across the two namespaces. + + Scale names never appear in the source checkpoint, so they are resolved by the + weight they belong to rather than judged absent on their own. + """ + rename = _SourceAwareRenamer(_MIXED_RULE_SHAPES[shape], _MIXED_SOURCE_KEYS) + sparse = "language_model.model.layers.1.mlp.experts.0.w1" + dense = "language_model.model.layers.0.mlp.down_proj" + assert ( + rename(sparse + suffix) + == "language_model.model.layers.1.block_sparse_moe.experts.0.w1" + suffix + ) + assert rename(dense + suffix) == dense + suffix + + +def test_mixed_namespace_config_names_track_the_weights(tmp_path): + """``exclude_modules`` / ``quantized_layers`` follow the same per-key decision.""" + pytest.importorskip("transformers.core_model_loading") + # Local import: transformers is an optional dependency for ModelOpt. + from transformers.core_model_loading import WeightRenaming + + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps( + {"weight_map": dict.fromkeys(_MIXED_SOURCE_KEYS, "model-00001-of-00001.safetensors")} + ) + ) + model = types.SimpleNamespace( + name_or_path=str(tmp_path), + config=None, + _weight_conversions=[WeightRenaming(".block_sparse_moe.", ".mlp.")], + ) + mapper = build_reverse_name_mapper(model) + assert ( + mapper("language_model.model.layers.1.mlp.experts.0.w1") + == "language_model.model.layers.1.block_sparse_moe.experts.0.w1" + ) + assert ( + mapper("language_model.model.layers.0.mlp.down_proj") + == "language_model.model.layers.0.mlp.down_proj" + ) + + +def test_model_without_a_conversion_mapping_is_a_no_op(): + """The empty-mapping early return must stay shape-compatible with its callers. + + ``build_reverse_name_mapper`` unpacks the result, so a short tuple here would raise + ``ValueError`` for every model that declares no conversion mapping at all. + """ + model = types.SimpleNamespace(name_or_path=None, config=None, _weight_conversions=[]) + split_rules, merge_rules, rename_rules, expert_leaves, source_keys = _build_reverse_rules(model) + assert (split_rules, merge_rules, rename_rules, expert_leaves, source_keys) == ( + [], + [], + [], + [], + None, + ) + assert build_reverse_name_mapper(model) is None + state = {"model.embeddings.weight": torch.ones(2, 2)} + assert revert_weight_conversion_quant_aware(model, state) is state + + +def test_source_checkpoint_keys_reads_sharded_index(tmp_path): + """The guard's input comes from the checkpoint the model was loaded from.""" + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": dict.fromkeys(_SOURCE_KEYS, "model-00001-of-00001.safetensors")}) + ) + model = types.SimpleNamespace(name_or_path=str(tmp_path), config=None) + assert _source_checkpoint_keys(model) == _SOURCE_KEYS + assert _source_checkpoint_keys(types.SimpleNamespace(name_or_path="", config=None)) is None + + +def test_source_checkpoint_keys_reads_single_file(tmp_path): + """The single-file ``model.safetensors`` layout is read through ``safe_open``.""" + # A fresh tensor per key: ``save_file`` rejects names that share storage. + save_file({key: torch.ones(1) for key in _SOURCE_KEYS}, tmp_path / "model.safetensors") + model = types.SimpleNamespace(name_or_path=str(tmp_path), config=None) + assert _source_checkpoint_keys(model) == _SOURCE_KEYS + + +def test_source_checkpoint_keys_tolerates_a_malformed_index(tmp_path): + """An unreadable checkpoint is "unknown", not an export failure.""" + (tmp_path / "model.safetensors.index.json").write_text("{ not json") + model = types.SimpleNamespace(name_or_path=str(tmp_path), config=None) + assert _source_checkpoint_keys(model) is None + + (tmp_path / "model.safetensors.index.json").write_text(json.dumps({"metadata": {}})) + assert _source_checkpoint_keys(model) is None + + +def test_source_checkpoint_keys_resolves_a_cached_hub_id(tmp_path, monkeypatch): + """A hub-id ``name_or_path`` resolves through the local HF cache, offline. + + ``from_pretrained("nvidia/...")`` leaves ``name_or_path`` as the repo id, which is + not a directory. Without cache resolution the guard would silently do nothing for + the most common way ``hf_ptq.py`` is invoked. + """ + snapshot = tmp_path / "snapshots" / "abc123" + snapshot.mkdir(parents=True) + index = snapshot / "model.safetensors.index.json" + index.write_text( + json.dumps({"weight_map": dict.fromkeys(_SOURCE_KEYS, "model-00001-of-00001.safetensors")}) + ) + seen = {} + + def _fake_cache(repo_id, filename, revision=None): + seen["revision"] = revision + return str(index) if filename == "model.safetensors.index.json" else None + + monkeypatch.setattr(quant_aware_conversion, "try_to_load_from_cache", _fake_cache) + model = types.SimpleNamespace(name_or_path="nvidia/Some-Model", config=None) + assert _source_checkpoint_keys(model) == _SOURCE_KEYS + assert seen["revision"] is None + + +def test_cached_hub_id_lookup_uses_the_loaded_revision(tmp_path, monkeypatch): + """A model pinned to a revision must not read ``main``'s snapshot. + + ``try_to_load_from_cache`` defaults to ``main``. For a model loaded from a pinned + commit that would read a different cached checkpoint and filter the wrong rules, + so the commit ``transformers`` recorded on the config is passed through. + """ + index = tmp_path / "model.safetensors.index.json" + index.write_text( + json.dumps({"weight_map": dict.fromkeys(_SOURCE_KEYS, "model-00001-of-00001.safetensors")}) + ) + seen = {} + + def _fake_cache(repo_id, filename, revision=None): + seen["revision"] = revision + return str(index) if filename == "model.safetensors.index.json" else None + + monkeypatch.setattr(quant_aware_conversion, "try_to_load_from_cache", _fake_cache) + model = types.SimpleNamespace( + name_or_path="nvidia/Some-Model", + config=types.SimpleNamespace(_commit_hash="deadbeef", _name_or_path=None), + ) + assert _source_checkpoint_keys(model) == _SOURCE_KEYS + assert seen["revision"] == "deadbeef" + + +def test_source_aware_guard_applies_through_the_export_paths(tmp_path): + """The guard reaches both the tensor rename and the quant-config name mapping. + + Reproduces the NemotronH conversion mapping: a ``backbone.`` -> ``model.`` prefix + rename that genuinely restores source names, plus a backwards-compatibility + ``embedding.weight`` -> ``embeddings.weight`` load rename whose reverse invents a + name the published checkpoint never carried. Only the second must be suppressed. + """ + pytest.importorskip("transformers.core_model_loading") + # Local import: transformers is an optional dependency for ModelOpt. + from transformers.core_model_loading import WeightRenaming + + source_keys = {"backbone.embeddings.weight", "backbone.layers.0.mixer.in_proj.weight"} + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": dict.fromkeys(source_keys, "model-00001-of-00001.safetensors")}) + ) + model = types.SimpleNamespace( + name_or_path=str(tmp_path), + config=None, + _weight_conversions=[ + WeightRenaming("backbone.", "model."), + WeightRenaming("embedding.weight", "embeddings.weight"), + ], + ) + state = _nvfp4_linear("model.layers.0.mixer.in_proj", 8, 16) + state["model.embeddings.weight"] = torch.ones(8, 16) + + written = revert_weight_conversion_quant_aware(model, state) + # The ``backbone.`` prefix reverse still applies; the embedding name stays canonical. + assert "backbone.embeddings.weight" in written + assert "backbone.embedding.weight" not in written + assert "backbone.layers.0.mixer.in_proj.weight" in written + + config = { + "quantized_layers": {"model.layers.0.mixer.in_proj": {"quant_algo": "NVFP4"}}, + "exclude_modules": ["model.embeddings"], + } + revert_quant_config_names(config, build_reverse_name_mapper(model)) + assert config["quantized_layers"] == { + "backbone.layers.0.mixer.in_proj": {"quant_algo": "NVFP4"} + } + assert config["exclude_modules"] == ["backbone.embeddings"] + + def test_rename_carries_scale_siblings(): """A module rename rewrites weight + all scale siblings with identical values.""" sd = _nvfp4_linear("model.language_model.layers.10.mlp.experts.40.gate_proj", 8, 16) diff --git a/tests/unit/torch/export/test_quantized_weight_state.py b/tests/unit/torch/export/test_quantized_weight_state.py new file mode 100644 index 00000000000..22ad3cac977 --- /dev/null +++ b/tests/unit/torch/export/test_quantized_weight_state.py @@ -0,0 +1,353 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import copy + +import pytest +import torch +import torch.nn as nn + +import modelopt.torch.export.quantized_weight_export as weight_export +import modelopt.torch.quantization as mtq +from modelopt.torch.export.convert_hf_config import convert_hf_quant_config_format +from modelopt.torch.export.quantized_weight_export import ( + capture_quantized_weight_export_state, + get_quantized_weight_export_spec, +) +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.nn import GroupedQuantizer, NVFP4StaticQuantizer, TensorQuantizer + + +def _fp8_linear() -> nn.Linear: + module = nn.Linear(4, 4, bias=False) + with torch.no_grad(): + module.weight.copy_(torch.arange(16, dtype=torch.float32).reshape(4, 4) / 8 - 1) + return mtq.quantize(module, copy.deepcopy(mtq.FP8_DEFAULT_CFG), lambda m: m(torch.ones(2, 4))) + + +def _static_w4a16_linear( + weight: torch.Tensor, + per_block_amax: torch.Tensor, + global_amax: torch.Tensor, +) -> nn.Linear: + module = nn.Linear(weight.shape[1], weight.shape[0], bias=False) + module.weight.data.copy_(weight) + cfg = QuantizerAttributeConfig( + num_bits=(2, 1), + block_sizes={-1: 16, "type": "static", "scale_bits": (4, 3)}, + ) + quantizer = NVFP4StaticQuantizer(quant_attribute_cfg=cfg) + quantizer.amax = per_block_amax.clone() + quantizer.global_amax = global_amax.clone() + module.weight_quantizer = quantizer + module.input_quantizer = TensorQuantizer() + module.input_quantizer.disable() + return module + + +class _SquareTransposedExperts(nn.Module): + def __init__(self): + super().__init__() + self.gate_up_proj = nn.Parameter(torch.arange(512, dtype=torch.float32).reshape(2, 16, 16)) + self.down_proj = nn.Parameter(torch.arange(512, dtype=torch.float32).reshape(2, 16, 16)) + + fp8_cfg = QuantizerAttributeConfig(num_bits=(4, 3)) + self.gate_up_proj_weight_quantizer = TensorQuantizer(fp8_cfg) + self.gate_up_proj_weight_quantizer._amax = torch.tensor(2.0) + self.gate_up_proj_input_quantizer = TensorQuantizer(fp8_cfg) + self.gate_up_proj_input_quantizer._amax = torch.tensor(1.0) + + nvfp4_cfg = QuantizerAttributeConfig( + num_bits=(2, 1), + block_sizes={-1: 16, "type": "static", "scale_bits": (4, 3)}, + ) + self.down_proj_weight_quantizer = NVFP4StaticQuantizer(quant_attribute_cfg=nvfp4_cfg) + self.down_proj_weight_quantizer.amax = torch.arange(32, dtype=torch.float32) + 1 + self.down_proj_weight_quantizer.global_amax = torch.tensor(32.0) + self.down_proj_input_quantizer = TensorQuantizer() + self.down_proj_input_quantizer.disable() + + def iter_weights_for_calibration(self): + for name in ("gate_up_proj", "down_proj"): + yield getattr(self, name).transpose(-1, -2), getattr(self, f"{name}_weight_quantizer") + + +class _GroupedWeights(nn.Module): + """Minimal TEGroupedLinear-style numbered-weight layout.""" + + def __init__(self): + super().__init__() + self.weight0 = nn.Parameter(torch.arange(16, dtype=torch.float32).reshape(4, 4)) + self.weight1 = nn.Parameter(torch.arange(16, 32, dtype=torch.float32).reshape(4, 4)) + cfg = QuantizerAttributeConfig(num_bits=(4, 3)) + self.quantizers = nn.ModuleList([TensorQuantizer(cfg), TensorQuantizer(cfg)]) + for index, quantizer in enumerate(self.quantizers, start=1): + quantizer._amax = torch.tensor(float(index)) + self.input_quantizer = TensorQuantizer() + self.input_quantizer.disable() + + def iter_weights_for_calibration(self): + yield self.weight0, self.quantizers[0] + yield self.weight1, self.quantizers[1] + + +def test_capture_does_not_modify_zero_amax(): + module = _fp8_linear() + module.weight_quantizer._amax.zero_() + module.input_quantizer._amax.zero_() + weight_amax = module.weight_quantizer._amax.clone() + input_amax = module.input_quantizer._amax.clone() + + capture_quantized_weight_export_state(module) + + torch.testing.assert_close(module.weight_quantizer._amax, weight_amax) + torch.testing.assert_close(module.input_quantizer._amax, input_amax) + + +_CAPTURE_FORMAT_CONFIGS = { + "fp8_pc_pt": {"num_bits": (4, 3), "axis": 0}, + "fp8_pb_wo": {"num_bits": (4, 3), "block_sizes": {-2: 128, -1: 128}}, + "nvfp4": { + "num_bits": (2, 1), + "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, + }, + "mxfp8": { + "num_bits": (4, 3), + "block_sizes": {-1: 32, "type": "dynamic", "scale_bits": (8, 0)}, + }, + "mxfp4": { + "num_bits": (2, 1), + "block_sizes": {-1: 32, "type": "dynamic", "scale_bits": (8, 0)}, + }, +} + + +def _capture_linear(quantization_format): + module = nn.Linear(256, 128, bias=False) + cfg = _CAPTURE_FORMAT_CONFIGS[quantization_format] + module.weight_quantizer = TensorQuantizer(QuantizerAttributeConfig(**cfg)) + module.input_quantizer = TensorQuantizer( + QuantizerAttributeConfig( + **({"num_bits": (4, 3), "type": "dynamic"} if "axis" in cfg else cfg) + ) + ) + if quantization_format == "fp8_pb_wo": + module.input_quantizer.disable() + return module + + +@pytest.mark.parametrize("quantization_format", _CAPTURE_FORMAT_CONFIGS) +def test_capture_preserves_format_state_and_detaches_tensors(quantization_format): + module = _capture_linear(quantization_format) + cfg = _CAPTURE_FORMAT_CONFIGS[quantization_format] + expected = {} + if quantization_format == "fp8_pc_pt": + scale = torch.arange(1, 129, dtype=torch.float32) + module.weight_quantizer._amax = (448 * scale).reshape(128, 1) + expected["weight_scale"] = (scale, (0,), (1,)) + elif quantization_format == "fp8_pb_wo": + scale = torch.tensor([[1.0, 2.0]]) + module.weight_quantizer._amax = 448 * scale + expected["weight_scale"] = (scale, (0, 1), (128, 128)) + elif quantization_format == "nvfp4": + module.weight_quantizer._amax = torch.tensor(2 * 6 * 448.0) + module.input_quantizer._amax = torch.tensor(3 * 6 * 448.0) + expected["weight_scale_2"] = (torch.tensor(2.0), (), ()) + expected["input_scale"] = (torch.tensor(3.0), (), ()) + elif quantization_format == "mxfp8": + scale = torch.full((128, 8), 127, dtype=torch.uint8) + module.weight_quantizer._scale = scale.clone() + expected["weight_scale"] = (scale, (0, 1), (1, 32)) + for quantizer in (module.weight_quantizer, module.input_quantizer): + if hasattr(quantizer, "_amax"): + quantizer._amax.requires_grad_() + + state = capture_quantized_weight_export_state(module) + + assert state.quantization_format == quantization_format + assert state.weight_shape == (128, 256) + assert state.packing_permutation == (0, 1) + assert state.block_size == (cfg.get("block_sizes") or {}).get(-1, 0) + assert state.input_quantizer_enabled is module.input_quantizer.is_enabled + with torch.no_grad(): + for quantizer in (module.weight_quantizer, module.input_quantizer): + for name in ("_amax", "_scale"): + if (value := getattr(quantizer, name, None)) is not None: + value.zero_() + assert {record.name for record in state.tensors} == set(expected) + for record in state.tensors: + value, axes, block_sizes = expected[record.name] + torch.testing.assert_close(record.value, value) + assert record.axes == axes + assert record.block_sizes == block_sizes + assert record.value.device.type == "cpu" + assert not record.value.requires_grad + + +@pytest.mark.parametrize("quantization_format", _CAPTURE_FORMAT_CONFIGS) +def test_capture_requires_calibration_only_for_static_scales(quantization_format): + module = _capture_linear(quantization_format) + + if quantization_format in {"mxfp8", "mxfp4"}: + state = capture_quantized_weight_export_state(module) + assert state.quantization_format == quantization_format + assert not state.tensors + else: + error = "amax" if quantization_format == "nvfp4" else "Missing calibrated" + with pytest.raises((RuntimeError, AssertionError), match=error): + capture_quantized_weight_export_state(module) + + +def test_capture_resolves_numbered_grouped_weights_by_storage(): + module = _GroupedWeights() + + state0 = capture_quantized_weight_export_state(module, "weight0") + state1 = capture_quantized_weight_export_state(module, "weight1") + + assert state0 is not None and state1 is not None + assert state0.quantization_format == state1.quantization_format == "fp8" + assert not state0.input_quantizer_enabled and not state1.input_quantizer_enabled + assert state0.tensors[0].value.item() != state1.tensors[0].value.item() + + +@pytest.mark.parametrize("has_iterator", [True, False]) +def test_capture_rejects_unresolved_numbered_weight(has_iterator): + module = _GroupedWeights() + module.weight_quantizer = GroupedQuantizer(*module.quantizers) + module.iter_weights_for_calibration = ( + (lambda: iter(((module.weight0, module.quantizers[0]),))) if has_iterator else None + ) + + with pytest.raises(NotImplementedError, match="weight1"): + capture_quantized_weight_export_state(module, "weight1") + + +def test_capture_rejects_full_fused_weight_with_per_expert_quantizers(): + module = _GroupedWeights() + module.weight = nn.Parameter(torch.stack((module.weight0, module.weight1))) + module.weight_quantizer = GroupedQuantizer(*module.quantizers) + module.iter_weights_for_calibration = lambda: ( + (module.weight[index], quantizer) for index, quantizer in enumerate(module.quantizers) + ) + + with pytest.raises(NotImplementedError, match="matching full weight view"): + capture_quantized_weight_export_state(module) + + +@pytest.mark.parametrize("parameter_name", ["bias", "layer_norm_weight"]) +def test_unquantized_sibling_has_no_export_state_or_spec(parameter_name): + module = _fp8_linear() + module.register_parameter(parameter_name, nn.Parameter(torch.ones(4))) + + assert capture_quantized_weight_export_state(module, parameter_name) is None + assert get_quantized_weight_export_spec(module, parameter_name) is None + + +def test_unquantized_weight_has_no_export_state_or_spec(): + module = nn.Linear(4, 4, bias=False) + + assert capture_quantized_weight_export_state(module) is None + assert get_quantized_weight_export_spec(module) is None + + +@pytest.mark.parametrize( + ("features", "quantizer_cfg", "error"), + [ + pytest.param(4, {"num_bits": 8}, "int8_wo", id="int8-weight-only"), + pytest.param( + 32, + { + "num_bits": (2, 1), + "block_sizes": {-1: 32, "type": "dynamic", "scale_bits": (8, 0)}, + }, + "weight-only MXFP4", + id="mxfp4-weight-only", + ), + pytest.param( + 128, + {"num_bits": (4, 3), "block_sizes": {-2: 64, -1: 128}}, + "only with 128x128 blocks", + id="fp8-rectangular-blocks", + ), + ], +) +def test_export_spec_rejects_unsupported_quantization(features, quantizer_cfg, error): + module = nn.Linear(features, features, bias=False) + module.weight_quantizer = TensorQuantizer(QuantizerAttributeConfig(**quantizer_cfg)) + module.input_quantizer = TensorQuantizer() + module.input_quantizer.disable() + + with pytest.raises(NotImplementedError, match=error): + get_quantized_weight_export_spec(module) + + +@pytest.mark.parametrize( + ("quantization_format", "block_size", "input_enabled"), + [("fp8", 0, True), ("fp8", 0, False), ("w4a16_nvfp4", 16, False)], +) +def test_export_spec_does_not_materialize_tensor_state( + quantization_format, block_size, input_enabled, monkeypatch +): + if quantization_format == "fp8": + module = _fp8_linear() + else: + weight = torch.arange(32, dtype=torch.float32).reshape(2, 16) / 32 + module = _static_w4a16_linear( + weight, weight.abs().amax(dim=1, keepdim=True), torch.tensor(1.0) + ) + monkeypatch.setattr( + weight_export, + "_state_tensor", + lambda *args, **kwargs: pytest.fail("export spec materialized tensor state"), + ) + + spec = get_quantized_weight_export_spec(module) + + assert spec is not None + assert spec.quantization_format == quantization_format + assert spec.block_size == block_size + if not input_enabled and module.input_quantizer.is_enabled: + module.input_quantizer.disable() + updated_spec = get_quantized_weight_export_spec(module) + assert updated_spec != spec + assert updated_spec.quantization_format == spec.quantization_format + assert updated_spec.block_size == spec.block_size + spec = updated_spec + assert spec.input_quantizer_enabled is input_enabled + + +def test_capture_uses_the_requested_weight_format_and_transposed_layout(): + module = _SquareTransposedExperts() + gate_state = capture_quantized_weight_export_state(module, "gate_up_proj") + down_state = capture_quantized_weight_export_state(module, "down_proj") + + assert gate_state.quantization_format == "fp8" + assert down_state.quantization_format == "w4a16_nvfp4" + assert gate_state.weight_shape == down_state.weight_shape == (2, 16, 16) + assert gate_state.packing_permutation == down_state.packing_permutation == (0, 2, 1) + + +def test_converted_hf_config_preserves_nvfp4_group_size(): + config = convert_hf_quant_config_format( + { + "quantization": { + "quant_algo": "W4A16_NVFP4", + "group_size": 32, + } + } + ) + + assert config["group_size"] == 32 + assert config["config_groups"]["group_0"]["weights"]["group_size"] == 32 diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py index 9c9467ba17e..1c747186a1e 100644 --- a/tests/unit/torch/export/test_unified_export_hf.py +++ b/tests/unit/torch/export/test_unified_export_hf.py @@ -38,7 +38,11 @@ postprocess_state_dict, sync_tied_input_amax, ) -from modelopt.torch.export.unified_export_hf import _resolve_export_dtype, read_unplaced_weights +from modelopt.torch.export.unified_export_hf import ( + _resolve_export_dtype, + _source_checkpoint, + read_unplaced_weights, +) from modelopt.torch.quantization.nn import TensorQuantizer @@ -1075,3 +1079,9 @@ def test_union_degrades_to_the_recorded_set_without_the_loader(tmp_path, monkeyp carried = read_unplaced_weights(model) assert key in carried, "recorded keys must still carry when the structural pass is unavailable" + + +def test_source_checkpoint_is_none_for_a_model_built_from_config(): + """``from_config`` leaves ``_name_or_path`` empty; it must not resolve to the working directory, + or every export would copy the cwd as the "source checkpoint".""" + assert _source_checkpoint(SimpleNamespace(config=SimpleNamespace(_name_or_path=""))) is None diff --git a/tests/unit/torch/quantization/plugins/test_moe_linear.py b/tests/unit/torch/models/test_step3p5_modeling_ptq.py similarity index 99% rename from tests/unit/torch/quantization/plugins/test_moe_linear.py rename to tests/unit/torch/models/test_step3p5_modeling_ptq.py index cceaacd44d5..66b52325c24 100644 --- a/tests/unit/torch/quantization/plugins/test_moe_linear.py +++ b/tests/unit/torch/models/test_step3p5_modeling_ptq.py @@ -27,13 +27,13 @@ import modelopt.torch.quantization as mtq from modelopt.torch.export.hf_export_handlers import _export_moe_linear from modelopt.torch.export.registry import ExportContext, ExportModuleRegistry -from modelopt.torch.quantization.nn import QuantModuleRegistry -from modelopt.torch.quantization.plugins.huggingface import ( +from modelopt.torch.models.step3p5.modeling_ptq import ( _is_expert_indexed_moe_linear, _QuantMoELinear, _reconstruct_fused_moe_linear, register_moe_linear_on_the_fly, ) +from modelopt.torch.quantization.nn import QuantModuleRegistry NUM_EXPERTS = 4 HIDDEN_SIZE = 32 diff --git a/tests/unit/torch/nas/test_full_algorithms_vision.py b/tests/unit/torch/nas/test_full_algorithms_vision.py index d217b79909f..f6c30c1da52 100644 --- a/tests/unit/torch/nas/test_full_algorithms_vision.py +++ b/tests/unit/torch/nas/test_full_algorithms_vision.py @@ -111,7 +111,7 @@ def run_search(): @pytest.mark.parametrize( ("flops_only", "bounded_latency"), [ - (False, True), # fast representative kept enabled for sanity + pytest.param(False, True, marks=pytest.mark.manual), pytest.param(True, True, marks=pytest.mark.manual), pytest.param(True, False, marks=pytest.mark.manual), pytest.param(False, False, marks=pytest.mark.manual), diff --git a/tests/unit/torch/nas/test_nas.py b/tests/unit/torch/nas/test_nas.py index 3d40b3a51d5..46b44a9ce3a 100644 --- a/tests/unit/torch/nas/test_nas.py +++ b/tests/unit/torch/nas/test_nas.py @@ -55,6 +55,7 @@ def get_data_loader(num_batches): yield torch.rand(2, 3, 56, 56), int(torch.randint(1000, (1,))) +@pytest.mark.manual # AutoNAS is not actively maintained; ~3s def test_automode(test_config): # Use model class to create a model inside mtn.convert model = mtn.convert((TinyMobileNetFeatures, (), {}), mode=[("autonas", test_config)]) @@ -118,6 +119,7 @@ def test_automode(test_config): assert not torch.allclose(out_dyn, out2) +@pytest.mark.manual # AutoNAS is not actively maintained; ~3s each @pytest.mark.parametrize("native_eval", [False, True]) def test_prep_for_eval(native_eval, test_config): """Test all four ways of calling prep_for_eval. @@ -288,6 +290,7 @@ def _get_current_bn_stats(): assert max_config == mtn.get_subnet_config(model) +@pytest.mark.manual # AutoNAS is not actively maintained; ~4s def test_set_auto_mode(test_config): """Test whether we can correctly set automode on/off.""" # setup model in autonas mode and training mode @@ -412,15 +415,24 @@ def _decorated_compare(): torch.randn(1, 16, 8, 8), ["autonas", "export_nas"], ), - ( + pytest.param( TinyMobileNetFeatures, (), "", False, torch.randn(1, 3, 64, 64), ["autonas", "export_nas"], + marks=pytest.mark.manual, # 5s; the InvertedResidual cases cover the same paths + ), + pytest.param( + TinyMobileNetFeatures, + (), + "", + False, + torch.randn(1, 3, 64, 64), + ["autonas"], + marks=pytest.mark.manual, ), - (TinyMobileNetFeatures, (), "", False, torch.randn(1, 3, 64, 64), ["autonas"]), (TinyMobileNetFeatures, (), "", False, torch.randn(1, 3, 64, 64), []), ], ) diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 6e7375dd695..4f224e3b7c0 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# 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"); @@ -15,6 +15,10 @@ """Tests for _QuantFusedExperts: generic fused MoE quantization and export.""" +import copy +import warnings +from contextlib import nullcontext + import pytest import torch import torch.nn as nn @@ -23,16 +27,18 @@ pytest.importorskip("transformers") from _test_utils.torch.quantization.tied_modules import tie_fused_experts_3d_params +from transformers import PretrainedConfig +from transformers.integrations.finegrained_fp8 import FP8Experts import modelopt.torch.quantization as mtq import modelopt.torch.quantization.nn.modules.tensor_quantizer as tensor_quantizer_module from modelopt.torch.export.model_utils import _release_exported_tensors -from modelopt.torch.export.moe_utils import _export_fused_experts +from modelopt.torch.export.moe_utils import _export_fused_experts, _fallback_amax from modelopt.torch.export.quant_utils import get_quant_config, get_quantization_format from modelopt.torch.quantization.config import QuantizerAttributeConfig from modelopt.torch.quantization.conversion import _normalize_fused_experts_quantizer_name from modelopt.torch.quantization.model_calib import local_hessian_calibrate -from modelopt.torch.quantization.nn import QuantModuleRegistry, TensorQuantizer +from modelopt.torch.quantization.nn import QuantModuleRegistry, SequentialQuantizer, TensorQuantizer from modelopt.torch.quantization.plugins.huggingface import ( _fused_experts_wrapper_class, _is_fused_experts_module, @@ -45,6 +51,71 @@ ) from modelopt.torch.quantization.utils.core_utils import weight_attr_names + +def test_native_fp8_experts_dequantized_before_nvfp4_export(): + """Calibrate and pack scaled FP8 gate/up and down weights, preserving their BF16 values.""" + generator = torch.Generator().manual_seed(42) + experts = FP8Experts( + PretrainedConfig(hidden_size=256, intermediate_size=256, num_experts=2, hidden_act="silu"), + block_size=(128, 128), + ) + reference = {} + for name, scale_range in (("gate_up_proj", (0.01, 0.08)), ("down_proj", (0.15, 0.3))): + weight = getattr(experts, name) + scales = getattr(experts, f"{name}_scale_inv") + assert weight.dtype == torch.float8_e4m3fn + with torch.no_grad(): + weight.copy_(torch.randn(weight.shape, generator=generator).to(weight.dtype)) + scales.copy_(torch.linspace(*scale_range, scales.numel()).reshape(scales.shape)) + # Build only the reference here; ModelOpt must consume the untouched FP8 source. + blocks = weight.float().reshape(2, weight.shape[1] // 128, 128, 2, 128) + reference[name] = (blocks * scales[:, :, None, :, None]).reshape(weight.shape).bfloat16() + + registered_before = QuantModuleRegistry.get(FP8Experts) + try: + mtq.quantize(experts, mtq.NVFP4_DEFAULT_CFG) + for name, expected in reference.items(): + assert getattr(experts, name).dtype == torch.float8_e4m3fn + quantizers = getattr(experts, f"{name}_weight_quantizers") + for idx, quantizer in enumerate(quantizers): + assert quantizer.is_enabled and quantizer.num_bits == (2, 1) + torch.testing.assert_close( + quantizer.amax.float(), + expected[idx].float().abs().amax(), + rtol=0.004, + atol=0, + ) + + _export_fused_experts(experts, torch.bfloat16) + levels = torch.tensor([0, 0.5, 1, 1.5, 2, 3, 4, 6, 0, -0.5, -1, -1.5, -2, -3, -4, -6]) + for idx in range(2): + exported = getattr(experts, str(idx)) + torch.testing.assert_close( + exported.gate_proj.weight_scale_2, exported.up_proj.weight_scale_2 + ) + for name, expected in ( + ("gate_proj", reference["gate_up_proj"][idx, :256]), + ("up_proj", reference["gate_up_proj"][idx, 256:]), + ("down_proj", reference["down_proj"][idx]), + ): + projection = getattr(exported, name) + packed = projection.weight + assert packed.dtype == torch.uint8 + codes = torch.stack((packed & 15, packed >> 4), dim=-1).reshape(expected.shape) + decoded = ( + levels[codes.long()] + * projection.weight_scale.float().repeat_interleave(16, dim=-1) + * projection.weight_scale_2.float() + ) + relative_l2 = torch.linalg.vector_norm( + decoded - expected.float() + ) / torch.linalg.vector_norm(expected.float()) + assert relative_l2 < 0.15, f"Expert {idx} {name}: relative L2 {relative_l2}" + finally: + if registered_before is None and QuantModuleRegistry.get(FP8Experts) is not None: + QuantModuleRegistry.unregister(FP8Experts) + + # --------------------------------------------------------------------------- # Synthetic fused expert module matching the HF transformers 5.0+ pattern # --------------------------------------------------------------------------- @@ -696,6 +767,156 @@ def _spy_export(wrapper, dtype, **_kwargs): if QuantModuleRegistry.get(expert_type) is not None: QuantModuleRegistry.unregister(expert_type) + @pytest.mark.parametrize("amax_state", ["calibrated", "zero", "missing"]) + def test_w4a8_sequential_weight_quantizer_exports_split_block_amax( + self, monkeypatch, amax_state + ): + """W4A8 AWQ weights use a SequentialQuantizer: an INT4 block stage (flat per-block + ``_amax`` of shape ``(rows * blocks, 1)``, as block-quant calibration records it) + followed by a per-tensor FP8 stage. Export must treat each stage on its own: slice the + block amax into the gate/up halves (updating the stale ``_amax_shape_for_export``), and + for experts that saw no calibration tokens derive a per-block fallback instead of + forcing a scalar into the block stage (``Changing shape when setting amax``). + """ + block = 16 + experts = _SyntheticFusedExperts() + expert_type = type(experts) + if QuantModuleRegistry.get(expert_type) is None: + QuantModuleRegistry.register({expert_type: "test.SyntheticFusedExperts"})( + _QuantFusedExperts + ) + try: + converted = QuantModuleRegistry.convert(experts) + + def _w4a8_quantizer(rows, fused): + int4 = TensorQuantizer( + QuantizerAttributeConfig(num_bits=4, block_sizes={-1: block, "type": "static"}) + ) + fp8 = TensorQuantizer(QuantizerAttributeConfig(num_bits=(4, 3))) + # gate_up_proj reduces over HIDDEN_DIM, down_proj over INTERMEDIATE_DIM. + blocks = (HIDDEN_DIM if fused else INTERMEDIATE_DIM) // block + int4._amax_shape_for_export = (rows, -1) + if amax_state != "missing": + # Row r of gate_up carries the marker r + 1, so a wrong slice is visible. + marker = torch.arange(1, rows + 1).float().repeat_interleave(blocks)[:, None] + int4.amax = ( + marker.clone() if amax_state == "calibrated" else torch.zeros_like(marker) + ) + fp8.amax = torch.tensor(2.0 if amax_state == "calibrated" else 0.0) + return SequentialQuantizer(int4, fp8) + + for idx in range(NUM_EXPERTS): + converted.gate_up_proj_weight_quantizers[idx] = _w4a8_quantizer( + 2 * INTERMEDIATE_DIM, fused=True + ) + converted.down_proj_weight_quantizers[idx] = _w4a8_quantizer( + HIDDEN_DIM, fused=False + ) + + seen = [] + + def _spy_export(wrapper, dtype, **_kwargs): + int4, fp8 = wrapper.weight_quantizer + seen.append( + ( + wrapper.weight.shape[0], + int4.export_amax().detach().clone(), + fp8._amax.detach().clone(), + ) + ) + + monkeypatch.setattr( + "modelopt.torch.export.unified_export_hf._export_quantized_weight", + _spy_export, + ) + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + _export_fused_experts(converted, torch.float16) + fallback_warned = any("was not calibrated" in str(w.message) for w in caught) + assert fallback_warned == (amax_state != "calibrated") + + assert len(seen) == 3 * NUM_EXPERTS + for rows, int4_amax, fp8_amax in seen: + assert int4_amax.shape[0] == rows # sliced to the projection, not the fused rows + assert torch.all(int4_amax > 0) + assert torch.all(fp8_amax > 0) + # Export order per expert is gate, up, down; gate and up must share the FP8 stage. + for gate, up in zip(seen[0::3], seen[1::3]): + assert torch.equal(gate[2], up[2]) + if amax_state == "calibrated": + gate_up_rows = {int4_amax[0, 0].item() for rows, int4_amax, _ in seen if rows == 16} + # gate starts at row marker 1, up at INTERMEDIATE_DIM + 1. + assert gate_up_rows == {1.0, float(INTERMEDIATE_DIM + 1)} + finally: + if QuantModuleRegistry.get(expert_type) is not None: + QuantModuleRegistry.unregister(expert_type) + + @pytest.mark.parametrize( + ("amax_layout", "expected_gate_shape", "expected_down_shape"), + [ + ("per_tensor", (1,), (1,)), + ("per_channel", (INTERMEDIATE_DIM, 1), (HIDDEN_DIM, 1)), + ("per_row_1d", (INTERMEDIATE_DIM,), (HIDDEN_DIM,)), + ("per_block", (INTERMEDIATE_DIM, 2), (HIDDEN_DIM, 2)), + # What block-quant calibration records: (rows * blocks, 1). + ("flat_block", (INTERMEDIATE_DIM, 2), (HIDDEN_DIM, 2)), + ], + ) + def test_uncalibrated_zero_amax_fallback_keeps_layout( + self, monkeypatch, amax_layout, expected_gate_shape, expected_down_shape + ): + """The weight-derived fallback must reproduce the layout of the quantizer's zero amax.""" + experts = _SyntheticFusedExperts() + expert_type = type(experts) + if QuantModuleRegistry.get(expert_type) is None: + QuantModuleRegistry.register({expert_type: "test.SyntheticFusedExperts"})( + _QuantFusedExperts + ) + try: + converted = QuantModuleRegistry.convert(experts) + layouts = { + "per_tensor": lambda rows: (1,), + "per_channel": lambda rows: (rows, 1), + "per_row_1d": lambda rows: (rows,), + "per_block": lambda rows: (rows, 2), + "flat_block": lambda rows: (rows * 2, 1), + } + for idx in range(NUM_EXPERTS): + for wq, rows in ( + (converted.gate_up_proj_weight_quantizers[idx], 2 * INTERMEDIATE_DIM), + (converted.down_proj_weight_quantizers[idx], HIDDEN_DIM), + ): + wq._disabled = False + wq.amax = torch.zeros(*layouts[amax_layout](rows)) + + seen = [] + + def _spy_export(wrapper, dtype, **_kwargs): + seen.append(wrapper.weight_quantizer._amax.detach().clone()) + + monkeypatch.setattr( + "modelopt.torch.export.unified_export_hf._export_quantized_weight", + _spy_export, + ) + with pytest.warns(UserWarning, match="was not calibrated"): + _export_fused_experts(converted, torch.float16) + + assert all(torch.all(a > 0) for a in seen) + # Export order per expert is gate, up, down. + for gate, up, down in zip(seen[0::3], seen[1::3], seen[2::3]): + assert gate.shape == up.shape == expected_gate_shape + assert down.shape == expected_down_shape + finally: + if QuantModuleRegistry.get(expert_type) is not None: + QuantModuleRegistry.unregister(expert_type) + + def test_fallback_amax_refuses_layouts_it_cannot_group(self): + quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=8, axis=1)) + quantizer.amax = torch.zeros(1, HIDDEN_DIM) # per input channel: not a (rows, blocks) grid + with pytest.raises(ValueError, match="Cannot derive amax"): + _fallback_amax(torch.randn(INTERMEDIATE_DIM, HIDDEN_DIM), quantizer) + # --------------------------------------------------------------------------- # Tests for tied-experts dedup in _export_fused_experts @@ -1574,3 +1795,191 @@ def test_quantize_collects_amax(self): finally: if registered_before is None and QuantModuleRegistry.get(experts_type) is not None: QuantModuleRegistry.unregister(experts_type) + + +@pytest.mark.parametrize("activation_scheme", ["dynamic", "static"]) +@pytest.mark.parametrize("fold", [True, False]) +@pytest.mark.parametrize("has_gate", [True, False]) +@pytest.mark.parametrize("block_size", [(128, 128), (64, 128), None]) +def test_native_fp8_experts_scales(has_gate, block_size, fold, activation_scheme): + torch.manual_seed(42) + config = PretrainedConfig( + hidden_size=256, intermediate_size=256, num_experts=2, hidden_act="silu" + ) + experts = FP8Experts( + config, block_size=block_size, has_gate=has_gate, activation_scheme=activation_scheme + ) + first = "gate_up_proj" if has_gate else "up_proj" + expected = {} + for name in (first, "down_proj"): + weight = getattr(experts, name) + scales = getattr(experts, f"{name}_scale_inv") + with torch.no_grad(): + weight.copy_(torch.randn(weight.shape).to(weight.dtype)) + scales.copy_(torch.linspace(1e-8, 0.08, scales.numel()).reshape(scales.shape)) + rows, cols = block_size or (weight.shape[1] // scales.shape[1], weight.shape[2]) + expanded = scales.float().repeat_interleave(rows, 1).repeat_interleave(cols, 2) + expected[name] = weight.float() * expanded[:, : weight.shape[1], : weight.shape[2]] + + reference = copy.deepcopy(experts) + for name, weight in expected.items(): + setattr(reference, name, nn.Parameter(weight, requires_grad=False)) + mtq.quantize( + experts, {"quant_cfg": [{"quantizer_name": "*", "enable": False}], "algorithm": None} + ) + for (weight, _), expected_weight in zip( + experts.iter_weights_for_calibration(), + [w for name in (first, "down_proj") for w in expected[name]], + strict=True, + ): + torch.testing.assert_close(weight.float(), expected_weight) + + for name in (first, "down_proj"): + for idx in range(2): + for dtype in (torch.float16, torch.bfloat16): + torch.testing.assert_close( + experts.get_expert_weight(name, idx, dtype), + expected[name][idx].to(dtype), + rtol=0, + atol=0, + ) + + inputs = torch.randn(2, 256) + indices = torch.tensor([[0], [1]]) + routing = torch.ones(2, 1) + torch.testing.assert_close( + experts(inputs, indices, routing), reference(inputs, indices, routing) + ) + for _, quantizer in experts.iter_weights_for_calibration(): + quantizer.enable() + mtq.calibrate(experts, algorithm="max", forward_loop=None) + for weight, quantizer in experts.iter_weights_for_calibration(): + torch.testing.assert_close(quantizer.amax, weight.abs().amax()) + quantizer.disable() + if fold: + experts.unpack_weight() + experts.fold_weight() + torch.testing.assert_close( + experts(inputs.bfloat16(), indices, routing), + copy.deepcopy(reference).bfloat16()(inputs.bfloat16(), indices, routing), + ) + torch.testing.assert_close( + experts(inputs, indices, routing), reference(inputs, indices, routing) + ) + _export_fused_experts(experts, torch.float32) + for idx in range(2): + exported = getattr(experts, str(idx)) + if has_gate: + torch.testing.assert_close(exported.gate_proj.weight, expected[first][idx, :256]) + torch.testing.assert_close(exported.up_proj.weight, expected[first][idx, 256:]) + else: + torch.testing.assert_close(exported.up_proj.weight, expected[first][idx]) + torch.testing.assert_close(exported.down_proj.weight, expected["down_proj"][idx]) + assert not any("scale_inv" in key or "activation_scale" in key for key in experts.state_dict()) + + +@pytest.mark.parametrize("has_gate", [True, False]) +@pytest.mark.parametrize("unpacked_projection", [None, "first", "down_proj"]) +@pytest.mark.parametrize("fold_api", ["direct", "permanent", "temporary"]) +def test_native_fp8_experts_reject_packed_fold(has_gate, unpacked_projection, fold_api): + torch.manual_seed(42) + experts = FP8Experts( + PretrainedConfig(hidden_size=32, intermediate_size=32, num_experts=2, hidden_act="silu"), + block_size=None, + has_gate=has_gate, + ) + first = "gate_up_proj" if has_gate else "up_proj" + for name in (first, "down_proj"): + weight = getattr(experts, name) + with torch.no_grad(): + weight.copy_(torch.randn(weight.shape).to(weight.dtype)) + getattr(experts, f"{name}_scale_inv").fill_(0.1) + # Put a foldable module first to check model-wide validation precedes all mutation. + model = nn.Sequential(nn.Linear(32, 32), experts) + mtq.quantize(model, mtq.INT8_DEFAULT_CFG, forward_loop=None) + if unpacked_projection is not None: + name = first if unpacked_projection == "first" else unpacked_projection + weight = torch.stack([experts.get_expert_weight(name, idx) for idx in range(2)]) + setattr(experts, name, nn.Parameter(weight, requires_grad=False)) + params = dict(model.named_parameters()) + state = copy.deepcopy(model.state_dict()) + quantizers = [q for q in model.modules() if isinstance(q, TensorQuantizer)] + enabled = [q.is_enabled for q in quantizers] + inputs = torch.randn(2, 32) + indices = torch.tensor([[0], [1]]) + routing = torch.ones(2, 1) + expected = experts(inputs, indices, routing) + + with pytest.raises(ValueError, match=r"Call unpack_weight\(\) before folding"): + if fold_api == "direct": + experts.fold_weight() + elif fold_api == "permanent": + mtq.fold_weight(model) + else: + with mtq.temporarily_fold_weights(model): + pytest.fail("Packed FP8 weights must be rejected before entering the context") + + assert [q.is_enabled for q in quantizers] == enabled + for name, param in model.named_parameters(): + assert param is params[name] + assert param.dtype == state[name].dtype + # Compare FP8 payloads in FP32 because CPU close/equality kernels do not support FP8. + torch.testing.assert_close( + {name: value.float() for name, value in model.state_dict().items()}, + {name: value.float() for name, value in state.items()}, + rtol=0, + atol=0, + ) + torch.testing.assert_close(experts(inputs, indices, routing), expected) + + +@pytest.mark.parametrize("has_gate", [True, False]) +@pytest.mark.parametrize("exception", [True, False]) +@pytest.mark.parametrize("snapshot_device", [None, "cpu"]) +def test_native_fp8_experts_unpacked_temporary_fold(has_gate, exception, snapshot_device): + torch.manual_seed(42) + config = PretrainedConfig( + hidden_size=32, intermediate_size=32, num_experts=2, hidden_act="silu" + ) + experts = FP8Experts(config, block_size=None, has_gate=has_gate) + first = "gate_up_proj" if has_gate else "up_proj" + for name in (first, "down_proj"): + weight = getattr(experts, name) + scales = getattr(experts, f"{name}_scale_inv") + with torch.no_grad(): + weight.copy_(torch.randn(weight.shape).to(weight.dtype)) + scales.copy_(torch.linspace(0.01, 0.08, scales.numel()).reshape(scales.shape)) + mtq.quantize( + experts, + { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*weight_quantizer", "cfg": {"num_bits": 4, "axis": None}}, + ], + "algorithm": "max", + }, + ) + experts.unpack_weight() + params = {name: getattr(experts, name) for name in (first, "down_proj")} + state = copy.deepcopy(experts.state_dict()) + quantizers = [q for _, q in experts.iter_weights_for_calibration()] + inputs = torch.randn(2, 32) + indices = torch.tensor([[0], [1]]) + routing = torch.ones(2, 1) + expected = experts(inputs, indices, routing) + expected_exit = pytest.raises(RuntimeError, match="fold exit") if exception else nullcontext() + with expected_exit, mtq.temporarily_fold_weights(experts, snapshot_device=snapshot_device): + assert all(not q.is_enabled for q in quantizers) + for name, param in params.items(): + assert getattr(experts, name) is param + assert param.dtype == torch.get_default_dtype() + assert not torch.equal(param, state[name]) + torch.testing.assert_close(experts(inputs, indices, routing), expected) + if exception: + raise RuntimeError("fold exit") + for name, param in params.items(): + assert getattr(experts, name) is param + assert param.dtype == torch.get_default_dtype() + assert all(q.is_enabled for q in quantizers) + torch.testing.assert_close(experts.state_dict(), state, rtol=0, atol=0) + torch.testing.assert_close(experts(inputs, indices, routing), expected) diff --git a/tests/unit/torch/quantization/test_autoquant.py b/tests/unit/torch/quantization/test_autoquant.py index f61cbf7b893..5ee55d81081 100644 --- a/tests/unit/torch/quantization/test_autoquant.py +++ b/tests/unit/torch/quantization/test_autoquant.py @@ -21,7 +21,6 @@ import pytest import torch -from _test_utils.torch.distributed.utils import spawn_multiprocess_job from _test_utils.torch.quantization.models import SimpleConv, SimpleConvLinear, SimpleLinear import modelopt.torch.opt as mto @@ -793,24 +792,29 @@ def test_active_moe_search_prefers_budget_lower_bound(): } +_BLOCKWISE_INT8 = [mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG, mtq.INT8_DEFAULT_CFG] +_AWQ_SMOOTHQUANT = [mtq.INT4_AWQ_CFG, mtq.INT8_SMOOTHQUANT_CFG] +_AWQ_CUSTOM = [mtq.INT4_AWQ_CFG, INT8_CUSTOM_QUANT_TEST_CFG] +_SMOOTHQUANT_ONLY = [mtq.INT8_SMOOTHQUANT_CFG] +_NONE_SMOOTHQUANT = [None, mtq.INT8_SMOOTHQUANT_CFG] + + +# A covering set rather than the full model x formats x method product: every model, format +# list and method appears at least once, and the methods only differ in how candidates are scored. @pytest.mark.parametrize( - "model_cls", - [SimpleConv, SimpleConvLinear, SimpleLinear, TransformerBlock], -) -@pytest.mark.parametrize( - ("search_formats", "min_bits", "search_bits"), + ("model_cls", "search_formats", "min_bits", "search_bits", "method"), [ - ([mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG, mtq.INT8_DEFAULT_CFG], 4.0, 6.0), - ([mtq.INT4_AWQ_CFG, mtq.INT8_SMOOTHQUANT_CFG], 4.0, 6.0), - ([mtq.INT4_AWQ_CFG, INT8_CUSTOM_QUANT_TEST_CFG], 4.0, 6.0), - ([mtq.INT8_SMOOTHQUANT_CFG], 8.0, 11.0), - ([None, mtq.INT8_SMOOTHQUANT_CFG], 8.0, 11.0), + (SimpleConv, _BLOCKWISE_INT8, 4.0, 6.0, "gradient"), + (SimpleConvLinear, _AWQ_SMOOTHQUANT, 4.0, 6.0, "gradient"), + (TransformerBlock, _NONE_SMOOTHQUANT, 8.0, 11.0, "gradient"), + (SimpleLinear, _AWQ_CUSTOM, 4.0, 6.0, "kl_div"), + (TransformerBlock, _SMOOTHQUANT_ONLY, 8.0, 11.0, "kl_div"), + (SimpleConvLinear, _BLOCKWISE_INT8, 4.0, 6.0, "kl_div"), + (SimpleConv, _AWQ_SMOOTHQUANT, 4.0, 6.0, "aumann_shapley"), + (SimpleLinear, _NONE_SMOOTHQUANT, 8.0, 11.0, "aumann_shapley"), + (TransformerBlock, _AWQ_CUSTOM, 4.0, 6.0, "aumann_shapley"), ], ) -@pytest.mark.parametrize( - "method", - ["gradient", "kl_div", "aumann_shapley"], -) def test_auto_quantize(model_cls, search_formats, min_bits, search_bits, method): model = model_cls() @@ -973,9 +977,8 @@ def _test_data_parallel_auto_quantize(rank, size): assert search_history["best"]["is_satisfied"] -def test_data_parallel_auto_quantize(skip_on_windows): - # 2 ranks fully exercise the cross-rank sync the test asserts; more just adds spawn overhead. - spawn_multiprocess_job(2, _test_data_parallel_auto_quantize, backend="gloo") +def test_data_parallel_auto_quantize(cpu_dist_workers): + cpu_dist_workers.run(_test_data_parallel_auto_quantize) def _test_data_parallel_moe_score_module(rank, size): @@ -1036,8 +1039,8 @@ def _test_data_parallel_moe_score_module(rank, size): assert scores == rank_zero_scores -def test_data_parallel_moe_score_module(skip_on_windows): - spawn_multiprocess_job(2, _test_data_parallel_moe_score_module, backend="gloo") +def test_data_parallel_moe_score_module(cpu_dist_workers): + cpu_dist_workers.run(_test_data_parallel_moe_score_module) def test_score_hparam_registration_preserves_order(): diff --git a/tests/unit/torch/quantization/test_autoquant_shapley.py b/tests/unit/torch/quantization/test_autoquant_shapley.py index b2a61fb4b1b..6752fcee65c 100644 --- a/tests/unit/torch/quantization/test_autoquant_shapley.py +++ b/tests/unit/torch/quantization/test_autoquant_shapley.py @@ -22,7 +22,6 @@ """ import copy -import functools import itertools import math from collections import namedtuple @@ -31,7 +30,6 @@ import pytest import torch import torch.nn.functional as F -from _test_utils.torch.distributed.utils import spawn_multiprocess_job import modelopt.torch.quantization as mtq import modelopt.torch.quantization.model_quant as _model_quant @@ -775,7 +773,7 @@ def _shapley_data_parallel(rank, size, baseline): assert got == pytest.approx(scores, rel=1e-3, abs=1e-9), f"{name}: {got} vs {scores}" -def test_data_parallel_aumann_shapley(skip_on_windows): +def test_data_parallel_aumann_shapley(cpu_dist_workers): """Data parallel aumann shapley.""" _model, single = _search( _Block(seed=0), method={"method": "aumann_shapley", "num_path_nodes": 2} @@ -785,9 +783,7 @@ def test_data_parallel_aumann_shapley(skip_on_windows): "f_corner": single["damage_model"]["f_corner"], "scores": {name: stat["scores"] for name, stat in single["candidate_stats"].items()}, } - spawn_multiprocess_job( - 2, functools.partial(_shapley_data_parallel, baseline=baseline), backend="gloo" - ) + cpu_dist_workers.run(_shapley_data_parallel, baseline=baseline) def test_anchor_ceiling_rejects_non_finite_measurements(): diff --git a/tests/unit/torch/quantization/test_dist.py b/tests/unit/torch/quantization/test_dist.py index 4130b8ce82a..0eb5329805e 100644 --- a/tests/unit/torch/quantization/test_dist.py +++ b/tests/unit/torch/quantization/test_dist.py @@ -15,7 +15,6 @@ import torch import torch.distributed as dist -from _test_utils.torch.distributed.utils import spawn_multiprocess_job from _test_utils.torch.quantization.models import SimpleLinear from torch.nn.parallel import DistributedDataParallel @@ -40,8 +39,7 @@ def forward_loop(model): amax = module.amax.clone() dist.all_reduce(amax, op=dist.ReduceOp.MAX) assert torch.allclose(amax, module.amax) - dist.destroy_process_group() -def test_data_parallel(skip_on_windows): - spawn_multiprocess_job(2, _test_data_parallel_helper, backend="gloo") +def test_data_parallel(cpu_dist_workers): + cpu_dist_workers.run(_test_data_parallel_helper) diff --git a/tests/unit/torch/quantization/test_dynamo_onnx_export.py b/tests/unit/torch/quantization/test_dynamo_onnx_export.py new file mode 100644 index 00000000000..1c0b64b7210 --- /dev/null +++ b/tests/unit/torch/quantization/test_dynamo_onnx_export.py @@ -0,0 +1,226 @@ +# 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. + +"""CPU contract tests for the supported Dynamo ONNX export path.""" + +import inspect + +import pytest +import torch +from torch import nn + +onnx = pytest.importorskip("onnx") +pytest.importorskip("onnxscript") + +from modelopt.torch.quantization._dynamo_onnx import _get_dynamo_onnx_translation_table + +_NVFP4_ARGS = (4, 2, 8, 4, "Float") +_MXFP8_ARGS = (8, 4, 9, 8, "Float") + + +class _StrictQuantOp(nn.Module): + def __init__(self, quant_format, axis=0): + super().__init__() + self.quant_format = quant_format + self.axis = axis + + def forward(self, x, amax): + if self.quant_format.startswith("fp8"): + return torch.ops.tensorrt.quantize_op.default( + x, + None if self.quant_format == "fp8_none" else amax, + 8, + 4, + False, + False, + "Float", + ) + if self.quant_format in {"int8", "uint8"}: + return torch.ops.tensorrt.quantize_op.default( + x, amax, 8, 0, self.quant_format == "uint8", False, "Float", None, self.axis + ) + if self.quant_format == "int4": + return torch.ops.tensorrt.quantize_op.default( + x, amax, 4, 0, False, True, "Float", 16, -1 + ) + if self.quant_format.startswith("nvfp4"): + return torch.ops.tensorrt.dynamic_block_quantize_op.default( + x, 16, amax, *_NVFP4_ARGS, self.quant_format.removeprefix("nvfp4_") + ) + return torch.ops.tensorrt.dynamic_block_quantize_op.overload( + x, 32, None, *_MXFP8_ARGS, self.quant_format.removeprefix("mxfp8_") + ) + + +def _assert_valid_without_custom_ops(model): + onnx.checker.check_model(model, full_check=True) + assert not model.functions + assert not any( + node.domain == "tensorrt" or node.op_type in {"quantize_op", "dynamic_block_quantize_op"} + for node in model.graph.node + ) + return model + + +def _raw_export(model, inputs, path, opset=23): + compatibility_options = ( + {"fallback": False} if "fallback" in inspect.signature(torch.onnx.export).parameters else {} + ) + torch.onnx.export( + model, + inputs, + path, + dynamo=True, + opset_version=opset, + custom_translation_table=_get_dynamo_onnx_translation_table(), + **compatibility_options, + ) + exported = onnx.load(path) + return _assert_valid_without_custom_ops(exported) + + +def _attribute(node, name): + return onnx.helper.get_attribute_value( + next(attr for attr in node.attribute if attr.name == name) + ) + + +def _constant(model, name): + initializer = next((item for item in model.graph.initializer if item.name == name), None) + if initializer is not None: + return onnx.numpy_helper.to_array(initializer) + node = next( + item for item in model.graph.node if item.op_type == "Constant" and name in item.output + ) + return onnx.numpy_helper.to_array(_attribute(node, "value")) + + +_STRICT_CASES = [ + ("fp8_none", ("trt", "TRT_FP8QuantizeLinear")), + ("fp8", ("trt", "TRT_FP8QuantizeLinear")), + ("int8", ("", "QuantizeLinear")), + ("uint8", ("", "QuantizeLinear")), + ("int4", ("trt", "DequantizeLinear")), + ("nvfp4_static", ("trt", "TRT_FP4QDQ")), + ("mxfp8_static", ("trt", "TRT_MXFP8DequantizeLinear")), + ("mxfp8_dynamic", ("trt", "TRT_MXFP8DynamicQuantize")), +] + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) +@pytest.mark.parametrize(("quant_format", "expected_node"), _STRICT_CASES) +def test_private_table_strict_formats(tmp_path, quant_format, expected_node, dtype): + sample_input = torch.randn(4, 32, dtype=dtype) + if quant_format in {"int8", "uint8"}: + amax = torch.ones(4, 1) + elif quant_format == "int4": + amax = torch.ones(4, 2) + else: + amax = torch.tensor(1.0) + exported_program = torch.export.export( + _StrictQuantOp(quant_format), (sample_input, amax), strict=True + ) + exported = _raw_export(exported_program, (), tmp_path / f"{quant_format}.onnx") + assert expected_node in {(node.domain, node.op_type) for node in exported.graph.node} + expected_dtype = onnx.TensorProto.FLOAT if dtype == torch.float32 else onnx.TensorProto.FLOAT16 + assert exported.graph.output[0].type.tensor_type.elem_type == expected_dtype + tensor_types = { + value.name: value.type.tensor_type.elem_type + for value in (*exported.graph.input, *exported.graph.value_info, *exported.graph.output) + } + if quant_format in {"int8", "uint8"}: + quantizer = next(node for node in exported.graph.node if node.op_type == "QuantizeLinear") + assert _attribute(quantizer, "axis") == 0 + unsigned = quant_format == "uint8" + assert tensor_types[quantizer.input[2]] == ( + onnx.TensorProto.UINT8 if unsigned else onnx.TensorProto.INT8 + ) + divisor = next(node for node in exported.graph.node if node.op_type == "Div").input[1] + assert _constant(exported, divisor).item() == (255.0 if unsigned else 127.0) + elif quant_format == "int4": + dequantizer = next( + node for node in exported.graph.node if node.op_type == "DequantizeLinear" + ) + assert _attribute(dequantizer, "axis") == -1 + assert _attribute(dequantizer, "block_size") == 16 + assert tensor_types[dequantizer.output[0]] == onnx.TensorProto.FLOAT + divisor = next(node for node in exported.graph.node if node.op_type == "Div").input[1] + assert _constant(exported, divisor).item() == 7.0 + elif quant_format == "fp8": + divisor = next(node for node in exported.graph.node if node.op_type == "Div").input[1] + assert _constant(exported, divisor).item() == 448.0 + + +@pytest.mark.parametrize("amax_shape", [(), (1,), (1, 1)]) +@pytest.mark.parametrize("axis", [0, -1]) +def test_private_table_int8_singleton_amax_is_per_tensor(tmp_path, amax_shape, axis): + inputs = torch.randn(4, 32) + exported_program = torch.export.export( + _StrictQuantOp("int8", axis), (inputs, torch.ones(amax_shape)), strict=True + ) + exported = _raw_export(exported_program, (), tmp_path / "int8_per_tensor.onnx") + tensor_shapes = { + value.name: value.type.tensor_type.shape + for value in (*exported.graph.input, *exported.graph.value_info, *exported.graph.output) + } + for node in exported.graph.node: + if node.op_type in {"QuantizeLinear", "DequantizeLinear"}: + assert len(tensor_shapes[node.input[1]].dim) == 0 + assert len(tensor_shapes[node.input[2]].dim) == 0 + + +def test_private_table_rejects_int4_without_amax(tmp_path): + exported_program = torch.export.export( + _StrictQuantOp("int4"), (torch.randn(4, 32), None), strict=True + ) + with pytest.raises(torch.onnx.OnnxExporterError, match="INT4 ONNX export requires amax"): + _raw_export(exported_program, (), tmp_path / "int4_without_amax.onnx") + + +def test_private_table_preserves_nvfp4_carrier_dtype(tmp_path): + inputs = torch.randn(4, 32, dtype=torch.float16) + exported_program = torch.export.export( + _StrictQuantOp("nvfp4_dynamic"), (inputs, torch.tensor(1.0)), strict=True + ) + + exported = _raw_export(exported_program, (), tmp_path / "nvfp4_dtype.onnx", 23) + + assert exported.graph.output[0].type.tensor_type.elem_type == onnx.TensorProto.FLOAT16 + assert {node.domain for node in exported.graph.node if node.op_type == "DequantizeLinear"} == { + "" + } + + +def test_private_table_rejects_int8_axis_disagreement(tmp_path): + inputs = torch.randn(2, 3, 4) + amax = torch.ones(1, 3, 1) + exported_program = torch.export.export(_StrictQuantOp("int8"), (inputs, amax), strict=True) + + with pytest.raises(torch.onnx.OnnxExporterError, match="axis 0 does not match amax axis 1"): + _raw_export(exported_program, (), tmp_path / "int8_axis.onnx", 23) + + +@pytest.mark.parametrize( + ("fmt", "inputs", "amax", "message"), + [ + ("fp8", torch.randn(2, 4), torch.ones(2, 1), "scalar amax"), + ("int8", torch.randn(2, 3, 4), torch.ones(2, 1, 4), "multi-axis"), + ("nvfp4_dynamic", torch.randn(1, 2, 3, 32), torch.tensor(1.0), "rank 2 or 3"), + ], +) +def test_private_table_rejects_shapes(tmp_path, fmt, inputs, amax, message): + exported_program = torch.export.export(_StrictQuantOp(fmt), (inputs, amax), strict=True) + with pytest.raises(torch.onnx.OnnxExporterError, match=message): + _raw_export(exported_program, (), tmp_path / f"{fmt}.onnx") diff --git a/tests/unit/torch/quantization/test_gptq.py b/tests/unit/torch/quantization/test_gptq.py index 345bd167112..1c55c089943 100644 --- a/tests/unit/torch/quantization/test_gptq.py +++ b/tests/unit/torch/quantization/test_gptq.py @@ -15,9 +15,13 @@ """CPU unit tests for GPTQ utilities.""" +from functools import partial + import torch +from _test_utils.torch.distributed.utils import spawn_multiprocess_job -from modelopt.torch.quantization.utils.calib_utils import update_hessian +from modelopt.torch.quantization.utils.calib_utils import GPTQHelper, update_hessian +from modelopt.torch.utils.distributed import ParallelState def test_update_hessian(): @@ -117,3 +121,25 @@ def test_update_hessian_zero_token_input_noops(): assert updated_hessian is hessian assert new_n_samples == n_samples torch.testing.assert_close(hessian, expected_hessian) + + +def _gptq_hessian_dp_sync_helper(rank, size, token_counts): + torch.manual_seed(0) + inputs = [torch.randn(n, 8) for n in token_counts] + linear = torch.nn.Linear(8, 4) + linear.parallel_state = ParallelState(data_parallel_group=None) + helper = GPTQHelper(linear, "linear") + helper.hessian, helper.n_samples = update_hessian(inputs[rank], helper.hessian, 0) + + helper.sync_hessian_across_dp() + + ref, n_ref = update_hessian(torch.cat(inputs), torch.zeros(8, 8), 0) + assert helper.n_samples == n_ref + torch.testing.assert_close(helper.hessian, ref) + + +def test_gptq_hessian_dp_sync(skip_on_windows): + """Token-weighted merge over a gloo (CPU) group, with one rank that saw no tokens.""" + spawn_multiprocess_job( + 2, partial(_gptq_hessian_dp_sync_helper, token_counts=(5, 0)), backend="gloo" + ) diff --git a/tests/unit/torch/quantization/test_int8_tensor.py b/tests/unit/torch/quantization/test_int8_tensor.py new file mode 100644 index 00000000000..47e190fe0bd --- /dev/null +++ b/tests/unit/torch/quantization/test_int8_tensor.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for INT8QTensor real quantization of small FP16/BF16 channels on CPU.""" + +import copy + +import pytest +import torch + +import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.qtensor import INT8QTensor + +# Rows: amax 2^-20 (its FP16 scale amax / 127 rounds to zero), amax 1e-5 (FP16 subnormal +# scale, rounded down by ~25%), all zeros, and a normal channel. +WEIGHT = [ + [2**-20, -(2**-20), 0.0, 2**-21], + [1e-5, -5e-6, 2e-6, 0.0], + [0.0, 0.0, 0.0, 0.0], + [1.0, -1.0, 0.0, 0.5], +] + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize( + ("layout", "weight"), + [ + ({"axis": 0}, WEIGHT), + ({"block_sizes": {-1: 2}}, WEIGHT), + # A per-tensor scale is 0-dim and does not promote `weight / scale` out of the dtype. + ({}, [list(range(8))]), + ], +) +def test_int8_quantize_matches_fp32(dtype, layout, weight): + """Real-quantizing low-precision weights must give the FP32 result, cast to that dtype.""" + weight = torch.tensor(weight, dtype=dtype) + block_sizes = layout.get("block_sizes") + + def round_trip(w): + qtensor, scale = INT8QTensor.quantize(w, **layout) + return qtensor.dequantize(dtype=dtype, scale=scale, block_sizes=block_sizes) + + dequantized = round_trip(weight) + + assert torch.equal(dequantized, round_trip(weight.float())) + # No nonzero value here is small enough to round to zero against its own amax. + assert torch.equal(dequantized == 0, weight == 0) + + +def test_int8_compress_keeps_small_fp16_channel(): + """mtq.compress must not zero a small FP16 channel that fake quantization keeps.""" + model = torch.nn.Sequential(torch.nn.Linear(4, 2, bias=False)).half() + with torch.no_grad(): + model[0].weight.copy_(torch.tensor([WEIGHT[0], WEIGHT[3]])) + x = torch.tensor([[1.0, 0.0, 0.0, 0.0]], dtype=torch.float16) + + mtq.quantize(model, copy.deepcopy(mtq.INT8_DEFAULT_CFG), lambda m: m(x)) + with torch.no_grad(): + fake_quant_output = model(x) + mtq.compress(model) + real_quant_output = model(x) + + assert fake_quant_output[0, 0] == 2**-20 + assert torch.equal(real_quant_output, fake_quant_output) diff --git a/tests/unit/torch/quantization/test_lsq.py b/tests/unit/torch/quantization/test_lsq.py index 1a1136d6ece..623b92342f1 100644 --- a/tests/unit/torch/quantization/test_lsq.py +++ b/tests/unit/torch/quantization/test_lsq.py @@ -40,6 +40,7 @@ _amax_to_scale, ) from modelopt.torch.quantization.tensor_quant import int_cast_ste +from modelopt.torch.quantization.utils.numeric_utils import E4M3_MAX from modelopt.torch.quantization.utils.shared_input import SharedWeightGlobalAmaxState from modelopt.torch.utils import to_empty_if_meta_device @@ -426,6 +427,7 @@ def spy_block_scale(amax, quantize): assert quantize_flags == [True, False] def test_skip_pre_scale_quantization_uses_raw_scale_floor(self, monkeypatch): + """Keep the raw pre-scale floor unchanged while normalizing the FP8 post-scale floor.""" q = self._make_lsq_quantizer() q._quantize_scales = True q._quantize_pre_scale = False @@ -446,7 +448,7 @@ def fake_amax_to_scale(amax, maxbound, min_value=1e-8): out = q._fake_quantize(torch.randn(4, 16)) assert out.shape == (4, 16) - assert torch.equal(min_values[0], torch.tensor([_FP8_E4M3_MIN_POSITIVE])) + assert torch.equal(min_values[0], torch.tensor([_FP8_E4M3_MIN_POSITIVE / E4M3_MAX])) assert min_values[1] == 1e-8 @@ -484,6 +486,7 @@ def _make_tied_lsq_group(self, global_amax=3.0, n_members=3): return members def test_block_scale_tracks_shared_update(self): + """Check each member's FP8 scale floor against the updated shared global amax.""" members = self._make_tied_lsq_group(global_amax=3.0) new_value = 5.0 # Mutate the shared buffer in place, mimicking export unification. @@ -493,7 +496,7 @@ def test_block_scale_tracks_shared_update(self): expected = _amax_to_scale(torch.tensor(new_value), member._quant_max_bound) scale = member._block_scale_from_amax(member.amax_post, quantize=True) assert scale.shape == member.amax_post.shape - assert torch.all(scale >= _FP8_E4M3_MIN_POSITIVE * expected) + assert torch.all(scale >= (_FP8_E4M3_MIN_POSITIVE / E4M3_MAX) * expected) def test_members_produce_identical_output_after_shared_update(self): members = self._make_tied_lsq_group(global_amax=3.0) diff --git a/tests/unit/torch/quantization/test_quantize_cpu.py b/tests/unit/torch/quantization/test_quantize_cpu.py index 0c663875fb1..07d5fc309a9 100644 --- a/tests/unit/torch/quantization/test_quantize_cpu.py +++ b/tests/unit/torch/quantization/test_quantize_cpu.py @@ -21,7 +21,6 @@ import pytest import torch -from _test_utils.torch.distributed.utils import spawn_multiprocess_job from _test_utils.torch.quantization.models import SimpleConv, SimpleConvLinear, SimpleLinear from _test_utils.torch.quantization.quantize_common import ( INT4_AWQ_CLIP_CFG, @@ -722,11 +721,10 @@ def _test_indexer_check_covers_all_ranks(rank, size): mtq.quantize(_ToyStage([rank]), config, lambda m: None) finally: mtq.unregister(_ToyIndexer) - torch.distributed.destroy_process_group() -def test_indexer_k_check_covers_all_ranks(skip_on_windows): - spawn_multiprocess_job(2, _test_indexer_check_covers_all_ranks, backend="gloo") +def test_indexer_k_check_covers_all_ranks(cpu_dist_workers): + cpu_dist_workers.run(_test_indexer_check_covers_all_ranks) def test_indexer_q_patterns_need_the_query_quantizer(): diff --git a/tests/unit/torch/quantization/test_tensor_quant_cpu.py b/tests/unit/torch/quantization/test_tensor_quant_cpu.py index d72bdca6e1c..26f048d0d1d 100644 --- a/tests/unit/torch/quantization/test_tensor_quant_cpu.py +++ b/tests/unit/torch/quantization/test_tensor_quant_cpu.py @@ -23,7 +23,7 @@ import modelopt.torch.quantization as mtq import modelopt.torch.quantization.nn.modules.tensor_quantizer as tensor_quantizer_module -from modelopt.torch.quantization import QuantModuleRegistry +from modelopt.torch.quantization import QuantModuleRegistry, tensor_quant from modelopt.torch.quantization.config import QuantizerAttributeConfig, RotateConfig from modelopt.torch.quantization.nn import ( SequentialQuantizer, @@ -37,6 +37,70 @@ class TestFakeTensorQuantCPU(FakeTensorQuantTester): device = "cpu" +def test_custom_op_schemas_keep_legacy_positional_calls(): + inputs = torch.randn(4, 32) + amax = torch.tensor(1.0) + fp8 = torch.ops.tensorrt.quantize_op(inputs, amax, 8, 4, False, False) + mxfp8 = torch.ops.tensorrt.dynamic_block_quantize_op.overload(inputs, 32, None, 8, 4, 9, 8) + + assert (fp8.shape, fp8.dtype) == (inputs.shape, inputs.dtype) + assert (mxfp8.shape, mxfp8.dtype) == (inputs.shape, inputs.dtype) + + +@pytest.mark.parametrize("num_bits", [8, 4]) +def test_cpu_integer_carrier_strict_capture_preserves_custom_op(num_bits): + class QuantizeCarrier(torch.nn.Module): + def forward(self, inputs, amax): + return torch.ops.tensorrt.quantize_op(inputs, amax, num_bits, 0, False, False, "Float") + + inputs = torch.randn(4, 32) + amax = torch.tensor(1.0) + exported = torch.export.export(QuantizeCarrier(), (inputs, amax), strict=True) + + assert any( + node.target == torch.ops.tensorrt.quantize_op.default for node in exported.graph.nodes + ) + expected = tensor_quant.fake_tensor_quant(inputs, amax, None, num_bits, False, False) + torch.testing.assert_close(exported.module()(inputs, amax), expected) + + +def test_cpu_fp8_quantizer_strict_capture_preserves_custom_op(): + quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=(4, 3), narrow_range=False)) + quantizer.amax = torch.tensor(1.0) + inputs = (torch.randn(4, 32),) + + exported = torch.export.export(quantizer, inputs, strict=True) + + assert any( + node.target == torch.ops.tensorrt.quantize_op.default for node in exported.graph.nodes + ) + torch.testing.assert_close(exported.module()(*inputs), quantizer(*inputs)) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) +@pytest.mark.parametrize( + ("amax_shape", "num_bits", "block_size", "axis"), + [([], 8, None, None), ([2, 1], 8, None, 0), ([4, 1], 4, 2, 1)], + ids=["scalar", "per_axis", "block"], +) +def test_registered_cpu_quantize_matches_eager(amax_shape, num_bits, block_size, axis, dtype): + inputs = torch.linspace(-3, 3, 8, dtype=dtype).reshape(2, 4) + if block_size is not None: + inputs = inputs.reshape(-1, block_size) + amax = torch.full(amax_shape, 2.0, dtype=dtype) + if amax.numel() > 1: + amax[0] = 0 + expected = tensor_quant.fake_tensor_quant(inputs, amax, None, num_bits, False, True) + + legacy = torch.ops.tensorrt.quantize_op(inputs, amax, num_bits, 0, False, True) + with_metadata = torch.ops.tensorrt.quantize_op( + inputs, amax, num_bits, 0, False, True, "Float", block_size, axis + ) + + torch.testing.assert_close(legacy, expected, rtol=0, atol=0) + torch.testing.assert_close(with_metadata, expected, rtol=0, atol=0) + + class TestQuantizerAttributeConfig: def test_scaled_mode(self): num_bits = np.random.randint(1, 16) diff --git a/tests/unit/torch/utils/test_dataset_utils.py b/tests/unit/torch/utils/test_dataset_utils.py index 35543f8e9b9..1666308e1d4 100644 --- a/tests/unit/torch/utils/test_dataset_utils.py +++ b/tests/unit/torch/utils/test_dataset_utils.py @@ -18,6 +18,7 @@ import pytest import torch +from _test_utils.torch.transformers_models import get_tiny_tokenizer from torch.utils.data import DataLoader from modelopt.torch.utils import dataset_utils @@ -59,6 +60,33 @@ def test_prepare_messages_for_chat_template(): ) +@pytest.mark.parametrize("include_labels", [False, True]) +def test_calibration_padding_avoids_masked_gradient_rows(monkeypatch, include_labels): + tokenizer = get_tiny_tokenizer() + monkeypatch.setattr( + dataset_utils, "get_dataset_samples", lambda *args, **kwargs: ["a", "a b c"] + ) + batch = next( + iter( + get_dataset_dataloader( + tokenizer=tokenizer, + batch_size=2, + num_samples=2, + include_labels=include_labels, + device="cpu", + ) + ) + ) + mask = batch["attention_mask"].bool() + assert tokenizer.padding_side == "left" + if include_labels: + causal_mask = torch.ones(mask.shape[-1], mask.shape[-1], dtype=torch.bool).tril() + assert (causal_mask & mask[:, None, :]).any(dim=-1).all() + assert torch.equal(batch["labels"][~mask], torch.full_like(batch["labels"][~mask], -100)) + else: + assert not mask[0, 0] and mask[0, -1] + + def setup_test_data(): # Create sample batch data batch_data = { diff --git a/tests/unit/torch/utils/test_distributed.py b/tests/unit/torch/utils/test_distributed.py index 7b630f93d50..f88377e2a5f 100644 --- a/tests/unit/torch/utils/test_distributed.py +++ b/tests/unit/torch/utils/test_distributed.py @@ -15,11 +15,15 @@ """Pure-function tests for ``modelopt.torch.utils.distributed``.""" +import os +import pickle + import pytest import torch import torch.nn as nn -from modelopt.torch.utils.distributed import _off_dtype_params +from modelopt.torch.nas.hparams.container import DepthHparam +from modelopt.torch.utils.distributed import _deserialize, _off_dtype_params, _serialize def _model(*sizes_and_dtypes) -> nn.Module: @@ -61,3 +65,24 @@ def test_off_dtype_params_dominant_is_by_numel(params, expected): off = _off_dtype_params(model) kept = {n for n, p in model.named_parameters() if p not in off} assert kept == {expected} + + +class _Gadget: + def __reduce__(self): + return (os.system, ("exit 1",)) + + +@pytest.mark.parametrize("force_no_weights_only", [False, True]) +def test_deserialize_rejects_pickle_gadget(monkeypatch, force_no_weights_only): + # The override that trusts local checkpoints must not also trust peer payloads. + if force_no_weights_only: + monkeypatch.setenv("TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD", "1") + with pytest.raises(pickle.UnpicklingError): + _deserialize(_serialize(_Gadget())) + + +def test_deserialize_traced_hp_subclass(): + # The allowlist matches exact classes, so each TracedHp subclass must be registered too. + hp = _deserialize(_serialize(DepthHparam([1, 2]))) + assert type(hp) is DepthHparam + assert hp.choices == [1, 2] diff --git a/tests/unit/torch/utils/test_hf_checkpoint_utils.py b/tests/unit/torch/utils/test_hf_checkpoint_utils.py index 490ca79d7ed..093e8e1cf66 100644 --- a/tests/unit/torch/utils/test_hf_checkpoint_utils.py +++ b/tests/unit/torch/utils/test_hf_checkpoint_utils.py @@ -16,18 +16,14 @@ """Tests for modelopt/torch/utils/plugins/hf_checkpoint_utils.py""" import json -import os import sys import warnings -from functools import partial from pathlib import Path from types import SimpleNamespace from unittest.mock import patch import pytest import torch -import torch.distributed as dist -from _test_utils.torch.distributed.utils import spawn_multiprocess_job from safetensors.torch import save_file from modelopt.torch.utils.plugins.hf_checkpoint_utils import ( @@ -227,6 +223,78 @@ def test_sanitize_hf_config_for_deployment_trims_nextn_layer_types(): assert config_data["layer_types"] == hidden_layer_types +@pytest.mark.parametrize("native", [False, True]) +@pytest.mark.parametrize("keep_tokenizer", [False, True]) +def test_export_config_reloads_nemotron_h_with_source_auto_map(tmp_path, native, keep_tokenizer): + # Native NemotronH is optional on the minimum-Transformers test lane. + transformers = pytest.importorskip("transformers", minversion="5.5") + from modelopt.torch.export.unified_export_hf import ( + _write_hf_export_config, + save_non_weight_artifacts, + ) + + source, export = tmp_path / "source", tmp_path / "export" + source.mkdir() + (source / "configuration_legacy.py").write_text( + "from transformers import PretrainedConfig\n" + "class LegacyConfig(PretrainedConfig):\n" + " model_type = 'nemotron_h'\n" + " def __init__(self, hybrid_override_pattern='M*', **kwargs):\n" + " self.hybrid_override_pattern = hybrid_override_pattern\n" + " super().__init__(**kwargs)\n" + " @property\n" + " def layers_block_type(self):\n" + " return ['mamba', 'attention']\n" + ) + auto_map = { + "AutoConfig": "configuration_legacy.LegacyConfig", + "AutoModel": "modeling_legacy.LegacyModel", + "AutoModelForCausalLM": "modeling_legacy.LegacyModel", + } + if keep_tokenizer: + auto_map["AutoTokenizer"] = [None, "tokenization_legacy.LegacyTokenizer"] + original_auto_map = dict(auto_map) + (source / "config.json").write_text( + json.dumps({"model_type": "nemotron_h", "auto_map": auto_map}) + ) + config = ( + transformers.NemotronHConfig(layers_block_type=["mamba", "attention"], auto_map=auto_map) + if native + else transformers.AutoConfig.from_pretrained( + source, trust_remote_code=True, local_files_only=True + ) + ) + model = torch.nn.Module() + model.config = config + model._modelopt_source_checkpoint = str(source) + quant_config = {"quantization": {"quant_algo": "FP8", "kv_cache_quant_algo": "FP8"}} + + save_non_weight_artifacts(model, export) + _write_hf_export_config(model, quant_config, export) + reloaded = transformers.AutoConfig.from_pretrained( + export, trust_remote_code=True, local_files_only=True + ) + + assert isinstance(reloaded, transformers.NemotronHConfig) == native + assert reloaded.layers_block_type == config.layers_block_type + assert reloaded.quantization_config["quant_algo"] == "FP8" + assert reloaded.quantization_config["kv_cache_scheme"] == { + "dynamic": False, + "num_bits": 8, + "type": "float", + } + assert config.auto_map == original_auto_map + expected_map = { + key: value + for key, value in original_auto_map.items() + if not native or key == "AutoTokenizer" + } + assert getattr(reloaded, "auto_map", {}) == expected_map + assert (export / "configuration_legacy.py").read_text() == ( + source / "configuration_legacy.py" + ).read_text() + + def test_sanitize_hf_config_for_deployment_adds_rope_theta_to_llama3_rope_parameters(): """Transformers 5.x requires rope_theta inside llama3 rope_parameters.""" config_data = { @@ -905,74 +973,3 @@ def test_copy_non_model_files_follows_hub_snapshot_links_into_blobs(tmp_path): assert copied == ["tokenizer.json"] assert (tmp_path / "export" / "tokenizer.json").read_text() == "{}" - - -def _ensure_local_checkpoint_job(tmp_path, fail, group_ranks, rank, size): - utils = hf_checkpoint_utils - snapshot = tmp_path / "snapshot" - - def fake_snapshot_download(repo_id, **kwargs): - (tmp_path / f"downloaded-by-rank{rank}").write_text(repo_id) - if fail: - raise OSError("hub unreachable") - snapshot.mkdir(exist_ok=True) - return str(snapshot) - - utils.snapshot_download = fake_snapshot_download - utils.HfApi = _FakeHfApi - os.environ.pop("HF_HUB_OFFLINE", None) # inherited from tests/unit/conftest.py - # new_group is collective over every rank, members or not. - group = None if group_ranks is None else dist.new_group(group_ranks, backend="gloo") - if group_ranks is not None and rank not in group_ranks: - return - if fail: - with pytest.raises(RuntimeError, match="hub unreachable"): - utils.ensure_local_checkpoint("org/model", group=group) - else: - result = utils.ensure_local_checkpoint("org/model", group=group) - assert result == ("org/model", str(snapshot)) - - -def _ensure_local_checkpoint_local_path_job(tmp_path, hidden_from_rank1, rank, size): - utils = hf_checkpoint_utils - - def fail(*args, **kwargs): - raise AssertionError("a local path must not be downloaded") - - utils.snapshot_download = fail - checkpoint = tmp_path / "checkpoint" - checkpoint.mkdir(exist_ok=True) - if hidden_from_rank1 and rank == 1: - # As if the directory sat on rank 0's node-local disk only. - isdir = os.path.isdir - utils.os.path.isdir = lambda path: path != str(checkpoint) and isdir(path) - with pytest.raises(RuntimeError, match="this rank cannot see"): - utils.ensure_local_checkpoint(str(checkpoint)) - else: - assert utils.ensure_local_checkpoint(str(checkpoint)) == (None, str(checkpoint)) - - -@pytest.mark.parametrize("hidden_from_rank1", [False, True]) -def test_ensure_local_checkpoint_lets_rank_0_resolve_a_local_path(tmp_path, hidden_from_rank1): - """Ranks that each checked a local path could disagree about it and strand one another in the - collective; rank 0 decides, and a rank that cannot see its answer raises instead.""" - job = partial(_ensure_local_checkpoint_local_path_job, tmp_path, hidden_from_rank1) - spawn_multiprocess_job(2, job, backend="gloo") - - -@pytest.mark.parametrize( - ("fail", "group_ranks", "downloader"), - [ - (False, None, "rank0"), - (True, None, "rank0"), - (False, [0, 1], "rank0"), - # Rank 0 of the group, not of the world: global rank 1. - (False, [1], "rank1"), - ], -) -def test_ensure_local_checkpoint_downloads_once_per_group(tmp_path, fail, group_ranks, downloader): - """Rank 0 of the group downloads once and every member gets its result -- including its - failure, so the waiting ranks raise instead of hanging.""" - job = partial(_ensure_local_checkpoint_job, tmp_path, fail, group_ranks) - spawn_multiprocess_job(2, job, backend="gloo") - assert [p.name for p in tmp_path.glob("downloaded-by-*")] == [f"downloaded-by-{downloader}"]