bridge.training.post_training.dspark.loss

DSpark draft training objective and acceptance metrics.

DSpark (arXiv:2607.05147) trains the draft with three position-weighted terms:

• ce: cross-entropy of the draft logits against the target’s next tokens;

• l1: the raw probability L1 distance ||p_draft - p_target||_1 to the target’s next-token distribution (paper Eq. 10, weight 0.9). This is the dominant term. The factor 1/2 that turns L1 into a total-variation distance belongs to the acceptance label of Eq. 8 only, never to this training term;

• confidence: binary cross-entropy training the confidence head against the analytical acceptance label 1 - 0.5 * L1.

Each block position k is weighted by exp(-k / loss_decay_gamma). Acceptance is measured analytically as accept_rate = 1 - 0.5 * L1 and the expected accepted prefix length of a block as tau = sum_k cumprod(accept_rate)_k + 1.

Return contract

func:

dspark_loss returns the Megatron-Core 3-tuple (loss, num_tokens, report), matching :func:megatron.bridge.training.losses.masked_next_token_loss:

• loss is the unnormalized weighted numerator summed over the micro-batch;

• num_tokens is this micro-batch’s supervised token count;

• report maps each metric to a [numerator, denominator] pair so the training loop can sum both over the log window and the data-parallel group and form the global ratio once.

Nothing here divides by a local denominator, because dividing per micro-batch and averaging afterwards computes a mean of micro-batch means, which is biased whenever micro-batches carry different supervised-token counts. The unbiased global mean requires calculate_per_token_loss=True on the model config, which makes Megatron-Core accumulate num_tokens across micro-batches and scale the gradient once by the data-parallel total in finalize_model_grads_func. With the default calculate_per_token_loss=False, Megatron-Core divides each micro-batch by its own num_tokens and re-introduces exactly that bias, so the integration that calls this function must set the flag; see

func:

assert_per_token_loss_config.

This module depends only on torch. Wiring it to a Megatron-Core forward step is a separate integration concern.

Module Contents

Classes

DSparkForwardOutput

Outputs of one DSpark training forward, consumed by :func:dspark_loss.

Functions

assert_per_token_loss_config	Fail fast unless model_config.calculate_per_token_loss is set.
_loss_weight_mask	eval_mask (float) scaled by the per-position decay exp(-k / gamma).
_chunk_terms	Per-row cross-entropy and probability L1 over one [rows, vocab] chunk.
_ce_and_l1_per_token	Per-position cross-entropy and exact FP32 probability L1, without full-vocab temporaries.
_pair	Pack a reducible [numerator, denominator] reporting pair.
dspark_loss	Compute the DSpark draft loss and acceptance metrics.
_acceptance_report	Analytical acceptance diagnostics as reducible [numerator, denominator] pairs.

Data

_PROBABILITY_CHUNK_TOKENS

API

bridge.training.post_training.dspark.loss._PROBABILITY_CHUNK_TOKENS

128

class bridge.training.post_training.dspark.loss.DSparkForwardOutput

Outputs of one DSpark training forward, consumed by :func:dspark_loss.

Shape symbols: batch, num_blocks (sampled anchor blocks per sample), block_size (draft positions per anchor), vocab.

.. attribute:: draft_logits

Draft logits [batch, num_blocks, block_size, vocab].

.. attribute:: target_ids

Teacher-forced next token per position [batch, num_blocks, block_size].

.. attribute:: eval_mask

Bool/float supervision mask [batch, num_blocks, block_size] (a block is a contiguous, in-bounds, loss-enabled prefix).

.. attribute:: block_keep_mask

Kept-anchor mask [batch, num_blocks].

.. attribute:: confidence_pred

Optional per-position acceptance logit [batch, num_blocks, block_size].

.. attribute:: aligned_target_logits

Optional target next-token logits [batch, num_blocks, block_size, vocab] (the L1 / confidence teacher).

draft_logits: torch.Tensor

None

target_ids: torch.Tensor

None

eval_mask: torch.Tensor

None

block_keep_mask: torch.Tensor

None

confidence_pred: torch.Tensor | None

None

aligned_target_logits: torch.Tensor | None

None

bridge.training.post_training.dspark.loss.assert_per_token_loss_config(model_config: Any) → None

Fail fast unless model_config.calculate_per_token_loss is set.

Func:

dspark_loss returns an unnormalized numerator; only the per-token-loss path sums num_tokens across micro-batches and scales the gradient once by the global total. Under the default path Megatron-Core normalizes each micro-batch by its own token count, which silently trains a mean of micro-batch means.

When the forward-step wiring lands, this belongs in ConfigContainer.validate alongside the existing context-parallel calculate_per_token_loss rule, so it fires before training rather than at the first loss call. It lives here for now because this package has no config field to key that rule off yet. The parameter is untyped because this module deliberately imports nothing beyond torch; the attribute is read without a default, so passing the wrong object raises AttributeError rather than reporting a config error that is not there.

Parameters:

model_config – The Megatron-Core model config used for the run.

Raises:

• AttributeError – If model_config has no calculate_per_token_loss.

• ValueError – If calculate_per_token_loss is not enabled.

bridge.training.post_training.dspark.loss._loss_weight_mask(

eval_mask: torch.Tensor,

loss_decay_gamma: float | None,

) → torch.Tensor

eval_mask (float) scaled by the per-position decay exp(-k / gamma).

bridge.training.post_training.dspark.loss._chunk_terms(

draft_logits: torch.Tensor,

target_ids: torch.Tensor,

target_logits: torch.Tensor | None,

) → tuple[torch.Tensor, torch.Tensor]

Per-row cross-entropy and probability L1 over one [rows, vocab] chunk.

Both terms are derived from a single FP32 log_softmax of the draft logits, so the draft distribution is normalized once rather than once per term.

bridge.training.post_training.dspark.loss._ce_and_l1_per_token(

draft_logits: torch.Tensor,

target_ids: torch.Tensor,

target_logits: torch.Tensor | None,

chunk_tokens: int = _PROBABILITY_CHUNK_TOKENS,

) → tuple[torch.Tensor, torch.Tensor]

Per-position cross-entropy and exact FP32 probability L1, without full-vocab temporaries.

Both terms need the draft’s normalized distribution over the whole vocabulary. A single [batch, num_blocks, block_size, vocab] FP32 tensor is 2.03 GiB at reference scale (num_blocks=512, block_size=7, vocab≈150k), and computing the two terms separately keeps several of them live at once, since the cross-entropy’s own log_softmax is saved for its backward. Measured on an 80 GB card, an unchunked pass costs 10.15 GiB of forward-plus-backward overhead above its inputs, against 2.06 GiB here.

Chunking the flattened rows caps the temporaries at chunk_tokens * vocab elements, and each chunk is recomputed in the backward (:func:torch.utils.checkpoint.checkpoint) so no chunk keeps its distribution alive for autograd. Chunking does not change the result: every row is reduced independently over the vocab axis.

Parameters:

• draft_logits – [..., vocab].

• target_ids – Teacher-forced next token per position [...].

• target_logits – [..., vocab] already detached by the caller, or None to compute cross-entropy only (the L1 output is then zeros).

• chunk_tokens – Rows of the flattened logits per chunk.

Returns:

(ce, l1), both [...]. l1 is twice the total-variation distance.

bridge.training.post_training.dspark.loss._pair(

numerator: torch.Tensor,

denominator: torch.Tensor,

) → torch.Tensor

Pack a reducible [numerator, denominator] reporting pair.

bridge.training.post_training.dspark.loss.dspark_loss(

outputs: bridge.training.post_training.dspark.loss.DSparkForwardOutput,

*,

ce_alpha: float = 0.1,

l1_alpha: float = 0.9,

confidence_alpha: float = 1.0,

loss_decay_gamma: float | None = 4.0,

chunk_tokens: int = _PROBABILITY_CHUNK_TOKENS,

) → tuple[torch.Tensor, torch.Tensor, dict[str, torch.Tensor]]

Compute the DSpark draft loss and acceptance metrics.

Follows the Megatron-Core 3-tuple loss contract: the returned loss is the unnormalized micro-batch numerator, and normalization happens once against the globally reduced num_tokens. See the module docstring for why, and

Func:

assert_per_token_loss_config for the config this requires.

Parameters:

• outputs – The draft forward outputs; see :class:DSparkForwardOutput. All tensors are [batch, num_blocks, block_size(, vocab)].

• ce_alpha – Cross-entropy weight.

• l1_alpha – Weight of the raw probability-L1 term (paper Eq. 10; no 1/2).

• confidence_alpha – Confidence-BCE weight (used only when outputs.confidence_pred is set).

• loss_decay_gamma – Per-position decay; None disables it.

• chunk_tokens –

  Rows per chunk of the FP32 probability L1; see

  func:

  _l1_distance_per_token.

Returns:

Tuple (loss, num_tokens, report). loss is the unnormalized weighted numerator for this micro-batch; num_tokens is its supervised token count as an int tensor; report maps each metric name to a detached [numerator, denominator] pair that the caller sums over the log window and the data-parallel group before dividing once.

Raises:

ValueError – If the L1 term or the confidence head is requested without aligned_target_logits.

bridge.training.post_training.dspark.loss._acceptance_report(

outputs: bridge.training.post_training.dspark.loss.DSparkForwardOutput,

eval_mask: torch.Tensor,

accept_rate: torch.Tensor,

) → dict[str, torch.Tensor]

Analytical acceptance diagnostics as reducible [numerator, denominator] pairs.

accept_rate (= 1 - 0.5 * L1) is the per-position acceptance probability. A draft token survives only if every earlier token in its block is accepted, hence tau is the running product over the block, + 1 for the verified anchor token.

Parameters:

• outputs – The forward outputs (for block_keep_mask).

• eval_mask – The float supervision mask [batch, num_blocks, block_size].

• accept_rate – Per-position acceptance [batch, num_blocks, block_size].

Returns:

{"dspark accept rate": pair, "dspark tau": pair}.