How JAX Allocates Memory

Intro At my previous job I ran into a problem: our JAX trainer in an RL job would OOM after running for a while, which is pretty rare. Most of my OOM debugging experience came from Torch, where there are basically two moves: The careful way: use the Torch Memory Visualizer to break down where memory goes and optimize from there. Often it’s just a tensor you forgot to free, or an allocation that…