Error when trying to finetune `lmsys/vicuna-7b-v1.5` with 6 A100 40G GPUs

Nobody has claimed this yet.

script used to finetune `lmsys/vicuna-7b-v1.5`

```
CUDA_VISIBLE_DEVICES="7,6,5,4,3,2" torchrun --nproc_per_node=4 --master_port=20001 fastchat/train/train_mem.py \
--model_name_or_path lmsys/vicuna-7b-v1.5 \
--data_path data/vicuna_dummy_train.json \
--bf16 True \
--output_dir output_vicuna \
--num_train_epochs 1 \
--per_device_train_batch_size 1 \
--per_device_eval_batch_size 1 \
--gradient_accumulation_steps 16 \
--evaluation_strategy "no" \
--save_strategy "steps" \
--save_steps 1200 \
--save_total_limit 10 \
--learning_rate 2e-5 \
--weight_decay 0. \
--warmup_ratio 0.03 \
--lr_scheduler_type "cosine" \
--logging_steps 1 \
--fsdp "full_shard auto_wrap" \
--fsdp_transformer_layer_cls_to_wrap 'LlamaDecoderLayer' \
--tf32 True \
--model_max_length 258 \
--gradient_checkpointing True \
--lazy_preprocess True
```
Got the following error, which is identical to a previous issue [#1428](https://github.com/lm-sys/FastChat/issues/1428)

```
/home/kun/anaconda3/envs/llm-fine-tune/lib/python3.9/site-packages/torch/distributed/fsdp/_state_dict_utils.py:312: UserWarning: Failed to clone() tensor with name lm_head.weight on rank 2. This may mean that this state_dict entry could point to invalid memory regions after returning from state_dict() call if this parameter is managed by FSDP. Please check clone implementation of lm_head.weight. Error: CUDA out of memory. Tried to allocate 500.00 MiB (GPU 2; 39.39 GiB total capacity; 37.52 GiB already allocated; 58.00 MiB free; 38.13 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
warnings.warn(
Traceback (most recent call last):
File "/home/kun/llm_fine_tune/vicuna-finetune/fast-chat-training/FastChat/fastchat/train/train_mem.py", line 13, in
train()
File "/home/kun/anaconda3/envs/llm-fine-tune/lib/python3.9/site-packages/fastchat/train/train.py", line 284, in train
safe_save_model_for_hf_trainer(trainer=trainer, output_dir=training_args.output_dir)
File "/home/kun/anaconda3/envs/llm-fine-tune/lib/python3.9/site-packages/fastchat/train/train.py", line 76, in safe_save_model_for_hf_trainer
cpu_state_dict = {key: value.cpu() for key, value in state_dict.items()}
File "/home/kun/anaconda3/envs/llm-fine-tune/lib/python3.9/site-packages/fastchat/train/train.py", line 76, in
cpu_state_dict = {key: value.cpu() for key, value in state_dict.items()}
RuntimeError: CUDA error: invalid argument
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
```

while checking the cuda usage, it shows that the memory usage distribution is highly skrew, GPU 7 (in my case) reached full capacity while the other GPUs are not fully used. Any ideas?

Start with fastchat/train/train_mem.py and fastchat/train/train.py, especially safe_save_model_for_hf_trainer and the reported FSDP state-dict failure. Reproduce the command with the six A100 GPUs and inspect the memory distribution during checkpoint saving. Done means the cause of the uneven usage and CUDA failure is identified, with a documented fix or confirmed workaround.

Written by the indexing model from the issue text.