fix(training): adopt the Megatron-Core loss contract for DSpark

Return the unnormalized micro-batch numerator plus num_tokens and reducible
[numerator, denominator] reporting pairs, instead of dividing by a locally
reduced denominator. Megatron-Core calls the loss once per micro-batch, so a
locally normalized value accumulates a mean of micro-batch means and biases the
objective whenever micro-batches carry different supervised-token counts. The
unbiased global mean needs calculate_per_token_loss=True, which
assert_per_token_loss_config now enforces at the integration point.

Chunk the FP32 probability L1 so the softmax temporaries are bounded by
chunk_tokens * vocab rather than the full [batch, num_blocks, block_size, vocab]
tensor, recomputing each chunk in the backward. Three such tensors at reference
scale are several GiB of forward temporaries before autograd storage.

Correct the objective's documentation: the training term is the raw weighted L1
of paper Eq. 10, and the factor 1/2 belongs only to the Eq. 8 acceptance label.
The implementation already did this; the docstring and the design page did not. A
disjoint-distribution regression now asserts l1_loss == 2 with accept_rate == 0
so the two cannot drift apart again.

Rewrite the 2-GPU functional test around the new contract: two ranks, two
micro-batches each, uneven supervised-token counts, asserting both that the
global-token gradient matches the single-process reference and that the
per-micro-batch normalization does not. Add a dedicated L0 launch script so CI
actually selects the test, using a single pytest coordinator because the test
self-spawns its ranks.

Signed-off-by: khazic <khazzz1c@gmail.com>

Author: khazic

src/megatron/bridge/training/post_training/dspark/__init__.py

@@ -16,7 +16,7 @@
 
 This package holds the pure-``torch`` pieces of DSpark draft training: the
 sequential (Markov / RNN) and confidence heads (:mod:`heads`) and the three-term
-CE / TV / confidence objective with its acceptance metrics (:mod:`loss`). The
+CE / L1 / confidence objective with its acceptance metrics (:mod:`loss`). The
 parallel-aware draft backbone and the Megatron-Core forward-step / provider wiring
 are separate, GPU-validated concerns. See ``docs/training/dspark-speculative-decoding.md``.
 """
@@ -28,7 +28,11 @@
     VanillaMarkov,
     build_markov_head,
 )
-from megatron.bridge.training.post_training.dspark.loss import DSparkForwardOutput, dspark_loss
+from megatron.bridge.training.post_training.dspark.loss import (
+    DSparkForwardOutput,
+    assert_per_token_loss_config,
+    dspark_loss,
+)
 
 
 __all__ = [
@@ -37,6 +41,7 @@
     "GatedMarkovHead",
     "RNNHead",
     "VanillaMarkov",
+    "assert_per_token_loss_config",
     "build_markov_head",
     "dspark_loss",
 ]