Repository navigation
Gather per-example tensors before reducing logged means - #7045
Closed
behroozazarkhalili wants to merge 11 commits into
Closed
behroozazarkhalili wants to merge 11 commits into
behroozazarkhalili wants to merge 11 commits into
Conversation
Several trainers reduced a metric to a rank-local scalar and then passed it to gather_for_metrics. That defeats the de-duplication the gather exists for: accelerate trims the padded final batch by slicing the batch dimension, and a scalar has none left, so on a dataset that does not divide evenly across ranks the duplicated example stays in the mean. Every affected site now hands the per-example tensor to the gather and reduces afterwards: RewardTrainer accuracy, margin and reward range; XPO log-probs, rewards, KL, entropy, margins and accuracy; Online DPO score margin, scores, KL, non-score reward, RLHF reward and entropy (per prompt, over its two completions); PPO objective KL, entropy, non-score reward, RLHF reward and scores; CPO and ORPO logits means; TPO entropy (per-token entropy and mask gathered together). Each trainer gains an oracle test built on a shared stand-in gather that drops the final example: the fixed code reports the mean over the remaining examples, the previous code reports the mean over all of them. Fixes #7010
|
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. |
…st round missed DPO and KTO reduced entropy and token counts to scalars before gather_for_metrics, TPO gathered logits and accuracy as scalar sums and its rewards, margins and log-probs with plain gather, Online DPO never gathered val/contain_eos_token, and CPO and ORPO logged nll_loss from the scalar mean of nn.CrossEntropyLoss. None of those can drop the example accelerate duplicates on the last batch. Every site now hands one row per example to gather_for_metrics and reduces afterwards; the CPO and ORPO forward returns per-example NLL sums and token counts and the training loss is still their total-sum over total-count. gather_for_metrics de-duplicates by slicing the leading dimension, and the DPO and TPO forward batches stack chosen and rejected rows (and a reference row for TPO), as Online DPO stacks two completions per prompt, so those tensors are folded to one row per example before the gather. Each site has an oracle test that fails with its change reverted, and PPO gains the regression test the earlier round lacked. The training-token counters keep their existing pattern.
…o fix/7010-gather-before-mean # Conflicts: # tests/experimental/test_online_dpo_trainer.py
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
TPO copies DPO's metric block and received the per-example fix in the earlier round while DPO and KTO kept the old shape: logits and token accuracy were summed to a rank-local scalar before gather_for_metrics, and DPO's rewards, accuracies, margins and log-probs went through plain gather, which never removes the duplicated last-batch padding at all. DPO now hands per-example logit sums, token counts, correct-token counts, rewards, accuracies, margins and log-probs to gather_for_metrics in both the standard and the Liger block. KTO does the same for the KL term, the token count, the chosen and rejected counts (one boolean per row), the logits and the rewards and log-probs, which are scattered into a zero-filled row-aligned tensor so the opposite label contributes zero to the sum and to the count. The Liger logits in both trainers stay scalar: liger_kernel returns chosen_logits_mean and rejected_logits_mean as 0-d sums, and the KTO fused loss returns its logits, log-probs and rewards the same way, so no per-example quantity exists to gather. The PR body lists this as a limit. The DPO and KTO oracle tests now assert every changed metric against the mean over the first N-1 examples and no longer stub accelerator.gather, so a leftover gather call fails loudly. On the previous commit the DPO oracle fails on that stub and both KTO cases keep the duplicated KL.
The previous commit moved KTO's num_tokens counter to a per-row gather. The PR keeps the attention_mask.sum() counters unchanged in every trainer as a separate convention, and DPO still uses the scalar form, so KTO follows it too. The oracle no longer asserts the token count.
A batch that holds only one label class gets the empty float32 default for the other class's rewards, and index_copy refuses a float32 source into the bfloat16 row tensor under mixed precision. Cast both reward tensors to the target dtype first. The new test runs a bfloat16 batch with a single label class each way; on the previous commit it raises the index_copy dtype error.
PPOTrainer left main in #7020, so the PPO objective-metric changes and their oracle test have no target. The testing_utils conflict keeps drop_last_for_metrics and takes main's removal of the bitsandbytes helpers.
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit 4ee8527. Configure here.
Datasets can carry 0/1 integer labels. On an integer tensor `~` is a bitwise inversion, so the rejected count and the rejected logits mask became -1 and -2 rows. The loss path already casts the field to bool; the metric path now does the same. The oracle runs over bool and int64 labels; on the previous commit the int64 rows fail with an IndexError.
Collaborator
Author
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Several trainers reduced metrics before gathering them, preventing removal of duplicated examples in the final distributed batch. I changed the affected paths to gather per-example values before reduction and corrected related padding and label handling. Tests that simulate duplicate removal pass and fail when the changes are reverted, but I did not run multiple processes.
Fixes #7010.
Note
Medium Risk
Touches metric logging across many RLHF trainers; training losses are mostly unchanged but logged dashboards and early-stopping on metrics could shift on uneven distributed datasets.
Overview
Fixes biased training/eval metrics when dataset size does not divide evenly across distributed ranks.
gather_for_metricsonly strips Accelerate’s duplicated last-batch examples along the batch dimension, so metrics that were averaged (or summed) before the gather kept the extra padding example baked in.Across DPO, KTO, Reward, CPO, ORPO, TPO, XPO, Online DPO (and related paths), logging now passes one value per example (or per prompt where batches stack chosen/rejected/reference or two completions), gathers, then reduces. Entropy and average logits use per-row token sums and counts; stacked batches are folded so one row maps to one example before gather. CPO/ORPO expose per-example NLL sum and token count from forward; the training loss is unchanged but
nll_lossis computed after gather. Online DPO aggregates KL, entropy, scores, and EOS flags per prompt and masks padding in KL/entropy. KTO casts labels to bool, gathers KL per row, and scatters chosen/rejected rewards/logps into full-batch tensors so single-class batches and bfloat16 dtypes work with gather. Reward and XPO log accuracy, margins, and rewards from per-pair tensors instead of pre-reduced scalars.Tests add
drop_last_for_metricsintesting_utils(mimics Accelerate de-duplication) and oracle tests that assert logged values match the mean over non-outlier examples when the last batch row is an outlier.Reviewed by Cursor Bugbot for commit 000fc99. Bugbot is set up for automated code reviews on this repo. Configure here.