#!/usr/bin/env python3 """ MLX port of nvidia/Nemotron-Labs-TwoTower-30B-A3B (NemotronHTwoTowerForCausalLM): a block-wise autoregressive *diffusion* LM. context_tower — frozen AR NemotronH backbone (KV + Mamba states for the prompt) denoiser_tower — diffusion decoder, adaLN-conditioned on the diffusion timestep mask-diffusion generation: per block, iteratively denoise a fully-masked block, commit high-confidence tokens, remask the rest, then commit the block to context. Reuses mlx-lm's `nemotron_h` building blocks (NemotronHModel, mixers, caches). Only the two-tower orchestration + adaLN + diffusion sampling are new here. Reference (CUDA-only, cannot run on this Mac): modeling_nemotron_twotower.py. """ import math from typing import List, Optional import mlx.core as mx import mlx.nn as nn from mlx_lm.models.nemotron_h import ModelArgs, NemotronHModel from mlx_lm.models.cache import ArraysCache, KVCache from mlx_lm.models.base import scaled_dot_product_attention # --------------------------------------------------------------------------- # Time conditioning (PixArt-alpha adaLN-single) # --------------------------------------------------------------------------- def timestep_embedding(t, dim, max_period=10000): half = dim // 2 freqs = mx.exp(-math.log(max_period) * mx.arange(half, dtype=mx.float32) / half) args = t[:, None].astype(mx.float32) * freqs[None] emb = mx.concatenate([mx.cos(args), mx.sin(args)], axis=-1) if dim % 2: emb = mx.concatenate([emb, mx.zeros_like(emb[:, :1])], axis=-1) return emb class TimestepEmbedder(nn.Module): def __init__(self, hidden_size, frequency_embedding_size=256, max_period=1000): super().__init__() self.frequency_embedding_size = frequency_embedding_size self.max_period = max_period # keys: mlp.0 (Linear), mlp.1 (SiLU, no params), mlp.2 (Linear) self.mlp = [ nn.Linear(frequency_embedding_size, hidden_size, bias=True), nn.SiLU(), nn.Linear(hidden_size, hidden_size, bias=True), ] def __call__(self, t): t_scaled = t * self.max_period h = timestep_embedding(t_scaled, self.frequency_embedding_size) h = self.mlp[0](h) h = self.mlp[1](h) return self.mlp[2](h) def quantization_predicate(bits, group_size=64): """Mixed-precision scheme 'mixed_v1' for the diffusion two-tower. Diffusion compounds quantization error across denoising steps, so the tiny but critical timestep-conditioning MLPs stay bf16 and the embeddings / LM heads stay >=8-bit; the bulk (MoE experts, attention & Mamba projections) uses the target bits. Used identically at pack time and load time so the quantized module structure matches the saved tensors.""" def pred(path, module): if not hasattr(module, "to_quantized"): return False if "t_embedder" in path or "t_block" in path: return False # keep bf16 if "embeddings" in path or "lm_head" in path: return {"group_size": group_size, "bits": max(bits, 8)} return {"group_size": group_size, "bits": bits} return pred def modulate(x, shift, scale): # x:(B,L,D) shift/scale:(B,D) -> x*(1+scale) + shift return x * (1.0 + scale[:, None, :]) + shift[:, None, :] def get_mod_params(t_emb, table): # t_emb:(B,3D) table:(3,D) -> shift,scale,gate each (B,D) B = t_emb.shape[0] D = table.shape[1] combined = table[None] + t_emb.reshape(B, 3, D) shift = combined[:, 0, :] scale = combined[:, 1, :] gate = combined[:, 2, :] return shift, scale, gate # --------------------------------------------------------------------------- # Two-Tower diffusion model # --------------------------------------------------------------------------- class TwoTowerModel(nn.Module): def __init__(self, args: ModelArgs): super().__init__() self.args = args H = args.hidden_size N = args.num_hidden_layers self.context_tower = NemotronHModel(args) self.context_lm_head = nn.Linear(H, args.vocab_size, bias=False) self.denoiser_tower = NemotronHModel(args) self.lm_head = nn.Linear(H, args.vocab_size, bias=False) # time conditioning self.t_embedder = TimestepEmbedder(H) self.t_block = [nn.SiLU(), nn.Linear(H, 3 * H, bias=True)] # t_block.0/.1 self.scale_shift_tables = [ mx.zeros((3, H)) for _ in range(N) ] # single-char block-type pattern, e.g. "M","E","*" self.pattern = list(args.hybrid_override_pattern) # NemotronHModel keeps a COMPACT cache list (one slot per M/* layer only). # Map each layer index -> its slot in that compact list (None for E/-). self.cache_index = [] c = 0 for t in self.pattern: if t in ("M", "*"): self.cache_index.append(c) c += 1 else: self.cache_index.append(None) # -- weight loading ----------------------------------------------------- def sanitize(self, weights): # conv1d weight layout + per-tower expert stacking (mirror nemotron_h) out = {} for k, v in weights.items(): if "conv1d.weight" in k and v.shape[-1] != 1: v = v.moveaxis(2, 1) out[k] = v weights = out for tower in ("context_tower", "denoiser_tower"): for l in range(self.args.num_hidden_layers): prefix = f"{tower}.layers.{l}.mixer" for m, n in [("down_proj", "fc2"), ("up_proj", "fc1")]: k0 = f"{prefix}.experts.0.{m}.weight" if k0 in weights: stack = [ weights.pop(f"{prefix}.experts.{e}.{m}.weight") for e in range(self.args.n_routed_experts) ] weights[f"{prefix}.switch_mlp.{n}.weight"] = mx.stack(stack) return weights # -- context tower cache (prefill / extend) ----------------------------- def _make_ctx_cache(self): # COMPACT: one entry per M/* layer, matching NemotronHModel.__call__. caches = [] for t in self.pattern: if t == "M": caches.append(ArraysCache(size=2)) elif t == "*": caches.append(KVCache()) return caches def build_context_cache(self, prompt_ids): """Prefill the whole prompt through the context tower, populating per-layer Mamba (conv+ssm) and attention (KV) caches. Returns the cache list.""" caches = self._make_ctx_cache() # NemotronHModel.__call__ populates the caches in place; return is unused. self.context_tower(prompt_ids, cache=caches) return caches def extend_context_cache(self, block_ids, caches): """Advance the context cache by a committed block (standard causal forward).""" self.context_tower(block_ids, cache=caches) return caches # -- denoiser bidirectional attention ----------------------------------- @staticmethod def _repeat_kv(x, n_rep): if n_rep == 1: return x b, h, l, d = x.shape x = mx.broadcast_to(x[:, :, None, :, :], (b, h, n_rep, l, d)) return x.reshape(b, h * n_rep, l, d) def _denoiser_attention(self, mixer, hidden, ctx_k, ctx_v): """Bidirectional self-attention over [context_KV | block_KV] (NoPE, is_causal=False).""" B, L, _ = hidden.shape nH = mixer.num_heads nKV = mixer.num_key_value_heads hd = mixer.head_dim q = mixer.q_proj(hidden).reshape(B, L, nH, hd).transpose(0, 2, 1, 3) k = mixer.k_proj(hidden).reshape(B, L, nKV, hd).transpose(0, 2, 1, 3) v = mixer.v_proj(hidden).reshape(B, L, nKV, hd).transpose(0, 2, 1, 3) if ctx_k is not None and ctx_k.shape[2] > 0: k = mx.concatenate([ctx_k.astype(k.dtype), k], axis=2) v = mx.concatenate([ctx_v.astype(v.dtype), v], axis=2) n_rep = nH // nKV k = self._repeat_kv(k, n_rep) v = self._repeat_kv(v, n_rep) # full (non-causal) attention, no mask o = mx.fast.scaled_dot_product_attention(q, k, v, scale=mixer.scale, mask=None) o = o.transpose(0, 2, 1, 3).reshape(B, L, nH * hd) return mixer.o_proj(o) def _denoiser_mamba(self, mixer, hidden, ctx_conv, ctx_ssm): """Forward-only chunk-scan of the block seeded from the context Mamba state. Reuses the mixer's own conv+ssm via a fresh seeded cache (context state is immutable; the fresh cache absorbs the writes).""" seed = ArraysCache(size=2) seed[0] = ctx_conv seed[1] = ctx_ssm return mixer(hidden, mask=None, cache=seed) # -- one diffusion denoiser forward over a full (masked) block ---------- def denoiser_forward(self, block_ids, ctx_caches, t): tower = self.denoiser_tower t_repr = self.t_embedder(t.astype(mx.float32)) t_emb = self.t_block[1](self.t_block[0](t_repr)) # (B,3H) hidden = tower.embeddings(block_ids) for i, block in enumerate(tower.layers): residual = hidden shift, scale, gate = get_mod_params(t_emb, self.scale_shift_tables[i]) bt = block.block_type if bt in ("M", "*"): # norm is fused after modulate in mcore -> modulate THEN norm h = modulate(hidden, shift, scale) h = block.norm(h.astype(block.norm.weight.dtype)) else: # E / - : separate pre-norm -> norm THEN modulate h = block.norm(hidden.astype(block.norm.weight.dtype)) h = modulate(h, shift, scale) if bt == "M": c = ctx_caches[self.cache_index[i]] h = self._denoiser_mamba(block.mixer, h, c[0], c[1]) elif bt == "*": c = ctx_caches[self.cache_index[i]] ctx_k = c.keys[..., : c.offset, :] if c.offset > 0 else None ctx_v = c.values[..., : c.offset, :] if c.offset > 0 else None h = self._denoiser_attention(block.mixer, h, ctx_k, ctx_v) else: h = block.mixer(h) h = gate[:, None, :] * h hidden = residual + h hidden = tower.norm_f(hidden) return self.lm_head(hidden.astype(self.lm_head.weight.dtype)).astype(mx.float32) # -- mask-diffusion generation ------------------------------------------ @staticmethod def _mdlm_logprobs(logits, xt, mask_token_id): logits = mx.array(logits) neg = mx.full(logits.shape[-1:], -1e12) # mask token -> -inf logits = mx.concatenate([ logits[..., :mask_token_id], mx.full(logits[..., mask_token_id:mask_token_id + 1].shape, -1e12), logits[..., mask_token_id + 1:], ], axis=-1) log_probs = logits - mx.logsumexp(logits, axis=-1, keepdims=True) # unmasked positions predict themselves w.p. 1 unmasked = (xt != mask_token_id) # (B,L) onehot = mx.zeros(log_probs.shape) # scatter 0.0 at xt for unmasked, -1e12 elsewhere forced = mx.full(log_probs.shape, -1e12) idx = mx.clip(xt, 0, log_probs.shape[-1] - 1) forced = _set_at_last(forced, idx, 0.0) log_probs = mx.where(unmasked[..., None], forced, log_probs) return log_probs def generate_mask_diffusion( self, input_ids, max_new_tokens=128, block_size=16, steps_per_block=16, mask_token_id=3, confidence_threshold=0.9, eos_token_id=None, verbose=False, step_callback=None, ): assert max_new_tokens % block_size == 0 B = input_ids.shape[0] num_blocks = max_new_tokens // block_size caches = self.build_context_cache(input_ids) context_ids = input_ids nfe = 0 for blk in range(num_blocks): xt = mx.full((B, block_size), mask_token_id, dtype=mx.int32) if step_callback is not None: step_callback(blk, -1, xt, context_ids) for step in range(steps_per_block): is_masked = (xt == mask_token_id) n_masked = int(is_masked.sum().item()) if n_masked == 0: break t_model = is_masked.astype(mx.float32).mean() t_vec = mx.broadcast_to(t_model.reshape(1), (B,)) logits = self.denoiser_forward(xt, caches, t_vec) nfe += 1 log_xt = self._mdlm_logprobs(logits, xt, mask_token_id) x_theta = mx.exp(log_xt) predicted = mx.argmax(log_xt, axis=-1).astype(mx.int32) conf = mx.take_along_axis( x_theta, predicted[..., None].astype(mx.int32), axis=-1 )[..., 0] conf = mx.where(is_masked, conf, mx.array(float("inf"))) is_last = (step == steps_per_block - 1) n_masked_row = is_masked.sum(-1) # (B,) if is_last: commit = n_masked_row else: remaining = max(1, steps_per_block - step) num_above = ((conf > confidence_threshold) & is_masked).sum(-1) commit = mx.where(num_above > 0, num_above, mx.ones_like(num_above)) min_commit = mx.ceil(n_masked_row.astype(mx.float32) / remaining).astype(commit.dtype) commit = mx.minimum(mx.maximum(commit, min_commit), n_masked_row) output = mx.where(is_masked, predicted, xt) # remask lowest-confidence (per row) new_xt = [] for b in range(B): row = output[b] nrem = int((n_masked_row[b] - commit[b]).item()) if nrem > 0: cb = conf[b] order = mx.argsort(cb) # ascending; masked lowest first remask_idx = order[:nrem] row = _set_at_index(row, remask_idx, mask_token_id) new_xt.append(row) xt = mx.stack(new_xt) mx.eval(xt) if step_callback is not None: step_callback(blk, step, xt, context_ids) context_ids = mx.concatenate([context_ids, xt], axis=1) caches = self.extend_context_cache(xt, caches) if verbose: print(f"[block {blk+1}/{num_blocks}] nfe={nfe}") if eos_token_id is not None: eos = [eos_token_id] if isinstance(eos_token_id, int) else list(eos_token_id) if any(bool((xt == e).any().item()) for e in eos): break self.last_nfe = nfe return context_ids # helpers for scatter on last axis / by index (mlx has no in-place scatter sugar) def _set_at_last(arr, idx, value): """arr:(...,V) idx:(...) -> set arr[...,idx]=value along last axis.""" oh = nn.losses._make_one_hot if False else None V = arr.shape[-1] onehot = (mx.arange(V) == idx[..., None]) return mx.where(onehot, mx.array(value, dtype=arr.dtype), arr) def _set_at_index(vec, indices, value): """vec:(L,) set vec[indices]=value.""" L = vec.shape[0] mask = mx.zeros((L,), dtype=mx.bool_) mask = mask + (mx.arange(L)[:, None] == indices[None, :]).any(axis=1) return mx.where(mask, mx.array(value, dtype=vec.dtype), vec)