First I had an OOM during loading.
Ok, so there must be buggy behaviour in the archive when compiling with 12.0, and compiling with 12.1 fixes it.
The OOM while training seems to be back
If you restart the runtime and run the whole notebook once, is the OOM behaviour consistent or is it random?
I think I got too optimistic in my previous response.
While using my custom-built version it OOM on epoch two, here it OOM on epoch 4. So I only watched until epoch 3 and concluded that it will not OOM. Conclusion: It’s later but still OOM
I see, thanks for all the details! So the OOM during training is not version-specific for the most part, for the initial one I will investigate more and probably update the precompile versions : )
Do you have anything to suggest regarding the OOM training, you said the model is quite small so it should fit on a card with 24GB of ram.
Is it the
preallocated temp allocation: 4.00GiB
That is too small?
I was talking about loading specifically and comparing the 0.5GB model to the available 24GB. Training can definitely be memory expensive. There are likely memory optimisations we could do in axon/bumblebee specific to this training, so far we’ve been mostly focused on inference and deployment in bumblebee.
The model loads fine when I build or use your suggested version of XLA. It’s on fine tuning that it runs out of memory.
Regarding the version: 8.9.4.25-1+cuda12.2
Yeah, I tried to reproduce using the usual XLA_TARGET=cuda120. Maybe there is something specific to the GPU model, driver, or else. I will probably just bump the precompiled/required versions to match jax.