feat(training): add DSpark draft heads and loss (algorithm core)

[khazic]

What

First implementation increment of DSpark speculative-decoding draft training
(feature #5230, design PR #5233): the architecture-agnostic, pure-torch
core. Opened as a draft because it is not yet wired to a trainer (see Scope).

Adds src/megatron/bridge/training/post_training/dspark/:

• heads.py: the DSpark sequential head that injects intra-block token
  dependency onto the parallel backbone's logits, as
  P(x | x0) = prod_k p_k(x_k | x0, x_1..x_{k-1}). Three variants:
  VanillaMarkov (low-rank token-to-token additive bias), GatedMarkovHead
  (gates the previous-token embedding by the backbone hidden), and RNNHead
  (carries recurrent state across the block). Plus the ConfidenceHead.
• loss.py: the three-term objective, position-weighted exp(-k/gamma):
  cross-entropy (0.1) + total-variation 0.5*||p_draft - p_target||_1 (0.9)
  • confidence BCE (1.0); and the analytical acceptance metrics
    accept_rate = 1 - TV and tau = sum cumprod(accept_rate) + 1. The
    data-parallel denominator reduction is an injectable dp_group, so the same
    math wires to a Megatron-Core forward step (and its DP group) in a follow-up.

Follows the DSpark paper (arXiv:2607.05147)
and the NeMo AutoModel reference implementation
(nemo_automodel/components/speculative/dspark/).

Scope (why draft)

This PR is deliberately limited to the pieces that are pure tensor math and
CPU-testable, so they can be reviewed on their own. The parts that need
Megatron-Core and multi-GPU validation are intentionally not in this PR:

• the parallel-aware draft backbone (the block-attention module that reads the
  frozen target's hidden states as context K/V), as an MCore-native or
  ModelOpt-injected module;
• the forward_step_dspark + provider hook + dspark.* HF weight mappings, and a
  recipe, following the ModelOpt distillation subsystem.

Tests

tests/unit_tests/training/post_training/test_dspark_heads.py and
test_dspark_loss.py (CPU): the Markov/RNN/confidence head math and shapes, and
the loss/acceptance math (accept_rate = 1 - TV, tau = block_size + 1 when the
draft matches the target, three-term combine, backward, masking, error paths).

Validation status

ruff clean. The core is pure torch, and its math was verified end-to-end
locally (all head and loss/acceptance assertions pass). It has no Megatron-Core
dependency, so the unit tests run in the CPU lane. The backbone and integration
follow-ups still require multi-GPU validation.

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

Added a multi-GPU functional test for the one piece of this PR that cannot be validated on a single process: the loss's dp_group denominator reduction.

tests/functional_tests/test_groups/training/test_dspark_loss_dp.py spawns 2 NCCL ranks, splits a batch across them, and asserts the DP-averaged gradient through a shared head equals the single-process full-batch gradient. The same commit carries a small correctness fix: dspark_loss(dp_group=None) is now local-only instead of reducing over the default process group, whose members are not the DP replicas under Megatron-Core TP/PP/CP, so the caller passes the data-parallel group explicitly.

This test needs a 2-GPU node and is not part of the CPU unit lane. I do not have permission to trigger GPU CI on this repo, so it needs a maintainer to run it (or /ok to test), and to wire it into an L1 launch script if it should run in CI. The pure-torch head and loss math is already covered by the CPU unit tests and was verified locally.

The remaining validation needs the multi-node cluster and is the Phase 1b/2 follow-up from the design (#5233), not part of this PR: the parallel-aware draft backbone (the block-attention module that reads the frozen target's hidden states as context K/V, with TP/PP sharding), the forward_step_dspark + provider hook + dspark.* weight mappings, and full-trainer acceptance-length (tau) convergence against the NeMo AutoModel reference.

[yaoyu-33]

Architecture/correctness disposition: changes requested at exact head 63ad70f3.

I independently checked the head against the DSpark paper, DeepSpec, NeMo AutoModel, current Bridge/MCore training contracts, issue #5230, RFC #5233, and the stacked #5241-#5244 changes. The head/teacher-forcing shapes, Markov and RNN recurrence, target detachment, 1 - 0.5 * L1 acceptance label, prefix-product tau + 1, explicit DP-group requirement, and single-call DP gradient algebra are sound. All 16 submitted CPU unit cases passed, and I ran the submitted test with exactly two GPUs: its NCCL gradient-equivalence assertion passed.

The blockers are the inline items below plus the unresolved subsystem ownership in #5233. This PR installs public, trainable modules with no provider/model/forward-step/recipe caller, so the actual causality mask, confidence input [h_k; embed(x_{k-1})], frozen/shared parameter ownership, TP/PP/CP placement, checkpoint schema, and MCore loss return contract cannot be verified here. Those are not separable details: the current per-call normalization is incompatible with uneven microbatches under MCore, and the full-vocabulary FP32 probability path is not viable at reference-scale anchor counts.

For the stack: #5241 should not be folded in as written—the primary objective trains raw weighted L1 and uses one-half only for the acceptance label. The mask/sentinel hardening in #5242/#5244 is valid if this standalone core is retained. #5243's missing-group policy should follow the approved trainer/process-group contract rather than precede it.

I am not issuing /ok to test yet: the GPU test is not wired into an active CI launch script, and the architecture/loss contract needs revision first. Please settle the owner/API boundary in #5233, then either move these modules to that owner or land them together with the smallest real Bridge integration that makes the distributed and checkpoint contracts testable.

[yaoyu-33]

The implementation below accumulates raw L1, not 0.5 * L1, and the implementation is the part that matches the primary algorithm. Paper Eq. 10 defines the training term as sum_k w_k ||p_d - p_t||_1; DeepSpec and NeMo AutoModel do the same. The factor 0.5 appears in Eq. 8 only when forming the acceptance/confidence label. With the published coefficients, changing this to half-L1 (as #5241 proposes) halves the dominant training term and changes the recipe. Please keep the raw-L1 computation, correct this docstring and the module/design text, and add a disjoint-distribution regression expecting l1_loss == 2 and accept_rate == 0 so the two contracts cannot drift again.

[yaoyu-33]

This denominator collective only spans DP ranks for one loss invocation; it does not avoid the claimed per-microbatch token-count bias. MCore invokes the loss function once per microbatch. Its 3-tuple path accumulates an unnormalized numerator plus num_tokens and scales gradients once by the global token total; its legacy 2-tuple path averages the already-returned losses across microbatches. Returning a normalized value here therefore computes an average of microbatch means when masks differ. A focused 1-token/3-token reproduction gives 1.1269 here versus the correct global token mean 0.6269; returning a token count later would instead double-normalize. Please define this together with the real MCore forward/loss return contract, return reducible numerator/denominator reporting values, and test uneven masks across both DP ranks and at least two microbatches. The current 2-GPU test validates only the single-call PyTorch-DDP algebra.

[yaoyu-33]

This materializes two full FP32 [B, N, L, V] probability tensors and then a same-sized subtraction/absolute-value temporary. At a reference-scale shape such as B=1, N=512, L=7, V≈150k, each FP32 tensor is about 2 GiB, so this helper alone needs over 6 GiB of forward temporaries before autograd storage. The tiny-vocabulary unit tests cannot expose that. The current NeMo AutoModel implementation chunks probability tokens (and checkpoints the chunks) for this reason. Please use a bounded chunked accumulation or establish and enforce a much smaller anchor-count contract with memory evidence before exposing this as a training core.

[yaoyu-33]

I ran this exact test with two GPUs and it passes, but CI will not: no active H100/GB200 launch script references this file. Functional tests in test_groups/ are selected by the launch scripts, not discovered repo-wide. Also, because this test already calls mp.spawn(..., nprocs=2), it must be launched by a single pytest coordinator with CUDA_VISIBLE_DEVICES=0,1; placing it under an existing two-rank torch.distributed.run launcher would create four workers and collide on devices/ports. Please add a dedicated active launch entry (at most two GPUs) using the prescribed uv run python -m pytest command.

[khazic]

Thanks @yaoyu-33, this is addressed in c8cf0bc. I checked each item against the sources you cited rather than taking them on faith, and all four held up. Numbers below are from a real 2-GPU run, not estimates.

1. Objective scaling. You are right, and the direction was the docstring's, not the implementation's. AutoModel's _compute_local_l1_term accumulates l1_dist_per_token * loss_weight_mask with no 1/2, and 0.5 appears only in _compute_accept_rate_3d; DeepSpec splits it the same way; Eq. 10 is the weighted raw L1 and Eq. 8 is where the 1/2 belongs. So the docstring and the design page were the wrong ones and are corrected, and test_disjoint_distributions_give_raw_l1_of_two_and_zero_acceptance asserts the pair together (l1_loss == 2 with accept_rate == 0) exactly as you asked, so the two cannot drift again. I closed #5241, which would have made precisely the change you flagged.

2. Loss return contract. Rewritten to the Megatron-Core 3-tuple: dspark_loss now returns the unnormalized micro-batch numerator, this micro-batch's num_tokens, and a report of reducible [numerator, denominator] pairs. Nothing divides locally, and dp_group is gone from the module entirely.

Tracing it through confirmed your point and sharpened it: the 3-tuple alone is not sufficient. At schedules.py:334 the calculate_per_token_loss=False path still does output_tensor /= torch.clamp(num_tokens, min=1) per micro-batch, so it reintroduces the same bias; only the True path defers to finalize_model_grads_func(..., total_num_tokens). assert_per_token_loss_config therefore fails fast when the flag is unset, and the requirement is written into the RFC. When the forward-step wiring lands, that check belongs in ConfigContainer.validate next to the existing CP>1 rule; it is a free function today only because this package has no config field to key it off yet.

Measured on 2 ranks x 2 micro-batches with uneven supervised-token counts ([6, 22] and [20, 8], 56 global):

[rank 0] global-token normalization    rel err vs single-process = 1.811e-07
[rank 0] per-micro-batch normalization rel err vs single-process = 3.723e-01
[rank 1] global-token normalization    rel err vs single-process = 1.811e-07
[rank 1] per-micro-batch normalization rel err vs single-process = 3.723e-01

A 37% gradient error, so the bias is well clear of noise. The test asserts both directions: the global-token path matches, and the per-micro-batch path must not, so it cannot silently stop discriminating. It also asserts the micro-batch token counts are actually unequal, since with equal counts the two normalizations coincide and the test would prove nothing.

3. Memory. Confirmed and larger than the estimate. Chunking the L1 alone was not enough: F.cross_entropy over the full [3584, 151936] logits saves its own full-size log_softmax for the backward, which is a 2.03 GiB tensor sitting right next to the term being chunked. Both terms now derive from a single FP32 log_softmax inside one chunked, checkpointed loop. On an A100-80GB at B=1, N=512, L=7, V=151936:

one FP32 [B,N,L,V] tensor = 2.03 GiB
unchunked (one full-vocab pass)   peak=12.18 GiB  overhead=10.15 GiB  l1=1.040990  ce=12.396630
chunked (chunk_tokens=128)        peak= 4.09 GiB  overhead= 2.06 GiB  l1=1.040990  ce=12.396630
reduction: 4.9x

Identical values for both terms, and a unit test pins that chunk size changes neither the values nor the gradients. Chunk size matches the AutoModel reference's 128.

4. CI launch. You were right that no active script referenced the file, and right about the four-worker collision. Rather than keep the bespoke launch mode, I dropped mp.spawn and moved the test onto the same torch.distributed.run --nproc_per_node=2 harness every other test in test_groups/training/ uses, via the shared initialize_distributed(). That removes the hand-rolled _free_port, the env preamble, and the teardown, and sidesteps mp.spawn under pytest plus coverage. L0_Launch_training_dspark.sh is the dedicated active entry you asked for, in the standard 2-GPU form and carrying a # CI_TIMEOUT=10 line, which the first version was missing.

On the stack. #5241 is closed. #5242 and #5244 are still valid but were stacked on it, so they carry the bad commit as their base and need rebasing onto this branch; I will do that once this one settles. #5243 is obsolete: dp_group no longer exists in the loss, and assert_per_token_loss_config covers what its guard was for.

On scope. Your point that this installs public modules with no caller stands, and the RFC now settles the owner question that gates it (Megatron-Core native, with the reasoning and the consequences written down). Say the word if you would rather these modules move to Megatron-Core now, or land here together with a minimal forward-step and recipe, and I will restructure accordingly.

How this was validated. The GPU runs above loaded loss.py and heads.py from this branch verbatim, with the megatron.bridge package __init__ chain stubbed out, since those two modules import only torch and my environment is outside the NeMo container. Unit tests: 24 passed. Functional: 1 passed under torchrun --nproc_per_node=2. CI exercises the packaged import path.

[kamran-nvidia]

/ok to test f376431

[kamran-nvidia]

/ok to test 91fac4a

[kamran-nvidia]

/ok to test 1f2ae35

[kamran-nvidia]

/ok to test c2878d0

[khazic]

This one is approved with all 96 checks green and shows as mergeable, but it carries both ready-to-merge and needs-more-tests, so I am not sure which one reflects the current state.

Is there specific coverage still missing? Happy to add it if so. Otherwise the needs-more-tests label may just be stale from an earlier round and can come off.