nemo_rl.algorithms.loss.draft#

Module Contents#

Classes#

DraftLossStats

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

_CachedVocabParallelSoftCE

_StreamingVocabParallelSoftCE

Functions#

_tile_log_normalizers

_tile_distributions

_cached_tile_distributions_and_cross_entropy

streaming_vocab_parallel_soft_ce

Compute raw soft-CE bins while bounding FP32 vocab work to one token tile.

Data#

DEFAULT_DRAFT_TOKEN_CHUNK_SIZE

API#

nemo_rl.algorithms.loss.draft.DEFAULT_DRAFT_TOKEN_CHUNK_SIZE#

4096

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._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]#
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.