Memory allocation on GPU

Question

Memory allocation on GPU

Closed this issue 4 months ago · 2 comments

I am loading the octo-small model on my GPU (NVIDIA GeForce RTX 4090) and checking nvidia-smi shows that it uses about 20GB of memory storage, which seems high to me. Upon loading the octo-base model, it still takes up about 20GB of storage, which doesn't make much sense. Does anyone know how I can decrease the amount of memory allocated to loading the pretrained octo model, or do these numbers sound about right to everyone?

Answer 1 · 2024-06-07T16:17:11.000Z

See https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html

Answer 2 · 2024-06-07T18:30:44.000Z

Exactly what I needed, thank you!