nemo_automodel.components.speculative.dspark.core

DSpark online training wrapper.

The DSpark draft is self-contained: it samples anchors, builds the block attention mask, runs the semi-autoregressive backbone + Markov head, and emits everything the objective needs. This module is therefore a thin wrapper that calls the draft with the target supervision and computes the three-term loss.

Module Contents

Classes

Name	Description
`DSparkStepMetrics`	Per-step training outputs for the DSpark draft.
`DSparkTrainerModule`	DSpark online training wrapper computing the three-term objective.

Data

__all__

API

class nemo_automodel.components.speculative.dspark.core.DSparkStepMetrics(
    loss: torch.Tensor,
    ce_loss: torch.Tensor,
    l1_loss: torch.Tensor,
    confidence_loss: torch.Tensor,
    accept_rate_per_pos_num: torch.Tensor,
    accept_rate_per_pos_den: torch.Tensor,
    tau_num: torch.Tensor,
    tau_den: torch.Tensor,
    confidence_abs_error_num: torch.Tensor,
    confidence_bias_num: torch.Tensor,
    confidence_cumprod_bias_num: torch.Tensor,
    confidence_diag_den: torch.Tensor
)

Dataclass

Per-step training outputs for the DSpark draft.

Beyond the loss and its three terms, this carries the acceptance diagnostics as unreduced (num, den) sums so the recipe can reduce both across the log window and the DP group and form the exact global ratio once: acceptance rate (the [block_size] per-position accept_rate@k sums, whose totals also give the aggregate accept_rate), tau (expected accepted block length), and the confidence-head calibration error/bias against the measured acceptance rate. A denominator is zero when the diagnostic was not computed (e.g. no confidence head, or no teacher signal), which the recipe uses to skip logging.

accept_rate_per_pos_den

Tensor

accept_rate_per_pos_num

Tensor

ce_loss

Tensor

confidence_abs_error_num

Tensor

confidence_bias_num

Tensor

confidence_cumprod_bias_num

Tensor

confidence_diag_den

Tensor

confidence_loss

Tensor

l1_loss

Tensor

loss

Tensor

tau_den

Tensor

tau_num

Tensor

class nemo_automodel.components.speculative.dspark.core.DSparkTrainerModule(
    draft_model: nemo_automodel.components.speculative.dspark.draft_qwen3.Qwen3DSparkModel,
    loss_decay_gamma: typing.Optional[float] = None,
    ce_loss_alpha: float = 0.1,
    l1_loss_alpha: float = 0.9,
    confidence_head_alpha: float = 1.0
)

Bases: Module

DSpark online training wrapper computing the three-term objective.

ce_loss_alpha

= float(ce_loss_alpha)

confidence_head_alpha

= float(confidence_head_alpha)

l1_loss_alpha

= float(l1_loss_alpha)

nemo_automodel.components.speculative.dspark.core.DSparkTrainerModule.forward(
    input_ids: torch.Tensor,
    target_hidden_states: torch.Tensor,
    loss_mask: torch.Tensor,
    target_last_hidden_states: typing.Optional[torch.Tensor] = None,
    position_ids: typing.Optional[torch.Tensor] = None,
    seq_lens: typing.Optional[torch.Tensor] = None,
    doc_remaining: typing.Optional[torch.Tensor] = None
) -> nemo_automodel.components.speculative.dspark.core.DSparkStepMetrics

Run the draft on the target supervision and compute the DSpark loss.

position_ids / seq_lens / doc_remaining (all None off the packing path) are forwarded to the draft, which keeps each anchor block inside one document (block-causal context, per-document positions, and document-truncated supervision).

nemo_automodel.components.speculative.dspark.core.__all__ = ['DSparkTrainerModule', 'DSparkStepMetrics']