From 4415fa5f3c1fc7ae1c98508ad35fc1f1eb93299d Mon Sep 17 00:00:00 2001 From: Jeronymous Date: Fri, 10 Jul 2026 15:55:22 +0200 Subject: [PATCH 1/4] Fixes #47246 : Fix inter-chunk recurrence in Zamba2/Nemotron-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 #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. --- src/transformers/models/nemotron_h/modeling_nemotron_h.py | 5 ++--- src/transformers/models/zamba2/modeling_zamba2.py | 5 ++--- src/transformers/models/zamba2/modular_zamba2.py | 5 ++--- 3 files changed, 6 insertions(+), 9 deletions(-) diff --git a/src/transformers/models/nemotron_h/modeling_nemotron_h.py b/src/transformers/models/nemotron_h/modeling_nemotron_h.py index 4f141b1628f5..f8dccf8ed654 100644 --- a/src/transformers/models/nemotron_h/modeling_nemotron_h.py +++ b/src/transformers/models/nemotron_h/modeling_nemotron_h.py @@ -551,9 +551,8 @@ def torch_forward(self, input_states, cache_params: Cache | None=None, attention states = torch.cat([previous_states, states], dim=1) decay_chunk = torch.exp(segment_sum(nn.functional.pad(A_cumsum[:, :, :, -1], (1, 0)))) - 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) + decay_chunk = decay_chunk.transpose(1, 3) + new_states = (decay_chunk[..., None, None] * states[:, :, None, ...]).sum(dim=1) states, ssm_state = new_states[:, :-1], new_states[:, -1] # Compute state -> output conversion per chunk diff --git a/src/transformers/models/zamba2/modeling_zamba2.py b/src/transformers/models/zamba2/modeling_zamba2.py index 35ab4bb3c677..639f67231f9a 100644 --- a/src/transformers/models/zamba2/modeling_zamba2.py +++ b/src/transformers/models/zamba2/modeling_zamba2.py @@ -839,9 +839,8 @@ def torch_forward(self, input_states, cache_params: Cache | None=None, attention states = torch.cat([previous_states, states], dim=1) decay_chunk = torch.exp(segment_sum(nn.functional.pad(A_cumsum[:, :, :, -1], (1, 0)))) - 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) + decay_chunk = decay_chunk.transpose(1, 3) + new_states = (decay_chunk[..., None, None] * states[:, :, None, ...]).sum(dim=1) states, ssm_state = new_states[:, :-1], new_states[:, -1] # Compute state -> output conversion per chunk diff --git a/src/transformers/models/zamba2/modular_zamba2.py b/src/transformers/models/zamba2/modular_zamba2.py index 6f4494313ac2..fc546171dc29 100644 --- a/src/transformers/models/zamba2/modular_zamba2.py +++ b/src/transformers/models/zamba2/modular_zamba2.py @@ -627,9 +627,8 @@ def torch_forward(self, input_states, cache_params: Cache | None=None, attention states = torch.cat([previous_states, states], dim=1) decay_chunk = torch.exp(segment_sum(nn.functional.pad(A_cumsum[:, :, :, -1], (1, 0)))) - 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) + decay_chunk = decay_chunk.transpose(1, 3) + new_states = (decay_chunk[..., None, None] * states[:, :, None, ...]).sum(dim=1) states, ssm_state = new_states[:, :-1], new_states[:, -1] # Compute state -> output conversion per chunk From c9288cd57f1e1e2841e34fe148bcf658b9973dd5 Mon Sep 17 00:00:00 2001 From: Jeronymous Date: Fri, 10 Jul 2026 20:19:08 +0200 Subject: [PATCH 2/4] Add regression tests for the Mamba2 slow-path inter-chunk recurrence 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. --- .../nemotron_h/test_modeling_nemotron_h.py | 42 ++++++++++ tests/models/zamba2/test_modeling_zamba2.py | 78 +++++++++++++++++++ 2 files changed, 120 insertions(+) diff --git a/tests/models/nemotron_h/test_modeling_nemotron_h.py b/tests/models/nemotron_h/test_modeling_nemotron_h.py index 67ec79c23856..cd49ad92ac86 100644 --- a/tests/models/nemotron_h/test_modeling_nemotron_h.py +++ b/tests/models/nemotron_h/test_modeling_nemotron_h.py @@ -520,6 +520,48 @@ def test_mamba2_slow_vs_fast_forward(self): 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): + """ + Regression test for the inter-chunk recurrence in the Mamba2 slow (torch) path. + + A single chunked forward over a multi-chunk sequence must match a token-by-token + recurrent decode. The input is deliberately large-magnitude so the SSM state is + O(1) and the inter-chunk contribution is not numerically negligible; with the + previous `.sum(dim=2)` reduction the two disagree by orders of magnitude. Runs on + CPU without the fast-path kernels. + """ + 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) + + with torch.no_grad(): + chunked = mixer.torch_forward(hidden_states) + cache = DynamicCache(config=config) + recurrent = torch.cat( + [mixer.torch_forward(hidden_states[:, t : t + 1], cache_params=cache) for t in range(seq_len)], + dim=1, + ) + + max_diff = (chunked - recurrent).abs().max().item() + self.assertLess(max_diff, 1e-3, f"slow-path chunked forward disagrees with recurrent decode: {max_diff}") + def test_attention_outputs(self): r""" Overriding the test_attention_outputs test as the NemotronH model outputs attention only for its attention layers diff --git a/tests/models/zamba2/test_modeling_zamba2.py b/tests/models/zamba2/test_modeling_zamba2.py index f5a570acfc06..9c013aca01d7 100644 --- a/tests/models/zamba2/test_modeling_zamba2.py +++ b/tests/models/zamba2/test_modeling_zamba2.py @@ -30,6 +30,7 @@ slow, torch_device, ) +from transformers.utils.import_utils import is_causal_conv1d_available, is_mamba_ssm_available from ...generation.test_utils import GenerationTesterMixin from ...test_configuration_common import ConfigTester @@ -298,6 +299,34 @@ def create_and_check_zamba2_chunked_prefill(self, config, input_ids, *args, devi 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): + """ + Test that cuda_kernels_forward and torch_forward produce consistent outputs for the + Mamba2 mixer, i.e. that the optimized CUDA kernel path and the pure PyTorch path are + equivalent. Guarded by the availability of the fast-path kernels and a CUDA device. + """ + if not (is_mamba_ssm_available() and is_causal_conv1d_available()): + self.parent.skipTest( + "This test needs the Mamba2 fast path. Skipping as the necessary packages have not been found." + ) + if torch_device != "cuda": + self.parent.skipTest("This test needs the Mamba2 fast path. Skipping as we need a cuda capable device.") + + model = Zamba2Model(config) + model.eval() + model.to(torch_device) + + # Find the first Mamba mixer in the model + mamba_mixer = next((layer.mamba for layer in model.layers if hasattr(layer, "mamba")), None) + if mamba_mixer is None: + self.parent.skipTest("No mamba layer found in the model configuration.") + + hidden_states = model.embed_tokens(input_ids.to(torch_device)) + + outputs_fast = mamba_mixer.cuda_kernels_forward(hidden_states) + outputs_slow = mamba_mixer.torch_forward(hidden_states) + self.parent.assertTrue(torch.allclose(outputs_fast, outputs_slow, atol=1e-3, rtol=1e-3)) + def prepare_config_and_inputs_for_common(self): config_and_inputs = self.prepare_config_and_inputs() ( @@ -351,6 +380,55 @@ def setUp(self): self.model_tester = Zamba2ModelTester(self) self.config_tester = ConfigTester(self, config_class=Zamba2Config, hidden_size=32) + def test_mamba2_slow_vs_fast_forward(self): + """ + Test that cuda_kernels_forward and torch_forward produce consistent outputs. + """ + config_and_inputs = self.model_tester.prepare_config_and_inputs() + self.model_tester.create_and_check_zamba2_slow_vs_fast_forward(*config_and_inputs) + + def test_mamba2_slow_path_multi_chunk(self): + """ + Regression test for the inter-chunk recurrence in the Mamba2 slow (torch) path. + + A single chunked forward over a multi-chunk sequence must match a token-by-token + recurrent decode. The input is deliberately large-magnitude so the SSM state is + O(1) and the inter-chunk contribution is not numerically negligible; with the + previous `.sum(dim=2)` reduction the two disagree by orders of magnitude. Runs on + CPU without the fast-path kernels. + """ + config = Zamba2Config( + vocab_size=99, + hidden_size=32, + mamba_d_state=16, + num_hidden_layers=1, + num_attention_heads=2, + n_mamba_heads=8, + intermediate_size=8, + chunk_size=8, + mamba_ngroups=1, + use_mamba_kernels=False, + layers_block_type=["mamba"], + num_mem_blocks=1, + use_mem_rope=True, + ) + torch.manual_seed(0) + mixer = Zamba2Model(config).eval().to(torch_device).layers[0].mamba + + seq_len = 5 * config.chunk_size + 3 + hidden_states = 50.0 * torch.randn(1, seq_len, config.hidden_size, device=torch_device) + + with torch.no_grad(): + chunked = mixer.torch_forward(hidden_states) + cache = DynamicCache(config=config) + recurrent = torch.cat( + [mixer.torch_forward(hidden_states[:, t : t + 1], cache_params=cache) for t in range(seq_len)], + dim=1, + ) + + max_diff = (chunked - recurrent).abs().max().item() + self.assertLess(max_diff, 1e-3, f"slow-path chunked forward disagrees with recurrent decode: {max_diff}") + @unittest.skip("We need at leat 3 layers to test weight tying!") def test_num_layers_is_small(self): pass From 10f37af759ee880737dbe5c9def8742bba4aee78 Mon Sep 17 00:00:00 2001 From: Jeronymous Date: Wed, 15 Jul 2026 12:13:26 +0200 Subject: [PATCH 3/4] Address review on the slow-path tests 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. --- .../nemotron_h/test_modeling_nemotron_h.py | 79 +++++++------ tests/models/zamba2/test_modeling_zamba2.py | 107 +++++++----------- 2 files changed, 83 insertions(+), 103 deletions(-) diff --git a/tests/models/nemotron_h/test_modeling_nemotron_h.py b/tests/models/nemotron_h/test_modeling_nemotron_h.py index cd49ad92ac86..a9743062c1cf 100644 --- a/tests/models/nemotron_h/test_modeling_nemotron_h.py +++ b/tests/models/nemotron_h/test_modeling_nemotron_h.py @@ -13,6 +13,7 @@ # limitations under the License. """Testing suite for the PyTorch NemotronH model.""" +import copy import tempfile import unittest @@ -356,6 +357,41 @@ def create_and_check_nemotron_h_chunked_prefill(self, config, input_ids, *args, 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): + """ + Regression test for the inter-chunk recurrence of the Mamba2 slow (torch) path: a + single chunked forward over a multi-chunk sequence must reproduce a token-by-token + recurrent decode. + + The mixer is fed a large-magnitude input so the SSM state is O(1). With a normal + activation scale the inter-chunk contribution is ~1e-7, i.e. below any tolerance, + which is why neither `create_and_check_mamba2_slow_vs_fast_forward` (single chunk, + needs the kernels) nor a plain `prepare_config_and_inputs` input surfaces the bug. + A small `chunk_size` lets a short sequence span several chunks. Runs on CPU without + the fast-path kernels. + """ + config = copy.deepcopy(config) + config.chunk_size = 8 + model = NemotronHModel(config).eval().to(torch_device) + mixer = next(layer.mixer for layer in model.layers if getattr(layer, "block_type", None) == "linear_attention") + + 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) + + with torch.no_grad(): + chunked = mixer.torch_forward(hidden_states) + cache = DynamicCache(config=config) + recurrent = torch.cat( + [mixer.torch_forward(hidden_states[:, t : t + 1], cache_params=cache) for t in range(seq_len)], + dim=1, + ) + + max_diff = (chunked - recurrent).abs().max().item() + self.parent.assertLess( + max_diff, 1e-3, f"slow-path chunked forward disagrees with recurrent decode: {max_diff}" + ) + def prepare_config_and_inputs_for_common(self): config_and_inputs = self.prepare_config_and_inputs() ( @@ -521,46 +557,9 @@ def test_mamba2_slow_vs_fast_forward(self): self.model_tester.create_and_check_mamba2_slow_vs_fast_forward(*config_and_inputs) def test_mamba2_slow_path_multi_chunk(self): - """ - Regression test for the inter-chunk recurrence in the Mamba2 slow (torch) path. - - A single chunked forward over a multi-chunk sequence must match a token-by-token - recurrent decode. The input is deliberately large-magnitude so the SSM state is - O(1) and the inter-chunk contribution is not numerically negligible; with the - previous `.sum(dim=2)` reduction the two disagree by orders of magnitude. Runs on - CPU without the fast-path kernels. - """ - 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) - - with torch.no_grad(): - chunked = mixer.torch_forward(hidden_states) - cache = DynamicCache(config=config) - recurrent = torch.cat( - [mixer.torch_forward(hidden_states[:, t : t + 1], cache_params=cache) for t in range(seq_len)], - dim=1, - ) - - max_diff = (chunked - recurrent).abs().max().item() - self.assertLess(max_diff, 1e-3, f"slow-path chunked forward disagrees with recurrent decode: {max_diff}") + """The Mamba2 slow path must reproduce the token-by-token recurrence across chunks.""" + config_and_inputs = self.model_tester.prepare_config_and_inputs() + self.model_tester.create_and_check_nemotron_h_slow_path_multi_chunk(*config_and_inputs) def test_attention_outputs(self): r""" diff --git a/tests/models/zamba2/test_modeling_zamba2.py b/tests/models/zamba2/test_modeling_zamba2.py index 9c013aca01d7..4eb90e00a2f5 100644 --- a/tests/models/zamba2/test_modeling_zamba2.py +++ b/tests/models/zamba2/test_modeling_zamba2.py @@ -13,6 +13,7 @@ # limitations under the License. """Testing suite for the PyTorch Zamba model.""" +import copy import tempfile import unittest @@ -30,7 +31,6 @@ slow, torch_device, ) -from transformers.utils.import_utils import is_causal_conv1d_available, is_mamba_ssm_available from ...generation.test_utils import GenerationTesterMixin from ...test_configuration_common import ConfigTester @@ -300,33 +300,52 @@ def create_and_check_zamba2_chunked_prefill(self, config, input_ids, *args, devi ) def create_and_check_zamba2_slow_vs_fast_forward(self, config, input_ids, *args): - """ - Test that cuda_kernels_forward and torch_forward produce consistent outputs for the - Mamba2 mixer, i.e. that the optimized CUDA kernel path and the pure PyTorch path are - equivalent. Guarded by the availability of the fast-path kernels and a CUDA device. - """ - if not (is_mamba_ssm_available() and is_causal_conv1d_available()): - self.parent.skipTest( - "This test needs the Mamba2 fast path. Skipping as the necessary packages have not been found." - ) - if torch_device != "cuda": - self.parent.skipTest("This test needs the Mamba2 fast path. Skipping as we need a cuda capable device.") - + """Slow vs fast path check guarded by require kernels to enable fast path""" model = Zamba2Model(config) model.eval() model.to(torch_device) - # Find the first Mamba mixer in the model - mamba_mixer = next((layer.mamba for layer in model.layers if hasattr(layer, "mamba")), None) - if mamba_mixer is None: - self.parent.skipTest("No mamba layer found in the model configuration.") - - hidden_states = model.embed_tokens(input_ids.to(torch_device)) - + mamba_mixer = next(layer.mamba for layer in model.layers if hasattr(layer, "mamba")) + hidden_states = model.embed_tokens(input_ids) outputs_fast = mamba_mixer.cuda_kernels_forward(hidden_states) outputs_slow = mamba_mixer.torch_forward(hidden_states) self.parent.assertTrue(torch.allclose(outputs_fast, outputs_slow, atol=1e-3, rtol=1e-3)) + def create_and_check_zamba2_slow_path_multi_chunk(self, config, input_ids, *args): + """ + Regression test for the inter-chunk recurrence of the Mamba2 slow (torch) path: a + single chunked forward over a multi-chunk sequence must reproduce a token-by-token + recurrent decode. + + The mixer is fed a large-magnitude input so the SSM state is O(1). With a normal + activation scale the inter-chunk contribution is ~1e-7, i.e. below any tolerance, + which is why neither `create_and_check_zamba2_slow_vs_fast_forward` (single chunk, + needs the kernels) nor a plain `prepare_config_and_inputs` input surfaces the bug. + A small `chunk_size` lets a short sequence span several chunks. Runs on CPU without + the fast-path kernels. + """ + config = copy.deepcopy(config) + config.chunk_size = 8 + model = Zamba2Model(config).eval().to(torch_device) + mixer = next(layer.mamba for layer in model.layers if hasattr(layer, "mamba")) + + 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) + + with torch.no_grad(): + chunked = mixer.torch_forward(hidden_states) + cache = DynamicCache(config=config) + recurrent = torch.cat( + [mixer.torch_forward(hidden_states[:, t : t + 1], cache_params=cache) for t in range(seq_len)], + dim=1, + ) + + max_diff = (chunked - recurrent).abs().max().item() + self.parent.assertLess( + max_diff, 1e-3, f"slow-path chunked forward disagrees with recurrent decode: {max_diff}" + ) + def prepare_config_and_inputs_for_common(self): config_and_inputs = self.prepare_config_and_inputs() ( @@ -380,54 +399,16 @@ def setUp(self): self.model_tester = Zamba2ModelTester(self) self.config_tester = ConfigTester(self, config_class=Zamba2Config, hidden_size=32) + @require_torch_accelerator + @require_kernels def test_mamba2_slow_vs_fast_forward(self): - """ - Test that cuda_kernels_forward and torch_forward produce consistent outputs. - """ config_and_inputs = self.model_tester.prepare_config_and_inputs() self.model_tester.create_and_check_zamba2_slow_vs_fast_forward(*config_and_inputs) def test_mamba2_slow_path_multi_chunk(self): - """ - Regression test for the inter-chunk recurrence in the Mamba2 slow (torch) path. - - A single chunked forward over a multi-chunk sequence must match a token-by-token - recurrent decode. The input is deliberately large-magnitude so the SSM state is - O(1) and the inter-chunk contribution is not numerically negligible; with the - previous `.sum(dim=2)` reduction the two disagree by orders of magnitude. Runs on - CPU without the fast-path kernels. - """ - config = Zamba2Config( - vocab_size=99, - hidden_size=32, - mamba_d_state=16, - num_hidden_layers=1, - num_attention_heads=2, - n_mamba_heads=8, - intermediate_size=8, - chunk_size=8, - mamba_ngroups=1, - use_mamba_kernels=False, - layers_block_type=["mamba"], - num_mem_blocks=1, - use_mem_rope=True, - ) - torch.manual_seed(0) - mixer = Zamba2Model(config).eval().to(torch_device).layers[0].mamba - - seq_len = 5 * config.chunk_size + 3 - hidden_states = 50.0 * torch.randn(1, seq_len, config.hidden_size, device=torch_device) - - with torch.no_grad(): - chunked = mixer.torch_forward(hidden_states) - cache = DynamicCache(config=config) - recurrent = torch.cat( - [mixer.torch_forward(hidden_states[:, t : t + 1], cache_params=cache) for t in range(seq_len)], - dim=1, - ) - - max_diff = (chunked - recurrent).abs().max().item() - self.assertLess(max_diff, 1e-3, f"slow-path chunked forward disagrees with recurrent decode: {max_diff}") + """The Mamba2 slow path must reproduce the token-by-token recurrence across chunks.""" + config_and_inputs = self.model_tester.prepare_config_and_inputs() + self.model_tester.create_and_check_zamba2_slow_path_multi_chunk(*config_and_inputs) @unittest.skip("We need at leat 3 layers to test weight tying!") def test_num_layers_is_small(self): From 65abbe6b5316222a275013debfb94c04174f8ee0 Mon Sep 17 00:00:00 2001 From: Jeronymous Date: Wed, 15 Jul 2026 16:53:48 +0200 Subject: [PATCH 4/4] Refine slow-path tests per review Drive the multi-chunk regression tests from the prepare_config_and_inputs input_ids and force CPU so the torch (slow) path is exercised. --- .../nemotron_h/test_modeling_nemotron_h.py | 29 ++++++++++--------- tests/models/zamba2/test_modeling_zamba2.py | 29 ++++++++++--------- 2 files changed, 30 insertions(+), 28 deletions(-) diff --git a/tests/models/nemotron_h/test_modeling_nemotron_h.py b/tests/models/nemotron_h/test_modeling_nemotron_h.py index a9743062c1cf..1086c202fed1 100644 --- a/tests/models/nemotron_h/test_modeling_nemotron_h.py +++ b/tests/models/nemotron_h/test_modeling_nemotron_h.py @@ -361,29 +361,30 @@ def create_and_check_nemotron_h_slow_path_multi_chunk(self, config, input_ids, * """ Regression test for the inter-chunk recurrence of the Mamba2 slow (torch) path: a single chunked forward over a multi-chunk sequence must reproduce a token-by-token - recurrent decode. - - The mixer is fed a large-magnitude input so the SSM state is O(1). With a normal - activation scale the inter-chunk contribution is ~1e-7, i.e. below any tolerance, - which is why neither `create_and_check_mamba2_slow_vs_fast_forward` (single chunk, - needs the kernels) nor a plain `prepare_config_and_inputs` input surfaces the bug. - A small `chunk_size` lets a short sequence span several chunks. Runs on CPU without - the fast-path kernels. + recurrent decode. Forced onto CPU so the slow (`torch_forward`) path is exercised. + + A small `chunk_size` lets the short `prepare_config_and_inputs` sequence span several + chunks, and the embedded input is rescaled so the SSM state is O(1) -- at the natural + activation scale the inter-chunk contribution is ~1e-7 (below any tolerance), which is + why the single-chunk `slow_vs_fast` check does not surface this regression. """ config = copy.deepcopy(config) - config.chunk_size = 8 - model = NemotronHModel(config).eval().to(torch_device) + config.chunk_size = 2 + torch.manual_seed(0) + model = NemotronHModel(config).eval().to("cpu") mixer = next(layer.mixer for layer in model.layers if getattr(layer, "block_type", None) == "linear_attention") - 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) + embeds = model.embeddings(input_ids[:1].to("cpu")) + hidden_states = 100.0 * embeds / embeds.std() with torch.no_grad(): chunked = mixer.torch_forward(hidden_states) cache = DynamicCache(config=config) recurrent = torch.cat( - [mixer.torch_forward(hidden_states[:, t : t + 1], cache_params=cache) for t in range(seq_len)], + [ + mixer.torch_forward(hidden_states[:, t : t + 1], cache_params=cache) + for t in range(hidden_states.shape[1]) + ], dim=1, ) diff --git a/tests/models/zamba2/test_modeling_zamba2.py b/tests/models/zamba2/test_modeling_zamba2.py index 4eb90e00a2f5..f9f239a94a4d 100644 --- a/tests/models/zamba2/test_modeling_zamba2.py +++ b/tests/models/zamba2/test_modeling_zamba2.py @@ -315,29 +315,30 @@ def create_and_check_zamba2_slow_path_multi_chunk(self, config, input_ids, *args """ Regression test for the inter-chunk recurrence of the Mamba2 slow (torch) path: a single chunked forward over a multi-chunk sequence must reproduce a token-by-token - recurrent decode. - - The mixer is fed a large-magnitude input so the SSM state is O(1). With a normal - activation scale the inter-chunk contribution is ~1e-7, i.e. below any tolerance, - which is why neither `create_and_check_zamba2_slow_vs_fast_forward` (single chunk, - needs the kernels) nor a plain `prepare_config_and_inputs` input surfaces the bug. - A small `chunk_size` lets a short sequence span several chunks. Runs on CPU without - the fast-path kernels. + recurrent decode. Forced onto CPU so the slow (`torch_forward`) path is exercised. + + A small `chunk_size` lets the short `prepare_config_and_inputs` sequence span several + chunks, and the embedded input is rescaled so the SSM state is O(1) -- at the natural + activation scale the inter-chunk contribution is ~1e-7 (below any tolerance), which is + why the single-chunk `slow_vs_fast` check does not surface this regression. """ config = copy.deepcopy(config) - config.chunk_size = 8 - model = Zamba2Model(config).eval().to(torch_device) + config.chunk_size = 2 + torch.manual_seed(0) + model = Zamba2Model(config).eval().to("cpu") mixer = next(layer.mamba for layer in model.layers if hasattr(layer, "mamba")) - 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) + embeds = model.embed_tokens(input_ids[:1].to("cpu")) + hidden_states = 100.0 * embeds / embeds.std() with torch.no_grad(): chunked = mixer.torch_forward(hidden_states) cache = DynamicCache(config=config) recurrent = torch.cat( - [mixer.torch_forward(hidden_states[:, t : t + 1], cache_params=cache) for t in range(seq_len)], + [ + mixer.torch_forward(hidden_states[:, t : t + 1], cache_params=cache) + for t in range(hidden_states.shape[1]) + ], dim=1, )