RockToken/gemma4_31b_to_e4b_onpolicy_code_5k

Gemma-4-E4B distilled from Gemma-4-31B — On-policy 5k (code)

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

Completes the Gemma-4 domain triplet alongside gemma4_31b_to_e4b_onpolicy_math_2k and gemma4_31b_to_e4b_onpolicy_science_5k — same recipe, different domain slice, run to check whether the Rock-Token findings hold outside the Qwen family. As in those runs there is no off-policy stage: the student starts from the released instruct checkpoint, so its whole training exposure is a single round:

1. On-policy KD (this run) on 5k code 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 == "code" slice
• 5,000 single-user-turn prompts; 81 exceeding prompt_max_len=2560 were filtered, leaving 4,919
• 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), 18 h 15 min wall-clock (~73 GPU-hours).

Key hyperparameters

Group	Value
Backend	fsdp2, bf16, gradient ckpt on
Epochs	1 (614 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	2560 (total max_len 14848)
Sampling	temperature 1.0, top-p 1.0
Teacher	TP=4, sleep/wakeup enabled

Two deviations from the math/science gemma runs, both memory-motivated: chunked_loss_size 64 instead of 128, and the SGLang rollout pool at 0.14 instead of 0.18 of GPU memory. The longer code prompts push max_len to 14,848 (vs 13,312 math / 13,568 science), and the original settings OOM'd in the student backward at step 257 of a first attempt. Neither knob changes the computed loss.

Training dynamics

Means over 100-step blocks (per-step values are noisy at batch 4):

steps	loss (reverse KL)	teacher–student top-4 overlap	mean gen length
1–100	2.96	0.630	2,654
101–200	2.44	0.650	2,596
201–300	2.34	0.650	2,399
301–400	2.22	0.653	2,372
401–500	2.16	0.656	2,401
501–600	2.16	0.655	2,290
601–614	2.13	0.658	2,312

Step 1 started at loss 3.90 / top-4 0.578.

No truncation. Generation lengths ran p50 = 2,390, p95 = 3,610, max = 5,679 against the 12,000 cap — no step ever hit it. This contrasts with the Qwen code companion (qwen3_30b_a3b_to_4b_onpolicy_code_5k), whose 8,000 cap truncated whole batches in 20% of steps; the two code runs are therefore not symmetric in supervision coverage, beyond their different starting points.

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; ~15k-token sequences were kept whole on 94 GB cards instead.

KDFlow required local patches. Gemma-4 breaks four 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.
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. Sharding is decided by comparing each table against lm_head.weight rather than by the config flag.
3. cuDNN attention backward. torch prefers the cuDNN kernel for sdpa, and its backward aborted with Expected mha_graph.execute(...).is_good() to be true, but got false. Gemma-4 alternates sliding-window and full attention, so mask shapes vary layer to layer. cuDNN is dropped from the sdpa candidate list.
4. 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: code (OpenThoughts-3 code split).

Limitations

• Trained end-to-end on code prompts only; not tuned for chat, safety, or other 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 a continual round.
• Requires transformers >= 5.6 for Gemma-4 support.