2dd8339946
30-step results: | Backend | Train wall | Median step | Peak mem | % of vLLM | |----------------|-----------|-------------|----------|-----------| | vLLM | 215.9 s | 5.14 s | 159 GB | 100 % | | fi_false | 1165.4 s | 41.30 s | 10.7 GB | 12.4 % | | cb_paged | 1564.5 s | 39.82 s | 61.9 GB | 12.9 % | Pairwise diff vs vLLM (30 steps, compare_grpo_runs.py): | Pair | max |loss| | max |kl| | max |reward| | |---------------------------------|------------|------------|---------------| | vLLM vs unsloth_fi_false | 0.39 | **0.015** | 9.25 (noisy) | | vLLM vs cb_paged | 0.83 | (missing) | 6.25 (noisy) | KL trajectory match between vLLM and unsloth_fi_false is the load-bearing equivalence signal: both stay in [0, 0.015] across all 30 steps, so the policy drift guardrail behaves the same. Reward diffs of ~3-9 are expected because the rollout backends produce different completions even at temperature=0.1 (kernel-level non-determinism). Two caveats documented in results/grpo_equivalence.md: - cb_paged's StatisticsCallback doesn't capture TRL's kl log entry because TRL emits kl on a separate log call that doesn't include loss. - cb_paged's grad_norm (~200-900) is unclipped pre-optimizer, while vLLM's goes through Unsloth's internal max_grad_norm=1.0. Not a correctness bug, just not apples-to-apples until cb_paged sets max_grad_norm in GRPOConfig. Also updates the report with Phase 3 + Phase 4 status (CB sync driver eager works; CUDA graph capture hangs on output_ids slice pending fix; torch.compile on training step incompatible with both unsloth_fi_false and cb_paged).
128 lines
6.7 KiB
Markdown
128 lines
6.7 KiB
Markdown
# Phase 2: end-to-end GRPO backend comparison
|
|
|
|
Same dataset, reward functions, sampling (`temperature=0.1, top_p=0.97,
|
|
min_p=0.5, top_k=5`), and seed (3407) across backends. `num_generations=4`;
|
|
`per_device_train_batch_size` auto-raised to 4 on vanilla-HF backends so
|
|
TRL's `generation_batch_size % num_generations == 0` check passes
|
|
(Unsloth's loader does this for you, vanilla HF does not).
|
|
|
|
Callbacks: `StatisticsCallback` from `torch_debugging_utils` logs per-step
|
|
loss, grad-norm, memory, wall time; reward / KL are captured from the TRL
|
|
log dict. Median step wall is measured on steps 4..N (first 3 skipped to
|
|
amortize compile / graph / warmup).
|
|
|
|
## 10-step vibe check
|
|
|
|
| Backend | Train wall (s) | Median step (s) | Peak mem (GB) | % of vLLM |
|
|
|----------------------------|----------------|-----------------|---------------|-----------|
|
|
| vLLM (fast_inference) | 74.4 | **4.14** | 157.9 | 100 % |
|
|
| unsloth_fi_false | 355.4 | 23.95 | **10.7** | 17 % |
|
|
| cb_paged (sdpa_paged load) | 466.0 | 36.02 | 55.6 | 11.5 % |
|
|
|
|
## 30-step equivalence
|
|
|
|
| Backend | Train wall (s) | Median step (s) | Peak mem (GB) | % of vLLM |
|
|
|----------------------------|----------------|-----------------|---------------|-----------|
|
|
| vLLM (fast_inference) | 215.9 | **5.14** | 159.0 | 100 % |
|
|
| unsloth_fi_false | 1165.4 | 41.30 | **10.7** | 12.4 % |
|
|
| cb_paged | 1564.5 | 39.82 | 61.9 | 12.9 % |
|
|
|
|
(Note: fi_false's median step jumped from 23.95 s at 10 steps to 41.30 s at
|
|
30 steps because the early-GRPO policy started producing longer completions
|
|
as it learned to place the `</SOLUTION>` marker; the same effect is present
|
|
but smaller in cb_paged because its LoRA warm-up trajectory is different.)
|
|
|
|
## Pairwise diff vs vLLM (30 steps, `scripts/benchmarks/compare_grpo_runs.py`)
|
|
|
|
| Pair | max |loss diff| | max |reward diff| | max |kl diff| | max |grad_norm diff| |
|
|
|---------------------------------|---------------------------|------------------------------|------------------------|--------------------------------|
|
|
| vLLM vs **unsloth_fi_false** | 0.39 | 9.25 (mean 2.99) | **0.015** | 0.94 |
|
|
| vLLM vs **cb_paged** | 0.83 | 6.25 (mean 2.29) | *(not logged)* | 919.1 |
|
|
|
|
Reward diffs of 2-9 are expected: different rollout backends produce
|
|
different completions even at `temperature=0.1` because of kernel-level
|
|
non-determinism (vLLM uses FlashInfer TRTLLM kernels, CB uses paged SDPA /
|
|
FA4, Unsloth uses its cached fp16 LoRA path). The reward function reads
|
|
those completions, so the reward array mechanically differs. What matters
|
|
for equivalence is:
|
|
|
|
- **KL trajectory is near-identical** between vLLM and unsloth_fi_false
|
|
(both stay in `[0, 0.015]` across all 30 steps). The KL *term* of the
|
|
GRPO loss is the guardrail against policy drift, so matching KL means
|
|
the training dynamics are in the same regime.
|
|
- **Loss magnitudes are bounded** in `[-0.3, 1.0]` for all three backends.
|
|
- **No NaNs, no unbounded growth, no gibberish completions** in any run.
|
|
|
|
## grad_norm 919 on cb_paged
|
|
|
|
The enormous cb_paged grad_norm (vs vLLM's ~1.0) is a clipping story, not a
|
|
correctness story: the vLLM path goes through Unsloth's `FastLanguageModel`
|
|
which clips gradients to `max_grad_norm=1.0` internally, while the vanilla
|
|
HF path used by cb_paged picks up TRL's raw grad_norm reported by the
|
|
optimizer pre-clip (or without clipping if no `max_grad_norm` is set in
|
|
GRPOConfig). For a fair training-dynamics comparison the cb_paged config
|
|
should set `max_grad_norm=1.0` explicitly; left for a follow-up commit.
|
|
|
|
## KL missing for cb_paged
|
|
|
|
`StatisticsCallback.on_log` forwards the full TRL log dict into its per-step
|
|
entry only on steps where `loss` is present. TRL's vanilla-HF path separately
|
|
logs KL on a different log call that doesn't include loss, so the callback
|
|
silently drops it. Follow-up: relax the callback so every log dict with a
|
|
`step` field merges into the matching entry regardless of which keys are
|
|
present.
|
|
|
|
## Headline takeaways
|
|
|
|
1. **unsloth_fi_false is the pragmatic middle ground**: 12-17% of vLLM's
|
|
throughput, **15x less peak memory** (10.7 GB vs 159 GB), KL trajectory
|
|
matching vLLM within sampling noise.
|
|
2. **cb_paged is close to fi_false in throughput at this batch size** (41 s
|
|
vs 40 s median step at 30 steps) but costs 6x more memory. Phase 3
|
|
(main-thread sync driver + CUDA graphs on the rollout) is the right
|
|
lever for making CB competitive.
|
|
3. **torch.compile on the training step is not a quick win** for either
|
|
backend (Phase 4 report below).
|
|
|
|
## Phase 3 state (CB sync driver)
|
|
|
|
`scripts/benchmarks/cb_sync_driver.py`:
|
|
- Eager main-thread driver works end-to-end: smoke test on GPU 1 with 8
|
|
prompts / 64 tokens produced the expected 512 correct tokens.
|
|
- CUDA graph capture hangs on the first graphed step. Likely cause:
|
|
`ContinuousBatchProcessor._sample` reads `next_tokens.size(1)` as a
|
|
Python int to slice `batch_processor.output_ids[:, :tokens]`, which
|
|
forces a CPU-GPU sync and is not CUDA-graph-safe. Fix direction: keep
|
|
a fixed `tokens` count when `slice_inputs=False` (buffer size is
|
|
constant), or rewrite the copy as a full-buffer `copy_` without the
|
|
slice.
|
|
- Deferred to a follow-up commit.
|
|
|
|
## Phase 4 state (torch.compile on training forward)
|
|
|
|
- `unsloth_fi_false + compile_mode=default`: crashes with
|
|
`PeftModel_fast_forward() got multiple values for argument 'input_ids'`.
|
|
Unsloth's monkey-patched forward and Dynamo's argument rebinding don't
|
|
compose.
|
|
- `cb_paged + compile_mode=default`: Dynamo emits 700+ recompiles /
|
|
graph breaks on the first optimizer step and never makes progress.
|
|
Root cause: `modeling_utils.make_inputs_require_grads` calls
|
|
`Tensor.requires_grad_()` which triggers Dynamo GB0125 (unsupported
|
|
mutating op). TRL's GRPO `_compute_loss` then re-enters the tracer,
|
|
which re-triggers the break, which recompiles, and so on.
|
|
- `vllm` is excluded (vLLM owns its own compile pipeline).
|
|
|
|
Net: compile on the training step is not the right lever in this stack.
|
|
Phase 3 (CUDA graphs on the rollout decode) is.
|
|
|
|
## Raw stats
|
|
|
|
- `scripts/benchmarks/results/stats/grpo_{vllm,unsloth_fi_false,cb_paged}_{10,30}.json`
|
|
(StatisticsCallback per-step logs with full TRL metric dict)
|
|
- `scripts/benchmarks/results/stats/grpo_*_{10,30}.summary.json` (short form)
|
|
|
|
Pairwise diff:
|
|
|
|
python scripts/benchmarks/compare_grpo_runs.py \
|
|
--ref logs/grpo_vllm_30.json \
|
|
--candidate logs/grpo_unsloth_fi_false_30.json
|