`nemo_rl.models.automodel.draft.common`#

Module Contents#

Classes#

DSparkForwardOutput

Outputs for one DSpark training forward.

Functions#

`build_anchor_candidate_mask`	Mask of anchor positions whose first predicted token is supervised.
`sample_anchor_positions`	Sample up to num_anchors valid anchor positions per sample.
`build_eval_mask`	Per-slot supervision mask for the anchored blocks.

Data#

__all__

API#

class nemo_rl.models.automodel.draft.common.DSparkForwardOutput#

Outputs for one DSpark training forward.

Shape symbols: batch_size: number of samples in the batch seq_len: source sequence length num_anchors: sampled anchor blocks per sample block_size: number of draft positions per anchor vocab_size: vocabulary size

The sampler keeps anchors whose first draft target is enabled by loss_mask. Later slots are supervised only while they remain inside seq_len and form a contiguous enabled prefix. Dummy anchors can still appear when a sample has too few valid anchors; they are masked out by block_keep_mask and eval_mask.

draft_logits: torch.Tensor#: None

target_ids: torch.Tensor#: None

eval_mask: torch.Tensor#: None

block_keep_mask: torch.Tensor#: None

confidence_pred: Optional[torch.Tensor]#: None

aligned_target_logits: Optional[torch.Tensor]#: None

first_supervised_slot: int#: 0

nemo_rl.models.automodel.draft.common.build_anchor_candidate_mask( *, seq_len: int, loss_mask: torch.Tensor, doc_remaining: Optional[torch.Tensor] = None, ) → torch.Tensor#

Mask of anchor positions whose first predicted token is supervised.

Parameters:

seq_len – source sequence length.
loss_mask – [batch_size, seq_len] per-token supervision mask.
doc_remaining – [batch_size, seq_len] tokens remaining in the current packed document, or None when unpacked.

Returns:

[batch_size, seq_len - 1] boolean mask; True at anchor position p when token p + 1 is supervised (and, when packing, stays inside the anchor’s document).

nemo_rl.models.automodel.draft.common.sample_anchor_positions( *, seq_len: int, loss_mask: torch.Tensor, num_anchors: int, device: torch.device, doc_remaining: Optional[torch.Tensor] = None, generator: Optional[torch.Generator] = None, ) → tuple[torch.Tensor, torch.Tensor]#

Sample up to num_anchors valid anchor positions per sample.

Parameters:

seq_len – source sequence length.
loss_mask – [batch_size, seq_len] per-token supervision mask.
num_anchors – maximum anchors to sample per sample.
device – device to allocate outputs on.
doc_remaining – packed-document remaining-token counts, or None.
generator – optional RNG for a reproducible draw (must be seeded identically across TP peers sharing the same draft replica).

Returns:

anchors is [batch_size, num_anchors] sorted anchor token positions (zero-padded past the valid count); keep_mask marks which of those are real vs. padding.

Return type:

(anchors, keep_mask)

nemo_rl.models.automodel.draft.common.build_eval_mask( *, seq_len: int, loss_mask: torch.Tensor, label_indices: torch.Tensor, safe_label_indices: torch.Tensor, block_keep_mask: torch.Tensor, doc_remaining: Optional[torch.Tensor] = None, anchor_positions: Optional[torch.Tensor] = None, supervised_from_slot: int = 0, ) → torch.Tensor#

Per-slot supervision mask for the anchored blocks.

supervised_from_slot marks the first supervised block slot: 0 when the anchor slot itself predicts (dspark next-token layout), 1 when the anchor is an unsupervised bonus token (dflash at-position layout). Leading unsupervised slots are treated as valid links for the trailing cumprod (their label — the anchor’s own token — may sit outside the loss mask without invalidating the rest of the block) and are zeroed afterwards.

nemo_rl.models.automodel.draft.common.__all__#: [‘DSparkForwardOutput’, ‘AcceptRatePredictor’, ‘extract_context_feature’, ‘validate_target_layer_ids…

nemo_rl.models.automodel.draft.common#