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. 1.On-policy KD (this run) on 5k code prompts → this checkpoint

Models

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

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):

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. 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. 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. 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. 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.