Upload 2 files

configuration_nemotron_h.py

@@ -239,4 +239,4 @@ class NemotronHConfig(PretrainedConfig):
         return [
             "mamba" if self.hybrid_override_pattern[i] == "M" else
             "attention" if self.hybrid_override_pattern[i] == "*" else "mlp"
-            for i in range(self.num_hidden_layers)]
+            for i in range(self.num_hidden_layers)]

modeling_nemotron_h.py

@@ -469,14 +469,12 @@ class NemotronHMamba2Mixer(nn.Module):
                         self.conv1d(hidden_states_B_C.transpose(1, 2))[..., :seq_len].transpose(1, 2)
                     )
                 else:
-                    # * 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 = 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,
@@ -485,23 +483,21 @@ class NemotronHMamba2Mixer(nn.Module):
                 )
 
                 # 3. SSM transformation
-                # * 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,
-                    )
+                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:
@@ -768,30 +764,31 @@ class NemotronHBlock(nn.Module):
         cache_position: Optional[torch.LongTensor] = None,
         attention_mask: Optional[torch.Tensor] = None,
     ):
-        residual = hidden_states
-        hidden_states = self.norm(hidden_states.to(dtype=self.norm.weight.dtype))
-        if self.residual_in_fp32:
-            residual = residual.to(torch.float32)
-
-        if self.block_type == "mamba":
-            hidden_states = self.mixer(
-                hidden_states, cache_params=cache_params, cache_position=cache_position #, attention_mask=attention_mask
-            )
-        elif self.block_type == "attention":
-            hidden_states = self.mixer(
-                hidden_states, cache_position=cache_position #, attention_mask=attention_mask
-            )
-            # hidden_states = (attn_output, attn_weights, past_key_value)
-            hidden_states = hidden_states[0]
-        elif self.block_type == "mlp":
-            hidden_states = self.mixer(
-                hidden_states
-            )
-        else:
-            raise ValueError(f"Invalid block_type: {self.block_type}")
+        with torch.cuda.stream(torch.cuda.default_stream(hidden_states.device)):
+            # * Use torch.cuda.stream() to avoid NaN issues when using multiple GPUs
+            residual = hidden_states
+            hidden_states = self.norm(hidden_states.to(dtype=self.norm.weight.dtype))
+            if self.residual_in_fp32:
+                residual = residual.to(torch.float32)
+
+            if self.block_type == "mamba":
+                hidden_states = self.mixer(
+                    hidden_states, cache_params=cache_params, cache_position=cache_position
+                )
+            elif self.block_type == "attention":
+                hidden_states = self.mixer(
+                    hidden_states, cache_position=cache_position
+                )
+                hidden_states = hidden_states[0]
+            elif self.block_type == "mlp":
+                hidden_states = self.mixer(
+                    hidden_states
+                )
+            else:
+                raise ValueError(f"Invalid block_type: {self.block_type}")
 
-        hidden_states = residual + hidden_states
-        return hidden_states
+            hidden_states = residual + hidden_states
+            return hidden_states
 
 
 # Copied from transformers.models.nemotron.modeling_nemotron Nemotron->NemotronH