Repository navigation
Context-parallel forward must return full-sequence outputs in caller order (ForwardOutput fields are currently rank-local shards) #911
Description
Summary
With context parallelism (the default TrainerRank topology for dense handlers when more than one GPU is visible: TP 1, CP = visible GPUs), every per-position field of ForwardOutput (hidden_states, target_logprobs, top_k, logits) is the rank's local shard of the sequence, in rank-plan order, not the caller's sequence. The API contract everywhere else (single GPU, DP, TP with sequence parallel) is that a ForwardOutput corresponds to the caller's ForwardInput in the caller's order. CP silently breaks that contract; callers that index outputs by position fail (or, worse, would compute on wrong positions if shapes happened to match).
Requested fix: make the CP forward return full-sequence outputs in source order, internally. Callers should not have to map positions or know the topology.
Observed
Caladan probe doing output.hidden_states[1:][mask] with a full-sequence mask on Qwen/Qwen3.8-27B, 20,436-token input:
- CP 2:
IndexError: The shape of the mask [20435] at index 0 does not match the shape of the indexed tensor [8659, 5120] - CP 4: same,
[6655, 5120]
Shards are uneven (prefix-tree rank plan), so nothing can be reconstructed caller-side without the plan.
Where it comes from
_prepare_context_parallel_forward dispatches arange(seq_len) through the rank plan and records positions_by_item (local rows) and source_positions_by_item (the source positions this rank owns). _project_head returns hidden_states = _select_positions(hidden_by_row, positions) and sizes target_logprobs[index] from labels.index_select(0, source_positions), i.e. local rows only. source_positions_by_item is used internally for labels and then discarded.
Proposed internal fix
Do for CP what _gather_sequence_parallel_hidden already does for TP sequence parallelism: gather each item's rows across the CP group and place them in source order before building ForwardOutput.
- For each item, all-gather
(source_positions, values)overps.get_context_parallel_group(). Shards are unequal, so pad to the rank plan'spad_multiple(or the max shard length) and drop the padding after the gather; ART already tracks the padded layout in the rank plan. - Scatter the gathered rows into a full-length tensor by
source_positions(index_copy_/scatter_). Apply tohidden_states,target_logprobs(with-100label positions handled as today),top_k, andlogits.logitsare the one large case (seq × padded vocab); gathering them is still correct, but it may be worth keeping them shard-local behind an explicit opt-out since nothing in caladan requests full logits under CP. - Use an autograd-aware gather (the same pattern as
gather_from_sequence_parallel_region, ortorch.distributed.nn.functional.all_gather) sohidden_states.requires_gradlosses computed on the gathered tensor route gradients back to the owning rank's rows. Because every CP rank then computes the same loss on the same full tensor, the reduce-scatter in the backward sumscp_sizeidentical contributions; the gather's backward should scale by1 / cp_size(or the loss reduction should), matching whatever convention the existing TP sequence-parallel hidden gather uses for the probe head. Worth confirming that convention at the same time, since TP +hidden_stateslosses have the same replication. - Once outputs are replicated across the CP group,
dp_reducemust not sum over CP ranks for quantities derived from full-sequence outputs. Today it reduces overget_data_parallel_group(with_context_parallel=True), which is what makes shard-local sum losses correct; after this change the natural contract isForwardOutput= full sequence anddp_reduce= pure DP (with_context_parallel=False). Any internal shard-sum path should reduce over the CP group itself before returning.
Net effect: CP becomes invisible to callers, as TP already is. This is the version we want; exposing positions and asking callers to map them is not acceptable as the fix.
Environment
ART bcfe13519 (caladan pins a3a248a8e), Megatron backend, dense qwen3_5 handler (Qwen/Qwen3.8-27B, Qwen/Qwen3.6-27B class), single node H200:2 / H200:4 via caladan Trainer. Interim workaround: ART_MEGATRON_CONTEXT_PARALLEL_SIZE=1 with TP = N (ART_MEGATRON_TENSOR_MODEL_PARALLEL_SIZE) or DP; measured on the 27B probe at 16 views/step: single GPU 118 s, TP 2 68 s, TP 4 40.5 s, DP 2 71 s.
Activity
[-]ForwardOutput per-position fields are context-parallel-local shards with no position mapping exposed[/-][+]Context-parallel forward must return full-sequence outputs in caller order (ForwardOutput fields are currently rank-local shards)[/+]bradhilton commentedon Sep 16, 2026
Updated the issue: the fix should be internal (gather to source order inside the CP forward, autograd-aware, with the dp_reduce contract adjusted), not a positions field that callers map themselves. See the revised proposal.
Return full source-order context-parallel trainer outputs (OpenPipe#912)