Skip to content

fix(engine): make ragged transport padding objective-safe - #1566

Open
EazyReal wants to merge 1 commit into
areal-project:mainfrom
EazyReal:vmax/ragged-transport-safety
Open

fix(engine): make ragged transport padding objective-safe#1566
EazyReal wants to merge 1 commit into
areal-project:mainfrom
EazyReal:vmax/ragged-transport-safety

Conversation

@EazyReal

@EazyReal EazyReal commented Jul 29, 2026

Copy link
Copy Markdown
Contributor

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=0 is 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

  • add transport-only microbatches only when synchronized execution requires them
  • bypass arbitrary objective weight and loss callbacks for transport rows
  • keep transport rows out of every loss denominator, metric, normalization population, and returned output
  • strip the marker before model execution and support ordinary and tree-packed paths across FSDP, Megatron, and Archon

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 main and 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 passed
  • pre-commit run --all-files
  • upstream hosted pre-commit and Ubuntu, macOS, SGLang-extra, and vLLM-extra installation checks passed; Docker was skipped by workflow policy
  • independent exact-range transport, Megatron, and Archon review of 08ab30fc..b0dbd4c4 — no actionable findings

A 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.

@gemini-code-assist

Copy link
Copy Markdown
Contributor

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

@gemini-code-assist

Copy link
Copy Markdown
Contributor

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant