nemo_rl.models.automodel.draft.loss#

Module Contents#

Functions#

_all_reduce_loss_denominators

_build_loss_weight_mask

Per-slot loss weights: eval mask times the positional decay.

_l1_probability_distance_chunk

_compute_l1_dist_per_token

Compute exact FP32 probability L1 distances without full-vocab temporaries.

_compute_accept_rate_3d

_compute_local_l1_term

_collect_acceptance_diagnostics

Per-batch numerator/denominator sums for the acceptance diagnostics.

_collect_local_terms

compute_dspark_loss

Compute the DSpark block-draft loss (CE + L1 + optional confidence).

Data#

_PROBABILITY_CHUNK_TOKENS

__all__

API#

nemo_rl.models.automodel.draft.loss._PROBABILITY_CHUNK_TOKENS#

128

nemo_rl.models.automodel.draft.loss._all_reduce_loss_denominators(
loss_terms: dict[str, torch.Tensor],
*,
process_group: Optional[torch.distributed.ProcessGroup],
) → dict[str, torch.Tensor]#
nemo_rl.models.automodel.draft.loss._build_loss_weight_mask(
*,
eval_mask: torch.Tensor,
block_size: int,
device: torch.device,
loss_decay_gamma: Optional[float],
first_supervised_slot: int = 0,
) → torch.Tensor#

Per-slot loss weights: eval mask times the positional decay.

Decay positions count PROPOSALS, not raw block slots: the k-th supervised slot gets exp(-k / gamma) so the first proposed token always has weight 1. Under the dflash bonus-anchor layout slot 0 is unsupervised (first_supervised_slot = 1); its clamped weight is irrelevant because the eval mask zeroes it.

nemo_rl.models.automodel.draft.loss._l1_probability_distance_chunk(
*,
draft_logits: torch.Tensor,
target_logits: torch.Tensor,
) → torch.Tensor#
nemo_rl.models.automodel.draft.loss._compute_l1_dist_per_token(
*,
draft_logits: torch.Tensor,
aligned_target_logits: torch.Tensor,
chunk_size: int = _PROBABILITY_CHUNK_TOKENS,
) → torch.Tensor#

Compute exact FP32 probability L1 distances without full-vocab temporaries.

nemo_rl.models.automodel.draft.loss._compute_accept_rate_3d(
l1_dist_per_token: Optional[torch.Tensor],
) → Optional[torch.Tensor]#
nemo_rl.models.automodel.draft.loss._compute_local_l1_term(
*,
l1_dist_per_token: Optional[torch.Tensor],
loss_weight_mask: torch.Tensor,
) → tuple[torch.Tensor, torch.Tensor]#
nemo_rl.models.automodel.draft.loss._collect_acceptance_diagnostics(
*,
outputs: nemo_rl.models.automodel.draft.common.DSparkForwardOutput,
accept_rate_3d: Optional[torch.Tensor],
loss_weight_mask: torch.Tensor,
has_confidence: bool,
) → dict[str, torch.Tensor]#

Per-batch numerator/denominator sums for the acceptance diagnostics.

accept_rate_3d is the TV-derived per-token acceptance probability. Every diagnostic is returned as an unreduced (num, den) sum; the recipe sums both across the log window and the data-parallel group and forms the global ratio once (sum(num) / sum(den)), so per-micro-batch token-count imbalance never biases the reported value. Returns zero sums when no teacher signal is available (accept_rate_3d is None).

nemo_rl.models.automodel.draft.loss._collect_local_terms(
*,
outputs: nemo_rl.models.automodel.draft.common.DSparkForwardOutput,
loss_decay_gamma: Optional[float],
l1_loss_alpha: float,
) → tuple[dict[str, torch.Tensor], bool]#
nemo_rl.models.automodel.draft.loss.compute_dspark_loss(
*,
outputs: nemo_rl.models.automodel.draft.common.DSparkForwardOutput,
loss_decay_gamma: Optional[float],
ce_loss_alpha: float,
l1_loss_alpha: float,
confidence_head_alpha: float,
process_group: Optional[torch.distributed.ProcessGroup] = None,
return_terms: bool = False,
) → torch.Tensor | tuple[torch.Tensor, dict[str, torch.Tensor]]#

Compute the DSpark block-draft loss (CE + L1 + optional confidence).

The returned backward_loss is built from global_denominators (all-reduced across process_group) so it is correct under FSDP’s gradient averaging. The per-term scalars in terms (when return_terms=True) are detached copies of the same global-denominator values that make up backward_loss – summing them across DP ranks (as the training loop’s metric reduction does) reproduces the backward-scale global mean, unlike a per-rank local-denominator mean, which would not.

Parameters:

  • outputs – forward-pass outputs from the DSpark draft model.

  • loss_decay_gamma – exponential per-block-position loss-decay factor, or None to disable decay.

  • ce_loss_alpha – weight on the cross-entropy term.

  • l1_loss_alpha – weight on the L1 (confidence-target) term.

  • confidence_head_alpha – weight on the confidence-head term.

  • process_group – process group to all-reduce loss denominators over (typically the DP group); None to skip the all-reduce.

  • return_terms – if True, also return a dict of detached diagnostic terms (global-denominator losses plus unreduced accept-rate/tau sums).

Returns:

backward_loss alone, or (backward_loss, terms) when return_terms is True.

nemo_rl.models.automodel.draft.loss.__all__#

[‘compute_dspark_loss’]