nemo_rl.algorithms.loss.utils

Module Contents

Functions

rescale_loss_metrics	Change global loss denominators while preserving raw counts and extrema.
map_teacher_logits_to_draft_vocab	Restrict full-vocab teacher logits to the draft vocabulary via d2t.
roll_packed_seq_dim	Left-shift a packed tensor by one along seq_dim within each segment.
pack_rolled_draft_token_mask	Build the packed draft-loss mask [1, T_packed] from unpacked masks.
_project_hidden_states_per_teacher	Project each sample’s hidden states through its own teacher’s LM head.
reconstruct_opd_full_teacher_logits	Turn the transported teacher payload into this rank’s teacher logit shard.
_pack_input_ids	Pack input_ids from [B, S] to [1, T_packed // CP] using sequence boundaries.

API

nemo_rl.algorithms.loss.utils.rescale_loss_metrics(

metrics: dict[str, Any],

normalizers: dict[str, nemo_rl.algorithms.loss.interfaces.MetricNormalizer],

*,

token_factor: float,

sequence_factor: float,

) → dict[str, Any]

Change global loss denominators while preserving raw counts and extrema.

Metrics absent from normalizers are returned unchanged. Unlike split-API normalization of raw sums, this only adjusts denominators explicitly advertised by the loss.

nemo_rl.algorithms.loss.utils.map_teacher_logits_to_draft_vocab(

teacher_logits: torch.Tensor,

d2t: Optional[torch.Tensor],

vocab_parallel_rank: Optional[int] = None,

vocab_parallel_group: Optional[torch.distributed.ProcessGroup] = None,

) → torch.Tensor

Restrict full-vocab teacher logits to the draft vocabulary via d2t.

d2t maps draft-vocab index i to target-vocab index i + d2t[i]. Under tensor parallelism the teacher logits arrive vocab-sharded, so they are gathered to the full vocab and re-sliced to this rank’s shard of the draft vocabulary (the draft output layer is sharded the same way). No-op when d2t is None (full-vocab drafts).

nemo_rl.algorithms.loss.utils.roll_packed_seq_dim(

tensor: torch.Tensor,

cu_seqlens_padded: torch.Tensor,

seq_dim: int,

) → torch.Tensor

Left-shift a packed tensor by one along seq_dim within each segment.

Equivalent to a per-sequence torch.roll(shifts=-1) over the packed layout: one global roll followed by zeroing each segment’s final slot (the only positions where the global roll would leak the next segment’s first row). Segment boundaries come from cu_seqlens_padded, the physical offsets of the packed layout.

nemo_rl.algorithms.loss.utils.pack_rolled_draft_token_mask(

token_mask: torch.Tensor,

sample_mask: torch.Tensor,

cu_seqlens: torch.Tensor,

cu_seqlens_padded: torch.Tensor,

) → torch.Tensor

Build the packed draft-loss mask [1, T_packed] from unpacked masks.

Mirrors the non-packed DRAFT prepare (token_mask left-shifted by one, scaled by sample_mask), laid out at each sequence’s padded offset. Each sequence’s last real slot (whose shifted target would cross the boundary) and all padding slots stay zero.

Reuses _pack_input_ids for the padded-offset layout (including its clamp for bin-alignment padding absorbed into the last sequence’s effective length) and roll_packed_seq_dim for the boundary-safe per-segment left shift.

nemo_rl.algorithms.loss.utils._project_hidden_states_per_teacher(

payload: torch.Tensor,

weight_by_index: dict[int, torch.Tensor],

teacher_index: torch.Tensor,

) → torch.Tensor

Project each sample’s hidden states through its own teacher’s LM head.

payload is [B, S, H]; teacher_index is [B] int, naming which loaded teacher LM head produced each row (see OPD_FULL_TEACHER_INDEX_FIELD). Grouped by teacher so this costs one matmul per distinct teacher present in the microbatch, not a per-sample loop – when sequence packing is on, each call already has B == 1 (one sequence at a time), so it takes the single-teacher path and never materializes the scatter buffer.

Raises:

ValueError – If teacher_index does not carry one entry per payload row, if a row’s teacher_index has no loaded LM head, or if an LM head disagrees with the payload or with the other teachers’.

nemo_rl.algorithms.loss.utils.reconstruct_opd_full_teacher_logits(

payload: torch.Tensor,

*,

teacher_payload: str,

student_logits: torch.Tensor,

vocab_parallel_rank: Optional[int],

context_parallel_group: Optional[torch.distributed.ProcessGroup],

teacher_output_layer_weight_by_index: Optional[dict[int, torch.Tensor]] = None,

teacher_index: Optional[torch.Tensor] = None,

) → torch.Tensor

Turn the transported teacher payload into this rank’s teacher logit shard.

The payload is canonical [B, S, D] (already TP/CP-gathered on the teacher side), so it is re-sharded onto the student’s own CP window here, exactly as from_parallel_logits_to_logprobs does for the student logits.

Parameters:

• payload – Teacher payload [B, S, D] – hidden states or full logits.

• teacher_payload – "hidden_states" or "logits".

• student_logits – Student logits [B, S_local, V_local], used for the target device, dtype-independent shapes, and CP geometry.

• vocab_parallel_rank – This rank’s vocabulary-parallel rank.

• context_parallel_group – Context-parallel process group, if any.

• teacher_output_layer_weight_by_index – [V_local, H_teacher] teacher LM-head shards for the hidden-state path, keyed by the stable index rows are tagged with. Routed by teacher_index whenever that column is present.

• teacher_index – [B] int, this microbatch’s per-row teacher index (see OPD_FULL_TEACHER_INDEX_FIELD). Required once teacher_output_layer_weight_by_index holds more than one entry, and validated against the loaded heads whenever present.

Returns:

Teacher logits [B, S_local, V_local] aligned with student_logits.

Raises:

ValueError – If the payload and student shard cannot be aligned, if the hidden-state path is used without a teacher LM-head shard, or if several shards are loaded without a teacher_index to route by.

nemo_rl.algorithms.loss.utils._pack_input_ids(

input_ids: torch.Tensor,

cu_seqlens_q: nemo_rl.utils.sequence_lengths.CpuIntTuple,

cu_seqlens_q_padded: nemo_rl.utils.sequence_lengths.CpuIntTuple,

cp_rank: int = 0,

cp_size: int = 1,

roll_shift: int = 0,

) → torch.Tensor

Pack input_ids from [B, S] to [1, T_packed // CP] using sequence boundaries.

Each sequence is individually padded to its padded length (from cu_seqlens_q_padded), optionally rolled, and CP-sharded at that padded length before being placed into the packed output. This matches how Megatron packs and CP-shards sequences in _pack_sequences_for_megatron.

Parameters:

• input_ids – Unpacked input IDs [B, S].

• cu_seqlens_q – Unpadded cumulative sequence lengths [B+1].

• cu_seqlens_q_padded – Padded cumulative sequence lengths [B+1].

• cp_rank – Context parallelism rank.

• cp_size – Context parallelism size.

• roll_shift – If non-zero, roll each padded sequence by this amount before CP-sharding. Use -1 to build shifted targets for next-token prediction.