fix(engine): make ragged transport padding objective-safe - #1566
Conversation
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
|
This makes transport-only padding rows safe for arbitrary objectives: they run through forward/backward for cross-rank and pipeline step parity, but are excluded from loss weight, metrics, and outputs. None of this exists on main yet (no |
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.