`nemo_rl.algorithms.loss.draft`#

Module Contents#

Classes#

`DraftLossStats`	Raw draft-loss reductions with one entry per pass or block slot.
`_CachedVocabParallelSoftCE`
`_StreamingVocabParallelSoftCE`
`_SumTensorParallelHiddenGradient`
`_StreamingProjectedVocabParallelSoftCE`

Functions#

`_dtype_header_code`
`_tile_log_normalizers`
`_tile_distributions`
`_float64_bits`	Encode a Python float as its signed IEEE-754 bit pattern.
`_tp_assert_projected_metadata_agreement`	Fail all TP ranks together before rank-local validation or CE collectives.
`_cached_tile_distributions_and_cross_entropy`
`_resolve_draft_bins`
`streaming_vocab_parallel_soft_ce`	Compute raw soft-CE bins while bounding FP32 vocab work to one token tile.
`projected_streaming_vocab_parallel_soft_ce`	Project hidden states against preselected teacher rows with TP-safe metadata.
`dflash_projected_vocab_parallel_soft_ce`	Map DFlash block slots to the live target head and indexed teacher rows.

Data#

`DEFAULT_DRAFT_TOKEN_CHUNK_SIZE`
`_DTYPE_HEADER_CODES`

API#

nemo_rl.algorithms.loss.draft.DEFAULT_DRAFT_TOKEN_CHUNK_SIZE#: 4096

nemo_rl.algorithms.loss.draft._DTYPE_HEADER_CODES: dict[torch.dtype, int]#: None

nemo_rl.algorithms.loss.draft._dtype_header_code(dtype: torch.dtype) → int#

class nemo_rl.algorithms.loss.draft.DraftLossStats#

Raw draft-loss reductions with one entry per pass or block slot.

numerators: torch.Tensor#: None

counts: torch.Tensor#: None

weights: torch.Tensor#: None

__post_init__() → None#

normalized(*, normalization_counts: torch.Tensor) → torch.Tensor#: Normalize local differentiable sums with externally reduced counts.

nemo_rl.algorithms.loss.draft._tile_log_normalizers( student_logits: torch.Tensor, teacher_logits: torch.Tensor, tp_group: torch.distributed.ProcessGroup | None, ) → torch.Tensor#

nemo_rl.algorithms.loss.draft._tile_distributions( student_logits: torch.Tensor, teacher_logits: torch.Tensor, log_normalizers: torch.Tensor, ) → tuple[torch.Tensor, torch.Tensor]#

nemo_rl.algorithms.loss.draft._float64_bits(value: float) → int#: Encode a Python float as its signed IEEE-754 bit pattern.

nemo_rl.algorithms.loss.draft._tp_assert_projected_metadata_agreement( *, tp_group: torch.distributed.ProcessGroup | None, reference: torch.Tensor, tensors: tuple[tuple[str, torch.Tensor | None], ...], scalars: tuple[tuple[str, int], ...], exact_tensors: tuple[tuple[str, torch.Tensor | None], ...], ) → None#: Fail all TP ranks together before rank-local validation or CE collectives.

nemo_rl.algorithms.loss.draft._cached_tile_distributions_and_cross_entropy( student_logits: torch.Tensor, teacher_logits: torch.Tensor, tp_group: torch.distributed.ProcessGroup | None, ) → tuple[torch.Tensor, torch.Tensor, torch.Tensor]#

class nemo_rl.algorithms.loss.draft._CachedVocabParallelSoftCE#

Bases: torch.autograd.Function

static forward( ctx: Any, student_logits: torch.Tensor, teacher_logits: torch.Tensor, mask: torch.Tensor, bin_ids: torch.Tensor, num_bins: int, tp_group: torch.distributed.ProcessGroup | None, ) → torch.Tensor#

static backward( ctx: Any, *grad_outputs: torch.Tensor, ) → tuple[torch.Tensor, None, None, None, None, None]#

class nemo_rl.algorithms.loss.draft._StreamingVocabParallelSoftCE#

Bases: torch.autograd.Function

static forward( ctx: Any, student_logits: torch.Tensor, teacher_logits: torch.Tensor, mask: torch.Tensor, bin_ids: torch.Tensor, num_bins: int, token_chunk_size: int, tp_group: torch.distributed.ProcessGroup | None, ) → torch.Tensor#

static backward( ctx: Any, *grad_outputs: torch.Tensor, ) → tuple[torch.Tensor, None, None, None, None, None, None]#

class nemo_rl.algorithms.loss.draft._SumTensorParallelHiddenGradient#

Bases: torch.autograd.Function

static forward( ctx: Any, student_hidden: torch.Tensor, tp_group: torch.distributed.ProcessGroup | None, ) → torch.Tensor#

static backward( ctx: Any, hidden_gradient: torch.Tensor, ) → tuple[torch.Tensor, None]#

class nemo_rl.algorithms.loss.draft._StreamingProjectedVocabParallelSoftCE#

Bases: torch.autograd.Function

static forward( ctx: Any, student_hidden: torch.Tensor, output_weight: torch.Tensor, selected_teacher_logits: torch.Tensor, mask: torch.Tensor, bin_ids: torch.Tensor, num_bins: int, token_chunk_size: int, tp_group: torch.distributed.ProcessGroup | None, ) → torch.Tensor#

static backward( ctx: Any, *grad_outputs: torch.Tensor, ) → tuple[torch.Tensor, None, None, None, None, None, None, None]#

nemo_rl.algorithms.loss.draft._resolve_draft_bins( mask: torch.Tensor, bin_ids: torch.Tensor | None, weights: torch.Tensor | None, ) → tuple[torch.Tensor, torch.Tensor, int]#

nemo_rl.algorithms.loss.draft.streaming_vocab_parallel_soft_ce( *, student_logits: torch.Tensor, teacher_logits: torch.Tensor, mask: torch.Tensor, token_chunk_size: int, tp_group: torch.distributed.ProcessGroup | None, bin_ids: torch.Tensor | None = None, weights: torch.Tensor | None = None, ) → nemo_rl.algorithms.loss.draft.DraftLossStats#: Compute raw soft-CE bins while bounding FP32 vocab work to one token tile.

nemo_rl.algorithms.loss.draft.projected_streaming_vocab_parallel_soft_ce( *, student_hidden: torch.Tensor, output_weight: torch.Tensor, selected_teacher_logits: torch.Tensor, mask: torch.Tensor, token_chunk_size: int, tp_group: torch.distributed.ProcessGroup | None, bin_ids: torch.Tensor | None = None, weights: torch.Tensor | None = None, _skip_tp_metadata_check: bool = False, ) → nemo_rl.algorithms.loss.draft.DraftLossStats#

Project hidden states against preselected teacher rows with TP-safe metadata.

TP ranks agree on a fixed structural header and exact mask/bin/weight metadata before validation or loss collectives, so malformed rank-local inputs fail together.

_skip_tp_metadata_check is a private escape hatch for callers that already ran the agreement check on the tensors these arguments are derived from (the DFlash adapter); it avoids doubling the collective traffic per step under TP.

nemo_rl.algorithms.loss.draft.dflash_projected_vocab_parallel_soft_ce( *, draft_hidden: torch.Tensor, output_weight: torch.Tensor, teacher_logits: torch.Tensor, sample_rows: torch.Tensor, label_positions: torch.Tensor, loss_mask: torch.Tensor, position_decay: float, token_chunk_size: int, tp_group: torch.distributed.ProcessGroup | None, ) → nemo_rl.algorithms.loss.draft.DraftLossStats#

Map DFlash block slots to the live target head and indexed teacher rows.

Called from the DFlash training path added in PR #3757.

Slot 0 is the verified anchor token and is always excluded from the loss: this function overrides loss_mask[:, 0] to False regardless of what the caller passed. The returned DraftLossStats.counts is a rank-local (gamma,) vector; callers must all-reduce it across DP/CP before passing it as normalization_counts to :meth:DraftLossStats.normalized, unlike the scalar global_valid_toks pattern used by the base draft cross-entropy path.

nemo_rl.algorithms.loss.draft#

Module Contents#

Classes#

Functions#

Data#

API#

`nemo_rl.algorithms.loss.draft`#