tensor size mismatch with larger gradient_accumulation_steps and fewer training data
Author: yyymeta
Pull request: no

### System Info

A100 Nvidia 80G GPU

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

it seems when I have fewer training examples (1000 or so), and when I use a larger gradient_accumulation_steps, (32)
I get tensor size mismatch on Adam gradient update:


```
\"/tmp/jetter.yrcudeja/torch/optim/optimizer.py\", line 33, in _use_grad\n    ret = func(self, *args, **kwargs)\n  File \"/tmp/jetter.yrcudeja/torch/optim/adamw.py\", line 173, in step\n    adamw(\n  File \"/tmp/jetter.yrcudeja/torch/optim/adamw.py\", line 323, in adamw\n    func(\n  File \"/tmp/jetter.yrcudeja/torch/optim/adamw.py\", line 502, in _multi_tensor_adamw\n    torch._foreach_add_(device_exp_avgs, device_grads, alpha=1 - beta1)\nRuntimeError: The size of tensor a (8192384) must match the size of tensor b (262156288) at non-singleton dimension 0\n", "errorTraits": null, "timestamp_us": 1692818123557766}
[4]:  File "/usr/local/fbcode/platform010/lib/python3.8/runpy.py", line 194, in _run_module_as_main
[4]:    return _run_code(code, main_globals, None,
[4]:  File "/usr/local/fbcode/platform010/lib/python3.8/runpy.py", line 87, in _run_code
[4]:    exec(code, run_globals)
[4]:  File "/tmp/jetter.12kzp8qf/aml/comment/llama_finetune/train.py", line 150, in <module>
[4]:    train()
[4]:  File "/tmp/jetter.12kzp8qf/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 346, in wrapper
[4]:    return f(*args, **kwargs)
[4]:  File "/tmp/jetter.12kzp8qf/aml/comment/llama_finetune/train.py", line 124, in train
[4]:    trainer.train()
[4]:  File "/tmp/jetter.12kzp8qf/transformers/trainer.py", line 1664, in train
[4]:    return inner_training_loop(
[4]:  File "/tmp/jetter.12kzp8qf/transformers/trainer.py", line 1998, in _inner_training_loop
[4]:    self.optimizer.step()
[4]:  File "/tmp/jetter.12kzp8qf/torch/optim/lr_scheduler.py", line 69, in wrapper
[4]:    return wrapped(*args, **kwargs)
[4]:  File "/tmp/jetter.12kzp8qf/torch/optim/optimizer.py", line 280, in wrapper
[4]:    out = func(*args, **kwargs)
[4]:  File "/tmp/jetter.12kzp8qf/torch/optim/optimizer.py", line 33, in _use_grad
[4]:    ret = func(self, *args, **kwargs)
[4]:  File "/tmp/jetter.12kzp8qf/torch/optim/adamw.py", line 173, in step
[4]:    adamw(
[4]:  File "/tmp/jetter.12kzp8qf/torch/optim/adamw.py", line 323, in adamw
[4]:    func(
[4]:  File "/tmp/jetter.12kzp8qf/torch/optim/adamw.py", line 502, in _multi_tensor_adamw
[4]:    torch._foreach_add_(device_exp_avgs, device_grads, alpha=1 - beta1)
[4]:RuntimeError: The size of tensor a (8192384) must match the size of tensor b (262156288) at non-singleton dimension 0
[7]:ERROR:aiplatform.error_reporting.error_reporting:Exception Found: The size of tensor a (8192384) must match the size of tensor b (262156288) at non-singleton dimension 0
```

using gradient_accumulation_steps=1 fixes it, but then it causes some impact on model quality


command was 

/packages/torchx_python/python

-m torch.distributed.run --rdzv_backend zeus --rdzv_id torchx-llama_finetune_train-k64xgt1h4dt52c --nnodes 4 --nproc_per_node 8 --tee 3 --role  -m aml.comment.llama_finetune.train --local_dir /tmp/users --model_manifold_bucket pi_adv_problems --model_manifold_dir tree/dpa_llama --input_model_filename 7B-converted --output_model_filename yytest__v7_instagram_basic_5e-6 --data_path manifold://pi_adv_problems/tree/appreview_llama/data/v7/train__instagram_basic.json --eval_data_path manifold://pi_adv_problems/tree/appreview_llama/data/v7/eval__instagram_basic.json --data_task generic --prompt_temp normal --processed True --model_max_length 1024 --num_train_epochs 30 --per_device_train_batch_size 2 --per_device_eval_batch_size 8 --gradient_accumulation_steps 32 --evaluation_strategy steps --eval_steps 10 --save_strategy steps --save_steps 200 --save_total_limit 1 --learning_rate 5e-6 --weight_decay 0. --warmup_ratio 0.03 --lr_scheduler_type cosine --logging_steps 50 --fsdp full_shard auto_wrap --fsdp_transformer_layer_cls_to_wrap LlamaDecoderLayer --bf16 True --tf32 True

### Expected behavior

note that the above error shows 

yyy@yyy-mbp ~ % echo '262156288/8192384'|bc -l
32.00000000000000000000

so somehow it seems only one gradient is obtained while maybe 32 are expected?

Comments:
Hey! If you want help from our trainer expert @pacman100 we're gonna need to have a look at the training script or at least have a reproducer.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

Hi, I have encountered the same problem, fewer training examples and adam optimizer. May I ask if you have resolved it? How was it resolved?

Same issue here. I get such an error when doing FSDP training. No error when using larger `gradient_accumulation_steps` without FSDP and also no error when gradient_accumulation_steps is 1.

I'd love to know if either of you solved it.

Same issue reported here: https://discuss.huggingface.co/t/errors-when-using-gradient-accumulation-with-fsdp-peft-lora-sfttrainer/105006

Thanks, will re-open and add to the list of tracked issues


Same thing happens here. 

same error happens with FSDP + PEFT QLoRA + SFTTrainer + adamw-schedulefree, batchsize=1/gpu * 8gpu * 3node, gradient_accumulation_steps=8.
when I resume the training from previous checkpoint, the error happens after the same iterations.
`dataloader_drop_last = True` does not help.
when I set `gradient_accumulation_steps=1`, the training raises no error.

I’m interested in working on this issue. Could you please assign it to me? @ArthurZucker @MekkCyber 

Hi @KevinCarpricorn, Thanks for your interest 😄, but the issue is fixed with `transformers>=4.47.0`. If you are interested, you can find more issues related to the `Trainer` here : https://github.com/huggingface/transformers/issues/33345

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.