PPO default yaml OOM

[yuki-97]

PPO default yaml OOM for both backend:

1. DTensor: examples/configs/ppo_math_1B.yaml
2. Megatron: examples/configs/ppo_math_1B_megatron.yaml
3. SC + Megatron: examples/configs/ppo_math_1B_megatron_single_controller.yaml (will be added after feat(sc): support PPO in single controller #3773)

run with uv run python examples/run_ppo.py --config examples/configs/ppo_math_1B.yaml.

error log:

ray.exceptions.RayTaskError(OutOfMemoryError): ray::DTensorPolicyWorkerV2.get_logprobs() (pid=2671339, ip=10.65.16.209, actor_id=d5836b60449efbd220b3c2f406000000, repr=DTensorPolicyWorkerV2[rank=0])
  File "/root/.local/share/uv/python/cpython-3.13.14-linux-x86_64-gnu/lib/python3.13/concurrent/futures/_base.py", line 453, in result
    return self.__get_result()
           ~~~~~~~~~~~~~~~~~^^
  File "/root/.local/share/uv/python/cpython-3.13.14-linux-x86_64-gnu/lib/python3.13/concurrent/futures/_base.py", line 402, in __get_result
    raise self._exception
  File "/lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_nemorl/users/yukih/reinforcer-backup/nemo_rl/utils/nsys.py", line 119, in wrapper
    ret = func(*args, **kwargs)
  File "/lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_nemorl/users/yukih/reinforcer-backup/nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py", line 653, in get_logprobs
    token_logprobs, _metrics, _ = forward_with_post_processing_fn(
                                  ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^
        model=self.model,
        ^^^^^^^^^^^^^^^^^
    ...<5 lines>...
        sequence_dim=sequence_dim,
        ^^^^^^^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "/lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_nemorl/users/yukih/reinforcer-backup/nemo_rl/models/automodel/train.py", line 401, in forward_with_post_processing_fn
    result = post_processing_fn(
        logits=logits,
    ...<4 lines>...
        sequence_dim=sequence_dim,
    )
  File "/lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_nemorl/users/yukih/reinforcer-backup/nemo_rl/models/automodel/train.py", line 755, in __call__
    token_logprobs = self._compute_local_logprobs(
        logits, processed_inputs.input_ids
    )
  File "/lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_nemorl/users/yukih/reinforcer-backup/nemo_rl/models/automodel/train.py", line 838, in _compute_local_logprobs
    log_probs = torch.nn.functional.log_softmax(logits, dim=-1)
  File "/opt/ray_venvs/nemo_rl.models.policy.workers.dtensor_policy_worker_v2.DTensorPolicyWorkerV2/lib/python3.13/site-packages/torch/nn/functional.py", line 2267, in log_softmax
    ret = input.log_softmax(dim)
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 37.09 GiB. GPU 0 has a total capacity of 79.11 GiB of which 10.58 GiB is free. Process 2669658 has 1.35 GiB memory in use. Including non-PyTorch memory, this process has 66.43 GiB memory in use. Process 2672344 has 720.00 MiB memory in use. Of the allocated memory 64.36 GiB is allocated by PyTorch, and 875.71 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://docs.pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-cuda-alloc-conf)

[fujial-code]

I'll follow up on this issue.