3 ms·
> I kept running out of memory (OOM) when trying to `jax.jit` the entire training step. Even gradual `jax.jit` didn't work. Were you using activation checkpoin
by logicchains 2y ago
> I kept running out of memory (OOM) when trying to `jax.jit` the entire training step. Even gradual `jax.jit` didn't work.
Were you using activation checkpointing? https://jax.readthedocs.io/en/latest/_autosummary/jax.checkpoint.html https://jax.readthedocs.io/en/latest/_autosummary/jax.checkp... is very important for keeping memory usage reasonable when training large models.