Hi Thomas,
Little updates:
From my last messages, I mainly worked with Adam optimizer using float32 precision, also I started doing batch accumulation (giving smoother training, but also slightly better accuracy in average). I did not got the NaN problem in f32, but I observed some times a peak 100x higher than usual in my loss after hours of training making the model loosing its progress. It’s possible that those jumps in the loss I observed in f32 training share the same root cause as the NaN problem I got with mixed precision training. I think it may still coming from v_hat going close to 0 sometimes, and I realized that this kind of problem should be tackled by using the amsgrad variant of adam. As with amsgrad v_hat is taken from past maximum rather than moving average. I launched an experiment with amsgrad 48h ago, the training goes well so far. I’ll try it later with mixed precision to check if it solve the original problem, I still plan to play with LARS latter also.
Best regards,
Thomas