Fine-tune LMs with GRPO Loss: A Training Loop Tutorial

In our last two lessons, we crafted the essential building blocks of our GRPO implementation: the function to compute group-relative advantages and the complete GRPO loss function. We now have the tools to evaluate generated responses and to calculate a gradient signal for model improvement. This final lesson brings everything together.

Our goal is to assemble a complete training loop that uses these components to fine-tune a language model. We will orchestrate the entire on-policy learning process, from generating experiences with the model to updating its weights. By the end of this lesson, you will have a runnable script that puts GRPO into action, marking the completion of your journey from the theory of GRPO to a full-fledged implementation.

The On-Policy Training Cycle

GRPO, like its predecessor PPO, is an on-policy algorithm. This means it follows a distinct, two-phase cycle: first, it collects a batch of experience using the current policy, and second, it uses that experience to perform one or more optimization steps. This cycle of "collect then train" is fundamental to its operation.

The video below provides a clear, high-level overview of this two-phase structure.

This cycle forms the backbone of our training loop. We will implement an outer loop that iterates for a set number of training steps. Inside this loop, we will execute the two key phases:

1. Experience Collection (Rollout): We use a fixed version of our policy (model_old) to interact with the environment (i.e., generate responses to prompts). We collect the prompts, the generated responses, the rewards, and the log probabilities of the responses under model_old.
2. Optimization: We use the collected data to compute the GRPO loss and update the weights of our main, trainable policy (model).

Anatomy of the Training Loop

Let's dissect the training loop into a sequence of concrete steps. We will structure our implementation around a single function, grpo_train_loop, which will execute one full iteration of the collect-and-train cycle.

The following resource provides a complete, well-structured implementation of a GRPO training loop. While it uses the Apple MLX framework, the logic is identical to what we will implement in PyTorch. We will use it as our primary guide.

Let's walk through these steps in detail.

Step 1: Prompt Sampling

The loop begins by sampling a mini-batch of prompts from our training dataset. This is standard practice in any deep learning training pipeline.

Step 2: Rollout Phase

This is the core of the experience collection phase. For each prompt in our batch, we generate a group of G responses. Crucially, we use model_old, a frozen copy of our main model, for this generation. This ensures that the policy used for generating data is stable during the collection of a full batch of experiences.

Step 3: Reward Calculation

For each generated response, we calculate a reward. For this implementation, we'll use a simple dummy reward function that checks if the correct answer is present in the generation. In a full RLHF pipeline, this would involve a complex, pre-trained reward model.

Step 4: Advantage & Old Log-Prob Calculation

With a group of rewards for each prompt, we can now call our compute_group_relative_advantage function from the first lesson. This standardizes the rewards, providing the $\hat{A}$ values.

Simultaneously, we need to calculate the log probabilities of the generated tokens under the policy that created them (model_old). These are the old_log_probs required for the probability ratio in our loss function. This is a forward pass through model_old with the generated sequences.

Step 5: The Optimization Step

This is the training phase. We now have all the necessary ingredients to calculate the GRPO loss:

• The generated sequences (prompts + responses).
• The advantages ($\hat{A}$).
• The log probabilities from the old policy (old_log_probs).

Inside the optimization step, we perform another forward pass, this time through the trainable model to get the current_log_probs. We also need a forward pass through the reference model model_ref to get ref_log_probs for the KL penalty.

We then pass all these tensors to our compute_grpo_loss function from the previous lesson. The returned loss is a scalar on which we can call .backward() to compute gradients, followed by optimizer.step() to update the weights of model.

The Complete Implementation

Let's now assemble these steps into a single, cohesive Python function using PyTorch. The function below orchestrates the entire process, using the helper functions we've built in previous lessons.

import torch
import torch.optim as optim
from tqdm import tqdm
import numpy as np




# Assume the following functions are defined from previous lessons:
# - compute_group_relative_advantage(rewards: list[list[float]]) -> torch.Tensor
# - compute_grpo_loss(...) -> torch.Tensor
# - generate(...) -> (str, torch.Tensor) # Returns decoded text and token IDs
# - calculate_log_probs(model, sequences, attention_mask, completion_mask) -> torch.Tensor

def grpo_train_loop(
    model: torch.nn.Module,
    model_old: torch.nn.Module,
    model_ref: torch.nn.Module,
    tokenizer,
    optimizer: optim.Optimizer,
    train_set,  # A list of dictionaries, e.g., [{"prompt": ..., "answer": ...}]
    iters: int = 200,
    group_size: int = 4,
    batch_size: int = 2,
    update_every: int = 10, # Sync old model every 10 iterations
    beta: float = 0.01,
    epsilon: float = 0.2,
):
    """The complete GRPO training loop."""
    model.train()
    model_old.eval()
    model_ref.eval()
    



    # Sync the old model initially
    model_old.load_state_dict(model.state_dict())
    
    losses, all_rewards = [], []
    pbar = tqdm(range(iters))

    for it in pbar:



        # --- 1. Experience Collection ---
        



        # Step 1: Sample a batch of prompts
        prompts = [train_set[i] for i in np.random.randint(0, len(train_set), batch_size)]
        
        rollout_data = {
            "sequences": [], "attention_masks": [], "completion_masks": [], "rewards": []
        }
        



        # Step 2 & 3: Rollout and Reward Calculation
        for p in prompts:
            prompt_tokens = tokenizer(p["prompt"], return_tensors="pt").to(model.device)
            group_rewards = []
            for _ in range(group_size):

```grasp
{
  "type": "exercise",
  "id": "b5c05fa2-8726-4da8-9d66-6f1a94c17cdd"
}

            # Generate a response using the OLD model
            response_text, response_tokens = generate(
                model_old, tokenizer, prompt_tokens.input_ids,
            )
            



            # Dummy reward: 1.0 if correct answer is in the response
            reward = 1.0 if p["answer"] in response_text else 0.0
            group_rewards.append(reward)
            



            # Store sequence data for training
            full_sequence = torch.cat([prompt_tokens.input_ids[0], response_tokens], dim=0)
            rollout_data["sequences"].append(full_sequence)



            # (Assume padding and mask creation happens later)
            
        rollout_data["rewards"].append(group_rewards)
        all_rewards.extend(group_rewards)
        



    # (Assume padding is applied to sequences to form a batch)
    # sequences_padded, attention_masks, completion_masks = pad_and_mask(rollout_data["sequences"])
    



    # Step 4: Compute Advantages & Old Log-Probs
    advantages = compute_group_relative_advantage(rollout_data["rewards"]) # shape: (B*G,)
    
    with torch.no_grad():
        old_log_probs = calculate_log_probs(
            model_old, sequences_padded, attention_masks, completion_masks
        )
        ref_log_probs = calculate_log_probs(
            model_ref, sequences_padded, attention_masks, completion_masks
        )




    # --- 2. Optimization ---
    



    # Step 5: Compute current log_probs and the final loss
    current_log_probs = calculate_log_probs(
        model, sequences_padded, attention_masks, completion_masks
    )
    
    loss = compute_grpo_loss(
        log_probs=current_log_probs,
        old_log_probs=old_log_probs,
        ref_log_probs=ref_log_probs,
        advantages=advantages,
        completion_mask=completion_masks,
        epsilon=epsilon,
        beta=beta,
    )
    



    # Backpropagation and optimizer step
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    losses.append(loss.item())
    pbar.set_description(f"Loss: {np.mean(losses[-10:]):.3f}, Mean Reward: {np.mean(all_rewards[-20:]):.3f}")




    # Sync old model weights periodically
    if (it + 1) % update_every == 0:
        model_old.load_state_dict(model.state_dict())
        print(f"\nIter {it+1}: Synced old model weights.")

return losses, all_rewards

```grasp
{
  "type": "exercise",
  "id": "58124453-9e17-4ef0-b061-4c89fde49928"
}

Model Management: model, model_old, and model_ref

A critical detail in the implementation is the management of three distinct model instances:

1. model: The main model whose weights are being actively trained.
2. model_old: A frozen copy of model. It is used to generate the experiences for a batch. Its weights are periodically updated to match model. This is the essence of the on-policy update rule.
3. model_ref: A frozen reference model, typically the base SFT model you started with. It is never updated. Its sole purpose is to provide a stable distribution for the KL-divergence penalty, preventing the model from drifting too far from its original capabilities.

The line model_old.load_state_dict(model.state_dict()) is where the synchronization happens, ensuring that the policy used for data collection doesn't get too "stale" relative to the policy being optimized.

Conclusion

Congratulations! You have now assembled a complete GRPO training loop. By orchestrating the flow of data from experience collection to optimization, you have created a functional implementation capable of fine-tuning a language model using group-relative rewards.

This concludes our course on GRPO. Let's recap the journey:

• You began by revisiting the fundamentals of PPO and RLHF to set the stage.
• You then delved into the theory of GRPO, understanding its critic-free approach and deriving its key components: the group-relative advantage and the full objective function.
• Finally, you translated this theory into practice, implementing the advantage calculation, the complete loss function, and now, the full training loop that brings it all together.

You have successfully met your goal of understanding both the mathematical theory and the practical implementation of GRPO, equipping you with the skills to apply and adapt this powerful algorithm.