Skip to main content

​
1. Model Introduction

Gemma-4 is Google’s multimodal model line. miles trains both released instruction-tuned sizes as language models, on the base VLM checkpoint directly. Both go through the HF to Megatron bridge (--megatron-to-hf-mode bridge), and on the rollout side sglang runs Gemma4ForConditionalGeneration, which loads Gemma-4’s hybrid head_dim weights correctly. There is no offline torch_dist conversion and no LLM-view rewrite of the checkpoint. Key highlights:
  • Two shapes, one recipe family: 26B-A4B is MoE (128 experts, top-8), 31B is dense. They differ mainly in expert parallelism and the token budget.
  • Bridge-mode load straight from the VLM checkpoint.
  • Tied embeddings: neither config passes --untie-embeddings-and-output-weights.
  • Single node: both recipes target 8 × H200.

​
2. Supported Variants

Both use GQA with kv_channels=256, RoPE base 1e6, and a 262144-token vocab. The 31B recipe requires the zhichen/gemma4-dense branch of radixark/Megatron-Bridge.

​
3. Environment Setup

​
3.1 Download model + datasets

--model-dir and --data-dir default to /root/models and /root/datasets. prepare performs these downloads for you.

​
3.2 No torch_dist conversion

The bridge reads the HF checkpoint directly, so --hf-checkpoint and --ref-load both point at the download:

​
4. Launch

Passing --num-nodes 1 puts the recipe into debug_minimal mode, which shortens --rollout-max-response-len to 256 for a quick smoke test. Multi-node runs use the full 8192.

​
5. Recipe Configuration

​
5.1 Parallelism

Activation checkpointing is on for the MoE recipe (--recompute-granularity full --recompute-method uniform --recompute-num-layers 1). The dense 31B runs a smaller token budget because its 60 dense layers at hidden 5376 cost more activation memory per token than the MoE’s 30 layers at 2816.

​
5.2 Algorithm

GRPO. The MoE recipe adds low-variance KL; the dense one runs without it:
Rollout batch 32 at 8 samples per prompt, global batch 256, --lr 1e-6. AIME evaluation every 20 steps is available behind --enable-eval and is off by default.

​
5.3 Rollout & SGLang

The MoE recipe pins sglang to conservative kernels:
--use-rollout-routing-replay replays the rollout’s expert routing during the training forward pass, so train log-probs match rollout log-probs. Every sigmoid- or softmax-routed MoE recipe in miles needs this; the dense 31B does not.

​
5.4 Notable quirks

  • Trained on the VLM checkpoint. miles does not strip the vision tower; the bridge and sglang both handle the multimodal config, and the RL recipe simply trains the language stack.
  • --attention-backend unfused on the training side for the MoE recipe.
  • Routing is softmax with seq_aux_loss balancing and the bias update rate at 0 (--moe-router-bias-update-rate 0 --moe-aux-loss-coeff 0), plus --moe-grouped-gemm and --moe-router-dtype fp32.
  • The 31B recipe needs the gemma4-dense branch of radixark/Megatron-Bridge, because the dense config is driven straight through Gemma4VLBridge.

​
6. Pairs Well With