Return full source-order context-parallel trainer outputs

[bradhilton]

Context-parallel forwards returned rank-local sequence shards, so caller expressions such as output.hidden_states[1:][mask] failed with a full-sequence mask. All ForwardOutput fields now contain the full flattened input sequence in source order on every TP/CP rank. No public positions field or caller-side mapping is needed.

Internally, each CP rank scatters its owned rows into source order and sums across the CP group. This handles uneven, reordered, and empty shards without padded all-gather buffers. Backward routes one copy of the replicated loss gradient to each owning row, matching the existing TP sequence-parallel convention. dp_reduce now combines only DP batches; registered custom heads average replicated CP gradients while model/LoRA shard gradients still sum. Memory admission accounts for local and gathered outputs coexisting.

The public docs explicitly require identical losses across TP/CP replicas. The native correctness harness now uses the public forward API. GPU CI includes the new distributed regression cases and the updated custom-parameter reduction oracle.

Validation:

• Distributed/head suite: 16 passed, 1 skipped. Covers CPU CP2/CP4/DP2×CP2 and two-H200 NCCL CP2; full outputs, masked readouts, nonlinear losses, decoder/head/probe gradients, ignored labels, empty shards, frozen decoders, no-grad, and DP metrics. The four-GPU case skips on this two-GPU machine.
• The new CP2 regression fails against the unmodified parent at the full-sequence shape assertion.
• Native pretrained Qwen3-0.6B (one layer), CP2, two LoRA slots: all 16 output combinations pass across no-sharing/full-sharing layouts; head-chunk and slot-backward parity pass.
• Ruff, formatting, lockfile checks, and changed library/test type checks pass. Required uv run prek run --all-files reports the same 23 existing type diagnostics as the unmodified parent (exact comparison).
• Broader regression suite: 429 passed. Native GPU integration: 4 passed, 1 skipped (DP/TP/CP custom-head reduction and the TP head/backward oracle); the combined TP2×CP2 case requires four GPUs. Total: 449 selected tests passed, 2 skipped.
• Native long-context probe canary passed on two H200s: Qwen3.8-27B architecture with four randomly initialized layers, CP2/TP1, two 20,436-token shared-prefix views, 5,120-wide hidden states. The original hidden_states[1:][mask] expression works; 13,624 masked tokens are counted once, custom-head gradients match the analytical full-sequence reference, backward completes, and the optimizer updates the registered probe. This is an architecture/path canary, not a full pretrained 27B run.

The review follow-up guards the optional Megatron dependency before spawning distributed test workers. The backend-only CI environment now passes the 12 head tests and skips its five runtime-dependent cases; the same file with Megatron and H200s still passes 16 cases with only the four-GPU case skipped. The unguarded version was reproduced failing with ModuleNotFoundError: megatron in the backend environment.

Fixes #911.

[bradhilton]

Review text removed on 2026-09-17; the consolidated review record is in the summary comment on this PR.

[bradhilton]

Review text removed on 2026-09-17; the consolidated review record is in the summary comment on this PR.

[bradhilton]

Review text removed on 2026-09-17; the consolidated review record is in the summary comment on this PR.

[bradhilton]

Consolidated review record (automated agent review traffic removed 2026-09-17)

Change. Context-parallel forwards returned rank-local shards, breaking full-sequence masks such as hidden_states[1:][mask]. All ForwardOutput fields now hold the full source-order sequence on every TP/CP rank via per-rank scatter plus CP-group sum; dp_reduce combines only DP batches and custom-head gradients average across CP. Fixes #911.

Review. Minsky and McCarthy independently cleared source correctness at 1e068ee (an earlier unnamed review at 855736f: LGTM pending GPU CI), both noting the CP semantics change is intentional, not behavior-neutral. Verified: CPU Gloo CP2/CP4/DP2xCP2 cases (64 passed, 2 skipped with Megatron; 12 passed, 5 skipped without), empty-shard and DP2 controls, negative control failing on the base runtime; two-H200 CI was pending at review time.

Findings addressed during review.

• CP tests failed with ModuleNotFoundError: megatron in the backend-only CI stage; pytest.importorskip("megatron.core") guards added.
• Docstring now requires identical losses on every TP/CP replica; comment distinguishes dp_reduce from internal CP-inclusive reductions.

Deferred / follow-ups.

• Four-GPU TP2xCP2 case unexercised; caladan 058 CP2/CP4 probe lanes to rerun.