Fix the gradient checkpointing bug of the llama model
Author: yqy2001
Pull request: yes

# What does this PR do?

The gradient checkpoint of the LLaMA model does not work. This PR fixed it following the [GPT-2 model's gradient checkpointing implementation](https://github.com/huggingface/transformers/blob/cf0af9a31beb84e8feec77af51f72d063ba905aa/src/transformers/models/gpt2/modeling_gpt2.py#L482).

### System Info

- `transformers` version: 4.28.0.dev0
- Platform: Linux-5.4.0-113-generic-x86_64-with-glibc2.10
- Python version: 3.8.8
- Huggingface_hub version: 0.11.1
- PyTorch version (GPU?): 2.1.0.dev20230317+cu118 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: Yes

### Reproducing the Bug

The bug is tested on 4 A100 40GB GPUs.

Please first cloning Stanford Alpaca's repo that finetunes the LLaMA model:
```shell
git clone git@github.com:tatsu-lab/stanford_alpaca.git
cd stanford_alpaca.git
pip install -r requirements.txt
```

CUDA out of memory error will raise if we run the training script with `per_device_train_batch_size=1`:
```sh
torchrun --nproc_per_node=4 --master_port=<your_random_port> train.py \
    --model_name_or_path <your_path_to_hf_converted_llama_ckpt_and_tokenizer> \
    --data_path ./alpaca_data.json \
    --bf16 True \
    --output_dir <your_output_dir> \
    --num_train_epochs 3 \
    --per_device_train_batch_size 1 \
    --per_device_eval_batch_size 4 \
    --gradient_accumulation_steps 8 \
    --evaluation_strategy "no" \
    --save_strategy "steps" \
    --save_steps 2000 \
    --save_total_limit 1 \
    --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
    --gradient_checkpointing
```

### Test after this PR:
After this PR, we can successfully train the LLaMA-7B models on 4 40GB GPUs with `per_device_train_batch_size=8`, using gradient checkpointing:
```sh
torchrun --nproc_per_node=4 --master_port=<your_random_port> train.py \
    --model_name_or_path <your_path_to_hf_converted_llama_ckpt_and_tokenizer> \
    --data_path ./alpaca_data.json \
    --bf16 True \
    --output_dir <your_output_dir> \
    --num_train_epochs 3 \
    --per_device_train_batch_size 8 \
    --per_device_eval_batch_size 4 \
    --gradient_accumulation_steps 8 \
    --evaluation_strategy "no" \
    --save_strategy "steps" \
    --save_steps 2000 \
    --save_total_limit 1 \
    --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
    --gradient_checkpointing
```

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [x] Did you write any new necessary tests?

## Who can review?

@ArthurZucker @zphang @sgugger




Comments:
_The documentation is not available anymore as the PR was closed or merged._