Release test failing with CUDA OOM at step 16/20

Release test failing off main branch with CUDA OOM at step 16/20

llm_grpo_gemma3_27b_it_8n8g_fsdp2tp8_actckpt_long

Error:

`torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 168.00 MiB. GPU 0 has a total capacity of 79.11 GiB of which 149.38 MiB is free. Process 1087837 has 5.05 GiB memory in use. Including non-PyTorch memory, this process has 73.81 GiB memory in use. Of the allocated memory 69.37 GiB is allocated by PyTorch, and 1.30 GiB 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)`