fix(training): fail fast when dspark_loss runs multi-rank without a dp_group

[khazic]

Stacked on #5233 and #5234: this fixes the DSpark core introduced there, so until those merge this PR also lists their commits. Only the last commit is new here; I will rebase once the base PRs land.

What

Commit a46c068 scoped the denominator all-reduce to the explicitly passed dp_group (correct: the default process group is not the DP replica set under TP/PP). But it left dp_group=None silently computing a locally normalized loss inside a multi-rank job. With uneven token counts across ranks, the DP-averaged gradient then no longer equals the gradient of the global position-weighted loss: ranks with few valid tokens are over-weighted, and training optimizes a biased objective with no visible failure.

Fix

dspark_loss now raises RuntimeError when torch.distributed is initialized with world size > 1 and dp_group is None. Single-process and world-size-1 usage is unchanged. Callers without data parallelism (e.g. TP-only jobs, or a reference path in tests) pass a size-1 group; this is a deliberate explicitness requirement, since the loss cannot know the DP degree on its own and guessing from the default group is exactly the bug a46c068 removed.

The 2-GPU functional test is updated accordingly: its single-process reference gradient now uses a per-rank size-1 group (identical numerics to the previous dp_group=None path), which the review caught as otherwise breaking under the new guard.

Tests

Two new unit tests via monkeypatched torch.distributed: multi-rank with dp_group=None raises; world-size-1 initialized still computes the local loss. The raise test fails on the parent commit and passes here. All 18 dspark unit tests pass; ruff is clean. The functional test needs a 2-GPU node and is unchanged in intent.

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

[khazic]

/ok to test c266402

[khazic]

Closing this as obsolete rather than rebasing it.

@yaoyu-33's note on #5234 was that this PR's missing-group policy should follow the approved trainer and process-group contract rather than precede it. Settling that contract removed the thing this guard protects: dspark_loss no longer takes a dp_group or performs any reduction of its own. It returns the Megatron-Core 3-tuple (unnormalized numerator, num_tokens, reducible reporting pairs) and lets the framework reduce once over the data-parallel group, so there is no group argument left to omit.

The failure mode this PR was guarding against is real and is still guarded, one level up. Returning the 3-tuple is not by itself sufficient: with calculate_per_token_loss=False, schedules.py divides each micro-batch by its own num_tokens, which reintroduces exactly the mean-of-micro-batch-means bias. #5234 adds assert_per_token_loss_config to fail fast on that, measured at a 37% gradient error against the full-batch reference on 2 ranks with uneven micro-batch token counts.

This PR was also stacked on the now-closed #5241, so its diff carries the incorrect 0.5 * L1 commit as a base. #5242 and #5244 remain valid and will be rebased onto the reworked #5234.