> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.speculative.dspark.loss

## Module Contents

### Functions

| Name                                                                                                                    | Description                                                                 |
| ----------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------- |
| [`_all_reduce_loss_denominators`](#nemo_automodel-components-speculative-dspark-loss-_all_reduce_loss_denominators)     | -                                                                           |
| [`_build_loss`](#nemo_automodel-components-speculative-dspark-loss-_build_loss)                                         | -                                                                           |
| [`_build_loss_weight_mask`](#nemo_automodel-components-speculative-dspark-loss-_build_loss_weight_mask)                 | -                                                                           |
| [`_collect_acceptance_diagnostics`](#nemo_automodel-components-speculative-dspark-loss-_collect_acceptance_diagnostics) | Per-batch numerator/denominator sums for the acceptance diagnostics.        |
| [`_collect_local_terms`](#nemo_automodel-components-speculative-dspark-loss-_collect_local_terms)                       | -                                                                           |
| [`_compute_accept_rate_3d`](#nemo_automodel-components-speculative-dspark-loss-_compute_accept_rate_3d)                 | -                                                                           |
| [`_compute_l1_dist_per_token`](#nemo_automodel-components-speculative-dspark-loss-_compute_l1_dist_per_token)           | Compute exact FP32 probability L1 distances without full-vocab temporaries. |
| [`_compute_local_l1_term`](#nemo_automodel-components-speculative-dspark-loss-_compute_local_l1_term)                   | -                                                                           |
| [`_l1_probability_distance_chunk`](#nemo_automodel-components-speculative-dspark-loss-_l1_probability_distance_chunk)   | -                                                                           |
| [`compute_dspark_loss`](#nemo_automodel-components-speculative-dspark-loss-compute_dspark_loss)                         | -                                                                           |

### Data

[`_PROBABILITY_CHUNK_TOKENS`](#nemo_automodel-components-speculative-dspark-loss-_PROBABILITY_CHUNK_TOKENS)

[`__all__`](#nemo_automodel-components-speculative-dspark-loss-__all__)

### API

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

```python
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
```

```python
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
```

```python
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`).

```python
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]
```

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

```python
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.

```python
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]
```

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

```python
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
)
```

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

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