qwen3-4b-mamba3-hybrid-stage2a-kd-bias0
Stage-2a endpoint, bias-0 arm -- forward-KL logit knowledge distillation against Qwen3-4B over the full 151,936-token vocabulary, teacher-forced at temperature 1, sequence length 4096, 600M tokens (287 steps of 2,097,152). All weights train, in two LR groups split on the parameter-name substring .linear_attn.: mixers at 2e-4, inherited weights at 2e-5, cosine to 0.1x with 3% warmup. 5 h 58 m on 4x B200 at micro-batch 8.
Final val KL 0.1297, top1-agreement 0.8580, wikitext-2 PPL 12.49, from a stage-1 starting point of 0.4951 / 0.7533 / 26.50. At this same stage the Qwen GDN arm reaches KL 0.1437 / top1 0.8515 / PPL 12.83 and the Qwen Mamba2 arm 0.1819 / 0.8295 / 12.60. Input to stage-2b long-context KD.
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)
| stage | bias 0 | bias 1 | Qwen GDN (dtfix) | Qwen Mamba2 (dtfix) |
|---|---|---|---|---|
| stage 1: val mean rel-MSE / worst layer / wikitext2 PPL | 0.00397 / L0 0.0371 / 26.54 | 0.00376 / L0 0.0357 / 25.35 | 0.0208 / - / - | 0.0175 / - / - |
| stage 2a (4k): val KL / top1 agreement / PPL | 0.1297 / 0.8580 / 12.49 | 0.1258 / 0.8601 / 12.41 | 0.1437 / 0.8515 / 12.83 | 0.1819 / 0.8295 / 12.60 |
| stage 2b (32k): val KL / top1 agreement / PPL | 0.2021 / 0.8056 / 11.93 | 0.1957 / 0.8089 / 11.81 | 0.2177 / 0.7985 / - | 0.2336 / 0.7875 / - |
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.