nemo_automodel.components.speculative.dspark.loss

Module Contents

Functions

Name	Description
`_all_reduce_loss_denominators`	-
`_build_loss`	-
`_build_loss_weight_mask`	-
`_collect_acceptance_diagnostics`	Per-batch numerator/denominator sums for the acceptance diagnostics.
`_collect_local_terms`	-
`_compute_accept_rate_3d`	-
`_compute_l1_dist_per_token`	Compute exact FP32 probability L1 distances without full-vocab temporaries.
`_compute_local_l1_term`	-
`_l1_probability_distance_chunk`	-
`compute_dspark_loss`	-

Data

_PROBABILITY_CHUNK_TOKENS

__all__

API

nemo_automodel.components.speculative.dspark.loss._all_reduce_loss_denominators(
    loss_terms: dict[str, torch.Tensor],
    world_size: int
) -> dict[str, torch.Tensor]

nemo_automodel.components.speculative.dspark.loss._build_loss(
    loss_terms: dict[str, torch.Tensor],
    global_denominators: dict[str, torch.Tensor],
    ce_loss_alpha: float,
    l1_loss_alpha: float,
    confidence_head_alpha: float,
    has_confidence: bool,
    world_size: int
) -> torch.Tensor

nemo_automodel.components.speculative.dspark.loss._build_loss_weight_mask(
    eval_mask: torch.Tensor,
    block_size: int,
    device: torch.device,
    loss_decay_gamma: typing.Optional[float]
) -> torch.Tensor

nemo_automodel.components.speculative.dspark.loss._collect_acceptance_diagnostics(
    outputs: nemo_automodel.components.speculative.dspark.common.DSparkForwardOutput,
    accept_rate_3d: typing.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_automodel.components.speculative.dspark.loss._collect_local_terms(
    outputs: nemo_automodel.components.speculative.dspark.common.DSparkForwardOutput,
    loss_decay_gamma: typing.Optional[float],
    l1_loss_alpha: float
) -> tuple[dict[str, torch.Tensor], bool]

nemo_automodel.components.speculative.dspark.loss._compute_accept_rate_3d(
    l1_dist_per_token: typing.Optional[torch.Tensor]
) -> typing.Optional[torch.Tensor]

nemo_automodel.components.speculative.dspark.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_automodel.components.speculative.dspark.loss._compute_local_l1_term(
    l1_dist_per_token: typing.Optional[torch.Tensor],
    loss_weight_mask: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]

nemo_automodel.components.speculative.dspark.loss._l1_probability_distance_chunk(
    draft_logits: torch.Tensor,
    target_logits: torch.Tensor
) -> torch.Tensor

nemo_automodel.components.speculative.dspark.loss.compute_dspark_loss(
    outputs: nemo_automodel.components.speculative.dspark.common.DSparkForwardOutput,
    loss_decay_gamma: typing.Optional[float],
    ce_loss_alpha: float,
    l1_loss_alpha: float,
    confidence_head_alpha: float,
    return_terms: bool = False
)

nemo_automodel.components.speculative.dspark.loss._PROBABILITY_CHUNK_TOKENS = 128

nemo_automodel.components.speculative.dspark.loss.__all__ = ['compute_dspark_loss']