Add onnxsim as an alternative ONNX simplification backend - #1
Closed
take-cheeze wants to merge 5 commits into
Closed
Add onnxsim as an alternative ONNX simplification backend#1take-cheeze wants to merge 5 commits into
take-cheeze wants to merge 5 commits into
Conversation
Re-organize Changelog for 0.46.0 into subsections Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
…runing (NVIDIA#1955) ### What does this PR do? Type of change: New feature Add Minitron pruning support for MoE models loaded with the efficient fused **grouped GEMM** experts (`TEGroupedMLP`), in addition to the existing `TESequentialMLP` path. - New `_DynamicTEGroupedLinear` + `_DynamicTEGroupedMLP` dynamic modules slice the per-expert grouped weights (`num_moe_experts` / `moe_ffn_hidden_size` / `hidden_size`) and reorder/drop experts via `num_gemms`, mirroring the SequentialMLP path with a minimal DynamicModule diff. - Since Minitron prunes homogeneously, a single shared `moe_ffn_hidden_size` is pruned across experts (SequentialMLP keeps one per expert). Same pruned width and independent per-expert weights; only the kept-channel index set is shared. - `examples/megatron_bridge/prune_minitron.py` uses grouped GEMM by default (`--no_moe_grouped_gemm` to fall back to SequentialMLP). ### Usage ```bash torchrun --nproc_per_node 2 prune_minitron.py \ --hf_model_name_or_path <moe-model> \ --prune_target_active_params 3e9 \ --output_hf_path /tmp/pruned # add --no_moe_grouped_gemm to use SequentialMLP ``` ### Testing Verified in an `nvcr.io/nvidia/nemo:26.06` + MBridge main (as of 23 Jul) mounted so it mimics nemo:26.08 behavior. - GPT MoE dynamic-module + pruning + parameter-sorting tests parametrized over both expert impls. - Mamba hybrid NAS metric tests: params-based now covers grouped GEMM, memory-based stays SequentialMLP (exact param counts / top-k / search-space-size assertions hold identically). - NemotronH end-to-end `test_prune_minitron` exercises real-forward NAS without grouped GEMM (to be enabled in 26.08 container) - Compared Nemotron-3-Nano-30B-A3B pruning: SequentialMLP vs GroupedGEMM on 4x B300 (accuracy + speed) | Metric | TESequentialMLP (old) | TEGroupedLinear (new) | |---|---|---| | Calibration time | 6.5 mins | 3.5 mins| | Time to evaluate Top-10 pruned candidates | 23 mins | 11 mins | Top-10 Pruned Candidates — MMLU Scores | # | export_config | params | TESequentialMLP (old) | TEGroupedLinear (new) | |---|---|---|---|---| | 1 | L52, h2688, mamba 56×48, 96 experts, ffn 1536, shared 3072 | 20.09B | 0.2713 | 0.2846 | | 2 | L52, h2688, mamba 48×56, 104 experts, ffn 1536, shared 3072 | 21.61B | 0.2580 | 0.2594 | | 3 | L52, h2560, mamba 48×64, 96 experts, ffn 1536, shared 3712 | 19.28B | 0.3951 | 0.3895 | | 4 | L52, h2304, mamba 64×64, 104 experts, ffn 1856, shared 3072 | 22.28B | 0.4951 | 0.4944 | | 5 | L52, h2560, mamba 48×48, 96 experts, ffn 1792, shared 3328 | 21.99B | 0.2685 | 0.2580 | | 6 | L48, h2560, mamba 56×56, 104 experts, ffn 1792, shared 3072 | 23.68B | 0.4741 | 0.4657 | | 7 | L46, h2560, mamba 64×56, 104 experts, ffn 1792, shared 3072 | 23.68B | 0.2385 | 0.2371 | | 8 | L52, h2688, mamba 48×56, 96 experts, ffn 1536, shared 3072 | 20.09B | 0.2587 | 0.2622 | | 9 | L52, h2304, mamba 64×64, 96 experts, ffn 1856, shared 3072 | 20.70B | 0.4888 | 0.4860 | | 10 | L50, h2560, mamba 48×48, 104 experts, ffn 1792, shared 3712 | 23.68B | 0.2517 | 0.2531 | ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ - Did you get Claude approval on this PR?: ✅ <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit ## Summary by CodeRabbit * **New Features** * Improved MoE pruning and NAS/search-space handling to support both grouped-GEMM and sequential expert layouts, including optimized grouped-GEMM execution for Minitron. * Added `--no_moe_grouped_gemm` to force the sequential expert path when required. * **Bug Fixes** * Tightened `--prune_score_func` parsing for MMLU to accept only `mmlu_<N>pct_bs<bs>`. * Added a safety check in Megatron prefill to prevent int32 indexing overflow by asking to reduce calibration batch size. * **Tests** * Expanded GPT and Mamba GPU pruning/search-space tests to cover both MoE modes. * **Documentation** * Updated release notes and bridge/pruning guidance for the grouped-GEMM behavior. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
### What does this PR do?
Type of change: new example + bug fixes
Adds a **MiniMax-M3 DSpark streaming training recipe** under
`tools/launcher/examples/MiniMaxAI/MiniMax-M3/`, plus the four fixes it
needs to actually run. Each fix addresses a failure mode that is silent
or misleading without it:
1. **Gemma-style final norm for the fake base**
(`modeling_final_norm.py`, `modeling_fakebase.py`): M3 uses a
gemma-style final RMSNorm (`(1 + weight)` scale, fp32
multiply-then-cast). Selecting the norm by `model_type` alone picks
plain `rmsnorm` — MiniMax's VL remote code coerces its `model_type`-less
`text_config` to **mixtral** — which silently drops the `+1` and
corrupts the distillation target. New `_FinalGemmaRMSNorm` + selection
by the explicit `use_gemma_norm` config flag (only MiniMax sets it).
2. **Loud failure for un-maskable `answer_only_loss`**
(`hf_streaming_dataset.py`): with a fast tokenizer whose chat template
has no `{% generation %}` tags,
`apply_chat_template(return_assistant_tokens_mask=True)` only warns and
returns an **all-zero mask** — training runs at zero loss on every
sample with no other symptom. Now raises at tokenization with an
actionable message.
3. **`SERVE_BLOCK_SIZE` knob** (`train_eagle_streaming.sh`): nemo_run
exports env values unquoted, so a multi-token
`SERVE_EXTRA_ARGS="--trust-remote-code --block-size 128"` loses
everything after the first token. M3's MSA sparse attention requires KV
block 128 (`ValueError: No common block size for 16` at engine init
otherwise), so `--block-size` gets a dedicated single-token knob.
4. **Relax the speculative_decoding `transformers` pin to `<5.13`**
(match `pyproject.toml`): the old `<5.4` pin downgrades recent vLLM
containers (e.g. transformers 5.12.1, which also provides in-tree
`minimax_m3_vl`) and breaks `vllm serve` (`ALLOWED_LAYER_TYPES` needs
>=5.5.3).
The example itself encodes the validated M3 specifics: generation-tagged
chat template copy (required for `answer_only_loss` — see fix 2), draft
dims + base `rope_theta=5e6` set explicitly (not inherited), mask token
200063 (reserved slot; added tokens end at 200060), `EAGLE_CAPTURE_IDS`
= draft default `target_layer_ids+1` + final layer, `trust_remote_code`
at serve/export, and AWS-EFA NIXL notes (UCX segfaults at agent init on
EFA nodes; LIBFABRIC required there).
### Usage
```bash
cd tools/launcher
export SLURM_HOST=localhost SLURM_ACCOUNT=<account> SLURM_PARTITION=<partition> \
SLURM_HF_LOCAL=<hf_models_dir> SLURM_JOB_DIR=<experiments_dir> NEMORUN_HOME=$PWD
uv run launch.py --yaml examples/MiniMaxAI/MiniMax-M3/hf_streaming_dspark_multi_node.yaml \
identity=$HOME/.ssh/id_ecdsa detach=True --yes
```
### Testing
- `tests/unit/torch/speculative/plugins/`: **129 passed** inside a
current vLLM x86_64 nightly container (transformers 5.12.1), including 3
new tests for the assistant-mask guard.
- Generation-tagged template verified against real corpus samples:
`input_ids` identical to the original template, mask covers exactly the
assistant turns (think prefix + content + eos), contiguous, no
user-prompt leak.
- The recipe is exercised end-to-end by a live M3 DSpark training run
(this yaml modulo cluster paths): streaming serve + NIXL transport +
resume all healthy; drafter MT-Bench AL exceeds our Kimi-K2.6 DSpark
reference by ~8k steps.
- `ruff check` / `ruff format` (0.15.20) clean.
### Before your PR is "*Ready for review*"
- Is this change backward compatible?: ✅ (norm selection only changes
models with `use_gemma_norm=True`, previously mis-normed; the guard
turns a silent zero-loss run into an error; `SERVE_BLOCK_SIZE` is
opt-in; the pin relax widens the allowed range)
- Did you write any new necessary tests?: ✅
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
❌ (can add if desired)
🤖 Generated with [Claude Code](https://claude.com/claude-code)
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
- **New Features**
- Added Gemma-style RMS normalization support for speculative decoding.
- Added MiniMax-M3 multi-node training configuration, including a full
Jinja chat template for tool calls, multimodal content, and thinking
modes.
- Added optional `SERVE_BLOCK_SIZE` support for vLLM serve launches.
- **Bug Fixes**
- Improved `answer_only_loss` masking validation: now fails fast with
clear errors when required `{% generation %}` markers are missing or
when using a non-fast tokenizer.
- **Compatibility**
- Expanded the supported Transformers version range for speculative
decoding examples.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
### What does this PR do?
Type of change: New feature (+ refactor/removal of the legacy multi-node
path) <!-- Use one of the following: Bug fix, new feature, new example,
new tests, documentation. -->
<!-- Details about the change. -->
Consolidates multi-node FSDP2 post-training quantization into the
standard hf_ptq.py entry point behind a --use_fsdp2 flag, and removes
the separate
Accelerate-based multinode_ptq.py script and fsdp2.yaml config. FSDP2
PTQ is now launched with torchrun and supports single-node multi-GPU and
multi-node out
of the box.
Highlights:
- --use_fsdp2 / --cpu_offload on hf_ptq.py — calibration runs under
PyTorch FSDP2; decoder layers are sharded (root stays replicated), with
optional CPU
offload of decoder shards between forwards.
- New modelopt/torch/utils/model_load_utils.py — parallel, round-robin
safetensors loading: each rank reads only the decoder layers it owns
from disk, then
broadcasts them so every rank can shard its slice. Non-decoder weights
(embed/lm_head/norm) are read on rank 0 and broadcast.
- New FSDP2 helpers in modelopt/torch/utils/distributed.py — fsdp2_wrap,
shard_dataloader, fsdp_aware_forward_loop, broadcast_state_dict.
- Distributed export (unified_export_hf.py) — replaces the Accelerate
get_state_dict gather with get_model_state_dict(full_state_dict,
cpu_offload,
broadcast_from_rank0); rank 0 writes files, other ranks sync on a
barrier.
- core_utils.py — relaxes the stale "root must be an FSDPModule" assert
(only decoder layers are wrapped), and adds a CPU↔GPU mirror so weight
access/writeback
works for CPU-offloaded shards.
- Docs — rewritten examples/llm_ptq/README.md FSDP2 section and the
SLURM PTQ reference, both using torchrun --use_fsdp2.
v1 scope: standard causal-LM checkpoints only. --use_fsdp2 raises
NotImplementedError for VILA, multimodal/VL,
pack-quantized/compressed-tensors, speculative decoding, MTP,
auto-quantize, and sparsity.
Design Notes
1. Why a custom parallel-safetensors loader instead of HF device_map /
accelerate
AutoModelForCausalLM.from_pretrained(device_map="auto" | "cpu") and
accelerate.load_checkpoint_in_model both load the full checkpoint on
every rank from disk (per-rank CPU peak ≈ full model size). For 70B
that's ~140 GiB/rank; for 200B+ it OOMs the node before sharding can
run. parallel_load_and_prepare_fsdp2 round-robins decoder layers across
ranks so each rank reads only ~model_size / world_size from disk, then
per-layer broadcasts to peers. Per-rank CPU peak is bounded by the
largest single layer + transient broadcast, not the full model. This is
what makes 200B+ FSDP2 PTQ feasible on commodity nodes; HF/accelerate's
existing entry points don't expose this composition.
2. Why a custom FSDP2 stack instead of keeping multinode_ptq.py +
accelerate launch
The deleted multinode_ptq.py ran on accelerate launch --config_file
fsdp2.yaml and duplicated hf_ptq.py's load → calibrate → export path.
Two consequences:
- Two divergent scripts for the same operation: CLI surface, calibration
loop, and export rewrites had to land twice. They had already drifted.
- Users had to know which script applied at which scale.
The new code path unifies under hf_ptq.py --use_fsdp2. Going direct to
FSDP2 primitives (fully_shard, CPUOffloadPolicy, MixedPrecisionPolicy)
instead of
routing through accelerate's wrappers buys:
- Direct control over mp_policy and offload_policy (accelerate's plugin
layer hides them).
- The parallel-read loader above (incompatible with accelerate's
per-rank full load).
- No YAML config file in the example dir.
-
The "custom stack" is intentionally thin: fsdp2_wrap is a 5-line wrapper
over fully_shard; the rest is pure torch.distributed. We're not
reimplementing FSDP2 —just composing the public PyTorch surface directly
rather than via accelerate's adapter.
3. fsdp_aware_forward_loop ↔ transformers_trainer.py:_quantize_model
duplication
Both implement the same trick: mtq.quantize unwraps the FSDP module
before calling the user's forward_loop, and forwarding through the
unwrapped module bypasses FSDP2's pre/post-forward hooks (no all-gather
→ broken calibration). Both capture the outer wrapped model and forward
through it instead. This PR extracts the pattern into
fsdp_aware_forward_loop (in distributed.py) as the canonical helper. The
QLoRA path (transformers_trainer.py:_quantize_model) keeps its inlined
version this PR because the QLoRA forward loop has training-specific
quirks (batch shape, loss accumulation, gradient flow) that need a
careful pass to share the helper cleanly. The TODO in the helper's
docstring marks the consolidation target.
### Usage
```python
# Single node, multiple GPUs
torchrun --standalone --nproc_per_node=<num_gpus> hf_ptq.py \
--pyt_ckpt_path <model> \
--qformat nvfp4 \
--export_path <out> \
--use_fsdp2
# Multi-node (run on each node)
torchrun \
--nnodes=<N> --node_rank=<rank> \
--master_addr=<node0_ip> --master_port=<port> \
--nproc_per_node=<gpus_per_node> \
hf_ptq.py \
--pyt_ckpt_path <model> --qformat nvfp4 \
--export_path <out> --use_fsdp2 --cpu_offload # --cpu_offload for very large models
```
### Testing
<!-- Mention how have you tested your change if applicable. -->
- tests/gpu/torch/quantization/test_fsdp2.py:
test_writeback_root_unwrapped (assert relaxation, root unwrapped) and
test_writeback_cpu_offload (CPU↔GPU mirror
round-trip under CPUOffloadPolicy).
- tests/unit/torch/utils/test_model_load_utils.py: pure-function tests
for weight_map_for (sharded / single-file / missing) and
read_safetensors_subset.
- Manual end-to-end PTQ runs on single- and multi-node torchrun (FP8 /
NVFP4 / NVFP4 layerwise), with and without --cpu_offload.
### Before your PR is "*Ready for review*"
Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)
and your commits are signed (`git commit -s -S`).
Make sure you read and follow the [Security Best
Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors)
(e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(...,
weights_only=False)`, `pickle`, etc.).
- Is this change backward compatible?: ❌ <!--- If ❌, explain why. -->
hf_ptq.py is fully backward compatible (FSDP2 is opt-in via a new flag),
but the legacy examples/llm_ptq/multinode_ptq.py script and fsdp2.yaml
are removed. Users of the old multi-node entry point must migrate to
hf_ptq.py --use_fsdp2
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A <!---
Mandatory -->
- Did you write any new necessary tests?: ✅ <!--- Mandatory for new
features or examples. -->
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
❌<!--- Only for new features, API changes, critical bug fixes or
backward incompatible changes. -->
- Did you get Claude approval on this PR?: ✅ / ❌ / N/A <!--- Run
`/claude review`. NVIDIA org members can self-trigger for complex
changes; orthogonal to CodeRabbit. -->
### Additional Information
<!-- E.g. related issue. -->
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
* **New Features**
* FSDP2-enabled PTQ: torchrun single-/multi-node execution, CPU offload,
and NVFP4 layerwise calibration; distributed model loading and export
support.
* **Documentation**
* Rewritten PTQ guide and SLURM notes with explicit torchrun/FSDP2
instructions.
* **Removed**
* Legacy Accelerate-based multinode PTQ workflow and YAML config.
* **Tests**
* Added FSDP2-focused tests for quantization writeback and safetensors
distributed loading.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
Signed-off-by: sugunav14 <178320438+sugunav14@users.noreply.github.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: realAsma <86726418+realAsma@users.noreply.github.com>
take-cheeze
force-pushed
the
claude/re-enable-onnxsim-c1wpem
branch
2 times, most recently
from
July 26, 2026 09:40
c2e56d6 to
345acfa
Compare
Keep onnxslim as the default simplifier and offer onnxsim as an equivalent, user-selectable option for modelopt.onnx.quantization.quantize(..., simplify=True) and the --simplify CLI flag. - quantize.py / __main__.py: add a simplify_backend argument (--simplify_backend) that selects between "onnxslim" (default, unchanged behavior) and "onnxsim". An unknown backend name or a missing onnxsim install fails loudly; genuine simplification failures still fall back gracefully to the original model. - pyproject.toml / uv.lock: add onnxsim>=0.7.0 to the onnx extra alongside onnxslim. onnxsim 0.7.0 ships wheels for Python 3.12+ and aarch64. - test_quantize_api.py: update the _preprocess_onnx mock signature for the new argument. - CHANGELOG.rst: document the new backend option under 0.47. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01PYET7LC6SwmWpkCSN18T1Z Signed-off-by: Claude <noreply@anthropic.com>
take-cheeze
force-pushed
the
claude/re-enable-onnxsim-c1wpem
branch
from
July 26, 2026 09:40
345acfa to
d98e6c9
Compare
Owner
Author
|
Moving to NVIDIA#2018 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Keep onnxslim as the default simplifier and offer onnxsim as an equivalent, user-selectable option for
modelopt.onnx.quantization.quantize(..., simplify=True) and the --simplify CLI flag.
Claude-Session: https://claude.ai/code/session_01PYET7LC6SwmWpkCSN18T1Z
What does this PR do?
Type of change: ?
Usage
# Add a code snippet demonstrating how to use thisTesting
Before your PR is "Ready for review"
Make sure you read and follow Contributor guidelines and your commits are signed (
git commit -s -S).Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded
trust_remote_code=True,torch.load(..., weights_only=False),pickle, etc.).CONTRIBUTING.md: ✅ / ❌ / N/AAdditional Information