Fix context-parallel normalization of masked training losses

[zupengwang]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

When a three-value loss callback returns CP-local loss sums and valid-token counts, the microbatch-averaged training path divides by each shard's local count. DDP then averages those gradients over DP × CP. Unequal masks therefore change the training objective: with one valid token, CP=2 halves the gradient and CP=4 reduces it to one quarter of the unsharded reference.

Reduce a cloned token count over the explicit CP group, normalize with the complete sequence's count, and compensate DDP's CP averaging. Pass that group through the non-pipelined, pipelined, and interleaved forward paths. Preserve the returned local count and metric aliases, the per-token/global-batch normalization path, and legacy two-value callbacks. Existing callers without an explicit group retain a documented MPU compatibility fallback.

This adds one scalar all-reduce per microbatch on the final pipeline stage when CP>1 and calculate_per_token_loss=False for a three-value loss callback. No performance improvement is claimed.

Issue tracking

Addresses finding 9 in #7452. The other findings are outside this change.

The MTP normalization work in #7498, MoE count fix in #7481, and evaluation metric reduction in #7207 are separate paths. Open DCP schedule PR #7047 does not change this normalization block.

Validation

Base: 3b1a08fc7e83c344653a4a54ae427c4f06a0d73d.

• Unmodified base, four NCCL ranks: 18 failed, 82 passed per rank on the same 100 training-step regressions. All failures are gradient comparisons in the masked, microbatch-averaged CP>1 cases.
• Patched tree: 104 passed per rank on four RTX 3090s, with zero skips.
• Expanded CI-style torch.distributed.run -> coverage run -> pytest, with PYTHONPATH unset and the dev marker filter: 114 passed per rank, including ten existing auxiliary-scaling/schedule cases. This is a local execution of the CI's Python invocation pattern, not an upstream CI run.
• Coverage: CP=1/2/4, PP=1/2, DP replicas, one/three accumulated microbatches, full/uneven/single-target/empty/varying masks, both normalization modes. Tests run the real schedule, Megatron DDP and finalize_model_grads, then compare every parameter gradient and an SGD update with an independent dense reference. Four additional cases verify CP subgroup isolation, local count aliases and the compatibility fallback.
• Black, isort, Ruff, Python syntax, whitespace and kernel determinism coverage checks passed.
• The autoformatter wrapper passes Black/isort, then stops because Pylint is unavailable. Mypy reports the same 22 diagnostics on the baseline and patch, with no added diagnostics; the new test file has none.

Environment: Python 3.12.13, PyTorch 2.13.0+cu130, pytest 9.1.1, four RTX 3090 GPUs. Transformer Engine and Apex are absent; the regression uses a small native token-wise two-projection model in FP32. The canonical container was unavailable due to Docker socket permissions. Full Transformer training/convergence, BF16, full-iteration CUDA graphs, MoE combined 1F1B and eight-GPU upstream CI were not run. Existing interleaved-schedule coverage is a mocked forward-only regression, not CP×VPP numerical equivalence.

Contribution process

• Relevant unit tests added
• Separate functional recipe added (the distributed regression is discovered by the existing pipeline-parallel unit-test bucket)
• Kernel determinism coverage checker passed; no GPU kernel implementation changed
• Added typing for the new optional group argument and updated loss-helper documentation
• Full autoformatter completed (Pylint unavailable, as described above)

AI-assisted implementation and validation. The human author has reviewed every line and authorized submission.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[Connor-XY]

@deepakn94, could you review the CP loss normalization and process-group handling here? The key check is that uneven masks produce the intended gradients across CP and pipeline schedules while preserving the existing per-token path. This addresses finding 9 of #7452.