Interaction of jax.checkpoint and jax.jit

[davisyoshida]

I've ported GPT-2 to JAX+Haiku, and found that I was able to do Adam steps with a batch size of 1 with a 382 token context size. I was hoping to bump this up higher by using jax.checkpoint on each layer self-attention block, but surprisingly I was getting OOM errors with the same context size which ran successfully without checkpointing. I'm guessing there's something going on with the interaction of jit and checkpoint, but I don't have any idea how to investigate this further. Any ideas on how I can prod at this to figure out what's happening?

Apologies for the lack of minimal example, but since I'm not sure what's going on, I couldn't figure out how to reproduce it in a small example.

[trevorcai]

If this is inside an hk.transform, you'll probably want to use hk.remat rather than jax.remat, as the blocks that you're transforming aren't pure functions since they interact with Haiku internal state.

https://dm-haiku.readthedocs.io/en/latest/notebooks/transforms.html

While I'm here, I'd recommend looking into donate_argnums to re-use your parameter buffers (as keeping two copies of 1.5B params around is a bit wasteful!).

[davisyoshida]

@trevorcai I am doing both of those things already. My understanding was that hk.remat is a pretty thin wrapper around the JAX version so I didn't think anything haiku specific was causing the issue.

[trevorcai]

In that case, I would roughly expect things to work out of the box - I personally use hk.remat on transformers on a regular basis on both GPU & TPU. There's probably something more subtle going on.

If you're able to make a minimal change to the Haiku example transformer that recreates this behavior, one of the currently-active Haiku maintainers (not me, unfortunately :P) would probably be able to take a look.

Otherwise, will rely on @mattjj to chime in on how OSS users should debug within-JIT memory problems.

[davisyoshida]

@trevorcai Thanks for the pointer, I'll try it out with that.

[davisyoshida]

For the record, I tried hk.remat with the example transformer, and it allowed me to increase the input size. I also tried remat with my GPT-2 implementation in the 345M parameter setting and it worked there as well. It's just the 1.5B parameter setting it seems to decrease the runnable context size in, which is really strange.