Lightning-AI
How to make use of NVIDIA GH200 Grace Hopper Superchip
I have access to a GH200 gpu and I'm trying to do model pretraining but when running the pretrain command i get Cuda out of memory error because litgpt isn't using the available unified memory of the chip.
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 224.00 MiB. GPU 0 has a total capacity of 94.88 GiB of which 172.19 MiB is free. Including non-PyTorch memory, this process has 94.69 GiB memory in use. Of the allocated memory 91.40 GiB is allocated by PyTorch, and 2.55 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://pytorch.org/docs/stable/notes/cuda.html#environment-variables)
Is there a way to use all the available memory?
Hi there. Based on the message you are getting, it looks like all the memory is used. Based on your nvidia-smi output, you have 97871 MiB, which is 97871 / 1024 ≈ 95 GiB.
This would match the statement
94.88 GiB of which 172.19 MiB is free.
For general strategies to reduce memory usage, you can try bf16 training, smaller batch and block sizes, etc. Pretraining can be a bit resource-hungry.
The question is how to use the unified memory. Apparently H200 has lots of CPU and GPU memory and a speedy connection between them, so these can be treated as a single memory. But LitGPT doesn't do anything specific to it. So the configuration has to be done on PyTorch side, like some env variables or torch. ... = ... Unfortunately, it's an uncharted territory for me.
Oh I see now, I didn't realize it was a unified memory between CPU and GPU (I thought it was somehow between multiple GPUs). In that case, I am actually not sure how to leverage that in PyTorch in general. My guess is that there will be larger transfer speeds between CPU and GPU memory, but not sure how to leverage that additional CPU RAM in PyTorch during training. Here, I am assuming the 95 Gb RAM are not already part of it, and there is an additional available memory on the CPU that you want to use.
The only way I could think of right now is to enable CPU offloading in FSDP/FSDPStrategy. I.e., by setting cpu_offload=True explicitly.
Looks like unified memory support is still work in progress on the PyTorch side: https://github.com/vllm-project/vllm/issues/10267 & https://github.com/pytorch/pytorch/issues/124807