[BUG]: Incompatible between colossalai_zero2 and LoRA tuning

[zhangliang-04]

🐛 Describe the bug

When i run this script:

torchrun --standalone --nproc_per_node=1 train_sft.py \
    --pretrain "/root/zl/download/pretrained/llama_7b_hf/" \
    --model 'llama' \
    --strategy colossalai_zero2 \
    --log_interval 10 \
    --save_path trained_models/Coati-7B \
    --dataset /root/zl/code/InstructionWild/data/instinwild_en.json \
    --batch_size 4 \
    --accimulation_steps 1 \
    --lr 2e-5 \
    --max_epochs 3 \
    --lora_rank 8 \
    --max_datasets_size 2

I got an error during backward:

Traceback (most recent call last):
  File "/data2/zl/code/ColossalAI/applications/Chat/train_sft.py", line 184, in <module>
    train(args)
  File "/data2/zl/code/ColossalAI/applications/Chat/train_sft.py", line 155, in train
    trainer.fit(logger=logger, log_interval=args.log_interval)
  File "/data2/zl/code/ColossalAI/applications/Chat/coati/trainer/sft.py", line 110, in fit
    self.strategy.optimizer_step(self.optimizer)
  File "/data2/zl/code/ColossalAI/applications/Chat/coati/trainer/strategies/colossalai.py", line 154, in optimizer_step
    optimizer.step()
  File "/opt/conda/lib/python3.9/site-packages/torch/optim/lr_scheduler.py", line 65, in wrapper
    return wrapped(*args, **kwargs)
  File "/opt/conda/lib/python3.9/site-packages/colossalai/zero/sharded_optim/low_level_optim.py", line 467, in step
    assert param_shape == flat_fp32_avg_grads.shape, \
AssertionError: fp32 param and grad have different shape torch.Size([20277248]) vs torch.Size([288768])

However, when i change the --strategy into ddp, it trains normally.
So is there a bug in implementing colossalai_zero2 , or it is incompatible with LoRA?

Environment

I use the docker image hpcaitech/colossalai:0.2.7

[zhangliang-04]

Additionally, when i use ddp and LoRA, there is also a bug when saving the checkpoint.
Script:

torchrun --standalone --nproc_per_node=1 train_sft.py \
    --pretrain "/root/zl/download/pretrained/llama_7b_hf/" \
    --model 'llama' \
    --strategy ddp \
    --log_interval 10 \
    --save_path trained_models/Coati-7B \
    --dataset /root/zl/code/InstructionWild/data/instinwild_en.json \
    --batch_size 4 \
    --accimulation_steps 1 \
    --lr 2e-5 \
    --max_epochs 3 \
    --lora_rank 8 \
    --max_datasets_size 10

Feedback:

Traceback (most recent call last):
  File "/data2/zl/code/ColossalAI/applications/Chat/train_sft.py", line 184, in <module>
    train(args)
  File "/data2/zl/code/ColossalAI/applications/Chat/train_sft.py", line 158, in train
    trainer.save_model(path=args.save_path, only_rank0=True, tokenizer=tokenizer)
  File "/data2/zl/code/ColossalAI/applications/Chat/coati/trainer/sft.py", line 158, in save_model
    self.strategy.save_model(model=self.model, path=path, only_rank0=only_rank0, tokenizer=tokenizer)
TypeError: save_model() got an unexpected keyword argument 'tokenizer'

[qinqinqaq]

I am having the same issue here

[YiAthena]

I am having the same issue here

[samulew]

met the same bug, and fixed it manually

[qinqinqaq]

met the same bug, and fixed it manually

hello,How did you manually solve it？

[alibabadoufu]

met the same bug, and fixed it manually

I am having the same issue too. Would you mind sharing your solution? Thanks a lot :D

[alibabadoufu]

Additionally, when i use ddp and LoRA, there is also a bug when saving the checkpoint. Script:

torchrun --standalone --nproc_per_node=1 train_sft.py \
    --pretrain "/root/zl/download/pretrained/llama_7b_hf/" \
    --model 'llama' \
    --strategy ddp \
    --log_interval 10 \
    --save_path trained_models/Coati-7B \
    --dataset /root/zl/code/InstructionWild/data/instinwild_en.json \
    --batch_size 4 \
    --accimulation_steps 1 \
    --lr 2e-5 \
    --max_epochs 3 \
    --lora_rank 8 \
    --max_datasets_size 10

Feedback:

Traceback (most recent call last):
  File "/data2/zl/code/ColossalAI/applications/Chat/train_sft.py", line 184, in <module>
    train(args)
  File "/data2/zl/code/ColossalAI/applications/Chat/train_sft.py", line 158, in train
    trainer.save_model(path=args.save_path, only_rank0=True, tokenizer=tokenizer)
  File "/data2/zl/code/ColossalAI/applications/Chat/coati/trainer/sft.py", line 158, in save_model
    self.strategy.save_model(model=self.model, path=path, only_rank0=only_rank0, tokenizer=tokenizer)
TypeError: save_model() got an unexpected keyword argument 'tokenizer'

I guess you can solve this one with this #3357 (comment)

[samulew]

Additionally, when i use ddp and LoRA, there is also a bug when saving the checkpoint. Script:

torchrun --standalone --nproc_per_node=1 train_sft.py \
    --pretrain "/root/zl/download/pretrained/llama_7b_hf/" \
    --model 'llama' \
    --strategy ddp \
    --log_interval 10 \
    --save_path trained_models/Coati-7B \
    --dataset /root/zl/code/InstructionWild/data/instinwild_en.json \
    --batch_size 4 \
    --accimulation_steps 1 \
    --lr 2e-5 \
    --max_epochs 3 \
    --lora_rank 8 \
    --max_datasets_size 10

Feedback:

Traceback (most recent call last):
  File "/data2/zl/code/ColossalAI/applications/Chat/train_sft.py", line 184, in <module>
    train(args)
  File "/data2/zl/code/ColossalAI/applications/Chat/train_sft.py", line 158, in train
    trainer.save_model(path=args.save_path, only_rank0=True, tokenizer=tokenizer)
  File "/data2/zl/code/ColossalAI/applications/Chat/coati/trainer/sft.py", line 158, in save_model
    self.strategy.save_model(model=self.model, path=path, only_rank0=only_rank0, tokenizer=tokenizer)
TypeError: save_model() got an unexpected keyword argument 'tokenizer'

I guess you can solve this one with this #3357 (comment)

according to the traceback message , i handled it simply :

ColossalAI/applications/Chat/train_sft.py", line 158
if args.strategy=="ddp":
trainer.save_model(path=args.save_path, only_rank0=True)
else:
trainer.save_model(path=args.save_path, only_rank0=True, tokenizer=tokenizer)

[qinqinqaq]

Additionally, when i use ddp and LoRA, there is also a bug when saving the checkpoint. Script:

torchrun --standalone --nproc_per_node=1 train_sft.py \
    --pretrain "/root/zl/download/pretrained/llama_7b_hf/" \
    --model 'llama' \
    --strategy ddp \
    --log_interval 10 \
    --save_path trained_models/Coati-7B \
    --dataset /root/zl/code/InstructionWild/data/instinwild_en.json \
    --batch_size 4 \
    --accimulation_steps 1 \
    --lr 2e-5 \
    --max_epochs 3 \
    --lora_rank 8 \
    --max_datasets_size 10

Feedback:

Traceback (most recent call last):
  File "/data2/zl/code/ColossalAI/applications/Chat/train_sft.py", line 184, in <module>
    train(args)
  File "/data2/zl/code/ColossalAI/applications/Chat/train_sft.py", line 158, in train
    trainer.save_model(path=args.save_path, only_rank0=True, tokenizer=tokenizer)
  File "/data2/zl/code/ColossalAI/applications/Chat/coati/trainer/sft.py", line 158, in save_model
    self.strategy.save_model(model=self.model, path=path, only_rank0=only_rank0, tokenizer=tokenizer)
TypeError: save_model() got an unexpected keyword argument 'tokenizer'

I guess you can solve this one with this #3357 (comment)

according to the traceback message , i handled it simply :

ColossalAI/applications/Chat/train_sft.py", line 158 if args.strategy=="ddp": trainer.save_model(path=args.save_path, only_rank0=True) else: trainer.save_model(path=args.save_path, only_rank0=True, tokenizer=tokenizer)

when i use dpp, the loss is always nan

[zhangliang-04]

Additionally, when i use ddp and LoRA, there is also a bug when saving the checkpoint. Script:

torchrun --standalone --nproc_per_node=1 train_sft.py \
    --pretrain "/root/zl/download/pretrained/llama_7b_hf/" \
    --model 'llama' \
    --strategy ddp \
    --log_interval 10 \
    --save_path trained_models/Coati-7B \
    --dataset /root/zl/code/InstructionWild/data/instinwild_en.json \
    --batch_size 4 \
    --accimulation_steps 1 \
    --lr 2e-5 \
    --max_epochs 3 \
    --lora_rank 8 \
    --max_datasets_size 10

Feedback:

Traceback (most recent call last):
  File "/data2/zl/code/ColossalAI/applications/Chat/train_sft.py", line 184, in <module>
    train(args)
  File "/data2/zl/code/ColossalAI/applications/Chat/train_sft.py", line 158, in train
    trainer.save_model(path=args.save_path, only_rank0=True, tokenizer=tokenizer)
  File "/data2/zl/code/ColossalAI/applications/Chat/coati/trainer/sft.py", line 158, in save_model
    self.strategy.save_model(model=self.model, path=path, only_rank0=only_rank0, tokenizer=tokenizer)
TypeError: save_model() got an unexpected keyword argument 'tokenizer'

I guess you can solve this one with this #3357 (comment)

according to the traceback message , i handled it simply :
ColossalAI/applications/Chat/train_sft.py", line 158 if args.strategy=="ddp": trainer.save_model(path=args.save_path, only_rank0=True) else: trainer.save_model(path=args.save_path, only_rank0=True, tokenizer=tokenizer)

when i use dpp, the loss is always nan

yes, I got the same issue. It seems that there is no parameters requires grad

[mynameischaos]

Additionally, when i use ddp and LoRA, there is also a bug when saving the checkpoint. Script:

torchrun --standalone --nproc_per_node=1 train_sft.py \
    --pretrain "/root/zl/download/pretrained/llama_7b_hf/" \
    --model 'llama' \
    --strategy ddp \
    --log_interval 10 \
    --save_path trained_models/Coati-7B \
    --dataset /root/zl/code/InstructionWild/data/instinwild_en.json \
    --batch_size 4 \
    --accimulation_steps 1 \
    --lr 2e-5 \
    --max_epochs 3 \
    --lora_rank 8 \
    --max_datasets_size 10

Feedback:

Traceback (most recent call last):
  File "/data2/zl/code/ColossalAI/applications/Chat/train_sft.py", line 184, in <module>
    train(args)
  File "/data2/zl/code/ColossalAI/applications/Chat/train_sft.py", line 158, in train
    trainer.save_model(path=args.save_path, only_rank0=True, tokenizer=tokenizer)
  File "/data2/zl/code/ColossalAI/applications/Chat/coati/trainer/sft.py", line 158, in save_model
    self.strategy.save_model(model=self.model, path=path, only_rank0=only_rank0, tokenizer=tokenizer)
TypeError: save_model() got an unexpected keyword argument 'tokenizer'

I guess you can solve this one with this #3357 (comment)

according to the traceback message , i handled it simply :
ColossalAI/applications/Chat/train_sft.py", line 158 if args.strategy=="ddp": trainer.save_model(path=args.save_path, only_rank0=True) else: trainer.save_model(path=args.save_path, only_rank0=True, tokenizer=tokenizer)

when i use dpp, the loss is always nan

yes, I got the same issue. It seems that there is no parameters requires grad

same question!

[chengeharrison]

Hi all, you can refer to #3436 for the unexpected keyword argument 'tokenizer' problem.

[alibabadoufu]

Additionally, when i use ddp and LoRA, there is also a bug when saving the checkpoint. Script:

torchrun --standalone --nproc_per_node=1 train_sft.py \
    --pretrain "/root/zl/download/pretrained/llama_7b_hf/" \
    --model 'llama' \
    --strategy ddp \
    --log_interval 10 \
    --save_path trained_models/Coati-7B \
    --dataset /root/zl/code/InstructionWild/data/instinwild_en.json \
    --batch_size 4 \
    --accimulation_steps 1 \
    --lr 2e-5 \
    --max_epochs 3 \
    --lora_rank 8 \
    --max_datasets_size 10

Feedback:

Traceback (most recent call last):
  File "/data2/zl/code/ColossalAI/applications/Chat/train_sft.py", line 184, in <module>
    train(args)
  File "/data2/zl/code/ColossalAI/applications/Chat/train_sft.py", line 158, in train
    trainer.save_model(path=args.save_path, only_rank0=True, tokenizer=tokenizer)
  File "/data2/zl/code/ColossalAI/applications/Chat/coati/trainer/sft.py", line 158, in save_model
    self.strategy.save_model(model=self.model, path=path, only_rank0=only_rank0, tokenizer=tokenizer)
TypeError: save_model() got an unexpected keyword argument 'tokenizer'

I guess you can solve this one with this #3357 (comment)

according to the traceback message , i handled it simply :
ColossalAI/applications/Chat/train_sft.py", line 158 if args.strategy=="ddp": trainer.save_model(path=args.save_path, only_rank0=True) else: trainer.save_model(path=args.save_path, only_rank0=True, tokenizer=tokenizer)

when i use dpp, the loss is always nan

Same for me

[puyuanOT]

Got the same issue.