Upload modeling_nemotron_h.py

modeling_nemotron_h.py

@@ -335,7 +335,7 @@ class NemotronHMamba2Mixer(nn.Module):
         A = torch.arange(1, self.num_heads + 1)
         self.A_log = nn.Parameter(torch.log(A))
         self.A_log._no_weight_decay = True
-        self.norm = MambaRMSNormGated(self.intermediate_size, eps=self.layer_norm_epsilon, group_size=self.intermediate_size)
+        self.norm = MambaRMSNormGated(self.intermediate_size, eps=self.layer_norm_epsilon, group_size=self.intermediate_size // self.n_groups)
         self.D = nn.Parameter(torch.ones(self.num_heads))
         self.D._no_weight_decay = True
 
@@ -469,13 +469,14 @@ class NemotronHMamba2Mixer(nn.Module):
                         self.conv1d(hidden_states_B_C.transpose(1, 2))[..., :seq_len].transpose(1, 2)
                     )
                 else:
-                    hidden_states_B_C = causal_conv1d_fn(
-                        x=hidden_states_B_C.transpose(1, 2),
-                        weight=self.conv1d.weight.squeeze(1),
-                        bias=self.conv1d.bias,
-                        activation=self.activation,
-                    ).transpose(1, 2)
-
+                    # * Use torch.cuda.stream() to avoid NaN issues when using multiple GPUs
+                    with torch.cuda.stream(torch.cuda.default_stream(hidden_states_B_C.device)):
+                        hidden_states_B_C = causal_conv1d_fn(
+                            x=hidden_states_B_C.transpose(1, 2),
+                            weight=self.conv1d.weight.squeeze(1),
+                            bias=self.conv1d.bias,
+                            activation=self.activation,
+                        ).transpose(1, 2)
                 hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask)
                 hidden_states, B, C = torch.split(
                     hidden_states_B_C,
@@ -484,21 +485,23 @@ class NemotronHMamba2Mixer(nn.Module):
                 )
 
                 # 3. SSM transformation
-                scan_output, ssm_state = mamba_chunk_scan_combined(
-                    hidden_states.view(batch_size, seq_len, -1, self.head_dim),
-                    dt,
-                    A,
-                    B.view(batch_size, seq_len, self.n_groups, -1),
-                    C.view(batch_size, seq_len, self.n_groups, -1),
-                    chunk_size=self.chunk_size,
-                    D=self.D,
-                    z=None,
-                    seq_idx=None,
-                    return_final_states=True,
-                    dt_bias=self.dt_bias,
-                    dt_softplus=True,
-                    **dt_limit_kwargs,
-                )
+                # * Use torch.cuda.stream() to avoid NaN issues when using multiple GPUs
+                with torch.cuda.stream(torch.cuda.default_stream(hidden_states_B_C.device)):
+                    scan_output, ssm_state = mamba_chunk_scan_combined(
+                        hidden_states.view(batch_size, seq_len, -1, self.head_dim),
+                        dt,
+                        A,
+                        B.view(batch_size, seq_len, self.n_groups, -1),
+                        C.view(batch_size, seq_len, self.n_groups, -1),
+                        chunk_size=self.chunk_size,
+                        D=self.D,
+                        z=None,
+                        seq_idx=None,
+                        return_final_states=True,
+                        dt_bias=self.dt_bias,
+                        dt_softplus=True,
+                        **dt_limit_kwargs,
+                    )
 
                 # Init cache
                 if ssm_state is not None and cache_params is not None: