perf(megatron): keep policy logits in model precision

[YolandaLyj]

Summary

• keep Megatron log-prob forward outputs in BF16/FP16 instead of materializing the complete FP32 logits tensor
• do the same for the normal policy-loss training forward while preserving FP32 outputs for value and other losses
• reuse the existing chunked log-prob path, which converts each logits chunk to FP32 immediately before fused vocab-parallel cross-entropy

Motivation

Float16Module.forward() defaults to fp32_output=True. For policy scoring with labels=None, that upcasts the complete per-rank [T, V/TP] logits tensor before log_probs_chunk_size can take effect. Large vocabularies and long packed trajectories can therefore OOM on the full FP32 output even though the loss already processes logits in chunks.

With this change, the normal path is:

BF16 logits [T, V/TP]
  -> BF16 chunk [C, V/TP]
  -> FP32 chunk [C, V/TP]
  -> fused vocab-parallel cross-entropy

Critic/value paths retain the existing FP32 output behavior.

The combined 1F1B schedule is not changed here because its Megatron PostProcessNode performs a separate unconditional float16_to_fp32; this PR only changes the normal model-forward paths.

Tests

Added regression coverage:

• actor log-prob scoring requests model-precision output with fp32_output=False
• BF16 logits are split before each chunk is converted to FP32

G2 GPU validation passed on 1x NVIDIA H200 with PyTorch 2.11.0+cu129 and Megatron TP=1:

• Training accuracy with loss checkpointing ([512, 8192], chunk size 128):
  • loss absolute difference: 0
  • BF16 gradient max / mean absolute difference: 0 / 0
• Training memory ([2048, 32768], chunk size 256):
  • old full-FP32 output + loss checkpoint: forward 412.01 MiB, forward+backward 768.00 MiB
  • BF16 output + loss checkpoint: forward 32.02 MiB, forward+backward 304.03 MiB
  • BF16 output without loss checkpoint: forward 256.04 MiB, forward+backward 304.03 MiB

The intended memory-saving training configuration therefore combines a positive --log-probs-chunk-size with --recompute-loss-function.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[YolandaLyj]

G2 validation (1x NVIDIA H200, PyTorch 2.11.0+cu129, Megatron TP=1) passed.

• Actor wiring: MegatronTrainRayActor.compute_log_prob() forwarded fp32_output=False.
• Chunk conversion: BF16 [5, 8], chunk size 2 reached fused CE as FP32 chunks [2, 8], [2, 8], [1, 8].
• Forward-only accuracy/memory ([4096, 32768], chunk size 256):
  • max/mean abs log-prob diff vs full-FP32 path: 9.5367e-7 / 2.8173e-8
  • CUDA peak delta: 1596.10 MiB -> 156.01 MiB (1440.09 MiB saved)
• Training accuracy with loss checkpoint ([512, 8192], chunk size 128):
  • loss abs diff: 0
  • BF16 gradient max/mean abs diff: 0 / 0
• Training memory ([2048, 32768], chunk size 256):
  • old full-FP32 output + loss checkpoint: forward 412.01 MiB, forward+backward 768.00 MiB
  • BF16 output + loss checkpoint: forward 32.02 MiB, forward+backward 304.03 MiB
  • BF16 output without loss checkpoint: forward 256.04 MiB, forward+backward 304.03 MiB

The training result confirms the intended combination is BF16 model output plus --recompute-loss-function and a positive --log-probs-chunk-size.

[yueming-yuan]

fixes bugs and added UT & snapshot test. approved now