Instructions to use RockToken/gemma4_31b_to_e4b_onpolicy_code_5k with libraries, inference providers, notebooks, and local apps. Follow these links to get started.

Libraries

How to use RockToken/gemma4_31b_to_e4b_onpolicy_code_5k with Transformers:

# Use a pipeline as a high-level helper
from transformers import pipeline

pipe = pipeline("text-generation", model="RockToken/gemma4_31b_to_e4b_onpolicy_code_5k")
messages = [
    {
        "role": "user",
        "content": [
            {"type": "image", "url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/p-blog/candy.JPG"},
            {"type": "text", "text": "What animal is on the candy?"}
        ]
    },
]
pipe(text=messages)

# Load model directly
from transformers import AutoProcessor, AutoModelForMultimodalLM

processor = AutoProcessor.from_pretrained("RockToken/gemma4_31b_to_e4b_onpolicy_code_5k")
model = AutoModelForMultimodalLM.from_pretrained("RockToken/gemma4_31b_to_e4b_onpolicy_code_5k", device_map="auto")
messages = [
    {
        "role": "user",
        "content": [
            {"type": "image", "url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/p-blog/candy.JPG"},
            {"type": "text", "text": "What animal is on the candy?"}
        ]
    },
]
inputs = processor.apply_chat_template(
	messages,
	add_generation_prompt=True,
	tokenize=True,
	return_dict=True,
	return_tensors="pt",
).to(model.device)

outputs = model.generate(**inputs, max_new_tokens=40)
print(processor.decode(outputs[0][inputs["input_ids"].shape[-1]:]))

Notebooks
Google Colab
Kaggle
Local Apps Settings

vLLM

How to use RockToken/gemma4_31b_to_e4b_onpolicy_code_5k with vLLM:

Install from pip and serve model

# Install vLLM from pip:
pip install vllm
# Start the vLLM server:
vllm serve "RockToken/gemma4_31b_to_e4b_onpolicy_code_5k"
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:8000/v1/chat/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "RockToken/gemma4_31b_to_e4b_onpolicy_code_5k",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

Use Docker

docker model run hf.co/RockToken/gemma4_31b_to_e4b_onpolicy_code_5k

SGLang

How to use RockToken/gemma4_31b_to_e4b_onpolicy_code_5k with SGLang:

Install from pip and serve model

# Install SGLang from pip:
pip install sglang
# Start the SGLang server:
python3 -m sglang.launch_server \
    --model-path "RockToken/gemma4_31b_to_e4b_onpolicy_code_5k" \
    --host 0.0.0.0 \
    --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/chat/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "RockToken/gemma4_31b_to_e4b_onpolicy_code_5k",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

Use Docker images

docker run --gpus all \
    --shm-size 32g \
    -p 30000:30000 \
    -v ~/.cache/huggingface:/root/.cache/huggingface \
    --env "HF_TOKEN=<secret>" \
    --ipc=host \
    lmsysorg/sglang:latest \
    python3 -m sglang.launch_server \
        --model-path "RockToken/gemma4_31b_to_e4b_onpolicy_code_5k" \
        --host 0.0.0.0 \
        --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/chat/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "RockToken/gemma4_31b_to_e4b_onpolicy_code_5k",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

Docker Model Runner
How to use RockToken/gemma4_31b_to_e4b_onpolicy_code_5k with Docker Model Runner:
```
docker model run hf.co/RockToken/gemma4_31b_to_e4b_onpolicy_code_5k
```

gemma4_31b_to_e4b_onpolicy_code_5k / README.md

Zhao666

Upload folder using huggingface_hub

0f3adca verified 7 days ago

preview code

Raw

History Blame Contribute Delete

6.82 kB

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