How to make use of NVIDIA GH200 Grace Hopper Superchip

[TheLukaDragar]

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?