Issue reproducing fine-tuning results – NaN loss after 40 epochs

[orspiegel]

Hello,
I’m trying to reproduce your fine-tuning results using the checkpoint multisubject_subj01_1024hid_nolow_300ep with subject01 data.
I’m training on a single L40S GPU with a batch size of 24.

However, after 40 epochs, I encountered the following error:

1173 Traceback (most recent call last):
1174   File "/net/mraid20/ifs/wisdom/userdata/orsp/repos/MindEyeV2/src/Train.py", line 834, in <module>
1175     utils.check_loss(loss)
1176   File "/net/mraid20/ifs/wisdom/userdata/orsp/repos/MindEyeV2/src/utils.py", line 208, in check_loss
1177     raise ValueError('NaN loss')
1178 ValueError: NaN loss

This occurs at:

utils.check_loss(loss)

I’m running the training (Train.ipynb) script using accel.slurm.

Do you have any insights on why this might be happening?
Thanks in advance!

[PaulScotti]

This is tricky because it likely depends on your specific cluster setup. I don't really have a great suggestion other than you could consider playing around with the learning rate, and you could try playing with your accelerate config (assuming you are using accelerate).

[orspiegel]

@PaulScotti Thank you for your quick response (and for the brilliant paper!).
I indeed managed to complete the training with very small learning rates.

Anyway, should this link still be available? https://stability.wandb.io/paul-scotti/mindeye?nw=nwuserpaulscotti
It seems to have expired.
It would be really helpful to look at your training logs for a sanity check.
:)

[PaulScotti]

Unfortunately this is because I left Stability AI :( so there's not anything I can do about these old wanbd links other than retraining the models with a separate wandb account

Here's a link to one wanbd run for fine-tuning MindEye2 model using 1 session worth of data, you can try to peek through the expiration pop-up to look at the plots: https://stability.wandb.io/paul-scotti/mindeye/runs/subj01_pretrained_1sess_1024hid_nolow_seed8?nw=nwuserpaulscotti

[orspiegel]

Oh I see, thank you a lot anyway!

[zjsndy]

Oh I see, thank you a lot anyway!
Have you solved this NaN issue? I’ve encountered the same problem. The code seems fine at first glance, but the error occurs after running for more than 20 epochs.