Repository navigation
Conversation
forward_backward and forward_only now reorder each slice_dp batch on the driver (Ray mode) so every DP rank gets a similar workload, reorder every batch-aligned argument with it, and return the outputs in the caller's order. TransformersModel(dp_balance=False) turns it off. - remote_function(dp_balance=True) asks the model's _dp_balance_call for a plan; calls with object refs, raw logits, unencoded inputs or lists of another length are dispatched unchanged. - Loss.dp_balance_group_size keeps paired samples together (preference and reward losses: 2, REAL: num_generations); embedding and listwise reranker losses, and plain callables, opt out. - calculate_loss rejects per-sample lists after a reordered forward_only. - dp_balance/* metrics are reported by calculate_metric. - The GRPO cookbook no longer calls plan_dp_balance itself. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com>
- Skip the blocking _dp_balance_spec round trip once a model reports dp_balance off; the flag is fixed at construction. - Do not reorder when the loss for the call cannot be resolved. - Leave the hidden-size default to estimate_workload. - Drop the unread wrapper._dp_balance and the one-line micro-batch config helper. - Document dp_balance in the TransformersModel docs (en/zh). Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com>
Each rank holds only its own slice, so the ranks all-gather their sequence lengths, plan the joined batch as the Ray driver would, swap samples when the plan reorders, and send per-sample outputs back to their rank. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com>
1 of 4 tasks
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
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.
PR type
PR information
Follow-up to #355. #355 added
plan_dp_balance, but the training loop had to call it and reorder every per-sample field itself. This PR builds it intoTransformersModeland turns it on by default.TransformersModel(dp_balance=True). In Ray mode, beforeforward_backward/forward_onlysplit a batch across DP ranks, the driver reorders the batch, and every per-sample argument with it (old_logps,advantages, ...), to even out the per-rank workload. It then returns the outputs in the caller's order. Passdp_balance=Falseto keep the original split.Loss.dp_balance_group_size):REALLosskeeps each group ofnum_generationstogether;return_logits=True;calculate_lossraises an error if it gets per-sample lists after a reorderedforward_only.dp_balance/*metrics are reported bymodel.calculate_metric.plan_dp_balance. Docs are updated (en/zh).Scope:
TransformersModel, Ray and local mode.Experiment results
The planner is the same as in #355.
Tests
tests/model/test_dp_balance.py(reordering, restoring outputs, keeping pairs together, skip cases, a two-process local-mode run on gloo),tests/model/test_micro_batch.py,tests/infra: 114 passed.