JVP/VJP materializes the full array

[neel04]

I'm trying to debug excessive memory use when training LMs. It seems that a 65M model with a batch size > 32 seems to OOM on a 16GB card. The situation is worse for bigger batches and models.

Profiling Graphs:

One Array I'm noticing seems to be this one, for the jvp which apears to be of shape (batch_size * seqlen, vocab_size) = (32 * 512, 50304). This is the biggest allocation on the heap, according to the profiler.

Surely this array shouldn't materialize, as for larger models it'd be severely memory inefficient? Is there something I'm missing here or is this expected?

[jakevdp]

It's certainly possible to write code for which a JVP would materialize an array of that size. Without seeing what code you're running, it's hard to say any more than that. Can you include a minimal reproduction of the issue?

[neel04]

Its just a standard LM training recipie. The code is a bit of a mess atm, but here's the full context if that helps. FWIW dlwh at Levanter faced the same issue.

Here's a MRE I expanded with a bit of help from claude. This isn't exactly what I obtain in the profiler, but appears close enough: colab (Check TensorBoard under memory_viewer -> jit_get_loss_and_grads)