[DM05 Training] CUDA OOM on 8× RTX 4090D in both VLA Arena and RoboTwin

[lyc-7]

Hi, I encountered CUDA OOM when training DM05 on 8 × NVIDIA GeForce RTX 4090D (24 GB) in both VLA Arena and RoboTwin.

1. VLA Arena:
   The error occurs during FSDP initialization before the first training step:
   accelerator.prepare(self.model) → FSDP(model) → handle.shard() → shard = chunk.clone()

   torch.OutOfMemoryError: CUDA out of memory.
   Tried to allocate 1.36 GiB.
   GPU 0 has a total capacity of 23.65 GiB of which 1.36 GiB is free.
   Process has 22.29 GiB memory in use.
   Of the allocated memory 21.73 GiB is allocated by PyTorch.
   
   

2. RoboTwin:
   I encountered almost exactly the same CUDA OOM during training:

   torch.OutOfMemoryError: CUDA out of memory.
   Tried to allocate 1.36 GiB.
   GPU 7 has a total capacity of 23.65 GiB of which 1.36 GiB is free.
   Process 282836 has 22.29 GiB memory in use.
   Of the allocated memory 21.73 GiB is allocated by PyTorch.
   

[hbzfeng]

Thanks for the report. The current VLA-Arena and RoboTwin default training entries use Full SFT. The traceback indicates that 24 GB RTX 4090D GPUs run out of memory during FSDP initialization, before the first training step.

For RTX 4090D training, we recommend adapting the existing LoRA examples. The current LIBERO and SO101 LoRA example configurations explicitly use eager for LLM attention and sdpa for vision and action attention.

For the current default Full SFT configuration, we recommend H20 or other GPUs with larger per-device memory.

[michaelmanly]

@lyc-7
are you specifically trying to keep this as full sft, or would lora work for what you're doing?

if you need full sft, i can spend a few bucks testing the training workload on a higher-vram gpu and see what setup actually gets through init + the first training steps.