Run RL fine-tuning for a 32B model on 256 GPUs

Kestrel Labs (fictional) trains a 32B coding model with reinforcement learning. The reward is whether the model's code passes unit tests. The research lead opens with: "Our runs are slow and the GPUs look idle half the time. Design the setup for the next run."

The interviewer gives these facts when asked. They are all listed here so you can check your math.

The model:

• 32 billion parameters, dense, BF16.
• 64 layers, hidden size 5,120, grouped-query attention with 8 KV heads of dimension 128.

One RL step:

• 512 coding prompts, 8 sampled completions each: 4,096 completions.
• Prompts average 1,200 tokens.
• Completions average 2,500 tokens. The longest hit the cap of 16,384 tokens.
• Each completion runs its unit tests in a CPU sandbox: 8 s on average. A separate pool of 1,024 CPU cores is available.
• The policy loss uses group-relative advantages (the 8 completions of a prompt are compared with each other) plus a KL penalty against a frozen reference model. The trainer splits the batch into 4 minibatches with one optimizer step each, so it needs the old policy's log-probs before the first update.

Hardware (assumptions):

• 32 nodes of 8 H100 GPUs. Per GPU: 80 GB, 3.35 TB/s, 989 TFLOPS dense BF16. NVLink at 900 GB/s inside a node.
• 8 × 400 Gb/s InfiniBand per node (50 GB/s per GPU).
• Shared storage reads and writes at 20 GB/s in total.
• Plan with 40% of peak compute for training and 75% of peak bandwidth for generation.
• A node fails about once every 20 days.

Target: 600 RL steps in 48 hours.

Asks:

1. Estimate the time per step for generation and for training. How do you split the 256 GPUs?
2. Design the system: rollout workers, reward sandboxes, trainer, weight sync, data flow.
3. What do the long completions do to the step time, and what do you change?
4. How do you checkpoint and recover from failures?

Write the 3 questions you would ask before you design anything. Then estimate the GPU-seconds one RL step needs for generation and for training. Which is bigger?