README.md

---
license: gemma
library_name: transformers
pipeline_tag: text-generation
language:
- en
base_model: google/gemma-4-E4B-it
datasets:
- open-thoughts/OpenThoughts3-1.2M
tags:
- knowledge-distillation
- on-policy-distillation
- code
- gemma-4
---

# 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`](https://huggingface.co/RockToken/gemma4_31b_to_e4b_onpolicy_math_2k) and [`gemma4_31b_to_e4b_onpolicy_science_5k`](https://huggingface.co/RockToken/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](https://github.com/songmzhang/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`](https://huggingface.co/RockToken/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.