Repository: huggingface/transformers
Pull request: Fix core dumped when `NemotronH` is torch compiled

# What does this PR do?

We had (flaky)

```bash
tests/models/nemotron_h/test_modeling_nemotron_h.py::NemotronHModelTest::test_sdpa_can_compile_dynamic Fatal Python error: Segmentation fault
```


`NemotronHBlock.forward` creates a temporary `torch.cuda.Stream` on every forward pass.

During `torch.compile`, `dynamo` stores a weakref to it — but since it's a local variable with no other references, CUDA frees the underlying C++ object before dynamo's `reset_user_object_tracking` cleanup runs, leaving a dangling pointer → SIGSEGV.

Fix: guard the stream creation with not `is_torchdynamo_compiling()`, same pattern already used in `NemotronHMamba2Mixer.forward`. The stream is only needed at runtime, not during tracing.


Changed files:
filename: src/transformers/models/nemotron_h/modeling_nemotron_h.py
patch: @@ -20,7 +20,6 @@
 # limitations under the License.
 
 
-import contextlib
 import math
 from collections.abc import Callable
 from typing import Any
@@ -222,6 +221,9 @@ def segment_sum(input_tensor):
     return tensor_segsum
 
 
+is_fast_path_available = False
+
+
 class NemotronHMamba2Mixer(nn.Module):
     """
     Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`.
@@ -677,7 +679,11 @@ def forward(
         **kwargs,
     ):
         if is_fast_path_available and "cuda" in self.in_proj.weight.device.type and not is_torchdynamo_compiling():
-            return self.cuda_kernels_forward(hidden_states, cache_params, attention_mask)
+            # Use cuda stream to avoid NaN when using multiple GPUs, which is caused by multi-GPU synchronization issue.
+            # Mamba might launch on the default cuda stream that not strictly respect the current Pytorch cuda stream.
+            # This leads to kernel reading uninitialized memory before the data transfer is complete.
+            with torch.cuda.stream(torch.cuda.default_stream(hidden_states.device)):
+                return self.cuda_kernels_forward(hidden_states, cache_params, attention_mask)
 
         return self.torch_forward(hidden_states, cache_params, attention_mask)
 
@@ -1036,34 +1042,26 @@ def forward(
         use_cache: bool | None = False,
         **kwargs: Unpack[TransformersKwargs],
     ):
-        if hidden_states.device.type == "cuda":
-            # Use cuda stream to avoid NaN when using multiple GPUs, which is caused by multi-GPU synchronization issue.
-            # Mamba might launch on the default cuda stream that not strictly respect the current Pytorch cuda stream.
-            # This leads to kernel reading uninitialized memory before the data transfer is complete.
-            stream_context = torch.cuda.stream(torch.cuda.default_stream(hidden_states.device))
+        residual = hidden_states
+        hidden_states = self.norm(hidden_states.to(dtype=self.norm.weight.dtype))
+
+        if self.block_type == "mamba":
+            hidden_states = self.mixer(hidden_states, cache_params=past_key_values, attention_mask=attention_mask)
+        elif self.block_type == "attention":
+            hidden_states, _ = self.mixer(
+                hidden_states=hidden_states,
+                past_key_values=past_key_values,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                user_cache=use_cache,
+                **kwargs,
+            )
         else:
-            stream_context = contextlib.nullcontext()
-
-        with stream_context:
-            residual = hidden_states
-            hidden_states = self.norm(hidden_states.to(dtype=self.norm.weight.dtype))
-
-            if self.block_type == "mamba":
-                hidden_states = self.mixer(hidden_states, cache_params=past_key_values, attention_mask=attention_mask)
-            elif self.block_type == "attention":
-                hidden_states, _ = self.mixer(
-                    hidden_states=hidden_states,
-                    past_key_values=past_key_values,
-                    attention_mask=attention_mask,
-                    position_ids=position_ids,
-                    user_cache=use_cache,
-                    **kwargs,
-                )
-            else:
-                hidden_states = self.mixer(hidden_states)
+            hidden_states = self.mixer(hidden_states)
 
-            hidden_states = residual + hidden_states
-            return hidden_states
+        hidden_states = residual + hidden_states
+
+        return hidden_states
 
 
 class NemotronHPreTrainedModel(PreTrainedModel):

filename: src/transformers/models/nemotron_h/modular_nemotron_h.py
patch: @@ -15,7 +15,6 @@
 
 from __future__ import annotations
 
-import contextlib
 import math
 
 import torch
@@ -35,14 +34,16 @@
 from ...models.zamba.modeling_zamba import ZambaForCausalLM
 from ...models.zamba2.modeling_zamba2 import Zamba2HybridDynamicCache, Zamba2MambaMixer, Zamba2RMSNormGated
 from ...processing_utils import Unpack
-from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, is_torchdynamo_compiling, logging
 from ...utils.generic import merge_with_config_defaults
 from ...utils.output_capturing import capture_outputs
 from .configuration_nemotron_h import NemotronHConfig
 
 
 logger = logging.get_logger(__name__)
 
+is_fast_path_available = False
+
 
 class NemotronHHybridDynamicCache(Zamba2HybridDynamicCache):
     def __init__(
@@ -129,6 +130,22 @@ def __init__(self, config: NemotronHConfig, layer_idx: int | None = None):
 
         self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.use_bias)
 
+    def forward(
+        self,
+        hidden_states,
+        cache_params: NemotronHHybridDynamicCache | None = None,
+        attention_mask: torch.Tensor | None = None,
+        **kwargs,
+    ):
+        if is_fast_path_available and "cuda" in self.in_proj.weight.device.type and not is_torchdynamo_compiling():
+            # Use cuda stream to avoid NaN when using multiple GPUs, which is caused by multi-GPU synchronization issue.
+            # Mamba might launch on the default cuda stream that not strictly respect the current Pytorch cuda stream.
+            # This leads to kernel reading uninitialized memory before the data transfer is complete.
+            with torch.cuda.stream(torch.cuda.default_stream(hidden_states.device)):
+                return self.cuda_kernels_forward(hidden_states, cache_params, attention_mask)
+
+        return self.torch_forward(hidden_states, cache_params, attention_mask)
+
 
 class NemotronHRMSNorm(LlamaRMSNorm):
     pass
@@ -307,34 +324,26 @@ def forward(
         use_cache: bool | None = False,
         **kwargs: Unpack[TransformersKwargs],
     ):
-        if hidden_states.device.type == "cuda":
-            # Use cuda stream to avoid NaN when using multiple GPUs, which is caused by multi-GPU synchronization issue.
-            # Mamba might launch on the default cuda stream that not strictly respect the current Pytorch cuda stream.
-            # This leads to kernel reading uninitialized memory before the data transfer is complete.
-            stream_context = torch.cuda.stream(torch.cuda.default_stream(hidden_states.device))
+        residual = hidden_states
+        hidden_states = self.norm(hidden_states.to(dtype=self.norm.weight.dtype))
+
+        if self.block_type == "mamba":
+            hidden_states = self.mixer(hidden_states, cache_params=past_key_values, attention_mask=attention_mask)
+        elif self.block_type == "attention":
+            hidden_states, _ = self.mixer(
+                hidden_states=hidden_states,
+                past_key_values=past_key_values,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                user_cache=use_cache,
+                **kwargs,
+            )
         else:
-            stream_context = contextlib.nullcontext()
-
-        with stream_context:
-            residual = hidden_states
-            hidden_states = self.norm(hidden_states.to(dtype=self.norm.weight.dtype))
-
-            if self.block_type == "mamba":
-                hidden_states = self.mixer(hidden_states, cache_params=past_key_values, attention_mask=attention_mask)
-            elif self.block_type == "attention":
-                hidden_states, _ = self.mixer(
-                    hidden_states=hidden_states,
-                    past_key_values=past_key_values,
-                    attention_mask=attention_mask,
-                    position_ids=position_ids,
-                    user_cache=use_cache,
-                    **kwargs,
-                )
-            else:
-                hidden_states = self.mixer(hidden_states)
+            hidden_states = self.mixer(hidden_states)
 
-            hidden_states = residual + hidden_states
-            return hidden_states
+        hidden_states = residual + hidden_states
+
+        return hidden_states
 
 
 class NemotronHPreTrainedModel(PreTrainedModel):