Fix inter-chunk recurrence in the Zamba2 / Nemotron-H Mamba2 slow path by Jeronymous · Pull Request #47250 · huggingface/transformers

Jeronymous · 2026-07-10T14:00:07Z

What does this PR do?

Fixes an incorrect reduction dimension in the inter-chunk SSM recurrence of the pure-PyTorch (no-kernel) torch_forward path shared by Zamba2MambaMixer and NemotronHMamba2Mixer.

The current code

states_permuted = states.permute(0, 2, 1, 3, 4)
result = (decay_chunk[..., None, None] * states_permuted[:, :, None, ...]).sum(dim=2)
new_states = result.permute(0, 2, 1, 3, 4)

sums over the target-chunk axis while states is broadcast over it, so states factors out of the sum. This collapses the intended recurrence

new_state[j] = Σ_i decay_chunk[j, i] · state[i]

into

new_state[q] = state[q] · Σ_p decay_chunk[p, q]

i.e. the per-chunk states are no longer propagated across chunk boundaries.

This is the same bug that was reported for mamba2 in #34817 and fixed in #35154; the fix was never propagated to zamba2, and therefore to nemotron_h, whose mixer inherits torch_forward from Zamba2MambaMixer. All other Mamba2 descendants (bamba, falcon_h1, granitemoehybrid) already use the corrected form. The change aligns these two models with mamba2:

decay_chunk = decay_chunk.transpose(1, 3)
new_states = (decay_chunk[..., None, None] * states[:, :, None, ...]).sum(dim=1)

It was applied to modular_zamba2.py and mirrored into the generated modeling_zamba2.py and modeling_nemotron_h.py (identical lines, no name substitution), so check_modular_conversion stays green.

Scope of impact

Only the slow path is affected (no mamba-ssm / causal-conv1d kernels — e.g. CPU, or use_mamba_kernels=False); the CUDA kernel path is correct. The bug is masked for a single chunk with an empty cache (seq_len <= chunk_size, fresh prefill), which is why existing short-sequence tests pass. It produces wrong outputs and wrong cached SSM states when:

the input is longer than chunk_size (prefill spanning ≥ 2 chunks), or
generation is continued from a populated SSM cache.

Reproduction

Standalone numerical check (no model download)

import torch


def segment_sum(input_tensor):  # verbatim from modeling_*.py
    chunk_size = input_tensor.size(-1)
    input_tensor = input_tensor[..., None].expand(*input_tensor.size(), chunk_size)
    mask = torch.tril(torch.ones(chunk_size, chunk_size, dtype=torch.bool), diagonal=-1)
    input_tensor = input_tensor.masked_fill(~mask, 0)
    tensor_segsum = torch.cumsum(input_tensor, dim=-2)
    mask = torch.tril(torch.ones(chunk_size, chunk_size, dtype=torch.bool), diagonal=0)
    return tensor_segsum.masked_fill(~mask, -torch.inf)


def run(num_chunks, zero_previous_state):
    b, h, d, n = 1, 2, 4, 5
    A_cumsum_last = torch.randn(b, h, num_chunks)          # == A_cumsum[:, :, :, -1]
    states = torch.randn(b, num_chunks + 1, h, d, n)       # cat([previous_states, per-chunk states])
    if zero_previous_state:
        states[:, :1] = 0                                  # fresh prefill => no cached state
    decay_chunk = torch.exp(segment_sum(torch.nn.functional.pad(A_cumsum_last, (1, 0))))

    # current (buggy): sum over dim=2
    sp = states.permute(0, 2, 1, 3, 4)
    buggy = (decay_chunk[..., None, None] * sp[:, :, None, ...]).sum(dim=2).permute(0, 2, 1, 3, 4)
    # fixed (mamba2 form): transpose(1, 3) + sum over dim=1
    dc_t = decay_chunk.transpose(1, 3)
    fixed = (dc_t[..., None, None] * states[:, :, None, ...]).sum(dim=1)
    # brute-force reference: new[j] = sum_i decay_chunk[h, j, i] * states[i]
    ref = torch.einsum("bhji,bihdn->bjhdn", decay_chunk, states)
    return (buggy - ref).abs().max().item(), (fixed - ref).abs().max().item()


torch.manual_seed(0)
for label, c, zero in [
    ("seq <= chunk_size, fresh prefill (empty cache)", 1, True),
    ("seq <= chunk_size, continuing from SSM cache",   1, False),
    ("seq  > chunk_size, fresh prefill",               3, True),
    ("seq  > chunk_size, continuing from SSM cache",   3, False),
]:
    bug, ok = run(c, zero)
    print(f"{label:48s}  buggy_vs_ref={bug:9.5f}  fixed_vs_ref={ok:9.5f}")

Output:

seq <= chunk_size, fresh prefill (empty cache)    buggy_vs_ref=  0.00000  fixed_vs_ref=  0.00000
seq <= chunk_size, continuing from SSM cache      buggy_vs_ref=  1.92688  fixed_vs_ref=  0.00000
seq  > chunk_size, fresh prefill                  buggy_vs_ref=  5.26675  fixed_vs_ref=  0.00000
seq  > chunk_size, continuing from SSM cache      buggy_vs_ref= 10.68568  fixed_vs_ref=  0.00000

Fixes #47246

Who can review?

@vasqu (fixed the equivalent mamba2 bug in #35154)

…on-H Mamba2 slow path The pure-PyTorch (no-kernel) path of the Mamba2 mixer summed the inter-chunk recurrence over the wrong dimension, reducing new_state[j] = sum_i decay_chunk[j, i] * state[i] to new_state[q] = state[q] * sum_p decay_chunk[p, q] so the per-chunk states were no longer mixed across chunk boundaries. Align the implementation with mamba2 (transpose(1, 3) + sum over dim=1), matching the fix made for mamba2 in huggingface#35154. This affects Zamba2 and Nemotron-H (which inherits torch_forward from Zamba2MambaMixer). The bug is masked for a single chunk with an empty cache, but produces wrong outputs and wrong cached SSM states for sequences longer than chunk_size or when continuing generation from a populated cache.

vasqu

Thanks! Can we add test(s) similar to create_and_check_mamba2_slow_vs_fast_forward in the other PR that fixed this

Add a slow-path multi-chunk regression test for Zamba2 and Nemotron-H that checks a single chunked `torch_forward` over a multi-chunk sequence matches a token-by-token recurrent decode. A large-magnitude input is used so the SSM state is O(1) and the inter-chunk term is observable; with the previous reduction the two disagree by orders of magnitude. Runs on CPU without the fast-path kernels. Also add the Mamba2 slow-vs-fast (kernel vs torch) consistency test to Zamba2, mirroring the existing mamba2 and nemotron_h tests.

vasqu

This is more about the tests (design), the fix itself is good!

vasqu · 2026-07-13T13:36:14Z

        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_mamba2_slow_vs_fast_forward(*config_and_inputs)

+    def test_mamba2_slow_path_multi_chunk(self):


Hmm was create_and_check_mamba2_slow_vs_fast_forward not enough? Or why the whole new test design?

Thanks. I reworked into a create_and_check_* helper driven by prepare_config_and_inputs.

vasqu · 2026-07-13T13:37:15Z

+        config = NemotronHConfig(
+            vocab_size=99,
+            hidden_size=32,
+            mamba_num_heads=8,
+            mamba_head_dim=8,
+            ssm_state_size=16,
+            n_groups=1,
+            mamba_chunk_size=8,
+            num_attention_heads=2,
+            num_key_value_heads=2,
+            head_dim=8,
+            intermediate_size=32,
+            use_mamba_kernels=False,
+            layers_block_type=["mamba"],
+        )
+        torch.manual_seed(0)
+        mixer = NemotronHModel(config).eval().to(torch_device).layers[0].mixer
+
+        seq_len = 5 * config.chunk_size + 3
+        hidden_states = 50.0 * torch.randn(1, seq_len, config.hidden_size, device=torch_device)


I would like to use prepare config and inputs if possible (if this test is still really needed after my comment above)

vasqu · 2026-07-13T13:37:50Z

            msg=f"Max diff: {(ref_first - under_test_first).abs().max().item():.6f}",
        )

+    def create_and_check_zamba2_slow_vs_fast_forward(self, config, input_ids, *args):


I think it changes a bit on main re mamba2 especially re decorators can you check

Zamba2 slow_vs_fast now gated with @require_torch_accelerator + @require_kernels, helper trimmed to match mamba2.

Wire the multi-chunk regression tests through prepare_config_and_inputs as create_and_check_* helpers, and gate the Zamba2 slow-vs-fast test with require_torch_accelerator + require_kernels to match mamba2.

github-actions · 2026-07-15T10:21:51Z

CI recap

Dashboard: View test results in Grafana
Latest run: 29114117332:2
Result: success | Jobs: 5 | Tests: 646 | Failures: 0 | Duration: 5m 16s

vasqu

Test design questions again

vasqu · 2026-07-15T13:06:05Z

+        seq_len = 4 * config.chunk_size + 1
+        torch.manual_seed(0)
+        hidden_states = 100.0 * torch.randn(1, seq_len, config.hidden_size, device=torch_device)


I mean we are still manually creating here, no? The idea was to directly use the created input ids

OK, thanks for clarifying. It should be fixed now.
feeds model.embeddings(input_ids) (the prepare_config_and_inputs ids) instead of a manual tensor. It's only rescaled so the SSM state is O(1), otherwise the inter-chunk term is ~1e-7 and invisible.

vasqu · 2026-07-15T13:07:14Z

            msg=f"Max diff: {(ref_first - under_test_first).abs().max().item():.6f}",
        )

+    def create_and_check_nemotron_h_slow_path_multi_chunk(self, config, input_ids, *args):


But I'm still a bit confused - why do we have this new test and not just the slow vs fast path test? (which should catch the same regression)

slow_vs_fast can't catch this one: it runs seq_length=7 < chunk_size, i.e. a single chunk with an empty cache, so the inter-chunk recurrence is never reached (previous_states=0, nothing to mix). Buggy and fixed give identical outputs there. It passes unchanged on the buggy code. The regression only shows with more than 2 chunks, hence the separate test.

vasqu · 2026-07-15T13:08:35Z

+        """
+        config = copy.deepcopy(config)
+        config.chunk_size = 8
+        model = NemotronHModel(config).eval().to(torch_device)


I mean we need to guarantee that this uses the slow path and here it could result into the fast path. Wouldnt it make more sense to just force cpu here?

Done. The model is now .to("cpu") so it always takes the torch_forward slow path regardless of kernel availability.

Drive the multi-chunk regression tests from the prepare_config_and_inputs input_ids and force CPU so the torch (slow) path is exercised.

github-actions · 2026-07-15T14:54:52Z

Thank you for your contribution 🤗!

CI Security Gate — automatic approval blocked

This PR was not automatically approved for CI because the security gate failed.

Possible reasons:

The PR touches 50 or more files — only PRs with fewer than 50 changed files are automatically approved
A changed file is outside the allowed directories (src/, tests/, docs/, utils/), has a disallowed extension (only .py, .txt, .md permitted outside tests/ and docs/), or is not .md/.yml inside docs/
A new high-severity security issue was detected in the changed Python files (Bandit check)

See the workflow run for the exact violations.

A maintainer can review and manually approve CI if a finding is a false positive.

github-actions · 2026-07-15T14:55:26Z

[For maintainers] Suggested jobs to run (before merge)

run-slow: nemotron_h, zamba2

vasqu

Gotcha thanks, I will check tomorrow again. Could you ping me then? I think we are mostly fine to merge

vasqu · 2026-07-15T16:21:55Z

run-slow: nemotron_h, zamba2

github-actions · 2026-07-15T16:23:34Z

Workflow Run ⚙️

This comment contains run-slow, running the specified jobs:

models: ["models/nemotron_h", "models/zamba2"]
quantizations: []

github-actions · 2026-07-15T16:44:41Z

CI Results

Workflow Run ⚙️

Commit Info

Context	Commit	Description
RUN	4dbe8928	workflow commit (merge commit)
PR	65abbe6b	branch commit (from PR)
main	e52d0fd6	base commit (on `main`)

⚠️ Model CI failed to report results

The test failure analysis could not be completed. Please check the workflow run for details.

Jeronymous · 2026-08-17T07:07:48Z

@vasqu is there anything still blocking this? It's approved, the slow tests were triggered (run-slow: nemotron_h, zamba2).

Also: the closed #47513 flagged a related slow-path issue in the same mixer (dt is clamped one-sided instead of two-sided against time_step_limit, unlike the kernel path). Want me to fold that fix into this PR so it fully closes #47246, or keep it separate?

vasqu · 2026-08-24T07:07:53Z

Hey @Jeronymous 👋 sorry I refactored all the mamba / lin attn models in #47630 so they all should work now, could you recheck? I lost track and was out of office for a bit sorry about that

Jeronymous · 2026-08-24T08:57:48Z

Rechecked. Looks good.
The inter-chunk recurrence on main now uses the correct transpose(1, 3) + sum(dim=1) form in both zamba2 and nemotron_h, and dt is two-sided clamped now too. So the fix from this PR is effectively in via #47630.

I close this PR since the fix is redundant now. One thing worth keeping though: the test_mamba2_slow_path_multi_chunk regression test I added here : it catches exactly this bug (multi-chunk slow-path forward vs. token-by-token recurrent decode, which the single-chunk slow_vs_fast test misses).

vasqu · 2026-08-25T06:45:51Z

Hey, I don't mind to add the test into e.g. only mamba2 as a baseline wdyt?

Jeronymous · 2026-08-25T16:06:18Z

Yes @vasqu it sounds good

Jeronymous mentioned this pull request Jul 10, 2026

Incorrect inter-chunk recurrence in NemotronHMamba2Mixer.torch_forward (slow path). the Mamba2 fix from #35154 was never propagated to nemotron_h #47246

Closed

4 tasks

Jeronymous force-pushed the fix-mamba2-slow-path-inter-chunk-recurrence branch from b5b0aad to 4415fa5 Compare July 10, 2026 14:15

vasqu approved these changes Jul 10, 2026

View reviewed changes

vasqu reviewed Jul 13, 2026

View reviewed changes

Address review on the slow-path tests

10f37af

Wire the multi-chunk regression tests through prepare_config_and_inputs as create_and_check_* helpers, and gate the Zamba2 slow-vs-fast test with require_torch_accelerator + require_kernels to match mamba2.

vasqu reviewed Jul 15, 2026

View reviewed changes

Refine slow-path tests per review

65abbe6

Drive the multi-chunk regression tests from the prepare_config_and_inputs input_ids and force CPU so the torch (slow) path is exercised.

vasqu approved these changes Jul 15, 2026

View reviewed changes

Rocketknight1 mentioned this pull request Jul 24, 2026

fix(nemotron_h): correct inter-chunk SSM recurrence axis and dt clamping in torch_forward #47513

Closed

Jeronymous closed this Aug 24, 2026

Uh oh!

Conversation

Jeronymous commented Jul 10, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Scope of impact

Reproduction

Who can review?

Uh oh!

vasqu left a comment

Choose a reason for hiding this comment

Uh oh!

vasqu left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Jul 15, 2026

CI recap

Uh oh!

vasqu left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Jul 15, 2026

CI Security Gate — automatic approval blocked

Uh oh!

github-actions Bot commented Jul 15, 2026

Uh oh!

vasqu left a comment

Choose a reason for hiding this comment

Uh oh!

vasqu commented Jul 15, 2026

Uh oh!

github-actions Bot commented Jul 15, 2026

Uh oh!

github-actions Bot commented Jul 15, 2026

CI Results

Commit Info

Uh oh!

Jeronymous commented Aug 17, 2026

Uh oh!

vasqu commented Aug 24, 2026

Uh oh!

Jeronymous commented Aug 24, 2026

Uh oh!

vasqu commented Aug 25, 2026

Uh oh!

Jeronymous commented Aug 25, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Jeronymous commented Jul 10, 2026 •

edited by github-actions Bot

Loading