Repository: google/tunix
File: scripts/grpo_demo_llama3_qwen2.py (Python)

Code before:
# training.
TEMPERATURE = 0.9
TOP_P = 1.0  # implies we don't do nucleus sampling
TOP_K = 50
# The number of times the policy generates multiple responses for a given prompt
# within a single training step. This corresponds to `G` in Algorithm 1 in the
# paper. The "group" in GRPO comes from here.
NUM_GENERATIONS = 4

# === other GRPO configs ===
# The number of iterations per batch (𝜇 in GRPO algo 1).
NUM_ITERATIONS = 1
# The coefficient for the KL divergence penalty (𝛽) in the GRPO loss function.
# Important to keep a high enough value for this, otherwise, the KL divergence
# can increase unchecked.
BETA = 0.08
# Epsilon value for clipping (𝜀 in GRPO loss in paper). Similar to PPO, for
# stable updates.
EPSILON = 0.2

# ====== Training ======
# 2 is the max we can do on v5e-8 with llama3 8B model.
# 4 is the max we can do on v5e-8 with llama3 1B model.
GLOBAL_BATCH_SIZE = len(jax.devices())
MINI_BATCH_SIZE = GLOBAL_BATCH_SIZE
TRAIN_MICRO_BATCH_SIZE = GLOBAL_BATCH_SIZE // 2
# To speed up for quick workflow validation, we can change NUM_BATCHES to e.g. 2
NUM_BATCHES = min(args.num_batches, 7473 // GLOBAL_BATCH_SIZE)
# Keep `NUM_TEST_BATCHES` low so that evaluation runs quickly. It can be
# increased to a max. of 330 (if batch size is 4).
# To speed up for quick workflow validation, we can change it to e.g. 1
NUM_TEST_BATCHES = args.num_test_batches

EVAL_EVERY_N_STEPS = 10  # this doesn't matter if `TRAIN_FRACTION = 1.0`.
NUM_EPOCHS = 1  # can potentially train for more epochs

# Number of training steps.
MAX_STEPS = int(NUM_BATCHES * NUM_ITERATIONS * TRAIN_FRACTION * NUM_EPOCHS)

# === AdamW, warmup, cosine scheduler ===
LEARNING_RATE = 3e-6
B1 = 0.9  # Adam beta1
B2 = 0.99  # Adam beta2
WEIGHT_DECAY = 0.1
# == Cosine decay with warmup scheduler ==
# Linearly increase learning rate from 0. to 5e-6 in the first 10% training
# steps, and then gradually decrease the learning rate to 0 using cosine
# scheduler.
WARMUP_STEPS = 0.1 * MAX_STEPS
# == Grad clipping ==
# Grad clipping to prevent large gradients. Found this

Review comment (performance):
This would probably cause llama 3 8B to go out of memory on v5e-8. Assuming the previous v5e-8 config was to use TRAIN_MICRO_BATCH_SIZE=2

Code after:

# ====== Reproducibility ======
SEED = 42

# ====== LoRA ======
ENABLE_LORA = False
RANK = 64
ALPHA = 64.0

# ====== Sharding ======
if "Qwen2.5-0.5B-Instruct" in args.model_version:
  TOTAL_TPU_TO_USE = 2
elif "Qwen2.5-7B-Instruct" in args.model_version:
  TOTAL_TPU_TO_USE = 4
else:
  TOTAL_TPU_TO_USE = jax.device_count()

MESH = [(1, TOTAL_TPU_TO_USE), ("fsdp", "tp")]  # YY

# ====== GRPO ======
# === Generation during GRPO training ===
MAX_PROMPT_LENGTH = 256
TOTAL_GENERATION_STEPS = 768
# Important to keep a high-ish temperature for varied, diverse responses during
# training.
TEMPERATURE = 0.9
TOP_P = 1.0  # implies we don't do nucleus sampling
TOP_K = 50
# The number of times the policy generates multiple responses for a given prompt
# within a single training step. This corresponds to `G` in Algorithm 1 in the
# paper. The "group" in GRPO comes from here.
NUM_GENERATIONS = 4

# === other GRPO configs ===
# The number of iterations per batch (𝜇 in GRPO algo 1).
NUM_ITERATIONS = 1
# The coefficient for the KL divergence penalty (𝛽) in the GRPO loss function.
# Important to keep a high enough value for this, otherwise, the KL divergence
# can increase unchecked.
BETA = 0.08
# Epsilon value for clipping (𝜀 in GRPO loss in paper). Similar to PPO, for
# stable updates.
EPSILON = 0.2

# ====== Training ======
# To speed up for quick workflow validation, we can change NUM_BATCHES to e.g. 2
NUM_BATCHES = min(args.num_batches, 7473 // args.global_batch_size)
# Keep `NUM_TEST_BATCHES` low so that evaluation runs quickly. It can be
# increased to a max. of 330 (if batch size is 4).
# To speed up for quick workflow validation, we can change it to e.g. 1
NUM_TEST_BATCHES = args.num_test_batches