fix(megatron): honor policy.logprob_chunk_size in the training loss path

nemo_rl/algorithms/loss/utils.py

@@ -40,6 +40,7 @@ def prepare_loss_input(
     context_parallel_group: Optional[torch.distributed.ProcessGroup] = None,
     sampling_params: Optional[TrainingSamplingParams] = None,
     d2t: Optional[torch.Tensor] = None,
+    chunk_size: Optional[int] = None,
 ) -> tuple[dict[str, Any], BatchedDataDict[Any]]:
     """Prepare loss input for a loss function.
 
@@ -52,6 +53,9 @@ def prepare_loss_input(
         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.
 
     Notes:
         vocab_parallel_rank, vocab_parallel_group, context_parallel_group are only used for megatron policy worker.
@@ -80,6 +84,7 @@ def prepare_loss_input(
                 vocab_parallel_group=vocab_parallel_group,
                 context_parallel_group=context_parallel_group,
                 sampling_params=sampling_params,
+                chunk_size=chunk_size,
             )
 
         # handle top-k/top-p filtering for logprobs, only used for ClippedPGLossFn now
@@ -102,6 +107,8 @@ def prepare_loss_input(
                     vocab_parallel_group=vocab_parallel_group,
                     context_parallel_group=context_parallel_group,
                     sampling_params=None,  # no filtering
+                    # Only reachable with top-k/top-p sampling active that has its own kernel path so don't chunk here
+                    chunk_size=None,
                 )
 
         loss_input = {"next_token_logprobs": logprobs}
@@ -238,6 +245,7 @@ def prepare_packed_loss_input(
     vocab_parallel_group: Optional[torch.distributed.ProcessGroup] = None,
     context_parallel_group: Optional[torch.distributed.ProcessGroup] = None,
     sampling_params: Optional[TrainingSamplingParams] = None,
+    chunk_size: Optional[int] = None,
 ) -> tuple[dict[str, Any], BatchedDataDict[Any]]:
     """Prepare loss input from packed logits in a single fused pass.
 
@@ -257,6 +265,9 @@ def prepare_packed_loss_input(
         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)
@@ -296,8 +307,15 @@ def prepare_packed_loss_input(
         roll_shift=-1,
     )
 
+    # With chunking, keep logits in their original dtype: the chunked logprob
+    # kernel casts each chunk to float32 internally.
+    use_chunking = chunk_size is not None and not need_top_k_or_top_p_filtering(
+        sampling_params
+    )
+    logits_for_logprobs = logits if use_chunking else logits.to(torch.float32)
+
     logprobs = from_parallel_logits_to_logprobs_packed_sequences(
-        logits.to(torch.float32),
+        logits_for_logprobs,
         packed_rolled_targets,
         cu_seqlens_q_padded,
         unpacked_seqlen,
@@ -307,6 +325,7 @@ def prepare_packed_loss_input(
         inference_only=False,
         cp_group=context_parallel_group,
         sampling_params=sampling_params,
+        chunk_size=chunk_size if use_chunking else None,
         target_is_pre_rolled=True,
     )
 
@@ -322,7 +341,7 @@ def prepare_packed_loss_input(
         ):
             data["curr_logprobs_unfiltered"] = (
                 from_parallel_logits_to_logprobs_packed_sequences(
-                    logits.to(torch.float32),
+                    logits_for_logprobs,
                     packed_rolled_targets,
                     cu_seqlens_q_padded,
                     unpacked_seqlen,
@@ -332,6 +351,7 @@ def prepare_packed_loss_input(
                     inference_only=False,
                     cp_group=context_parallel_group,
                     sampling_params=None,
+                    chunk_size=chunk_size if use_chunking else None,
                     target_is_pre_rolled=True,
                 )
             )

nemo_rl/distributed/model_utils.py

@@ -326,9 +326,7 @@ def backward(
         seq_size = int(vocab_parallel_logits.shape[1])
         num_chunks = (seq_size + chunk_size - 1) // chunk_size
 
-        grad_input: torch.Tensor = torch.zeros_like(
-            vocab_parallel_logits, dtype=torch.float32
-        )
+        grad_input: torch.Tensor = torch.zeros_like(vocab_parallel_logits)
 
         for chunk_idx in range(num_chunks):
             chunk_start = chunk_idx * chunk_size
@@ -351,18 +349,14 @@ def backward(
                 num_classes=partition_vocab_size,
             )
 
-            # Inplace index into the preallocated grad_input tensor
-            grad_input_chunk = grad_input[:, chunk_start:chunk_end, :]
-
-            grad_input_chunk.copy_(
-                is_chosen.float().sub_(softmax_output)
-            )  # inplace copy
-            grad_input_chunk.mul_(
+            chunk_grad_fp32 = is_chosen.float().sub_(softmax_output)
+            chunk_grad_fp32.mul_(
                 grad_output[:, chunk_start:chunk_end].unsqueeze(dim=-1)
             )
+            grad_input[:, chunk_start:chunk_end, :].copy_(chunk_grad_fp32)
 
             # Explicitly free before next iteration allocates
-            del softmax_output, is_chosen, logits
+            del softmax_output, is_chosen, logits, chunk_grad_fp32
 
         # if you add an argument to the forward method, then you must add a corresponding None here
         return grad_input, None, None, None, None, None, None
@@ -1354,6 +1348,7 @@ def get_next_token_logprobs_from_logits(
     vocab_parallel_group: Optional[torch.distributed.ProcessGroup] = None,
     context_parallel_group: Optional[torch.distributed.ProcessGroup] = None,
     sampling_params: Optional[TrainingSamplingParams] = None,
+    chunk_size: Optional[int] = None,
 ) -> torch.Tensor:
     """Compute token log-probabilities from logits, handling parallel and non-parallel cases.
 
@@ -1370,11 +1365,20 @@ def get_next_token_logprobs_from_logits(
         vocab_parallel_group: Process group for vocab parallelism
         context_parallel_group: Process group for context parallelism
         sampling_params: Sampling parameters for top-k/top-p filtering
+        chunk_size: Sequence-dim chunk size for the vocab-parallel path; only
+            applied without top-k/top-p sampling.
 
     Returns:
         Token log-probabilities of shape [batch_size, seq_len - 1]
     """
-    next_token_logits = next_token_logits.to(torch.float32)
+    # ChunkedDistributedLogprob casts each chunk to float32 internally.
+    use_chunking = (
+        vocab_parallel_group is not None
+        and chunk_size is not None
+        and not need_top_k_or_top_p_filtering(sampling_params)
+    )
+    if not use_chunking:
+        next_token_logits = next_token_logits.to(torch.float32)
 
     if vocab_parallel_group is not None:
         assert vocab_parallel_rank is not None, (
@@ -1389,6 +1393,7 @@ def get_next_token_logprobs_from_logits(
             inference_only=False,
             cp_group=context_parallel_group,
             sampling_params=sampling_params,
+            chunk_size=chunk_size if use_chunking else None,
         )
         # slice off to the correct length to remove potential CP padding
         logprobs = logprobs[:, : input_ids.shape[1] - 1]

nemo_rl/models/megatron/train.py

@@ -414,8 +414,12 @@ def __call__(
             Callable: Function that takes output tensor and returns (loss, metrics) tuple
         """
         # wrap prepare_loss_input with sampling_params and optional d2t mapping
+        logprob_chunk_size = self.cfg.get("logprob_chunk_size", None)
         prepare_loss_input_wrapped = partial(
-            prepare_loss_input, sampling_params=self.sampling_params, d2t=self.d2t
+            prepare_loss_input,
+            sampling_params=self.sampling_params,
+            d2t=self.d2t,
+            chunk_size=logprob_chunk_size,
         )
 
         # wrap loss function with loss input preparation
@@ -425,7 +429,9 @@ def __call__(
             if fuse_loss:
                 wrapper_cls = SequencePackingFusionLossWrapper
                 prepare_fn = partial(
-                    prepare_packed_loss_input, sampling_params=self.sampling_params
+                    prepare_packed_loss_input,
+                    sampling_params=self.sampling_params,
+                    chunk_size=logprob_chunk_size,
                 )
             else:
                 wrapper_cls = SequencePackingLossWrapper

tests/unit/distributed/test_distributed_logprob.py

@@ -30,6 +30,7 @@
     ChunkedDistributedLogprob,
     DistributedLogprob,
     _compute_distributed_log_softmax,
+    get_next_token_logprobs_from_logits,
 )
 
 
@@ -403,3 +404,119 @@ def test_chunked_distributed_entropy(
         inference_only=inference_only,
     )
     distributed_test_runner(test_fn, world_size=tp_size)
+
+
+def _run_chunk_memory(rank, world_size, tp_size):
+    """Chunking must cut the vocab-parallel logprob call's peak memory by at
+    least one full-sequence fp32 logits copy (the eager cast it avoids)."""
+    tp_group = torch.distributed.new_group(ranks=list(range(tp_size)))
+
+    batch_size, seq_len, full_vocab_size = 1, 4096, 32768
+    vocab_part_size = full_vocab_size // tp_size
+    vocab_start_index = rank * vocab_part_size
+    vocab_end_index = (rank + 1) * vocab_part_size
+
+    torch.manual_seed(42)
+    full_logits = torch.randn(
+        batch_size, seq_len, full_vocab_size, device="cuda", dtype=torch.bfloat16
+    )
+    input_ids = torch.randint(0, full_vocab_size, (batch_size, seq_len), device="cuda")
+
+    def peak_bytes(chunk_size):
+        logits = (
+            full_logits[:, :, vocab_start_index:vocab_end_index]
+            .clone()
+            .detach()
+            .requires_grad_(True)
+        )
+        torch.cuda.synchronize()
+        torch.cuda.reset_peak_memory_stats()
+        floor = torch.cuda.memory_allocated()
+        logprobs = get_next_token_logprobs_from_logits(
+            input_ids=input_ids,
+            next_token_logits=logits,
+            vocab_parallel_rank=rank,
+            vocab_parallel_group=tp_group,
+            chunk_size=chunk_size,
+        )
+        logprobs.sum().backward()
+        torch.cuda.synchronize()
+        peak = torch.cuda.max_memory_allocated() - floor
+        del logits, logprobs
+        torch.cuda.empty_cache()
+        return peak
+
+    peak_full = peak_bytes(None)
+    peak_chunked = peak_bytes(512)
+    one_fp32_logits_copy = seq_len * vocab_part_size * 4
+    assert peak_full - peak_chunked > one_fp32_logits_copy, (
+        f"chunked={peak_chunked} full={peak_full} need_saving>{one_fp32_logits_copy}"
+    )
+
+
+@pytest.mark.parametrize("tp_size", [1, 2])
+def test_get_next_token_logprobs_chunking_reduces_memory(
+    distributed_test_runner, tp_size
+):
+    test_fn = functools.partial(_run_chunk_memory, tp_size=tp_size)
+    distributed_test_runner(test_fn, world_size=tp_size)
+
+
+def _run_chunk_equivalence(rank, world_size, tp_size, dtype):
+    """Forward logprobs and backward grad must match between chunk_size=None and chunk_size=32."""
+    tp_group = torch.distributed.new_group(ranks=list(range(tp_size)))
+
+    batch_size, seq_len, full_vocab_size = 2, 128, 2048
+    vocab_part_size = full_vocab_size // tp_size
+    vocab_start_index = rank * vocab_part_size
+    vocab_end_index = (rank + 1) * vocab_part_size
+
+    torch.manual_seed(42)
+    full_logits = torch.randn(
+        batch_size, seq_len, full_vocab_size, device="cuda", dtype=dtype
+    )
+    input_ids = torch.randint(0, full_vocab_size, (batch_size, seq_len), device="cuda")
+
+    def run(cs):
+        logits = (
+            full_logits[:, :, vocab_start_index:vocab_end_index]
+            .detach()
+            .clone()
+            .requires_grad_(True)
+        )
+        logprobs = get_next_token_logprobs_from_logits(
+            input_ids=input_ids,
+            next_token_logits=logits,
+            vocab_parallel_rank=rank,
+            vocab_parallel_group=tp_group,
+            chunk_size=cs,
+        )
+        logprobs.sum().backward()
+        return logprobs.detach(), logits.grad
+
+    logprobs_full, grad_full = run(None)
+    logprobs_chunked, grad_chunked = run(32)
+
+    torch.testing.assert_close(logprobs_chunked, logprobs_full, rtol=1e-6, atol=1e-6)
+    torch.testing.assert_close(
+        grad_chunked.to(torch.float32),
+        grad_full.to(torch.float32),
+        rtol=1e-6,
+        atol=1e-6 if dtype == torch.float32 else 5e-3,
+    )
+
+
+@pytest.mark.parametrize(
+    "tp_size, dtype",
+    [
+        (1, torch.float32),
+        (2, torch.float32),
+        (1, torch.bfloat16),
+        (2, torch.bfloat16),
+    ],
+)
+def test_get_next_token_logprobs_chunk_equivalence(
+    distributed_test_runner, tp_size, dtype
+):
+    test_fn = functools.partial(_run_chunk_equivalence, tp_size=tp_size, dtype=dtype)
+    distributed_test_runner(test_fn, world_size=tp_size)