Redirect the KTO liger loss through the distributed wrapper under DDP - #7247
Conversation
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e78ddd538b
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
`compute_loss` entered `_forward_redirection` only for ZeRO-3 and FSDP, so under plain DDP the loss ran on the unwrapped model. `DistributedDataParallel.forward()` never ran, its reducer was never armed, and gradients were never all-reduced: every rank kept its own. Measured on 2 GPUs with different data per rank, comparing every parameter's gradient across ranks: 310/310 parameters differed before, 0/310 after. DPO has the same bug but cannot take the same fix: its fused loss runs `torch.func.grad_and_value` internally, which raises `NotImplementedError: Cannot access storage of TensorWrapper` when run inside `DistributedDataParallel.forward()`. It is fixable once #7243 replaces that path. Ref #7245
2d95e70 to
74c134a
Compare
A ZeRO-3 model is a `DeepSpeedEngine`, which accelerate's `extract_model_from_parallel` unwraps, so `model is not unwrapped_model` already covers it. Only FSDP2 needs its own check, since it shards in place and preserves object identity. Verified on 2 GPUs with zero3, zero2 and ddp.
74c134a to
43e196e
Compare
qgallouedec: same block, same text as huggingface#7247 - drop the is_zero3 check (accelerate already unwraps DeepSpeedEngine) and use the FSDP2 identity check plus the DDP wrapper-identity check.
What does this PR do?
Partially addresses #7245, for KTO only.
compute_lossentered_forward_redirectiononly for ZeRO-3 and FSDP, and otherwise ran the loss on the unwrapped model. Under plain DDP that meansDistributedDataParallel.forward()never runs, so its reducer is never armed byprepare_for_backward(), the autograd hooks return early, and gradients are never all-reduced. Each rank keeps its own, silently.Measured on 2 GPUs, one step, different data per rank, comparing every parameter's gradient across ranks:
use_liger_kernel=True, beforeuse_liger_kernel=True, afterThe condition now matches the one
compute_ref_log_probsalready uses in this same file.Why KTO only
DPO has the identical bug and the identical guard, but cannot take this fix yet. Its fused loss calls
torch.func.grad_and_valueinternally, and running that insideDistributedDataParallel.forward()raises:on every rank. KTO is unaffected because it already computes log-probabilities with
_ChunkedLogProbFunction, a plainautograd.Function.Applying this same guard on top of #7243, which moves DPO onto the chunked path, gives 0/310 differing. So the DPO half belongs there.
Not covered
test_compute_ref_log_probs_redirects_wrapped_liger_model) only pins theifcondition and would pass againstif True:. What catches this is a 2-GPU gradient comparison, which does not fittests/distributed, whererun_commandonly asserts the process exits 0.is_zero3is left in place.DeepSpeedEngineis in accelerate's unwrap list, somodel is not unwrapped_modellikely covers ZeRO-3 and the term is redundant, but I could not verify that here and did not want to touch the path Fix DPO/KTOuse_liger_kernelunder DeepSpeed ZeRO-3 #6372 fixed.Before submitting
Note
Medium Risk
Changes distributed training loss path for KTO+Liger; fixes incorrect gradients under DDP but touches a critical training code path.
Overview
Fixes silent broken multi-GPU training for KTO when
use_liger_kernel=Trueunder plain DDP: the Liger chunked path ran loss on the unwrapped model, soDistributedDataParallel.forward()never ran, DDP’s reducer was never armed, and gradients were not all-reduced across ranks.compute_lossnow routes the Liger loss through_forward_redirectionwhenever the model is still wrapped (model is not unwrapped_model) or FSDP is on—the same guardcompute_ref_log_probsalready uses—instead of only checking DeepSpeed ZeRO-3 explicitly. Comments are updated to note DDP and FSDP2 behavior alongside ZeRO-3.Reviewed by Cursor Bugbot for commit 43e196e. Bugbot is set up for automated code reviews on this repo. Configure here.