Gemma4 Dense and MoE with GSM8K

Gemma4 Dense and MoE with GSM8K#

This example is a small model-support validation for the Gemma4 text models. It uses GSM8K because the purpose is to verify the Megatron model path, SGLang rollout load path, loss masking, backward pass, and live weight update without adding task-specific runtime variables.

Larger task-specific recipes should be layered on after this validation passes.

What to Run#

Run the dense and MoE variants separately on one 8-GPU node:

Model

Script

Megatron topology

SGLang topology

google/gemma-4-31B-it

scripts/run-gemma4-31B-gsm8k.sh

TP2 PP4 CP1

TP8

google/gemma-4-26B-A4B-it

scripts/run-gemma4-26B-A4B-gsm8k.sh

TP2 PP2 EP2 CP1

TP8

The scripts default to two rollouts with short responses. They are intended to prove that the model can train, not to report a meaningful GSM8K score. A small default --entropy-coef keeps the optimizer path active even when the tiny sample receives zero reward.

Use a fresh converted checkpoint directory for each model and topology. The default paths include TP/PP/EP/CP because Megatron distributed checkpoints are sharded by the conversion topology.

Prepare Checkpoints and Data#

cd /root
git clone https://github.com/THUDM/slime.git
cd slime
pip install -e . --no-deps

hf download google/gemma-4-31B-it --local-dir /root/gemma-4-31B-it
hf download google/gemma-4-26B-A4B-it --local-dir /root/gemma-4-26B-A4B-it
hf download --repo-type dataset zhuzilin/gsm8k --local-dir /root/datasets/gsm8k

Convert the dense checkpoint:

cd /root/slime
source scripts/models/gemma4-31B.sh
PYTHONPATH=/root/Megatron-LM torchrun --nproc-per-node 8 \
   tools/convert_hf_to_torch_dist.py \
   "${MODEL_ARGS[@]}" \
   --hf-checkpoint /root/gemma-4-31B-it \
   --tensor-model-parallel-size 2 \
   --pipeline-model-parallel-size 4 \
   --context-parallel-size 1 \
   --save /root/gemma-4-31B-it_tp2_pp4_cp1_torch_dist

Convert the MoE checkpoint:

cd /root/slime
source scripts/models/gemma4-26B-A4B.sh
PYTHONPATH=/root/Megatron-LM torchrun --nproc-per-node 8 \
   tools/convert_hf_to_torch_dist.py \
   "${MODEL_ARGS[@]}" \
   --hf-checkpoint /root/gemma-4-26B-A4B-it \
   --tensor-model-parallel-size 2 \
   --pipeline-model-parallel-size 2 \
   --expert-model-parallel-size 2 \
   --context-parallel-size 1 \
   --save /root/gemma-4-26B-A4B-it_tp2_pp2_ep2_cp1_torch_dist

Run Training#

cd /root/slime
bash scripts/run-gemma4-31B-gsm8k.sh
bash scripts/run-gemma4-26B-A4B-gsm8k.sh

To log the validation runs:

USE_WANDB=1 WANDB_PROJECT=slime-gemma4-gsm8k bash scripts/run-gemma4-31B-gsm8k.sh
USE_WANDB=1 WANDB_PROJECT=slime-gemma4-gsm8k bash scripts/run-gemma4-26B-A4B-gsm8k.sh

Expected Signal#

A successful run should show:

  • SGLang loading Gemma4ForConditionalGeneration.

  • At least one completed rollout and train step.

  • train/loss, train/grad_norm, and entropy metrics in stdout or W&B.

  • Successful raw update_weights from Megatron to SGLang.

For quality training, increase the rollout count, batch sizes, response length, and evaluation interval, and set ENTROPY_COEF=0.