diff --git a/mlx_lm/models/nemotron_h.py b/mlx_lm/models/nemotron_h.py index 353de36c9..86eae9352 100644 --- a/mlx_lm/models/nemotron_h.py +++ b/mlx_lm/models/nemotron_h.py @@ -1,8 +1,9 @@ # Copyright © 2025 Apple Inc. +from copy import copy from dataclasses import dataclass from functools import partial -from typing import Any, List, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple import mlx.core as mx import mlx.nn as nn @@ -56,13 +57,58 @@ class ModelArgs(BaseModelArgs): time_step_limit: Optional[Tuple[float, float]] = None time_step_min: Optional[float] = None time_step_max: Optional[float] = None + block_configs: Optional[List[Dict[str, Any]]] = None # Map from layers_block_type names to single-char pattern codes _block_type_to_char = {"mamba": "M", "attention": "*", "moe": "E", "mlp": "-"} + @classmethod + def from_dict(cls, params): + params = dict(params) + if params.get("num_hidden_layers") is None: + blocks = params.get("block_configs") or params.get("layers_block_type") + if blocks is not None: + params["num_hidden_layers"] = len(blocks) + return super().from_dict(params) + def __post_init__(self): if self.time_step_limit is None: - self.time_step_limit = (0.0, float("inf")) + lower_limit = ( + self.time_step_min + if self.model_type == "nemotron_h_puzzle" + and self.time_step_min is not None + else 0.0 + ) + self.time_step_limit = (lower_limit, float("inf")) + + if self.block_configs is not None: + if len(self.block_configs) != self.num_hidden_layers: + raise ValueError( + "block_configs must contain one entry per hidden layer: " + f"got {len(self.block_configs)} entries for " + f"{self.num_hidden_layers} layers" + ) + block_types = [] + for layer_idx, config in enumerate(self.block_configs): + block_type = config.get("block_type") + if block_type not in self._block_type_to_char: + raise ValueError( + f"Puzzle block {layer_idx} has unsupported block_type " + f"{block_type!r}; expected one of " + f"{sorted(self._block_type_to_char)}" + ) + block_types.append(block_type) + if block_type == "moe": + missing = { + "moe_intermediate_size", + "num_experts_per_tok", + } - config.keys() + if missing: + raise ValueError( + f"Puzzle MoE block {layer_idx} is missing " + f"{sorted(missing)}" + ) + self.layers_block_type = block_types # Normalize to hybrid_override_pattern (single-char list) if self.hybrid_override_pattern is None and self.layers_block_type is not None: @@ -72,6 +118,34 @@ def __post_init__(self): if self.hybrid_override_pattern is not None: self.num_hidden_layers = len(self.hybrid_override_pattern) + def for_layer(self, layer_idx: int): + """Return a shallow config copy with Puzzle blockwise values applied.""" + if self.block_configs is None: + return self + + layer_args = copy(self) + block_config = self.block_configs[layer_idx] + for name in ("moe_intermediate_size", "num_experts_per_tok"): + if name in block_config: + setattr(layer_args, name, block_config[name]) + return layer_args + + +class NemotronHRMSNorm(nn.Module): + """Reference-compatible RMSNorm with float32 reduction and scaling.""" + + def __init__(self, hidden_size: int, eps: float): + super().__init__() + self.eps = eps + self.weight = mx.ones(hidden_size) + + def __call__(self, x: mx.array) -> mx.array: + input_dtype = x.dtype + x = x.astype(mx.float32) + variance = mx.mean(mx.square(x), axis=-1, keepdims=True) + x = x * mx.rsqrt(variance + self.eps) + return (self.weight.astype(mx.float32) * x).astype(input_dtype) + class MambaRMSNormGated(nn.Module): def __init__(self, hidden_size: int, eps: float, group_size: int): @@ -112,9 +186,11 @@ def __init__(self, args: ModelArgs): bias=args.use_conv_bias, ) + self.is_puzzle = args.model_type == "nemotron_h_puzzle" + projection_bias = args.use_bias if self.is_puzzle else args.mamba_proj_bias projection_size = self.intermediate_size + self.conv_dim + self.num_heads self.in_proj = nn.Linear( - self.hidden_size, projection_size, bias=args.mamba_proj_bias + self.hidden_size, projection_size, bias=projection_bias ) self.dt_bias = mx.ones(self.num_heads) @@ -128,7 +204,7 @@ def __init__(self, args: ModelArgs): group_size=group_size, ) self.out_proj = nn.Linear( - self.intermediate_size, self.hidden_size, bias=args.mamba_proj_bias + self.intermediate_size, self.hidden_size, bias=projection_bias ) def _conv( @@ -175,6 +251,23 @@ def _ssm( mask: Optional[mx.array], ) -> mx.array: batch_size, seq_len, _ = hidden_states.shape + output_dtype = hidden_states.dtype + + if self.is_puzzle: + # NVIDIA promotes the state-space operands and A_log + # exponentiation to float32, but computes softplus(dt + dt_bias) + # in the projected activation dtype. Preserve that boundary to + # match the Puzzle reference implementation. + hidden_states = hidden_states.astype(mx.float32) + B = B.astype(mx.float32) + C = C.astype(mx.float32) + A_log = self.A_log.astype(mx.float32) + D = self.D.astype(mx.float32) + dt_bias = self.dt_bias.astype(dt.dtype) + else: + A_log = self.A_log + D = self.D.astype(hidden_states.dtype) + dt_bias = self.dt_bias hidden_states = hidden_states.reshape( batch_size, seq_len, self.num_heads, self.head_dim @@ -189,19 +282,22 @@ def _ssm( y, state = ssm_update( hidden_states, - self.A_log, + A_log, B, C, - self.D.astype(hidden_states.dtype), + D, dt, - self.dt_bias, + dt_bias, state, self.time_step_limit, - mask, + mask=mask, + promote_dt=not self.is_puzzle, ) if cache: cache[1] = state + if self.is_puzzle: + y = y.astype(output_dtype) return y.reshape(batch_size, seq_len, self.intermediate_size) def __call__( @@ -354,12 +450,19 @@ def __init__(self, config: ModelArgs): self.routed_scaling_factor = config.routed_scaling_factor self.n_group = config.n_group self.topk_group = config.topk_group + self.is_puzzle = config.model_type == "nemotron_h_puzzle" self.weight = mx.zeros((self.n_routed_experts, config.hidden_size)) self.e_score_correction_bias = mx.zeros((self.n_routed_experts,)) def __call__(self, x): + if self.is_puzzle: + # Puzzle selects among 512 experts and computes router logits in + # float32 before top-k selection. + router_logits = x.astype(mx.float32) @ self.weight.astype(mx.float32).T + else: + router_logits = x @ self.weight.T return group_expert_select( - x @ self.weight.T, + router_logits, self.e_score_correction_bias, self.top_k, self.n_group, @@ -427,7 +530,10 @@ def __call__(self, x): class NemotronHBlock(nn.Module): def __init__(self, args: ModelArgs, block_type: str): super().__init__() - self.norm = nn.RMSNorm(args.hidden_size, eps=args.layer_norm_epsilon) + norm_cls = ( + NemotronHRMSNorm if args.model_type == "nemotron_h_puzzle" else nn.RMSNorm + ) + self.norm = norm_cls(args.hidden_size, eps=args.layer_norm_epsilon) self.block_type = block_type @@ -460,10 +566,13 @@ def __init__(self, args: ModelArgs): super().__init__() self.embeddings = nn.Embedding(args.vocab_size, args.hidden_size) self.layers = [ - NemotronHBlock(args, block_type) - for block_type in args.hybrid_override_pattern + NemotronHBlock(args.for_layer(layer_idx), block_type) + for layer_idx, block_type in enumerate(args.hybrid_override_pattern or []) ] - self.norm_f = nn.RMSNorm(args.hidden_size, eps=args.layer_norm_epsilon) + norm_cls = ( + NemotronHRMSNorm if args.model_type == "nemotron_h_puzzle" else nn.RMSNorm + ) + self.norm_f = norm_cls(args.hidden_size, eps=args.layer_norm_epsilon) self.fa_idx = 0 self.ssm_idx = 0 for b in args.hybrid_override_pattern: @@ -535,7 +644,23 @@ def make_cache(self): caches.append(KVCache()) return caches + @property + def quant_predicate(self): + if self.model_type != "nemotron_h_puzzle": + return lambda _path, _module: True + # Puzzle's 131k-token output projection is unusually sensitive to + # low-bit affine quantization. Quantizing it at 4 bits produced + # unrelated and repetitive generations while the same checkpoint + # generated correctly with the BF16 head restored. + return lambda path, _: path != "lm_head" + def sanitize(self, weights): + # Official Hugging Face Puzzle checkpoints use ``model.*`` while MLX's + # Nemotron-H implementation names the same module ``backbone.*``. + weights = { + (f"backbone.{k[6:]}" if k.startswith("model.") else k): v + for k, v in weights.items() + } 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: diff --git a/mlx_lm/models/ssm.py b/mlx_lm/models/ssm.py index eb7199c96..d5c2e78fe 100644 --- a/mlx_lm/models/ssm.py +++ b/mlx_lm/models/ssm.py @@ -11,6 +11,12 @@ def compute_dt(dt, dt_bias, time_step_limit): return mx.clip(dt, time_step_limit[0], time_step_limit[1]) +@mx.compile +def compute_dt_native(dt, dt_bias, time_step_limit): + dt = nn.softplus(dt + dt_bias) + return mx.clip(dt, time_step_limit[0], time_step_limit[1]) + + def make_ssm_kernel(): if not mx.metal.is_available(): return None @@ -74,12 +80,14 @@ def ssm_update_kernel( dt_bias: mx.array, state: mx.array, time_step_limit: Tuple[float, float], + promote_dt: bool = True, ): 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) + dt_fn = compute_dt if promote_dt else compute_dt_native + dt = dt_fn(dt, dt_bias, time_step_limit) return _ssm_kernel( inputs=[hidden_states, A_log, B, C, D, dt, state], template=[ @@ -125,6 +133,7 @@ def ssm_attn( mask: Optional[mx.array] = None, lengths: Optional[mx.array] = None, step: int = 256, + promote_dt: bool = True, ) -> Tuple[mx.array, mx.array]: """SSD-SSM forward pass. @@ -148,9 +157,13 @@ def ssm_attn( b, l, h, dh = x.shape _, _, g, d = B.shape - dt = compute_dt(dt, dt_bias, time_step_limit) + dt_fn = compute_dt if promote_dt else compute_dt_native + dt = dt_fn(dt, dt_bias, time_step_limit) repeats = h // g - A = -mx.exp(A_log).astype(dt.dtype) + # In native-dt mode the state transition follows A_log's dtype. Callers + # that require an FP32 recurrence must therefore pass FP32 A_log. + A_dtype = dt.dtype if promote_dt else A_log.dtype + A = -mx.exp(A_log).astype(A_dtype) dtA = dt * A.reshape(1, 1, -1) dtx = dt.reshape(b, l, h, 1) * x @@ -226,6 +239,7 @@ def ssm_update( time_step_limit: Tuple[float, float] = (0.001, 100.0), mask: Optional[mx.array] = None, lengths: Optional[mx.array] = None, + promote_dt: bool = True, ): seq_len = hidden_states.shape[1] if ( @@ -246,6 +260,7 @@ def ssm_update( time_step_limit, mask=mask, lengths=lengths, + promote_dt=promote_dt, ) else: return ssm_update_kernel( @@ -258,4 +273,5 @@ def ssm_update( dt_bias, state, time_step_limit, + promote_dt, ) diff --git a/mlx_lm/utils.py b/mlx_lm/utils.py index 68990b89e..b4a9b2b31 100644 --- a/mlx_lm/utils.py +++ b/mlx_lm/utils.py @@ -53,6 +53,7 @@ "minimax_m2": "minimax", "iquestcoder": "llama", "gemma4_unified": "gemma4", # encoder-free multimodal variant; vision/audio weights stripped by sanitize() + "nemotron_h_puzzle": "nemotron_h", } MAX_FILE_SIZE_GB = 5 diff --git a/tests/test_nemotron_h_puzzle.py b/tests/test_nemotron_h_puzzle.py new file mode 100644 index 000000000..d5836002b --- /dev/null +++ b/tests/test_nemotron_h_puzzle.py @@ -0,0 +1,316 @@ +import json +from dataclasses import asdict, replace +from unittest.mock import patch + +import mlx.core as mx +import mlx.nn as nn + +from mlx_lm.models import nemotron_h, ssm +from mlx_lm.models.cache import ArraysCache +from mlx_lm.utils import _get_classes, load_model, quantize_model, save_model + + +def puzzle_args(): + return nemotron_h.ModelArgs( + model_type="nemotron_h_puzzle", + vocab_size=128, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + max_position_embeddings=128, + num_attention_heads=4, + num_key_value_heads=2, + attention_bias=False, + mamba_num_heads=4, + mamba_head_dim=8, + mamba_proj_bias=False, + ssm_state_size=8, + conv_kernel=3, + n_groups=1, + mlp_bias=False, + layer_norm_epsilon=1e-5, + use_bias=False, + use_conv_bias=True, + layers_block_type=["moe", "moe"], + block_configs=[ + { + "block_type": "moe", + "moe_intermediate_size": 32, + "num_experts_per_tok": 1, + }, + { + "block_type": "moe", + "moe_intermediate_size": 64, + "num_experts_per_tok": 2, + }, + ], + moe_shared_expert_intermediate_size=32, + moe_latent_size=32, + n_group=1, + n_routed_experts=4, + n_shared_experts=1, + topk_group=1, + norm_topk_prob=True, + routed_scaling_factor=1.0, + ) + + +def test_puzzle_model_type_remaps_to_nemotron_h(): + model_cls, args_cls = _get_classes({"model_type": "nemotron_h_puzzle"}) + assert model_cls is nemotron_h.Model + assert args_cls is nemotron_h.ModelArgs + + +def test_puzzle_quantization_preserves_output_head(): + args = puzzle_args() + model = nemotron_h.Model(args) + + quantize_model(model, asdict(args), group_size=32, bits=4) + + assert isinstance(model.lm_head, nn.Linear) + assert not hasattr(model.lm_head, "scales") + assert hasattr(model.layers[0].mixer.switch_mlp.fc1, "scales") + + +def test_official_config_derives_layer_count(): + config = asdict(puzzle_args()) + config.pop("num_hidden_layers") + + args = nemotron_h.ModelArgs.from_dict(config) + + assert args.num_hidden_layers == len(config["block_configs"]) + + +def test_puzzle_layers_use_heterogeneous_moe_dimensions_and_top_k(): + model = nemotron_h.Model(puzzle_args()) + + first, second = model.layers + assert first.mixer.switch_mlp.fc1.output_dims == 32 + assert second.mixer.switch_mlp.fc1.output_dims == 64 + assert first.mixer.gate.top_k == 1 + assert second.mixer.gate.top_k == 2 + + logits = model(mx.array([[1, 2, 3]])) + mx.eval(logits) + assert logits.shape == (1, 3, 128) + + +def test_mamba_prefill_matches_cached_decode(): + args = replace( + puzzle_args(), + num_hidden_layers=1, + hybrid_override_pattern=["M"], + layers_block_type=["mamba"], + block_configs=[{"block_type": "mamba"}], + ) + block = nemotron_h.NemotronHBlock(args.for_layer(0), "M") + mx.random.seed(7) + inputs = mx.random.normal((1, 4, args.hidden_size)).astype(mx.bfloat16) + + prefill = block(inputs) + cache = ArraysCache(size=2) + decoded = mx.concatenate( + [ + block(inputs[:, position : position + 1], cache=cache) + for position in range(4) + ], + axis=1, + ) + mx.eval(prefill, decoded) + + assert cache[1].dtype == mx.float32 + assert mx.allclose(prefill, decoded, rtol=2e-2, atol=2e-2).item() + + +def test_puzzle_preserves_timestep_activation_precision(): + args = replace( + puzzle_args(), + num_hidden_layers=1, + hybrid_override_pattern=["M"], + layers_block_type=["mamba"], + block_configs=[{"block_type": "mamba"}], + time_step_min=0.001, + time_step_limit=None, + ) + mixer = nemotron_h.NemotronHMamba2Mixer(args) + hidden = mx.zeros((1, 2, mixer.intermediate_size), dtype=mx.bfloat16) + B = mx.zeros((1, 2, mixer.n_groups * mixer.ssm_state_size), dtype=mx.bfloat16) + C = mx.zeros_like(B) + dt = mx.zeros((1, 2, mixer.num_heads), dtype=mx.bfloat16) + captured = {} + + def fake_ssm_update(*values, **kwargs): + captured["hidden_dtype"] = values[0].dtype + captured["B_dtype"] = values[2].dtype + captured["C_dtype"] = values[3].dtype + captured["dt_dtype"] = values[5].dtype + captured["dt_bias_dtype"] = values[6].dtype + captured["promote_dt"] = kwargs["promote_dt"] + captured["time_step_limit"] = values[8] + return values[0], None + + with patch.object(nemotron_h, "ssm_update", fake_ssm_update): + output = mixer._ssm(hidden, B, C, dt, cache=None, mask=None) + + assert output.dtype == mx.bfloat16 + assert captured == { + "hidden_dtype": mx.float32, + "B_dtype": mx.float32, + "C_dtype": mx.float32, + "dt_dtype": mx.bfloat16, + "dt_bias_dtype": mx.bfloat16, + "promote_dt": False, + "time_step_limit": (0.001, float("inf")), + } + + +def test_native_timestep_keeps_state_transition_in_float32(): + hidden = mx.array([[[[0.5]], [[-0.25]]]], dtype=mx.float32) + A_log = mx.array([1.234567], dtype=mx.float32) + B = mx.array([[[[0.75]], [[-0.5]]]], dtype=mx.float32) + C = mx.array([[[[0.25]], [[1.5]]]], dtype=mx.float32) + D = mx.array([0.125], dtype=mx.float32) + dt = mx.array([[[0.125], [0.25]]], dtype=mx.bfloat16) + dt_bias = mx.array([-0.1], dtype=mx.bfloat16) + initial_state = mx.ones((1, 1, 1, 1), dtype=mx.float32) + + output, final_state = ssm.ssm_attn( + hidden, + A_log, + B, + C, + D, + dt, + dt_bias, + initial_state, + (0.001, float("inf")), + promote_dt=False, + ) + + reference_dt = nn.softplus(dt + dt_bias).astype(mx.float32) + reference_A = -mx.exp(A_log) + reference_state = initial_state + reference_output = [] + for position in range(hidden.shape[1]): + delta = reference_dt[:, position] + reference_state = reference_state * mx.exp( + delta[:, :, None, None] * reference_A[None, :, None, None] + ) + reference_state = reference_state + ( + delta[:, :, None, None] + * hidden[:, position, :, :, None] + * B[:, position, :, None, :] + ) + current = (reference_state * C[:, position, :, None, :]).sum(axis=-1) + current = current + D[None, :, None] * hidden[:, position] + reference_output.append(current) + reference_output = mx.stack(reference_output, axis=1) + mx.eval(output, final_state, reference_output, reference_state) + + assert mx.allclose(output, reference_output, rtol=1e-5, atol=1e-5).item() + assert mx.allclose(final_state, reference_state, rtol=1e-5, atol=1e-5).item() + + +def test_explicit_time_step_limit_is_preserved(): + explicit_limit = (0.001, 100.0) + for model_type in ("nemotron_h", "nemotron_h_puzzle"): + args = replace( + puzzle_args(), + model_type=model_type, + time_step_min=0.002, + time_step_limit=explicit_limit, + ) + assert args.time_step_limit == explicit_limit + + +def test_puzzle_precision_changes_do_not_change_base_nemotron_h(): + common = dict( + num_hidden_layers=1, + hybrid_override_pattern=["M"], + layers_block_type=["mamba"], + block_configs=[{"block_type": "mamba"}], + use_bias=True, + mamba_proj_bias=False, + time_step_min=0.001, + time_step_limit=None, + ) + puzzle = replace(puzzle_args(), **common) + base = replace(puzzle_args(), model_type="nemotron_h", **common) + puzzle_mixer = nemotron_h.NemotronHMamba2Mixer(puzzle) + base_mixer = nemotron_h.NemotronHMamba2Mixer(base) + puzzle_block = nemotron_h.NemotronHBlock(puzzle, "M") + base_block = nemotron_h.NemotronHBlock(base, "M") + router_dtypes = [] + + def capture_router_dtype(gates, *_args): + router_dtypes.append(gates.dtype) + return None, None + + puzzle_gate = nemotron_h.MoEGate(puzzle) + base_gate = nemotron_h.MoEGate(base) + puzzle_gate.weight = puzzle_gate.weight.astype(mx.bfloat16) + base_gate.weight = base_gate.weight.astype(mx.bfloat16) + with patch.object(nemotron_h, "group_expert_select", capture_router_dtype): + puzzle_gate(mx.ones((1, puzzle.hidden_size), dtype=mx.bfloat16)) + base_gate(mx.ones((1, base.hidden_size), dtype=mx.bfloat16)) + + assert "bias" in puzzle_mixer.in_proj.parameters() + assert "bias" in puzzle_mixer.out_proj.parameters() + assert "bias" not in base_mixer.in_proj.parameters() + assert "bias" not in base_mixer.out_proj.parameters() + assert isinstance(puzzle_block.norm, nemotron_h.NemotronHRMSNorm) + assert isinstance(base_block.norm, nn.RMSNorm) + assert router_dtypes == [mx.float32, mx.bfloat16] + assert puzzle.time_step_limit == (0.001, float("inf")) + assert base.time_step_limit == (0.0, float("inf")) + + +def test_official_source_model_prefix_is_remapped_to_backbone(): + model = nemotron_h.Model(puzzle_args()) + weights = { + "model.embeddings.weight": mx.ones((128, 32)), + "lm_head.weight": mx.ones((128, 32)), + } + + sanitized = model.sanitize(weights) + + assert "model.embeddings.weight" not in sanitized + assert "backbone.embeddings.weight" in sanitized + + +def test_official_source_experts_are_remapped_and_stacked(): + model = nemotron_h.Model(puzzle_args()) + weights = {} + for expert in range(4): + prefix = f"model.layers.0.mixer.experts.{expert}" + weights[f"{prefix}.up_proj.weight"] = mx.full((32, 32), expert) + weights[f"{prefix}.down_proj.weight"] = mx.full((32, 32), expert) + + sanitized = model.sanitize(weights) + + fc1 = sanitized["backbone.layers.0.mixer.switch_mlp.fc1.weight"] + fc2 = sanitized["backbone.layers.0.mixer.switch_mlp.fc2.weight"] + assert fc1.shape == (4, 32, 32) + assert fc2.shape == (4, 32, 32) + assert mx.array_equal(fc1[:, 0, 0], mx.arange(4)).item() + assert not any("experts." in key for key in sanitized) + + +def test_strict_load_of_heterogeneous_quantized_checkpoint(tmp_path): + args = puzzle_args() + model = nemotron_h.Model(args) + nn.quantize(model, group_size=32, bits=4) + mx.eval(model.parameters()) + + save_model(tmp_path, model) + config = asdict(args) + config["quantization"] = {"group_size": 32, "bits": 4, "mode": "affine"} + (tmp_path / "config.json").write_text(json.dumps(config)) + + loaded, _ = load_model(tmp_path, strict=True) + logits = loaded(mx.array([[1, 2, 3]])) + mx.eval(logits) + + assert loaded.layers[0].mixer.switch_mlp.fc1.output_dims == 32 + assert loaded.layers[1].mixer.switch_mlp.fc1.output_dims == 64 + assert logits.shape == (1, 3, 128)