From 73c8550478a39a7a163e9aec9bef6435b4592062 Mon Sep 17 00:00:00 2001 From: Angelos Katharopoulos Date: Mon, 16 Mar 2026 10:59:14 -0700 Subject: [PATCH] Nemotron super support (#992) --- mlx_lm/models/nemotron_h.py | 54 ++++++++++++++++++++++++++++++------- mlx_lm/models/ssm.py | 17 +++++++++--- 2 files changed, 58 insertions(+), 13 deletions(-) diff --git a/mlx_lm/models/nemotron_h.py b/mlx_lm/models/nemotron_h.py index 5fb500e..0c391e4 100644 --- a/mlx_lm/models/nemotron_h.py +++ b/mlx_lm/models/nemotron_h.py @@ -40,10 +40,12 @@ class ModelArgs(BaseModelArgs): layer_norm_epsilon: float use_bias: bool use_conv_bias: bool - hybrid_override_pattern: List[str] + hybrid_override_pattern: Optional[List[str]] = None + layers_block_type: Optional[List[str]] = None head_dim: Optional[int] = None moe_intermediate_size: Optional[int] = None moe_shared_expert_intermediate_size: Optional[int] = None + moe_latent_size: Optional[int] = None n_group: Optional[int] = None n_routed_experts: Optional[int] = None n_shared_experts: Optional[int] = None @@ -55,13 +57,20 @@ class ModelArgs(BaseModelArgs): time_step_min: Optional[float] = None time_step_max: Optional[float] = None + # Map from layers_block_type names to single-char pattern codes + _block_type_to_char = {"mamba": "M", "attention": "*", "moe": "E", "mlp": "-"} + def __post_init__(self): - if ( - self.time_step_limit is None - and self.time_step_min is not None - and self.time_step_max is not None - ): - self.time_step_limit = (self.time_step_min, self.time_step_max) + if self.time_step_limit is None and self.time_step_min is not None: + self.time_step_limit = (self.time_step_min, float("inf")) + + # Normalize to hybrid_override_pattern (single-char list) + if self.hybrid_override_pattern is None and self.layers_block_type is not None: + self.hybrid_override_pattern = [ + self._block_type_to_char[t] for t in self.layers_block_type + ] + if self.hybrid_override_pattern is not None: + self.num_hidden_layers = len(self.hybrid_override_pattern) class MambaRMSNormGated(nn.Module): @@ -365,8 +374,16 @@ class NemotronHMoE(nn.Module): super().__init__() self.config = config self.num_experts_per_tok = config.num_experts_per_tok + self.moe_latent_size = config.moe_latent_size + + # When latent projection is used, experts operate on the latent dim + expert_input_dim = ( + config.moe_latent_size + if config.moe_latent_size is not None + else config.hidden_size + ) self.switch_mlp = SwitchMLP( - config.hidden_size, + expert_input_dim, config.moe_intermediate_size, config.n_routed_experts, activation=nn.ReLU2(), @@ -379,12 +396,30 @@ class NemotronHMoE(nn.Module): config, intermediate_size=intermediate_size ) + # Latent projection layers for dimensionality reduction before/after experts + if config.moe_latent_size is not None: + self.fc1_latent_proj = nn.Linear( + config.hidden_size, config.moe_latent_size, bias=config.mlp_bias + ) + self.fc2_latent_proj = nn.Linear( + config.moe_latent_size, config.hidden_size, bias=config.mlp_bias + ) + def __call__(self, x): + residuals = x inds, scores = self.gate(x) + + if self.moe_latent_size is not None: + x = self.fc1_latent_proj(x) + y = self.switch_mlp(x, inds) y = (y * scores[..., None]).sum(axis=-2).astype(y.dtype) + + if self.moe_latent_size is not None: + y = self.fc2_latent_proj(y) + if self.config.n_shared_experts is not None: - y = y + self.shared_experts(x) + y = y + self.shared_experts(residuals) return y @@ -501,6 +536,7 @@ class Model(nn.Module): return caches def sanitize(self, weights): + weights = {k: v for (k, v) in weights.items() if not k.startswith("mtp.")} for k, v in weights.items(): if "conv1d.weight" in k and v.shape[-1] != 1: weights[k] = v.moveaxis(2, 1) diff --git a/mlx_lm/models/ssm.py b/mlx_lm/models/ssm.py index aeb16bf..eb7199c 100644 --- a/mlx_lm/models/ssm.py +++ b/mlx_lm/models/ssm.py @@ -6,6 +6,7 @@ import mlx.nn as nn @mx.compile def compute_dt(dt, dt_bias, time_step_limit): + dt = dt.astype(mx.float32) dt = nn.softplus(dt + dt_bias) return mx.clip(dt, time_step_limit[0], time_step_limit[1]) @@ -44,7 +45,7 @@ def make_ssm_kernel(): auto idx = d_idx * Ds + s_idx; auto dB_by_x = x_ * dt_ * static_cast(B_[s_idx]); auto state = dA * i_state[idx] + dB_by_x; - o_state[idx] = static_cast(state); + o_state[idx] = static_cast(state); acc += state * C_[s_idx]; } acc = simd_sum(acc); @@ -76,15 +77,23 @@ def ssm_update_kernel( ): n, _, h, d = hidden_states.shape input_type = hidden_states.dtype + state_type = state.dtype hb, ds = B.shape[-2:] dt = compute_dt(dt, dt_bias, time_step_limit) return _ssm_kernel( inputs=[hidden_states, A_log, B, C, D, dt, state], - template=[("T", input_type), ("Dh", d), ("Ds", ds), ("H", h), ("G", h // hb)], + template=[ + ("T", input_type), + ("U", state_type), + ("Dh", d), + ("Ds", ds), + ("H", h), + ("G", h // hb), + ], grid=(32, d, h * n), threadgroup=(32, 8, 1), output_shapes=[(n, 1, h, d), state.shape], - output_dtypes=[input_type, input_type], + output_dtypes=[input_type, state_type], ) @@ -186,7 +195,7 @@ def ssm_attn( mx.expand_dims(lengths < 0, (1, 2, 3)), state, next_state ) - return y, next_state + return y.astype(x.dtype), next_state ys = [] for i in range(0, l, step):