[subhajitdchow]

Motivation

OnlineDSparkModel all-reduces its supervised-token denominator on every
micro-step and normalizes the loss inside the model. With gradient
accumulation this adds one blocking collective per micro-step (16 per optimizer
step at accumulation_steps: 16), each a cross-rank sync point.

DFlash, Domino, MTP and DSpine already return loss_terms and let
TrainerCore normalize gradients once at the optimizer boundary. DSpark is the
only DFlash-family strategy that does not.

Modifications

• OnlineDSparkModel no longer all-reduces per micro-step; it exports
  loss_terms = (α_ce·CE + α_l1·L1 + α_conf·conf, local token count), the same
  contract as DFlash. The returned loss is the rank-local mean.
• Remove the per-micro-step _assert_async(global_loss_den > 0): a rank may
  have no supervised tokens in one micro-step, and the trainer already rejects
  a non-positive window denominator.
• DSparkTrainStrategy passes loss_terms to StepOutput.

Gradients are now normalized by the token count of the whole optimizer window
instead of averaging per-micro-step means. This matches DFlash/Domino/EAGLE3
and is identical when micro-steps carry equal token counts.

Related Issues

N/A

Accuracy Test

• New tests in tests/test_utils/test_dflash_losses.py (CPU) and
  tests/test_runtime/test_dflash_host_syncs.py (CUDA) check that the DSpark
  forward issues no all_reduce, that a zero-token micro-step stays finite,
  that TrainerCore gradients equal Σnum/Σden over a window with unequal token
  counts, and that a DSpark micro-step is host-sync-free. The three CPU tests
  fail on main and pass with this change.
• The affected modules (118 tests) pass on 8× MI355X; the full suite shows no
  new failures relative to main.
• Qwen3-4B training (setup below): final loss at step 100 is 2.6598 vs 2.6551
  on main (median per-step difference 0.5%), accuracy identical. Runs are
  bit-deterministic with a warm compile cache, so the difference is the
  intended window weighting, not noise.

Benchmark & Profiling

Qwen3-4B DSpark, upstream offline/colocated/qwen3-4b-dspark.yaml with
nproc_per_node: 8 and accumulation_steps: 16, 1 node of 8× AMD Instinct
MI355X (ROCm 7.2, PyTorch 2.14), identical features, seed and data order. Optimizer
step time over steps 20–100:

	median	mean	p90
main	1.131 s	1.164 s	1.153 s
This PR	1.021 s	1.055 s	1.076 s

The median optimizer step time drops from 1.131 s to 1.021 s (−9.7%), and
training throughput improves by 10.8% (1.11×). A second pair of runs gave
the same result (−9.8% step time, +10.9% throughput).

Measured on a single node. Offline features were generated with a HuggingFace
forward pass (same layers and final hidden state) instead of the SGLang capture
plugin.

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel at https://sgl-fru7574.slack.com/archives/C09784E3EN6 to discuss your PR.

[subhajitdchow]

@jiapingW can you please take a look at this PR?