JAX 是怎么分配内存的

引言 在之前的工作中我遇到一个问题：在我们的RL Job里，Jax Trainer跑着跑着OOM了，这很罕见。 我之前大部分的OOM Debug经验都来自Torch——对付Torch的OOM，基本就是两个套路： 细致点的话，用 Torch Memory Visualizer 把显存占用的来源一点点分析出来，然后针对性优化，很多时候其实是忘记释放或者申请了没有必要的大内存； 懒一点的话，在疑似Memory Leak的角落加上 torch.cuda.empty_cache() 。 所以碰到Jax OOM这个问题的时候我想到了两个更基本的问题： 为什么我很少在Jax中遇到OOM？ 我能不能使用 jax.cuda.empty_cache() ？ 想回答这两个问题，就得先搞明白Jax到底是怎么分配内存的——于是就有了这篇blog。 这篇文章也同步分享在知乎： JAX 是怎么分配内存的 。 概览…