RockToken/gemma4_31b_to_e4b_onpolicy_math_2k

Gemma-4-E4B distilled from Gemma-4-31B — On-policy 2k (math)

On-policy KD for a Gemma-4-E4B student toward the Gemma-4-31B teacher, on 2,000 math prompts from OpenThoughts-3.

This is the Gemma-4 counterpart to the Qwen3 checkpoints in this organisation, run to test whether the Rock-Token findings hold outside the Qwen family. Unlike the Qwen chain, there is no off-policy stage: the student starts from the released instruct checkpoint, so the effective training exposure of this checkpoint is a single round:

1. On-policy KD (this run) on 2k math prompts → this checkpoint

Models

Role	Model
Student	google/gemma-4-E4B-it
Teacher	google/gemma-4-31B-it (dense)

enable_thinking=False throughout.

Training data

• Source: open-thoughts/OpenThoughts3-1.2M, domain == "math" slice
• 2,000 single-user-turn prompts, median length ~222 chars
• Only the prompts are used; on-policy KD never reads the dataset's reference answers

Training setup

Framework: KDFlow v0.2.0 — FSDP2 + SGLang rollout, Ray-orchestrated GPU co-location with sleep/wakeup.

Hardware: 1× node, 4× H100 (94 GB), 8 h 27 min wall-clock (~34 GPU-hours).

Key hyperparameters

Group	Value
Backend	fsdp2, bf16, gradient ckpt on
Epochs	1 (250 rollout iterations)
Train batch	4 (micro 1)
Learning rate	2e-6, cosine, warmup 5%
KD ratio	1.0
KD loss	reverse KL (rkl)
KD algorithm	vanilla_kd
Temperature (KD)	1.0
Rollout engine	SGLang, TP=2, 1 engine
Rollout batch	8 prompts × 4 samples/prompt
generate_max_len	12000
prompt_max_len	1024 (total max_len 13312)
Sampling	temperature 1.0, top-p 1.0
Teacher	TP=4, sleep/wakeup enabled

Training dynamics

	step 3	step 250
loss (reverse KL)	3.16	1.45
teacher–student top-4 overlap	0.616	0.709
gradient norm	43.8	~25–30
mean response length	3,103	5,238

Generated responses averaged 2,000–5,800 tokens against the 12,000 cap, so generation was effectively untruncated.

Deviations from the Qwen pipeline

No sequence parallelism. ring_flash_attn 0.1.8 imports is_flash_attn_greater_or_equal_2_10 from transformers.modeling_flash_attention_utils, which transformers 5.x removed, while Gemma-4 requires transformers ≥ 5.6. Ring attention is therefore unavailable for this model family; 13k-token sequences were kept whole on 94 GB cards instead.

KDFlow required local patches. Gemma-4 breaks three assumptions that hold for Qwen3:

1. Cross-layer KV sharing. Gemma-4 threads one mutable shared_kv_states dict through all 42 decoder layers (22–23 write, 24–41 read). fully_shard's forward wrapper rebuilds the containers in a layer's arguments, so a wrapped writer mutates a private copy and the readers raise KeyError: 22. The two writer layers are left unsharded; readers shard normally, since the rebuilt dict carries existing entries.
2. Per-layer embeddings. KDFlow skips embedding sharding whenever tie_word_embeddings is set. Gemma-4 ties only embed_tokens (1.34 GB) to lm_head, while embed_tokens_per_layer is a separate 5.64 GB table — 44% of the model would stay replicated on every rank. Sharding is now decided by comparing each table against lm_head.weight rather than by the config flag.
3. lm_head loading. load_only_lm_head materialised the whole 49.8 GB teacher shard to read one 2.8 GB tensor; safe_open reads only that tensor.

Patches 1 and 2 change how the model is sharded, so results are not bit-identical to what stock KDFlow would produce.

Intended use

Research on distillation dynamics and on the cross-family generality of the Rock-Token analysis. Domain: math (OpenThoughts-3 math split).

Limitations

• Trained end-to-end on math prompts only; not tuned for chat, safety, or non-math domains.
• enable_thinking=False — this student does not emit thinking traces.
• Single on-policy round from the base instruct model, with no off-policy warm-up, so it is not directly comparable to the Qwen chain checkpoints, which carry off-policy KD plus several continual rounds.
• Requires transformers >= 5.6 for Gemma-4 support.