scale training

[scott-vsi]

I have a large GPU. When I run the SMAX training, it only uses a few GB of GPU memory. Is there a way to scale that up? I tried upping num_env but that didn't affect GPU memory.

[sash-a]

I assume you are running an anakin system if training SMAX. How many envs are you using exactly? Increase it to 512-2048 and you should see a noticeable uptick in GPU use. But JAX pre-allocates memory by default so you will need to set a flag to see the true GPU usage. Likely something like:

XLA_PYTHON_CLIENT_PREALLOCATE=false python path/to/system.py

Here's a description of the flags: https://docs.jax.dev/en/latest/gpu_memory_allocation.html

[scott-vsi]

Sorry for not including some of those details. Yes. I'm using the Anakin system.

Thank you for pointing out that Jax preallocates memory. Now I see the memory scaling.