nemo_rl.algorithms.loss.loss_input#

Builds the loss-function input dict from model logits.

Module Contents#

Functions#

_prepare_opd_full_loss_input

Build the full-vocabulary MOPD loss input from student logits + teacher payload.

prepare_loss_input

Prepare loss input for a loss function.

prepare_packed_loss_input

Prepare loss input from packed logits in a single fused pass.

API#

nemo_rl.algorithms.loss.loss_input._prepare_opd_full_loss_input(
logits: torch.Tensor,
data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any],
loss_fn: nemo_rl.algorithms.loss.interfaces.LossFunction,
*,
vocab_parallel_rank: Optional[int],
vocab_parallel_group: Optional[torch.distributed.ProcessGroup],
context_parallel_group: Optional[torch.distributed.ProcessGroup],
sampling_params: Optional[nemo_rl.algorithms.logits_sampling_utils.TrainingSamplingParams],
chunk_size: Optional[int],
teacher_output_layer_weight_by_index: Optional[dict[int, torch.Tensor]] = None,
) → dict[str, Any]#

Build the full-vocabulary MOPD loss input from student logits + teacher payload.

Runs the distributed reverse-KL kernel here rather than inside the loss so the loss stays free of process-group plumbing, mirroring the DISTILLATION branch.

Parameters:

  • logits – Student vocabulary-parallel logits [B, S_local, V_local].

  • data – Microbatch carrying the teacher payload column.

  • loss_fn – The opd_full-configured loss function.

  • vocab_parallel_rank – Vocabulary-parallel rank.

  • vocab_parallel_group – Vocabulary-parallel process group.

  • context_parallel_group – Context-parallel process group.

  • sampling_params – Training sampling params for the sampled-token logprobs.

  • chunk_size – Sequence-dim chunk size for the sampled-token logprobs.

  • teacher_output_layer_weight_by_index – Per-teacher LM-head shards keyed by index, for the hidden-state path.

Returns:

Loss input dict with the per-token divergence and, when requested, the entropy/cross-entropy decomposition.

Raises:

  • ValueError – If the teacher payload column is missing.

  • NotImplementedError – If no vocabulary-parallel group is available.

nemo_rl.algorithms.loss.loss_input.prepare_loss_input(
logits: torch.Tensor,
data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any],
loss_fn: nemo_rl.algorithms.loss.interfaces.LossFunction,
vocab_parallel_rank: Optional[int] = None,
vocab_parallel_group: Optional[torch.distributed.ProcessGroup] = None,
context_parallel_group: Optional[torch.distributed.ProcessGroup] = None,
sampling_params: Optional[nemo_rl.algorithms.logits_sampling_utils.TrainingSamplingParams] = None,
d2t: Optional[torch.Tensor] = None,
chunk_size: Optional[int] = None,
cp_sharder: Optional[nemo_automodel.components.distributed.context_parallel.ContextParallelSharder] = None,
teacher_output_layer_weight_by_index: Optional[dict[int, torch.Tensor]] = None,
) → tuple[dict[str, Any], nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any]]#

Prepare loss input for a loss function.

Parameters:

  • logits – Logits from the model.

  • data – Microbatch data. Will be updated if sampling_params is not None.

  • loss_fn – Loss function.

  • vocab_parallel_rank – Vocab parallel rank.

  • vocab_parallel_group – Vocab parallel group.

  • context_parallel_group – Context parallel group.

  • sampling_params – Sampling parameters.

  • d2t – Draft to target token mapping.

  • chunk_size – Sequence-dim chunk size for the vocab-parallel logprob computation (policy.logprob_chunk_size); avoids materializing full-size float32 logits during training.

  • cp_sharder – Automodel ContextParallelSharder owning this forward’s sequence layout; set only when cp_size > 1. logits are then this rank’s CP-local shard while data stays canonical.

  • teacher_output_layer_weight_by_index – This TP rank’s [V_local, H_teacher] teacher LM-head shards, keyed by the stable index rows are tagged with. The opd_full hidden-state path projects each row’s teacher payload into teacher logits with them.

.. rubric:: Notes

vocab_parallel_rank, vocab_parallel_group, context_parallel_group are only used for megatron policy worker. sampling_params is only used for LossInputType.LOGPROB, and currently only supported for ClippedPGLossFn. d2t is only used for LossInputType.DRAFT. teacher_output_layer_weight_by_index is only used for LossInputType.OPD_FULL.

Returns:

tuple(loss_input, maybe_updated_data)

nemo_rl.algorithms.loss.loss_input.prepare_packed_loss_input(
logits: torch.Tensor,
data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any],
loss_fn: nemo_rl.algorithms.loss.interfaces.LossFunction,
cu_seqlens_q: nemo_rl.utils.sequence_lengths.CpuIntTuple,
cu_seqlens_q_padded: nemo_rl.utils.sequence_lengths.CpuIntTuple,
vocab_parallel_rank: Optional[int] = None,
vocab_parallel_group: Optional[torch.distributed.ProcessGroup] = None,
context_parallel_group: Optional[torch.distributed.ProcessGroup] = None,
sampling_params: Optional[nemo_rl.algorithms.logits_sampling_utils.TrainingSamplingParams] = None,
chunk_size: Optional[int] = None,
) → tuple[dict[str, Any], nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any]]#

Prepare loss input from packed logits in a single fused pass.

Unlike prepare_loss_input which operates on a single (unpacked) sequence, this function computes log probabilities from packed logits across all sequences at once using from_parallel_logits_to_logprobs_packed_sequences.

Currently only supports LossInputType.LOGPROB.

Parameters:

  • logits – Packed logits from the model [1, T_packed // CP, V // TP].

  • data – Microbatch data (unpacked, [B, S]).

  • loss_fn – Loss function (must have input_type == LossInputType.LOGPROB).

  • cu_seqlens_q – Unpadded cumulative sequence lengths [B+1]. CPU-resident integer tuples are required at this host-side loss boundary.

  • cu_seqlens_q_padded – Padded cumulative sequence lengths [B+1]. CPU-resident integer tuples are required at this host-side loss boundary.

  • vocab_parallel_rank – Vocab parallel rank.

  • vocab_parallel_group – Vocab parallel group.

  • context_parallel_group – Context parallel group.

  • sampling_params – Sampling parameters.

  • chunk_size – Sequence-dim chunk size for the logprob computation (policy.logprob_chunk_size); avoids materializing full-size float32 logits during training.

Returns:

tuple(loss_input, maybe_updated_data)