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

Models

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

Training dynamics

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