`nemo_rl.algorithms.loss.wrapper`#

Module Contents#

Classes#

`SequencePackingLossWrapper`
`SequencePackingFusionLossWrapper`	Fused sequence packing loss wrapper that processes all sequences in one forward pass.
`DraftLossWrapper`	Combine policy loss with draft soft cross-entropy loss.
`DraftRuntimeLossWrapper`	Combine policy loss with a co-trained drafter loss (automodel path).

Functions#

wrap_loss_fn_with_input_preparation

Wraps a loss function to handle input preparation for megatron policy worker.

Data#

`Tensor`
`_SEQ_METRIC_MIN`
`_SEQ_METRIC_MAX`

API#

nemo_rl.algorithms.loss.wrapper.Tensor#: ‘TypeVar(…)’

nemo_rl.algorithms.loss.wrapper._SEQ_METRIC_MIN: frozenset[str]#: ‘frozenset(…)’

nemo_rl.algorithms.loss.wrapper._SEQ_METRIC_MAX: frozenset[str]#: ‘frozenset(…)’

class nemo_rl.algorithms.loss.wrapper.SequencePackingLossWrapper( loss_fn: nemo_rl.algorithms.loss.interfaces.LossFunction, prepare_fn: Callable[Any, Any], cu_seqlens_q: nemo_rl.utils.sequence_lengths.CpuIntTuple, cu_seqlens_q_padded: Optional[nemo_rl.utils.sequence_lengths.CpuIntTuple] = None, vocab_parallel_rank: Optional[int] = None, vocab_parallel_group: Optional[torch.distributed.ProcessGroup] = None, context_parallel_group: Optional[torch.distributed.ProcessGroup] = None, )#

Initialization

Wrap a loss function to handle sequence packing.

Parameters:

loss_fn – Loss function.
prepare_fn – Prepare function.
cu_seqlens_q – Unpadded cu seqlens q.
cu_seqlens_q_padded – Padded cu seqlens q.
vocab_parallel_rank – Vocab parallel rank.
vocab_parallel_group – Vocab parallel group.
context_parallel_group – Context parallel group.
vocab_parallel_rank
vocab_parallel_group
worker. (context_parallel_group are only used for megatron policy)

Returns:

Sequence packing loss wrapper.

__call__( next_token_logits: nemo_rl.algorithms.loss.wrapper.Tensor, data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any], global_valid_seqs: nemo_rl.algorithms.loss.wrapper.Tensor | None, global_valid_toks: nemo_rl.algorithms.loss.wrapper.Tensor | None, ) → tuple[nemo_rl.algorithms.loss.wrapper.Tensor, dict[str, Any]]#: Wraps a loss function to handle sequence packing by doing one sequence at a time to avoid excessive padding.

class nemo_rl.algorithms.loss.wrapper.SequencePackingFusionLossWrapper( loss_fn: nemo_rl.algorithms.loss.interfaces.LossFunction, prepare_fn: Callable[..., Any], cu_seqlens_q: nemo_rl.utils.sequence_lengths.CpuIntTuple, cu_seqlens_q_padded: Optional[nemo_rl.utils.sequence_lengths.CpuIntTuple] = None, vocab_parallel_rank: Optional[int] = None, vocab_parallel_group: Optional[torch.distributed.ProcessGroup] = None, context_parallel_group: Optional[torch.distributed.ProcessGroup] = None, )#

Fused sequence packing loss wrapper that processes all sequences in one forward pass.

Unlike SequencePackingLossWrapper which iterates over sequences one at a time, this wrapper calls prepare_fn once on the packed logits to compute log probabilities in a single shot, then calls the loss function once with the pre-computed result.

This avoids per-sequence kernel launches and TP/CP communication overhead while producing numerically identical results.

The prepare_fn should be prepare_packed_loss_input (from nemo_rl.algorithms.loss.loss_input), which currently only supports LossInputType.LOGPROB.

Initialization

__call__( next_token_logits: nemo_rl.algorithms.loss.wrapper.Tensor, data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any], global_valid_seqs: nemo_rl.algorithms.loss.wrapper.Tensor | None, global_valid_toks: nemo_rl.algorithms.loss.wrapper.Tensor | None, ) → tuple[nemo_rl.algorithms.loss.wrapper.Tensor, dict[str, Any]]#: Compute loss for all packed sequences in one forward pass.

class nemo_rl.algorithms.loss.wrapper.DraftLossWrapper( loss_fn: Callable[..., tuple[torch.Tensor, dict[str, Any]]], prepare_fn: Optional[Callable[Any, Any]], data_dict: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any], loss_weight: float = 1.0, vocab_parallel_rank: Optional[int] = None, vocab_parallel_group: Optional[torch.distributed.ProcessGroup] = None, context_parallel_group: Optional[torch.distributed.ProcessGroup] = None, defer_normalization: bool = False, cu_seqlens_q: Optional[torch.Tensor] = None, cu_seqlens_q_padded: Optional[torch.Tensor] = None, d2t: Optional[torch.Tensor] = None, student_logits: Optional[torch.Tensor] = None, token_chunk_size: int = DEFAULT_DRAFT_TOKEN_CHUNK_SIZE, )#

Combine policy loss with draft soft cross-entropy loss.

Two layouts are supported:

Unpacked (default): prepare_fn (prepare_loss_input with LossInputType.DRAFT) builds the shifted teacher logits and token mask from the [B, S] batch, and the student logits come from data["student_logits"].
Packed (cu_seqlens_q given): the draft cross-entropy is computed once over the packed [1, T_packed] layout. The teacher logits are left-shifted within each packed segment and restricted to the draft vocabulary, and the token mask is packed with the same per-sequence shift; student_logits must be passed explicitly (it is kept out of the data dict so the packing loss wrappers never slice it).

Initialization

_packed_draft_loss( next_token_logits: torch.Tensor, data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any], global_valid_seqs: torch.Tensor | None, global_valid_toks: torch.Tensor, ) → torch.Tensor#

__call__(

next_token_logits: torch.Tensor,

data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any],

global_valid_seqs: torch.Tensor | None,

global_valid_toks: torch.Tensor | None,

**kwargs: Any,

) → tuple[torch.Tensor, dict[str, Any]]#

nemo_rl.algorithms.loss.wrapper.wrap_loss_fn_with_input_preparation( next_token_logits: nemo_rl.algorithms.loss.wrapper.Tensor, data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any], global_valid_seqs: nemo_rl.algorithms.loss.wrapper.Tensor | None, global_valid_toks: nemo_rl.algorithms.loss.wrapper.Tensor | None, loss_fn: nemo_rl.algorithms.loss.interfaces.LossFunction, prepare_fn: Callable[Any, Any], vocab_parallel_rank: Optional[int] = None, vocab_parallel_group: Optional[torch.distributed.ProcessGroup] = None, context_parallel_group: Optional[torch.distributed.ProcessGroup] = None, ) → tuple[nemo_rl.algorithms.loss.wrapper.Tensor, dict[str, Any]]#: Wraps a loss function to handle input preparation for megatron policy worker.

class nemo_rl.algorithms.loss.wrapper.DraftRuntimeLossWrapper( loss_fn: Callable[..., tuple[torch.Tensor, dict[str, Any]]], prepare_fn: Callable[..., Any], draft_runtime: Any, )#

Combine policy loss with a co-trained drafter loss (automodel path).

Structural sibling of DraftLossWrapper: the policy loss is computed exactly as without a draft, then the draft runtime (dspark/dflash block drafter or eagle3 TTT drafter) runs the draft forward on the hidden states and raw teacher logits captured during the policy forward (both detached, so the policy trunk’s gradients are unchanged).

No context-parallel correction is needed here: the trainer’s outer backward multiplier (dp*cp/cp_gradient_fanout, see LossPostProcessor) is applied uniformly to the whole combined loss, and that multiplier is calibrated for a loss that’s replicated across CP ranks (each local model contributes to cp_size loss consumers, hence the /cp_gradient_fanout correction). The draft loss shares that same replicated shape under CP – its teacher/hidden inputs are allgathered and computed identically on every CP rank – so the policy-calibrated outer multiplier already gives it the correct per-replica gradient scale with no further adjustment.

Initialization

__call__( next_token_logits: torch.Tensor, data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any], global_valid_seqs: torch.Tensor | None, global_valid_toks: torch.Tensor | None, ) → tuple[torch.Tensor, dict[str, Any]]#

nemo_rl.algorithms.loss.wrapper#

Module Contents#

Classes#

Functions#

Data#

API#

`nemo_rl.algorithms.loss.wrapper`#