High GPU memory on simple import
Closed this issue · 2 comments
jens321 commented
When importing craftax functions such as the one below, my GPU memory immediately fills up for a large part. E.g. when I run the single line below
from craftax.craftax_env import make_craftax_env_from_name
and then check the gpu memory, I get that 60GB is already taken up on a H100 machine:
![Screenshot 2024-06-20 at 12 12 31 PM](https://private-user-images.githubusercontent.com/14240488/341487492-d7734629-0bd9-4ab1-8b15-be667ed7c469.png?jwt=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3MjU0NTIzMjgsIm5iZiI6MTcyNTQ1MjAyOCwicGF0aCI6Ii8xNDI0MDQ4OC8zNDE0ODc0OTItZDc3MzQ2MjktMGJkOS00YWIxLThiMTUtYmU2NjdlZDdjNDY5LnBuZz9YLUFtei1BbGdvcml0aG09QVdTNC1ITUFDLVNIQTI1NiZYLUFtei1DcmVkZW50aWFsPUFLSUFWQ09EWUxTQTUzUFFLNFpBJTJGMjAyNDA5MDQlMkZ1cy1lYXN0LTElMkZzMyUyRmF3czRfcmVxdWVzdCZYLUFtei1EYXRlPTIwMjQwOTA0VDEyMTM0OFomWC1BbXotRXhwaXJlcz0zMDAmWC1BbXotU2lnbmF0dXJlPWNhM2M0ZTZlNzZjOGZhM2Q5OTE3Yjg2NTMyYzUzYTIxMWZkYzA4Njk3MGQwNTVjZWFmMDBiNDQ1OWU0OGEyNWEmWC1BbXotU2lnbmVkSGVhZGVycz1ob3N0JmFjdG9yX2lkPTAma2V5X2lkPTAmcmVwb19pZD0wIn0.0eeC7WBS6h3Z29f_G8Awlk_-qLo-S8ZDTfOyl4RXFNE)
Is this normal? And if so, what's causing this huge memory allocation?
Thanks a lot for your help!
MichaelTMatthews commented
It is default JAX behaviour to preallocate 75% of GPU memory.
This behaviour can be controlled as detailed here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
jens321 commented
I had no idea, thanks for the pointer!