[bug] `policy.logprob_chunk_size` is ignored by the training loss path which leads to increased memory

[kaloyan-inherent]

Describe the bug

Observed OOM with Megatron backend and large context length in the train forward pass:

MegatronPolicyWorker.train
  → SequencePackingLossWrapper.__call__ → prepare_loss_input
  → get_next_token_logprobs_from_logits   (model_utils.py:1383)
  → from_parallel_logits_to_logprobs      (model_utils.py:1012, un-chunked branch)
  → DistributedLogprob.forward → _compute_distributed_log_softmax (model_utils.py:59)
      sum_exp_logits = vocab_parallel_logits.exp().sum(-1, keepdim=True).float()
torch.OutOfMemoryError: Tried to allocate 12.60 GiB. GPU had 12.44 GiB free of 79.19

Steps/Code to reproduce bug

Run GRPO training with the Megatron training backend.
policy.logprob_chunk_size is honoured by the inference logprob paths (prev/reference logprobs via LogprobsPostProcessor) but never reaches the training loss path. During the loss forward/backward, current-policy logprobs are always computed by the un-chunked DistributedLogprob on an eagerly fp32-cast, full-sequence logits tensor — even with logprob_chunk_size and defer_fp32_logits set ; which increases memory consumption unnecessarily.

Expected behavior

logprob_chunk_size applies during training.

Additional context

Addressed in #2872