Katkat
Back

GRPO Deep Dive: Critic-Free RL, Parallelism, MLA, and Reward Design for a Reasoning LLM

AmazonApplied Scientist · Technical Screen
HardApplied ScientistTechnical ScreenMachine Learning

Last updated

This question tests your deep understanding of Group Relative Policy Optimization (GRPO) for post-training reasoning models, including how critic-free optimization alters training dynamics, distributed memory scaling, KV-cache efficiency via Multi-Head Latent Attention (MLA), and reward hacking mitigations.

Scenario

You are leading the post-training team for an open-ended mathematical and code reasoning large language model (70B parameters). Standard Proximal Policy Optimization (PPO) incurs severe GPU memory pressure and infrastructure overhead because of maintaining four concurrent models: Actor, Critic, Reference, and Reward models. Your team decides to replace standard PPO with Group Relative Policy Optimization (GRPO) paired with an architecture utilizing Multi-Head Latent Attention (MLA).

Walk through the theoretical, architectural, and systems design decisions involved in this pipeline.

Questions to Address

1. Critic-Free Advantage Estimation and Mathematical Formulation

  • Detail how GRPO formulates advantage estimation without training an auxiliary Value/Critic network.
  • Given a prompt q and a set of G sampled candidate outputs {o_1, o_2, ..., o_G} with scalar rewards {r_1, r_2, ..., r_G}, state the advantage equation A_i for output o_i.
  • Explain how the clipped surrogate objective and the token-level KL divergence penalty are calculated in GRPO compared to standard PPO. What failure modes emerge when group size G is too small (e.g., G = 2) versus sufficiently large (e.g., G = 16)?

2. Distributed Training and Memory Footprint

  • Quantify the GPU memory savings achieved by eliminating the critic model during both generation (rollout) and backpropagation phases for a 70B parameter model.
  • Discuss how Multi-Head Latent Attention (MLA) impacts rollout efficiency compared to standard Grouped-Query Attention (GQA). Specifically, detail how compressed latent vectors for keys and values affect KV-cache memory limits, batch sizing, and maximum context length during group generation.

3. Reward Modeling and Failure Modes in Reasoning

  • Define a composite reward function R(o) = R_acc(o) + lambda * R_format(o) where R_acc is rule-based execution/verification (e.g., unit tests or symbolic math equality) and R_format checks structural constraints (e.g., ... tags).
  • Describe two critical training failure modes observed in reasoning models (such as "length gaming" or "entropy collapse" where the model outputs repetitive reasoning traces without reaching valid conclusions).
  • What concrete modifications to the reward function, advantage normalization, or KL regularization mitigate these failure modes?

Hints

Hint 1

Recall that GRPO standardizes rewards within the group of outputs generated from the exact same prompt, eliminating the baseline value function V(s).

Hint 2

MLA projects keys and values into a lower-dimensional latent space before caching, significantly lowering KV cache memory per token compared to standard GQA or MHA.

Hint 3

If rewards are sparse and accuracy is binary (0 or 1), a group where all responses fail yields zero variance, causing advantage normalization instability unless an epsilon term or group filtering is applied.

Follow-up questions

  • How does off-policy sample reuse (multiple gradient update steps per rollout batch) impact policy divergence and clipping thresholds when using GRPO?
  • If all G outputs for a hard prompt receive a reward of 0, how should the advantages be assigned to avoid noisy or misleading gradient signals?
  • How would you dynamically schedule the KL penalty coefficient or the format reward weight lambda across training steps to prevent premature reasoning collapse?