import compute
app = compute.App("rl-grpo")
image = compute.Image.rocm_pytorch().pip_install(
"transformers",
"peft",
"datasets",
"trl",
"accelerate",
"math_verify",
)
DEFAULT_MODEL = "Qwen/Qwen2.5-0.5B-Instruct"
DEFAULT_DATASET = "trl-lib/DeepMath-103K"
@app.function(gpu="MI300X", image=image, timeout=1800)
def train(
model_id: str = DEFAULT_MODEL,
dataset_id: str = DEFAULT_DATASET,
max_steps: int = 1,
sample_count: int = 16,
lr: float = 1e-5,
rank: int = 8,
num_generations: int = 4,
max_completion_length: int = 256,
max_prompt_length: int = 512,
) -> dict:
import time
import torch
from datasets import load_dataset
from peft import LoraConfig, TaskType
from trl import GRPOConfig, GRPOTrainer
from trl.rewards import accuracy_reward
if not torch.cuda.is_available():
raise RuntimeError("this entrypoint needs a GPU")
if num_generations < 2:
raise ValueError("num_generations must be >= 2")
if sample_count < num_generations:
raise ValueError("sample_count must be >= num_generations")
dataset = load_dataset(dataset_id, split=f"train[:{sample_count}]")
trainer = GRPOTrainer(
model=model_id,
reward_funcs=accuracy_reward,
train_dataset=dataset,
args=GRPOConfig(
output_dir="/tmp/lora-grpo",
max_steps=max_steps,
per_device_train_batch_size=num_generations,
gradient_accumulation_steps=1,
learning_rate=lr,
logging_steps=1,
bf16=True,
optim="adamw_torch",
report_to="none",
save_strategy="no",
gradient_checkpointing=True,
num_generations=num_generations,
max_completion_length=max_completion_length,
max_prompt_length=max_prompt_length,
remove_unused_columns=False,
model_init_kwargs={"torch_dtype": "bfloat16"},
),
peft_config=LoraConfig(
task_type=TaskType.CAUSAL_LM,
r=rank,
lora_alpha=rank * 2,
lora_dropout=0.05,
bias="none",
target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
),
)
t0 = time.perf_counter()
metrics = trainer.train().metrics
reward_keys = sorted(
key for key in metrics if key.startswith("rewards/") and key.endswith("/mean")
)
return {
"ok": True,
"method": "grpo",
"model_id": model_id,
"dataset_id": dataset_id,
"max_steps": max_steps,
"sample_count": sample_count,
"num_generations": num_generations,
"device": torch.cuda.get_device_name(0),
"train_loss": float(metrics.get("train_loss", 0.0)),
"reward_mean": float(metrics[reward_keys[0]]) if reward_keys else None,
"wall_s": round(time.perf_counter() - t0, 3),
}
