Nemotron-Labs-TwoTower-30B-A3B-mlx-bf16

Nemotron-Labs-TwoTower-30B-A3B — MLX bfloat16 (unquantized)

An MLX port of nvidia/Nemotron-Labs-TwoTower-30B-A3B-Base-BF16, a block-wise autoregressive diffusion language model, running on Apple Silicon.

Two NemotronH towers: a frozen context tower (prefills the prompt into KV + Mamba states) and a denoiser tower (adaLN-conditioned on the diffusion timestep). Generation is block-wise mask diffusion — each block starts fully masked and is iteratively denoised, committing high-confidence tokens and remasking the rest.

⚠️ Custom code required

Stock mlx-lm has no two-tower diffusion architecture, so this repo ships the MLX modeling code (nemotron_twotower_mlx.py) + a runner (run_twotower_mlx.py). The entry point is generate_mask_diffusion, not mlx_lm.generate.

📦 Usage guide, examples & benchmarks: https://github.com/PipeNetwork/nemotron-twotower-mlx

pip install mlx mlx-lm transformers
huggingface-cli download pipenetwork/Nemotron-Labs-TwoTower-30B-A3B-mlx-bf16 --local-dir tt-bf16
python tt-bf16/run_twotower_mlx.py --model tt-bf16 \
  --prompt "The capital of France is" --max-new-tokens 64 \
  --block-size 16 --steps-per-block 16 --mask-token-id 3

import sys; sys.path.insert(0, "tt-bf16")
from run_twotower_mlx import load
import mlx.core as mx
model, tok = load("tt-bf16")
ids = mx.array([tok("The capital of France is")["input_ids"]])
out = model.generate_mask_diffusion(ids, max_new_tokens=64, block_size=16,
        steps_per_block=16, mask_token_id=3, eos_token_id=tok.eos_token_id)
print(tok.decode(out[0].tolist()))

Token id 3 is the mask token (the model's training convention). Example output: "The capital of France is Paris, the capital of Germany is Berlin, the capital of Japan is Tokyo, ..."

Model family

Validation

The MLX conversion was checked against NVIDIA's reference implementation running on an NVIDIA GB10 (CUDA). Greedy decoding matched token-for-token: 120/120 tokens (100%) and 5/5 top-1 across the test prompts — e.g. both produce "George Washington. He was elected in 1789 and served two terms until 1797." (The AR/context tower is exercised directly; it is the shared backbone the diffusion denoiser also uses.)

Benchmarks

Measured on Apple M3 Ultra (512 GB unified memory), MLX 0.31 — steady-state (post-warmup). Peak RAM is the unified-memory high-water mark during generation.

AR / context tower — 128-token single-stream generation:

TwoTower diffusion — 64 new tokens, block size 16, ≤16 steps/block:

Diffusion runs steps_per_block denoiser passes per block, so it is slower per token than the AR tower — lower --steps-per-block trades quality for speed. Higher-precision builds tend to converge in fewer denoiser evaluations, but each pass moves more memory.

License

Governed by the NVIDIA Open Model License of the base model.