7852b3aed4
Phase 4 (torch.compile on training step) -- negative result documented:
- unsloth_fi_false + compile_mode=default: crashes immediately 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+ graph breaks /
recompiles during the first optimizer step and never makes progress
(killed after 10 minutes at step 0/10). Root cause trail:
`Tensor.requires_grad_()` inside `modeling_utils.make_inputs_require_grads`
triggers GB0125 (no Dynamo support), which propagates up through TRL's
`_compute_loss`. Fixing this would require restructuring GRPO's LoRA
gradient enablement path -- out of scope for this PR.
- vllm is excluded from Phase 4 because vLLM owns its own compile pipeline.
Net: torch.compile on the training step is not a quick win for the non-vLLM
paths in this stack. Phase 3 (CUDA graph on the rollout decode step) remains
the right lever for closing the CB <-> vLLM gap, and Phase 1 already
demonstrates that Unsloth's fast_inference=False path narrows the gap to
~14% of vLLM at 1/7th the peak memory without any compile.
Helper script `scripts/benchmarks/compare_grpo_runs.py`:
- Wraps `torch_debugging_utils.compare_training_runs` (loss/grad-norm diff)
- Adds reward/KL/time pairwise diffs with `max_abs` and `mean_abs`
- Reads the StatisticsCallback-emitted JSON written by
`qwen3_grpo_unified.py`
Sample output on Phase 2 vibe (10 steps):
vllm vs unsloth_fi_false:
max_loss_diff = 0.40, max_kl_diff = 0.009 (both tiny), reward_diff
mean 1.85 (different rollouts expected across backends at temp=0.1)
vllm vs cb_paged:
max_loss_diff = 0.29, max_grad_norm_diff = 715 (cb_paged has no
gradient clipping on the vanilla-HF path; vLLM path is clipped to
1.0 by Unsloth internally -- apples-to-oranges without matching
clipping, tracked for the 30-step run).
|
||
|---|---|---|
| .. | ||
| benchmarks | ||
| enforce_kwargs_spacing.py | ||
| install_gemma4_mlx.sh | ||
| install_qwen3_6_mlx.sh | ||
| run_ruff_format.py | ||