Some fixes for NemotronH - #110
Merged
Merged
Conversation
Refactor Mamba2Block to support arbitrary sequence lengths via an ONNX Scan op that iterates token-by-token over the conv + SSM recurrence. Key changes: - Use scan_input_axes/scan_output_axes to iterate over axis 1 directly, eliminating pre/post Transpose ops. - Build Scan body via GraphBuilder.subgraph() with a _scan_body method, using the child builder's scoped value names to avoid SSA collisions (replaces manual rename_subgraph_values). - Pre-realize conv1d/ssm parameters on the parent builder so the Scan body can reference them as implicit inputs. - Simplify op.Constant(value_ints=[...]) to plain [...] in Split, Slice, Reshape, and Concat arguments. - Remove explicit shapes from Scan body ir.Values (inference fills them). - Batch input projection (in_proj) and output projection (out_proj) over the full sequence outside the Scan loop. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam <grama@microsoft.com>
HuggingFace's _init_weights() unconditionally overwrites dt_bias with torch.rand() after loading checkpoint weights. Add _fix_nemotron_h_dt_bias() to reload the correct values from safetensors when running --compare-hf. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam <grama@microsoft.com>
Add MOBIUS_MAMBA_SCAN flag (default True). When False, Mamba2Block uses the original single-token forward pass without Scan, useful for debugging numerical divergence between the two paths. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam <grama@microsoft.com>
Implement the chunked Structured State Space Duality algorithm for parallel multi-token processing in Mamba2Block, gated behind the flags.mamba_scan feature flag (default True). Key changes: - Mamba2Block._forward_chunked_ssd: processes all tokens in parallel within chunks of chunk_size (default 256), propagating SSM state across chunk boundaries via cumulative matrix products. - dt clamp (time_step_min=0.001) after softplus in both chunked SSD and single-token Mamba2Scan paths, matching HF behavior. - NemotronHConfig: added mamba_time_step_min field. - Example script: config attribute name fixes for NemotronH. Standalone Mamba2Block verified to match PyTorch reference within ~4e-7 for all n_groups (1,2,4,8). Full-model prefill shows growing divergence (compound FP32 amplification across 21 Mamba layers), with top tokens matching 24/26 positions. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam <grama@microsoft.com>
Change flags.mamba_scan from bool to str with three modes: - 'chunked_ssd' (default): chunked SSD parallel algorithm - 'scan': ONNX Scan op token-by-token iteration - 'single': single-token path (seq_len must be 1) The Scan-based path (from rama/mamba branch) is integrated alongside the chunked SSD and single-token paths. All three are selectable via MOBIUS_MAMBA_SCAN env var. Backwards compatible: 1/true -> chunked_ssd, 0/false -> single. Add _env_str helper to _flags.py for string-valued flags with choices and boolean-alias backwards compatibility. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam <grama@microsoft.com>
HF's _init_weights also reinitializes out_proj.weight in every Mamba2 layer with kaiming_uniform_ scaled by 1/sqrt(num_hidden_layers) when rescale_prenorm_residual is True. This overwrites trained checkpoint weights, causing massive ONNX-vs-HF divergence. Extended _fix_nemotron_h_dt_bias() to reload both dt_bias and out_proj.weight from the safetensors checkpoint (42 params total: 21 dt_bias + 21 out_proj.weight). With both fixes, full-model comparison improves from 0/26 to 25/26 token matches (max diff 7.2, down from 25.0). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam <grama@microsoft.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam <grama@microsoft.com>
When no --prompt is given, the script now enters an interactive loop prompting for queries. Type 'quit' or Ctrl-D to stop. The ONNX session is reused across all queries. With --compare-hf, HF comparison runs after all interactive queries are collected. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam <grama@microsoft.com>
Tokens are now printed as they are generated instead of waiting until the end. Makes slow generation much more usable. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam <grama@microsoft.com>
Default 1.2 to prevent repetitive loops in greedy decoding. Uses the same algorithm as HF's RepetitionPenaltyLogitsProcessor: divide positive logits by penalty, multiply negative logits by penalty. Both ONNX and HF paths use the same --repetition-penalty value. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam <grama@microsoft.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam <grama@microsoft.com>
Performance Comparison
|
Contributor
There was a problem hiding this comment.
Pull request overview
Note
Copilot was unable to run its full agentic suite in this review.
Improves NemotronH/Mamba2 correctness and usability by aligning ONNX behavior more closely with HuggingFace, adding multi-token Mamba2 execution strategies, and refining the Nemotron example generation script.
Changes:
- Add
time_step_minsupport to Mamba2 SSM step to reduce numeric divergence vs HF. - Implement multi-token Mamba2Block paths (chunked-SSD and ONNX Scan) and dispatch via a new
MOBIUS_MAMBA_SCANflag. - Enhance the Nemotron text-generation example (interactive mode, repetition penalty, HF comparison fixes).
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/build_graph_test.py | Enables graph test coverage for mamba2 by removing it from the skip list. |
| src/mobius/models/nemotron_h.py | Plumbs mamba_time_step_min into the Mamba2Block construction. |
| src/mobius/components/_ssm.py | Adds time_step_min parameter and clamps dt accordingly. |
| src/mobius/components/_mamba_block.py | Adds chunked-SSD + Scan multi-token implementations and flag-based dispatch. |
| src/mobius/_flags.py | Introduces mamba_scan flag parsing from environment with validated choices. |
| src/mobius/_configs.py | Adds config field for mamba_time_step_min and HF config mapping. |
| examples/nemotron_3_nano_text_generation.py | Improves script UX and correctness (interactive prompts, repetition penalty, HF weight patching). |
🏗️ Architecture Diff
mamba (ssm-text-generation) / model — 21 change(s)Op summary: 103 → 99 nodes --- base
+++ head
@@ -2,13 +2,9 @@
RMSNormalization
Transpose
MatMul
-Constant
Split
Transpose
Concat
-Constant
-Constant
-Constant
Slice
Conv
UnsqueezeRemoved nodes:
Modified attributes:
Connectivity changes:
Initializer changes:
Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Replace the out_proj.weight reload workaround with a simpler config fix: set rescale_prenorm_residual=False before from_pretrained. This prevents HF's _init_weights from corrupting out_proj.weight with kaiming_uniform_ reinit (a training-time GPT-2 residual scaling flag, not needed for inference). The dt_bias reload is still needed (separate unconditional bug in _init_weights). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam <grama@microsoft.com>
Split the monolithic 1047-line Mamba2Block class into: - _mamba_block.py: Mamba2BlockBase (shared __init__), Mamba2BlockSingle (single-token path), and Mamba2Block factory function - _mamba_block_scan.py: Mamba2BlockScan (ONNX Scan token-by-token) - _mamba_block_chunked.py: Mamba2BlockChunkedSSD (parallel chunked SSD) The factory function reads flags.mamba_scan at construction time and returns the matching subclass. Zero caller changes needed — all 4 consumers (bamba, granitemoehybrid, mamba, nemotron_h) use the factory exactly like the old class. Design follows my/mamba-2-design.md Option B+D: - Subclass over boolean flags (no mode dispatch in forward()) - _realize_submodule is a single staticmethod on Mamba2BlockBase - Chunked SSD helpers (_segment_sum, _segment_sum_dynamic) move to the chunked file Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam <grama@microsoft.com>
- Change mamba_scan flag default from 'single' to 'chunked_ssd' and align the docstring to label chunked_ssd as the default mode - Update example docstring to reflect that multi-token Mamba2 modes now exist (removes outdated 'only support single-token decode') - Align _scan_body parameter names with subgraph input dict keys for consistency (conv_state_in→conv_state, xbc_in→xbc_t, etc.) The _realize_submodule duplicate issue was already resolved by the prior refactoring commit (single staticmethod on Mamba2BlockBase). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam <grama@microsoft.com>
Member
|
Tests are failing |
- Add noqa: N802 to Mamba2Block factory function (intentional PascalCase) - Add per-file-ignores for N803/N806 in _mamba_block_chunked.py (SSM math notation: A, B, C, D, H, N variables) - Fix E402 in example: move mobius imports above warnings.filterwarnings - Add noqa: N806 for _CORRUPTED_KEYS constant-like variable - Add noqa: RUF069 for intentional float equality check (penalty == 1.0) - Auto-fix import sorting Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam <grama@microsoft.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam <grama@microsoft.com>
The causal conv Pad in Mamba2BlockChunkedSSD used a float32 constant for the pad value, causing an ORT type error when the model is built with f16 or bf16. Fix by using CastLike to match the pad constant to the input tensor dtype. Internal chunked SSD Pads (on explicitly float32 tensors) keep explicit float32 constants. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam <grama@microsoft.com>
The default was mistakenly changed to 'chunked_ssd' when addressing PR review feedback. Restore 'single' as the intended default and update the docstring to list it first with the (default) label. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam <grama@microsoft.com>
Member
Can
|
| Mamba2 tensor | LinearAttention arg | Packed shape | Notes |
|---|---|---|---|
C_mat (readout) |
query |
(B, T, H_q × d_state) |
q_num_heads = num_heads |
B_mat (input matrix) |
key |
(B, T, H_kv × d_state) |
kv_num_heads = n_groups |
dt * x (discretized) |
value |
(B, T, H_kv × d_head) |
absorb dt into value |
A * dt (log-decay) |
decay |
(B, T, H_kv) |
per-head scalar |
ssm_state |
past_state |
(B, H, d_state, d_head) |
requires transpose (see below) |
Attributes: update_rule="gated", q_num_heads=num_heads, kv_num_heads=n_groups, scale=1.0
The n_groups > 1 case maps naturally onto LinearAttention's GQA grouping (q_num_heads / kv_num_heads = heads_per_group).
Three Items Outside the Op
-
dtdiscretization —dt = softplus(dt_proj(x) + dt_bias)is computed before calling LA;dt * xis passed asvalueandA * dtasdecay. No structural change needed. -
D skip connection —
y += D * xis a plainAddafter the LA output. One line. -
State shape transpose — Current Mamba2 state is
(B, H, d_head, d_state)but LA expects(B, H, d_k, d_v) = (B, H, d_state, d_head). NeedTranspose([0,1,3,2])onpast_stateinput andpresent_stateoutput. Minor.
What This Means for This PR
This PR implements three modes — Mamba2BlockSingle (T=1), Mamba2BlockScan (ONNX Scan, sequential), Mamba2BlockChunkedSSD (parallel within chunks) — because ORT doesn't have a native op that handles both decode and prefill. LinearAttention removes that need entirely. A single LA call handles both paths, and ORT's implementation is chunk-parallel for prefill internally.
The entire Mamba2BlockChunkedSSD._chunked_ssd() method (~180 ONNX ops: segment_sum, CumSum, Trilu, cross-chunk propagation, 5-stage SSD) would collapse to:
# Outside: compute dt, pack tensors
decay = op.Mul(a_2d, dt) # A * dt, shape (B,T,H_kv)
value = op.Mul(x_disc_packed, dt_packed) # dt * x, shape (B,T,H_kv*d_head)
# The op
output, new_state = op.LinearAttention(
C_packed, # query (B,T,H_q*d_state)
B_packed, # key (B,T,H_kv*d_state)
value, # value (B,T,H_kv*d_head)
past_state_transposed, # (B,H_kv,d_state,d_head)
decay, # (B,T,H_kv) — per-head scalar decay
domain="com.microsoft",
update_rule="gated",
q_num_heads=self.num_heads,
kv_num_heads=self.n_groups,
scale=1.0,
)
# Outside: D skip
y = op.Add(output, op.Mul(D_skip, x))Recommendation
Consider gating this PR on LinearAttention availability. Since LinearAttention is already in com.microsoft ContribOps (confirmed in ContribOperators.md), adding a rewrite rule that replaces the Mamba2BlockChunkedSSD ONNX subgraph with a single LinearAttention node (similar to how GQA rewrites work for attention) would:
- Remove ~300 lines of chunked SSD ops from the ONNX graph
- Let ORT choose the optimal chunk size at runtime
- Unify decode and prefill into one code path
If LinearAttention isn't available in the minimum supported ORT version, the current ONNX-native fallback in this PR is still valuable — they can coexist as primary and fallback paths.
Investigated by reading the PR diff, Mamba2Scan.forward() in _ssm.py, Mamba2BlockChunkedSSD._chunked_ssd(), and the LinearAttention spec in ContribOperators.md.
Switch NemotronH from create_attention_bias (float additive mask) to create_padding_mask (bool mask). ORT 1.24.4's opset 24 Attention op does not support bfloat16 for the mask/bias input, causing a type error: 'Type parameter (T) of Optype (Add) bound to different types (tensor(bfloat16) and tensor(float))'. The bool padding mask works with all dtypes (f32, f16, bf16) because: - Attention op already uses is_causal=1 for causal masking - The bool mask only encodes padding information - This also enables Flash Attention eligibility in ORT Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam <grama@microsoft.com>
NemotronH does not use rotary position embeddings — its HuggingFace reference (NemotronHAttention.forward) applies no rotary encoding. The model relies on Mamba layers' inherent position-awareness for sequence ordering. The ONNX model was incorrectly applying DefaultRope(theta=10000) because ArchitectureConfig.from_transformers() defaults rope_type to 'default' when the HF config doesn't specify one. This caused ~0.12 RMS divergence at the first attention layer (layer 12), growing to >1.0 by the final attention layer (layer 32) during prefill. After this fix, all 57 generated tokens match HF exactly with max logit diff of 5.7e-5 (pure FP32 rounding noise). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam <grama@microsoft.com>
Map --device to the corresponding EP name so the built model uses EP-specific contrib ops (e.g. GroupQueryAttention on CPU/CUDA). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam <grama@microsoft.com>
Signed-off-by: G Ramalingam <grama@microsoft.com> # Conflicts: # examples/nemotron_3_nano_text_generation.py
gramalingam
enabled auto-merge (squash)
April 10, 2026 14:53
|
The author of this PR, gramalingam, is not an activated member of this organization on Codecov. |
Collaborator
Author
Thanks for the suggestion. This version is working correctly at last. I suggest merging it. I will look at using LinearAttention as a follow-up PR. (Sometime copilots do a few dumb things that make it hard to get back to an earlier state. Better to get this version in, so that I can fallback to this easily.) A second follow-on relates to the issue that was causing the numeric-precision: this model doesn't use RoPE (in HF), its config doesn't have rope_type, but we end up using a default RoPE instead of no RoPE. Not sure if we need to change anything about the way the default works in this setting. The current approach (the model class explicitly makes that decision) might be fine, but worth thinking about. |
The mamba2 test config had intermediate_size=128 (from TINY_INTERMEDIATE default) but num_heads=4 * head_dim=16 = 64. Mamba2 requires d_inner = num_heads * head_dim, so the GatedRMSNorm received inputs with incompatible shapes ([batch,64] vs [batch,128]), causing shape inference to fail and the ONNX checker to report missing shape fields. Fix: add explicit intermediate_size=64 to the mamba2 test config. Also add __post_init__ validation to Mamba2Config to catch this constraint violation early. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam <grama@microsoft.com>
gramalingam
added a commit
that referenced
this pull request
Apr 11, 2026
- Change mamba_scan flag default from 'single' to 'chunked_ssd' and align the docstring to label chunked_ssd as the default mode - Update example docstring to reflect that multi-token Mamba2 modes now exist (removes outdated 'only support single-token decode') - Align _scan_body parameter names with subgraph input dict keys for consistency (conv_state_in→conv_state, xbc_in→xbc_t, etc.) The _realize_submodule duplicate issue was already resolved by the prior refactoring commit (single staticmethod on Mamba2BlockBase). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam <grama@microsoft.com>
gramalingam
disabled auto-merge
April 12, 2026 01:12
justinchuby
added a commit
that referenced
this pull request
Apr 22, 2026
`_extract_rope_config()` never returned `None`, so NoPE models (NemotronH, GraniteMoeHybrid, GPT-2/BERT/OPT families) silently received `rope_type='default'` and had spurious rotary ops injected into their ONNX graphs. The existing NemotronH fix (#110) is a model-class workaround; this PR fixes the root cause in the config system. ### Changes - **`_extract_rope_config()`** returns `None` unless the HF config has a real RoPE signal: `rope_parameters`, `rope_scaling`, or legacy `rotary_dim` / `rotary_pct` / `rotary_emb_base`. `rope_theta` alone is dead data (NemotronH carries it without using RoPE) and is deliberately not treated as a signal. - **`ArchitectureConfig.from_transformers()`** propagates `None` to both the `rope` sub-config and every flat RoPE field when the model is NoPE. `dataclasses.replace(rope_config, ...)` call sites now guard against `None`. - **`ArchitectureConfig.rope_type`** default is `None` (the structural NoPE signal). `rope_theta` / `partial_rotary_factor` keep inert numeric defaults so `ArchitectureConfig(rope_type="default", ...)` still works for direct construction in tests. - **`initialize_rope()`** returns `None` when `rope_type is None` and `mrope_section is None`. - **`TextModel.forward()`** guards `self.rotary_emb(...)` and passes `position_embeddings=None` down when RoPE is absent. - **`Attention.__init__`** tolerates `partial_rotary_factor=None` (treated as the inert 1.0) so NoPE-routed attention doesn't crash on `math.isclose(None, 1.0)`. ### Tests - New: `test_from_transformers_nope_model_has_none_rope`, `test_from_transformers_legacy_rotary_dim_enables_rope`, `test_rope_theta_alone_is_not_a_rope_signal`, `test_nope_returns_none` (initialize_rope). - Existing fakes in `_config_resolver_test.py` / `_configs_test.py` now include `rope_parameters={"rope_type": "default"}` to match how real HF `PretrainedConfig.__post_init__` populates the field. - `make_config()` test helper opts into `rope_type="default"` so component tests keep the RoPE code path. Direct `BambaConfig` / `JambaConfig` / `Gemma2Config` test constructors and `deepseek_ocr2.py`'s internal `ArchitectureConfig` opt in explicitly. ### Intentionally out of scope - Simplifying the NemotronH / GraniteMoeHybrid text-model workarounds — they use bespoke layer types (`NemotronHMambaLayer`, etc.) and can't drop straight onto `TextModel`. They remain correct and now coexist with a structural defense in the config layer. - Phase 2 (calling `validate()` in the build path, nullable MoE sub-config, deprecating flat RoPE fields) and Phase 3 (feature-group sub-configs, per-model config classes) from the issue. ### Example ```python # NemotronH config → no RoPE signal → structurally NoPE class FakeNemotronH: model_type = "nemotron_h" rope_theta = 10_000.0 # dead data, ignored # no rope_parameters, no rope_scaling, no rotary_* ... cfg = ArchitectureConfig.from_transformers(FakeNemotronH()) assert cfg.rope is None assert cfg.rope_type is None assert initialize_rope(cfg) is None # TextModel now skips RoPE automatically ``` --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com> Co-authored-by: Justin Chu <justinchuby@users.noreply.github.com>
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
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.
Uh oh!
There was an error while loading. Please reload this page.