Katkat
Back

Explain Pipeline Parallelism in GRPO Training

GoogleMachine Learning Engineer · Technical Screen
HardMachine Learning EngineerTechnical ScreenML System Design

Last updated

Design and analyze an efficient pipeline parallel training architecture for Group Relative Policy Optimization (GRPO) on large reasoning models. Tests deep understanding of pipeline scheduling, microbatching, distributed advantage normalization, and policy staleness trade-offs.

Scenario

You are an ML infrastructure engineer tasked with scaling Group Relative Policy Optimization (GRPO) to train a 70B-parameter reasoning model across a cluster of accelerator nodes (e.g., 64x 80GB GPUs/TPUs). GRPO removes the learned critic network found in standard PPO, instead generating a group of G candidate completions per prompt, computing scalar rewards for each, and normalizing advantages across the group (A_i = (R_i - mean(R)) / std(R)).

Because the 70B model cannot fit on a single device alongside activation memory and KV-cache, you must employ Pipeline Parallelism (PP) combined with Tensor Parallelism (TP).

Requirements

Functional

  • Support the two alternating phases of GRPO: Rollout Generation (autoregressive decoding of G outputs per prompt) and Policy Optimization (forward/backward passes with microbatching for policy gradient updates).
  • Formulate group advantage computation across microbatches and distributed stages without introducing synchronization deadlocks or excessive pipeline bubbles.
  • Handle policy divergence and staleness when generation and training are decoupled or pipelined asynchronously.

Non-Functional & Scale

  • Model: 70B parameters (~140GB in bfloat16 weights alone).
  • Topology: 8 nodes of 8 GPUs (64 GPUs total). Target config: TP = 4, PP = 8, DP = 2.
  • Workload: Group size G = 8 rollouts per prompt, prompt length 1K tokens, generation length up to 4K tokens.
  • Efficiency: Minimize pipeline bubble fraction (F_bubble) while keeping device memory safely under 80GB.

Discussion Points

Your design should address:

  1. Pipeline Architecture & Scheduling: Compare colocating generation and training on the same PP engine versus decoupling them into dedicated inference and training clusters. How does 1F1B (One-Forward-One-Backward) schedule interact with group-level loss computation?
  2. Group Advantage Synchronization: In standard PP, microbatches flow independently through pipeline stages. Where and when in the pipeline are group advantages computed, and how do you organize microbatches so that a prompt's G outputs do not fragment across incompatible boundaries?
  3. Memory & Activation Stash: Analyze activation memory during the backward pass in stage PP-1 versus earlier stages, and quantify the KV cache memory footprint required during the generation phase.
  4. Policy Staleness & Importance Sampling: If asynchronous pipelining is used between rollouts and updates, how does policy drift affect the PPO/GRPO clipping mechanism (clip(r_theta, 1 - eps, 1 + eps))? What mitigations maintain training stability?

Hints

Hint 1

Consider packing all G completions of a single prompt into the same microbatch or ensuring deterministic aggregation across microbatches before computing loss tokens.

Hint 2

Autoregressive generation exhibits different computational intensity and communication patterns than training forward/backward passes, making static 1F1B schedules sub-optimal if generation is inline.

Hint 3

Think about where the reference policy logits and token-level log probabilities are evaluated, and whether they can be computed concurrently with the target policy backward pass to minimize stored activations.

Follow-up questions

  • How would you adapt the pipeline schedule if you switched from synchronized batched generation to dynamic-length speculative decoding during the rollout phase?
  • If group size G is larger than a single microbatch's maximum token capacity, how do you handle cross-microbatch mean and variance synchronization for rewards without stalling the pipeline?
  • How does vLLM-style paged attention integration into the rollout workers affect weight synchronization latency between the PP training engine and the inference engine?