qwen3-4b-mamba2-hybrid-stage2a-kd
Stage-2a endpoint -- 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 (601,882,624 consumed). Unlike stage 1, 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.
Final val KL 0.1819, top1-agreement 0.8295, wikitext2 PPL 12.60, from a stage-1 starting point of 1.0832 / 0.6063 / 44.22. 6 h 29 m on 4x H100.
Lineage note: the identical run from the broken dt init reached KL 0.2060 / top1 0.8187 / PPL 12.61 -- so the dt repair bought teacher agreement, not perplexity. For reference the Qwen Gated-DeltaNet arm at this same stage reaches KL 0.1449 / top1 0.8508 / PPL 12.83: Mamba2 is the better standalone language model here (lower PPL) but the worse imitator of the teacher (higher KL). Validation KL was flat from step 260 of 287, so this converged inside its token budget.
RULER needle-in-a-haystack (5 seeds x 100 items pooled, n=500 per cell, greedy, thinking disabled, 128-token answer cap):
| task | 4K | 8K | 16K | 32K | macro |
|---|---|---|---|---|---|
| niah_single | 100.0 | 100.0 | 99.6 | 96.4 | 99.0 |
| niah_multikey | 94.4 | 72.6 | 51.8 | 27.6 | 61.6 |
| niah_multiquery | 92.0 | 73.6 | 59.3 | 24.8 | 62.4 |
Against the Qwen Gated-DeltaNet arm at the same stage (99.7 / 56.4 / 68.5 macro) Mamba2 is ahead on multikey at 4K, 8K and 16K and behind on multiquery, with the gap concentrated at 32K -- which is extrapolation for both arms, since stage 2a trains only at sequence length 4096. Format pathologies are gone at this stage: repetition loops fall to 0-1.4% of generations, and niah_single accuracy conditional on emitting an answer is 100.0% at every length, so the remaining misses are retrieval failures rather than formatting. The same run from the broken dt init reached only 23.1 multikey macro.
Commonsense and math batteries have not been run yet. Intermediate checkpoint, published for the ladder.
Loading
This is a custom container (model_type: qwen3_mamba2), not a stock
architecture, so trust_remote_code=True is required. The modeling code is
bundled in this repo 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)
Install mamba-ssm and causal-conv1d for the fast SSD kernels. Without them
transformers.models.bamba.modeling_bamba.is_fast_path_available is False and
every mixer silently falls back to torch_forward, which materializes a
[batch, n_chunks, heads, 256, 256] score tensor per layer. Correctness is
unaffected; throughput and memory are not. Gate on that flag in job scripts.
Never left-pad. apply_mask_to_padding_states is skipped on the training
branch, so pads reach the SSM state through conv1d.bias. Right padding is
causally harmless.
The study
Full-attention Qwen3-4B converted to a Mamba2 hybrid: 27 of 36 layers
have their attention replaced by a Mamba2 mixer (BambaMixer, n_groups=16,
expand=1.6, 32 heads x 128 head_dim, d_state=128), 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.
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 a Gated DeltaNet arm trained identically. Comparability is the deliverable.
Repaired dt initialization
These checkpoints are the dtfix lineage. transformers' Bamba initializer
sets dt_bias to 1.0 (modeling_bamba.py:1090), contradicting its own comment
two hundred lines above: softplus(1.0) = 1.3133, about 13x the intended
time_step_max of 0.1. Since dt sets both the per-token decay
exp(-dt*|A|) and the write strength, and this container's A = -arange(1, H+1)
floors |A| >= 1, the stock init leaves dt*|A| at a median of 19.38 --
the recurrent state is dead at step one. The repaired init samples
dt ~ exp(U(log 1e-3, log 1e-1)) and stores inv_softplus(dt), giving a median
dt*|A| of 0.12-0.13, inside the 0.020-0.574 band measured across nine
trained Mamba2-family checkpoints. The same defect is present in
granitemoehybrid and qwen3_next; mamba2 and zamba2 are correct.