GPT-OSS 20B Megatron SFT fails due to OOM after X steps

**Describe the bug**

Can someone please help me understand what could cause the job to fail after some steps and not at first step?

I've been trying to perform SFT on gpt-oss-20b with Megatron backend and openMathInstruct2 dataset. The job fails due to OOM after Step 683 consistently.

```
step 683/3710 =========================
▶ Preparing batch...
▶ Taking a training step...
Traceback (most recent call last):
File "/opt/nemo-rl/examples/run_sft.py", line 192, in
main()
File "/opt/nemo-rl/examples/run_sft.py", line 178, in main
sft_train(
File "/opt/nemo-rl/nemo_rl/algorithms/sft.py", line 455, in sft_train
train_results = policy.train(train_data, loss_fn)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/nemo-rl/nemo_rl/models/policy/lm_policy.py", line 532, in train
results = self.worker_group.get_all_worker_results(futures)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/nemo-rl/nemo_rl/distributed/worker_groups.py", line 953, in get_all_worker_results
return future_bundle.get_results(
^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/nemo-rl/nemo_rl/distributed/worker_groups.py", line 103, in get_results
all_results = ray.get(object_refs)
^^^^^^^^^^^^^^^^^^^^
File "/opt/nemo_rl_venv/lib/python3.12/site-packages/ray/_private/auto_init_hook.py", line 22, in auto_init_wrapper
return fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^
File "/opt/nemo_rl_venv/lib/python3.12/site-packages/ray/_private/client_mode_hook.py", line 104, in wrapper
return func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/opt/nemo_rl_venv/lib/python3.12/site-packages/ray/_private/worker.py", line 2882, in get
values, debugger_breakpoint = worker.get_objects(object_refs, timeout=timeout)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/nemo_rl_venv/lib/python3.12/site-packages/ray/_private/worker.py", line 968, in get_objects
raise value.as_instanceof_cause()
ray.exceptions.RayTaskError(OutOfMemoryError): ray::MegatronPolicyWorker.train() (pid=132277, ip=100.64.106.98, actor_id=02695ae32a93c4eaa4efbb7909000000, repr=MegatronPolicyWorker[rank=1])
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/nemo-rl/nemo_rl/utils/nsys.py", line 88, in wrapper
ret = func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/opt/nemo-rl/nemo_rl/models/policy/workers/megatron_policy_worker.py", line 1146, in train
losses_reduced = forward_backward_func(
^^^^^^^^^^^^^^^^^^^^^^
File "/opt/nemo-rl/3rdparty/Megatron-LM-workspace/Megatron-LM/megatron/core/pipeline_parallel/schedules.py", line 2257, in forward_backward_pipelining_without_interleaving
input_tensor_grad = backward_step(
^^^^^^^^^^^^^^
File "/opt/nemo-rl/3rdparty/Megatron-LM-workspace/Megatron-LM/megatron/core/pipeline_parallel/schedules.py", line 480, in backward_step
custom_backward(output_tensor[0], output_tensor_grad[0])
File "/opt/nemo-rl/3rdparty/Megatron-LM-workspace/Megatron-LM/megatron/core/pipeline_parallel/schedules.py", line 173, in custom_backward
Variable._execution_engine.run_backward(
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 2.38 GiB. GPU 0 has a total capacity of 139.81 GiB of which 2.33 GiB is free. Including non-PyTorch memory, this process has 137.47 GiB memory in use. Of the allocated memory 132.28 GiB is allocated by PyTorch, and 782.95 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://pytorch.org/docs/stable/notes/cuda.html#environment-variables)
```

**Steps/Code to reproduce bug**
1. NemoRL image created at e480f8f
2. Kick off SFT job with following overrides
Override config:
```
uv run examples/run_sft.py \
--config examples/configs/sft_openmathinstruct2_megatron.yaml \
+data.train.processor=sft_processor \
policy.model_name=openai/gpt-oss-20b \
policy.tokenizer.name=openai/gpt-oss-20b \
policy.train_global_batch_size=256 \
policy.train_micro_batch_size=1 \
policy.sequence_packing.enabled=false \
policy.dtensor_cfg.enabled=false \
policy.megatron_cfg.enabled=true \
policy.megatron_cfg.sequence_parallel=false \
policy.megatron_cfg.expert_model_parallel_size=8 \
policy.megatron_cfg.tensor_model_parallel_size=1 \
policy.megatron_cfg.pipeline_model_parallel_size=2 \
policy.make_sequence_length_divisible_by=1 \
policy.megatron_cfg.freeze_moe_router=true \
policy.megatron_cfg.moe_router_dtype=fp64 \
policy.megatron_cfg.moe_router_load_balancing_type=aux_loss \
policy.megatron_cfg.moe_router_bias_update_rate=1e-3 \
+policy.megatron_cfg.env_vars.NRL_MEGATRON_CHECKPOINT_DIR=/fsx/models/megatron/gpt-oss-20b \
++policy.megatron_cfg.env_vars.PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
policy.megatron_cfg.optimizer.use_distributed_optimizer=true \
policy.megatron_cfg.optimizer.use_precision_aware_optimizer=true \
policy.megatron_cfg.optimizer.params_dtype=float32 \
logger.wandb_enabled=false \
logger.tensorboard_enabled=false \
logger.mlflow_enabled=true \
logger.mlflow.experiment_name=nemo-rl-experiments \
logger.mlflow.run_name=gpt-oss-20b-test-sft-run \
+logger.mlflow.tracking_uri= \
cluster.gpus_per_node=8 \
cluster.num_nodes=2 \
checkpointing.checkpoint_dir=/fsx/ws/dpo-k8s/checkpoints-megatron/gpt-oss-sft-20b/nemo-sft-rayjob \
checkpointing.keep_top_k=5

```

**Expected behavior**

SFT job to finish running successfully.

**Additional context**

- The job failed at 21st step consistently, so I experimented with increasing `pipeline_model_parallel_size` to 2 from 1. That helped the job run longer, but now fails at Step 683.
- Similarly, reducing the `global_batch_size` from 512 to 256 helped postpone the failure, but not fix the issue.
- I tried setting `PYTORCH_CUDA_ALLOC_CONF` env variable to `expandable_segments:True`. That don't seem to help either.
- Instance type: p5en.48xlarge (H200 GPUs)
- GPU utilization of the pods don't seem to go beyond 90%.