CogView4-6B lora training: grad_norm=nan, loss=nan

[liux750]

CogKit/quickstart/scripts/t2i : sh start_train.sh

logs:
Some weights of GlmForCausalLM were not initialized from the model checkpoint at ./models--THUDM--CogView4-6B and are newly initialized: ['lm_head.weight']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.

Training steps: 0%| | 2/20000000000 [00:07<19850985:12:25, 3.57s/it, grad_norm=nan, loss=nan, lr=2e-5]

[Blcony]

I have the same problem. Have you solved this problem?

Thanks!

[Blcony]

@OleehyO Thanks for your wonderful work!

But I notice the same question as depicted above. Can you offer some suggestions to help us figure out the problem?

[Blcony]

@liux750 @OleehyO

I find that the batch_size>1 (such as batch_size=2) or enable_packing=False in the config.yaml will cause the loss nan, but the batch_size=1 and enable_packing=True will not cause the problem.

Can you give me some insights about this situation?

[Blcony]

@OleehyO

I find Line#75 in src/cogkit/finetune/diffusion/models/cogview/cogview4/lora_trainer.py replace_attn_processor(components.transformer, CogView4TrainingAttnProcessor()) is the key to cause the problem.

If I remove the Line#75, batch_size>1 will not cause the loss nan problem.

So is this line necessary for the single-resolution training experiments? If it is necessary, how to fix it to fit the batch_size>1? Or did I miss some key problems, something like package versions?