RockToken/gemma4_31b_to_e4b_onpolicy_science_5k

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

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

Companion to gemma4_31b_to_e4b_onpolicy_math_2k — same recipe, different domain slice, run to check whether the Rock-Token findings hold outside the Qwen family and outside math. As in the math run 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 science 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 == "science" slice
• 5,000 single-user-turn prompts
• 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), 5 h 55 min wall-clock (~24 GPU-hours).

Key hyperparameters

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

Only prompt_max_len differs from the math run (1280 vs 1024): science prompts run longer, p99 ≈ 1182 tokens.

Training dynamics

Per-step values are noisy at batch 4; these are means over 100-step blocks.

steps	loss (reverse KL)	teacher–student top-4 overlap
1–100	3.54	0.606
101–200	3.20	0.621
201–300	3.16	0.622
301–400	3.04	0.627
401–500	3.01	0.627
501–600	2.99	0.628
601–617	2.89	0.628

Both curves move monotonically, but they flatten early: most of the top-4 gain lands in the first 200 steps and the last 400 add 0.006. The math run over the same number of GPU-hours went 0.616 → 0.709.

The likely reason is token budget rather than the domain being harder. Mean response length here is 1,017 tokens (p95 1,307, max 2,233), against 2,000–5,200 for math. Science answers are simply shorter, so despite 617 iterations against math's 250, this run sees fewer generated tokens end to end — and in on-policy KD the loss is summed over the student's own tokens.

Generation was never truncated: the longest per-step mean was 2,233 tokens against the 12,000 cap, using under 19% of the budget.

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 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, 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. 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 six steps in. Gemma-4 alternates sliding-window and full attention, so mask shapes vary layer to layer. cuDNN is dropped from the sdpa candidate list; the flash / mem-efficient kernels handle the same shapes.
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: science (OpenThoughts-3 science split).

Limitations

• Trained end-to-end on science 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 several continual rounds.
• The top-4 curve is flat over the second half of training, so this checkpoint is unlikely to differ much from one taken at step 300.
• Requires transformers >= 5.6 for Gemma-4 support.