[Bug]SFT vicuna-7b-v1.3 with train_mem.py (with flash-attention) can not work

Nobody has claimed this yet.

script file: train_vicuna_7b.sh (which call the train_mem.py)
weight file: vicuna-7b-v1.3

Will meet the following error.

```
│ /opt/conda/lib/python3.8/site-packages/torch/optim/adamw.py:321 in adamw │
│ │
│ 318 │ else: │
│ 319 │ │ func = _single_tensor_adamw │
│ 320 │ │
│ ❱ 321 │ func( │
│ 322 │ │ params, │
│ 323 │ │ grads, │
│ 324 │ │ exp_avgs, │
│ │
│ /opt/conda/lib/python3.8/site-packages/torch/optim/adamw.py:389 in _single_tensor_adamw │
│ │
│ 386 │ │ param.mul_(1 - lr * weight_decay) │
│ 387 │ │ │
│ 388 │ │ # Decay the first and second moment running average coefficient │
│ ❱ 389 │ │ exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1) │
│ 390 │ │ exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2) │
│ 391 │ │ │
│ 392 │ │ if capturable or differentiable: │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
**RuntimeError: The size of tensor a (65537024) must match the size of tensor b (262148096) at non-singleton dimension 0**
WARNING:torch.distributed.elastic.multiprocessing.api:Sending process 7202 closing signal SIGTERM
ERROR:torch.distributed.elastic.multiprocessing.api:failed (exitcode: 1) local_rank: 1 (pid: 7203) of binary: /opt/conda/bin/python3
Traceback (most recent call last):
File "/opt/conda/bin/torchrun", line 33, in
sys.exit(load_entry_point('torch==2.0.1', 'console_scripts', 'torchrun')())
File "/opt/conda/lib/python3.8/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 346, in wrapper
return f(*args, **kwargs)
File "/opt/conda/lib/python3.8/site-packages/torch/distributed/run.py", line 794, in main
run(args)
File "/opt/conda/lib/python3.8/site-packages/torch/distributed/run.py", line 785, in run
elastic_launch(
File "/opt/conda/lib/python3.8/site-packages/torch/distributed/launcher/api.py", line 134, in __call__
return launch_agent(self._config, self._entrypoint, list(args))
File "/opt/conda/lib/python3.8/site-packages/torch/distributed/launcher/api.py", line 250, in launch_agent
raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError:
============================================================
../fastchat/train/train_mem.py FAILED
------------------------------------------------------------
Failures:
[1]:
time : 2023-07-04_05:09:53
host : 7a6a7980a2ea
rank : 2 (local_rank: 2)
exitcode : 1 (pid: 7204)
error_file:
traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
[2]:
time : 2023-07-04_05:09:53
host : 7a6a7980a2ea
rank : 3 (local_rank: 3)
exitcode : 1 (pid: 7205)
error_file:
traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
------------------------------------------------------------
Root Cause (first observed failure):
[0]:
time : 2023-07-04_05:09:53
host : 7a6a7980a2ea
rank : 1 (local_rank: 1)
exitcode : 1 (pid: 7203)
error_file:
traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
============================================================
wandb: 🚀 View run colorful-elevator-72 at: https://wandb.ai/ldwang/huggingface/runs/yko99tre0 MB deduped)
wandb: Synced 5 W&B file(s), 0 media file(s), 0 artifact file(s) and 0 other file(s)
wandb: Find logs at: ./wandb/run-20230704_050928-yko99tre/logs

```

Start with train_vicuna_7b.sh and the train_mem.py entry point, then reproduce the reported vicuna-7b-v1.3 training run with flash-attention. Inspect the AdamW tensor-size mismatch shown in the traceback; the work is done when this configuration completes training without the reported runtime error.

Written by the indexing model from the issue text.