qwen3-4b-mamba3-hybrid-stage2b-kd-bias0

Stage-2b endpoint, bias-0 arm -- long-context forward-KL logit distillation against Qwen3-4B at sequence length 32,768, 294M tokens (141 steps of 2,097,152; micro-batch 1 x 4 ranks x 16 accumulation), from the stage-2a endpoint of this arm. LR halved from stage 2a: mixers 1e-4, inherited weights 1e-5, cosine to 0.1x with 3% warmup; all weights train. 4x B200.

Final val KL at 32k 0.2021, top1-agreement 0.8056, wikitext-2 PPL (4096 context) 11.93, from a starting point of 1.1727 / 0.7167 / 12.49. At this same stage the Qwen GDN arm reaches KL 0.2177 / top1 0.7985 and the Qwen Mamba2 arm 0.2336 / 0.7875; see the table below for the bias-1 arm. This is the input to stage-3 on-policy distillation, which has not been run. No downstream evaluation yet.

Loading

This is a custom container (model_type: qwen3_mamba3), not a stock architecture, so trust_remote_code=True is required. The modeling code is bundled in this repo (modeling_qwen3_mamba3.py) and needs no other checkout:

from transformers import AutoModelForCausalLM, AutoTokenizer
m = AutoModelForCausalLM.from_pretrained(REPO, trust_remote_code=True, dtype="bfloat16")
tok = AutoTokenizer.from_pretrained(REPO)

Kernel. The mixers call upstream's fused Mamba-3 MIMO kernel, mamba_ssm.ops.tilelang.mamba3.mamba3_mimo, from state-spaces/mamba at commit e9594ce (TileLang + Triton; the pip mamba_ssm 2.2.6.post3 predates Mamba-3, so install from source). If that import fails the model warns once and every mixer runs mamba3_mimo_reference, a sequential fp32 torch transcription of the kernel's semantics: numerically equivalent (fwd rel err 6e-3 vs the kernel in bf16), and slow -- fine for a decode-parity check, not for throughput. Set Qwen3Mamba3Mixer.force_reference = True to opt into it explicitly.

Training with the kernel. Upstream's mamba3_mimo.py backward at e9594ce reads ctx.saved_tensors twice, which PyTorch's non-reentrant activation checkpointing (use_reentrant=False, HF's default) rejects with CheckpointError: Unpack is being triggered for a tensor that was already unpacked once. The fix is a one-line change (read it once); the training runs below used a local commit with that change. Inference is unaffected.

The study

Full-attention Qwen3-4B converted to a Mamba-3 hybrid: 27 of 36 layers have their attention replaced by a Mamba-3 MIMO mixer, with every 4th layer (indices 3, 7, 11 ... 35) retained as full attention -- the u4 pattern. Retained layers are literal Qwen3DecoderLayers, so retention is bit-exact. Mixer geometry: 32 heads x 128 head_dim, 8 B/C groups, d_state 128, MIMO rank 4, expand 1.6, chunk 8, RoPE on half of the state dims, the pre-gate grouped RMSNorm fused into the kernel. Parameter layout and math are upstream's exactly (in_proj rows [z | x | B | C | dd_dt | dd_A | trap | angles], rank-major (r g n) B/C), so the upstream kernel is called unchanged. 4,741,428,416 params, of which the 27 mixers are 1,426,756,032 (30.1%).

Capability is recovered by staged distillation rather than pretraining: surgery/init -> stage-1 per-layer hidden-state alignment (100M tok) -> stage-2a forward-KL KD @4k (600M) -> stage-2b long-context KD @32k (294M) -> stage-3 on-policy distillation. Teacher: Qwen/Qwen3-4B. This arm exists to isolate the mixer: it shares the teacher, the packed data, the token budgets and the schedules with the Gated DeltaNet and Mamba2 arms trained identically.

Inherited init

The mixer starts as the teacher's attention head, not as a random SSM (inherited fraction 99.95%): per converted layer, x head h <- v_proj head h // 4; B[r, g] <- k_proj head g; C[r, g] <- q_proj head 4g + r; out_proj <- o_proj; C_norm / B_norm <- q_norm / k_norm. The dd_dt, dd_A, trapezoid and angle rows of in_proj are zero (per-head constant dt, A = -1, trapezoid 1/2, no rotation; all keep gradient). dt_bias is the softplus-inverse of exp(U(log 1e-3, log 1e-1)) (the repaired dt init of the Mamba2 arm; horizons of 10-1000 tokens). mimo_x is a random simplex over ranks per (head, dim) so the ranks get distinct gradients while the init state equals upstream's uniform 1/R; mimo_o selects each head's own rank; mimo_z, D and the norm weights are 1. At init a converted head equals the teacher's head with softmax replaced by causal decay (tested to 3e-8).

B/C bias. Upstream initializes the B/C biases to 1.0. On an inherited head that adds sum(q) + sum(k) + 128 to every score, which looked like it would swamp the teacher's structure, so the primary arm used 0 and a second arm kept upstream's 1.0. Measured: bias 1 led at every stage-1 and stage-2a evaluation (below). Both arms are published; bias 1 is now the default.

Results so far (trainer metrics only)

Val KL is the forward KL to the teacher over the full vocabulary on held-out packed text at the stage's sequence length; PPL is wikitext-2 at 4096 context. Mamba-3 enters stage 2a at KL 0.50 (Mamba2: 1.08) and passes Mamba2's final KL by step 20 of 287.

Not yet evaluated. No downstream battery (commonsense, RULER needle-in-a-haystack, GSM8K / MATH / AIME, MMLU) has been run on any Mamba-3 checkpoint; the numbers above are the trainers' own validation metrics. Stage 2b (32k KD) is complete for both arms. Stage 3 (on-policy distillation) has not been run. Intermediate checkpoints, published for the ladder.