Context-parallel forward must return full-sequence outputs in caller order (ForwardOutput fields are currently rank-local shards)

[bradhilton]

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.

1. For each item, all-gather (source_positions, values) over ps.get_context_parallel_group(). Shards are unequal, so pad to the rank plan's pad_multiple (or the max shard length) and drop the padding after the gather; ART already tracks the padded layout in the rank plan.
2. Scatter the gathered rows into a full-length tensor by source_positions (index_copy_ / scatter_). Apply to hidden_states, target_logprobs (with -100 label positions handled as today), top_k, and logits. logits are 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.
3. Use an autograd-aware gather (the same pattern as gather_from_sequence_parallel_region, or torch.distributed.nn.functional.all_gather) so hidden_states.requires_grad losses 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 sums cp_size identical contributions; the gather's backward should scale by 1 / 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_states losses have the same replication.
4. Once outputs are replicated across the CP group, dp_reduce must not sum over CP ranks for quantities derived from full-sequence outputs. Today it reduces over get_data_parallel_group(with_context_parallel=True), which is what makes shard-local sum losses correct; after this change the natural contract is ForwardOutput = full sequence and dp_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.

[bradhilton]

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.