Training gives nan after the first iteration

[Ping-C]

Hello, when I try to run training with the command

python scripts/train.py --config configs/base.py:base

It seems like the train loss becomes nan immediately after the first iteration. Is this something that you have encountered before?

[Ping-C]

It seems like I was able to fix the error by uncommenting this line in train. Maybe this is something that you may want to consider updating in your repo?

# seems like remat is actually enabled by default -- this disables it
    # @partial(jax.checkpoint, policy=jax.checkpoint_policies.everything_saveable)

I still don't know why it gave an error though. It seems like this is something you encountered before. Would you mind sharing with me what this fix meant?

[kvablack]

I ran into this issue before with Jax 0.4.13 and fixed it by downgrading to 0.4.11. Very interesting that disabling remat fixes it --- I never figured that out. The commented-out line is just a remnant from when I was experimenting with memory usage. You might want to try different Jax versions instead of disabling remat, since (at least in theory) remat can give significant speed and memory savings.

[Ping-C]

I see. Thanks for getting back to me! Will update that :) Very cool paper by the way!