fix(engine): make ragged transport padding objective-safe - #1566
Open
EazyReal wants to merge 1 commit into
Open
Conversation
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
This was referenced Jul 29, 2026
EazyReal
marked this pull request as ready for review
July 31, 2026 02:33
EazyReal
requested review from
PrometheusComing,
fishcrap,
garrett4wade,
geshi001,
nuzant,
rchardx and
sitabulaixizawaluduo
as code owners
July 31, 2026 02:33
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
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.
Problem
After semantic compaction, token-aware splitting can give one data-parallel rank three real microbatches and another only one. FSDP collectives and Megatron/Archon pipeline schedules still require every rank and stage to execute the same number of steps; a shorter rank cannot simply stop without hanging its peers.
Why transport-only padding
Collection should first compact and backfill real work (#1563), but that cannot guarantee equal physical microbatch counts. Truncating to the smallest rank drops valid samples, and changing the loss denominator does not satisfy distributed collectives. Duplicating a real sample changes the objective, zero-length model inputs are not generally valid, and
loss_mask=0is insufficient for DPO, reward-model, or custom objectives whose weights come from sequence structure.The smallest shared backend mechanism is therefore a model-valid transport row with an explicit marker: it participates in forward/backward and communication, but the engine bypasses objective callbacks and assigns it zero loss weight.
Solution
Scope boundary
#1563 owns semantic membership, compaction, group boundaries, and normalization. This PR acts only after those decisions, at the engine scheduling boundary; transport rows never become rollout members.
Review stack
GitHub cannot base an upstream PR on a contributor-fork branch, so #1563 remains temporarily cumulative against
mainand shrinks to its core diff after this foundation lands.Part of #1559.
Verification
pytest tests/test_utils.py tests/test_tree_transport.py tests/test_eval_dispatch.py -q— 86 passedpre-commit run --all-files08ab30fc..b0dbd4c4— no actionable findingsA real multi-GPU Megatron/Archon pipeline canary was not available on this host; this is static, CPU/distributed-test, and hosted-install readiness rather than hardware certification.