diff --git a/nemo_rl/algorithms/loss/utils.py b/nemo_rl/algorithms/loss/utils.py index a811bdf0db9..98921429d36 100644 --- a/nemo_rl/algorithms/loss/utils.py +++ b/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, ) ) diff --git a/nemo_rl/distributed/model_utils.py b/nemo_rl/distributed/model_utils.py index 2e7791da4fc..e8df18ab87a 100644 --- a/nemo_rl/distributed/model_utils.py +++ b/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] diff --git a/nemo_rl/models/megatron/train.py b/nemo_rl/models/megatron/train.py index 5ee01e422dc..d9a1ca6dc70 100644 --- a/nemo_rl/models/megatron/train.py +++ b/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 diff --git a/tests/unit/distributed/test_distributed_logprob.py b/tests/unit/distributed/test_distributed_logprob.py index a4cae57a394..67b8390a2d0 100644 --- a/tests/unit/distributed/test_distributed_logprob.py +++ b/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)