Training instability encountered during training 1024 model.

[SunzeY]

I have trained a Pixart-512 on my own data and it works fine. However, when I extend the resolution to 1024. Severe training stability is encountered. I double checked to make sure that batch size and learning rate are exactly the same as 512 version? Does anyone have faced the similar problem? Or same thing happened to the official training process?

2024-04-07 19:05:43,000 - PixArt - INFO - Step/Epoch [1885/6][25/372]:total_eta: 10:52:57, epoch_eta:0:13:33, time_all:2.029, time_data:0.018, lr:1.000e-04, s:(64, 64), train_loss:0.0810, grad_norm:0.1405
2024-04-07 19:06:34,638 - PixArt - INFO - Step/Epoch [1910/6][50/372]:total_eta: 10:50:58, epoch_eta:0:12:33, time_all:2.065, time_data:0.019, lr:1.000e-04, s:(64, 64), train_loss:0.0689, grad_norm:nan
2024-04-07 19:07:26,034 - PixArt - INFO - Step/Epoch [1935/6][75/372]:total_eta: 10:48:58, epoch_eta:0:11:33, time_all:2.056, time_data:0.019, lr:1.000e-04, s:(64, 64), train_loss:0.2009, grad_norm:nan
2024-04-07 19:08:17,354 - PixArt - INFO - Step/Epoch [1960/6][100/372]:total_eta: 10:46:59, epoch_eta:0:10:34, time_all:2.053, time_data:0.019, lr:1.000e-04, s:(64, 64), train_loss:-1.1976, grad_norm:inf

[lawrence-cj]

I hope you are using the train_scripts/train.py so that you can set the fp32_attention = True in your config file to solve the NaN problem as we mentioned here:

PixArt-alpha/configs/pixart_config/PixArt_xl2_img1024_lcm.py

Line 10 in 82c8559

fp32_attention = False # Set to True if you got NaN loss

[SunzeY]

I have tested smaller lr. but this cause the model cannot learn good enough as 512 did...

[SunzeY]

Do you pretrain Pixart with linear warm-up then constant lr until finish? I believe for other models, it is better to use cosine lr scheduler to train with lr decreasing. Is constant lr a normal choice for diffusion finetuning/training?

[lawrence-cj]

Yes, I use lr warm-up and the same constant lr with the original DiT.

[lawrence-cj]

I have tested smaller lr. but this cause the model cannot learn good enough as 512 did...

Well, I think you can try to resume from the ckpt before the loss becomes nan and at the same time change the lr smaller. So that we will know if the problem is caused by the lr.

[SunzeY]

OK, I will try this. Thank you for this open-source work and kind reply!