add deepspeed grad ckpt
Author: SeunghyunSEO
Pull request: yes
Labels: name: DeepSpeed

hi, @younesbelkada , sorry for the late PR.
its been busy days ;-)

i wrote some draft lines for [deepspeed gradient checkpointing](https://github.com/huggingface/transformers/issues/29648),
i tried to change original code as little as i can, but 2 things were inevitable

- 1. add `gradient_checkpointing_kwargs` as member variables for parsing `num_checkpoints` and `use_deepspeed_grad_ckpt` keys to check use advanced grad ckpt or not
- 2. second thing is not about code but deepspeed checkpoint function force me to use deepspeed.init because of [this line](https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/runtime/activation_checkpointing/checkpointing.py#L539-L545).

here is my code for sanity checking advanced activation checkpointing

```python
import os
import copy
import random
from pdb import set_trace as Tra

import torch
import deepspeed
from deepspeed import get_accelerator
from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments


def _reset_seeds(seed=1234):
    torch.manual_seed(seed)
    random.seed(seed)

def _get_dummy_inputs(B=2, T=10, seed=1234):
    _reset_seeds(seed)
    return {
        'input_ids': torch.rand(B, T).long().cuda(),
        'attention_mask' : torch.ones(B, T).cuda(),
    }

def _get_optimizer(model):
    return torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

if __name__ == "__main__":

    ## set distributed arguments (because deepspeed ckpt check rank for printing...)
    local_rank = int(os.environ["LOCAL_RANK"])
    get_accelerator().set_device(local_rank)
    device = torch.device(get_accelerator().device_name(), local_rank)
    deepspeed.init_distributed()

    ## get model, tokenizer and dummy optimizer
    model_path = "lmsys/vicuna-7b-v1.5"
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoModelForCausalLM.from_pretrained(model_path).cuda().train() # fp16 can cause nan 
    optimizer = _get_optimizer(model)

    ## vanilla gradient checkpointing
    model.gradient_checkpointing_enable()
    inputs = _get_dummy_inputs()
    optimizer.zero_grad()
    output1 = model(**inputs)
    output1.logits.sum().backward()
    tmp_grad1 = copy.deepcopy(model.model.layers[0].self_attn.q_proj.weight.grad.cpu())

    ## deepspeed CPU offloading, selective gradient checkpointing
    gradient_checkpointing_kwargs = {
        "use_deepspeed_grad_ckpt" : True,
        "num_checkpoints" : 4,
        "checkpoint_in_cpu" : True,
    }
    model.gradient_checkpointing_enable(gradient_checkpointing_kwargs)
    inputs = _get_dummy_inputs()
    optimizer.zero_grad()
    output2 = model(**inputs)
    output2.logits.sum().backward()
    tmp_grad2 = copy.deepcopy(model.model.layers[0].self_attn.q_proj.weight.grad.cpu())

    grad_allclose = torch.allclose(tmp_grad1, tmp_grad2, rtol=1e-05, atol=1e-08)
    assert grad_allclose
    print(f'''
    grad_allclose : {grad_allclose}
    tmp_grad1     : {tmp_grad1}
    tmp_grad2     : {tmp_grad2}
    ''')
```

and the result was like

```python
/path/to/dir/transformers$ python -m torch.distributed.launch test_act_ckpt.py
/path/to/dir/venv/transformers_pr/lib/python3.10/site-packages/torch/distributed/launch.py:183: FutureWarning: The module torch.distributed.launch is deprecated
and will be removed in future. Use torchrun.
Note that --use-env is set by default in torchrun.
If your script expects `--local-rank` argument to be set, please
change it to read from `os.environ['LOCAL_RANK']` instead. See 
https://pytorch.org/docs/stable/distributed.html#launch-utility for 
further instructions

  warnings.warn(
[2024-04-13 08:46:25,432] [INFO] [real_accelerator.py:191:get_accelerator] Setting ds_accelerator to cuda (auto detect)
[2024-04-13 08:46:26,405] [INFO] [comm.py:637:init_distributed] cdb=None
[2024-04-13 08:46:26,405] [INFO] [comm.py:668:init_distributed] Initializing TorchBackend in DeepSpeed with backend nccl
Loading checkpoint shards: 100%|████████████████████████████████████████████████████████████████| 2/2 [00:02<00:00,  1.50s/it]
/path/to/dir/transformers/src/transformers/generation/configuration_utils.py:497: UserWarning: `do_sample` is set to `False`. However, `top_p` is set to `0.6` -- this flag is only used in sample-based generation modes. You should set `do_sample=True` or unset `top_p`.
  warnings.warn(
`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`.
[2024-04-13 08:46:38,841] [INFO] [checkpointing.py:539:forward] Activation Checkpointing Information
[2024-04-13 08:46:38,841] [INFO] [checkpointing.py:540:forward] ----Partition Activations False, CPU CHECKPOINTING True
[2024-04-13 08:46:38,841] [INFO] [checkpointing.py:541:forward] ----contiguous Memory Checkpointing False with 4 total layers
[2024-04-13 08:46:38,841] [INFO] [checkpointing.py:543:forward] ----Synchronization False
[2024-04-13 08:46:38,841] [INFO] [checkpointing.py:544:forward] ----Profiling time in checkpointing False

    grad_allclose : True
    tmp_grad1     : tensor([[-2.6913e-12,  1.5395e-12,  8.9857e-13,  ..., -8.0048e-13,
          5.5549e-12, -4.5321e-13],
        [ 4.7248e-12, -2.7027e-12, -1.5775e-12,  ...,  1.4053e-12,
         -9.7519e-12,  7.9564e-13],
        [-4.4621e-12,  2.5525e-12,  1.4898e-12,  ..., -1.3271e-12,
          9.2097e-12, -7.5140e-13],
        ...,
        [-1.0189e-12,  5.8284e-13,  3.4018e-13,  ..., -3.0304e-13,
          2.1030e-12, -1.7158e-13],
        [-2.7976e-12,  1.6004e-12,  9.3407e-13,  ..., -8.3210e-13,
          5.7743e-12, -4.7112e-13],
        [ 7.3264e-12, -4.1910e-12, -2.4461e-12,  ...,  2.1791e-12,
         -1.5122e-11,  1.2337e-12]])
    tmp_grad2     : tensor([[-2.6913e-12,  1.5395e-12,  8.9857e-13,  ..., -8.0048e-13,
          5.5549e-12, -4.5321e-13],
        [ 4.7248e-12, -2.7027e-12, -1.5775e-12,  ...,  1.4053e-12,
         -9.7519e-12,  7.9564e-13],
        [-4.4621e-12,  2.5525e-12,  1.4898e-12,  ..., -1.3271e-12,
          9.2097e-12, -7.5140e-13],
        ...,
        [-1.0189e-12,  5.8284e-13,  3.4018e-13,  ..., -3.0304e-13,
          2.1030e-12, -1.7158e-13],
        [-2.7976e-12,  1.6004e-12,  9.3407e-13,  ..., -8.3210e-13,
          5.7743e-12, -4.7112e-13],
        [ 7.3264e-12, -4.1910e-12, -2.4461e-12,  ...,  2.1791e-12,
         -1.5122e-11,  1.2337e-12]])
```

my environment was

```
python -c "import torch; print(torch.__version__); \
import transformers; print(transformers.__version__); \
import deepspeed; print(deepspeed.__version__)"
```

```python
2.2.2+cu118
4.40.0.dev0
[2024-04-13 08:44:58,468] [INFO] [real_accelerator.py:191:get_accelerator] Setting ds_accelerator to cuda (auto detect)
0.14.0
```

Comments:
cc @younesbelkada 

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_30233). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

I'm glad to see this PR, it's very helpful for me. When will this feature be merged into the main branch and released?

cc @muellerzr @SunMarc 

@glowwormX Sorry for the late feedback. I almost forgot about this PR because there hasn't been an update in a while..!
i will check @younesbelkada 's comments ASAP

> @glowwormX抱歉反馈晚了。我差点忘了这条 PR，因为已经有一段时间没有更新了……！ 我会检查一下@younesbelkada尽快发表评论

Thank you very much for your reply, I would also like to know if there is a solution for the potential issue you mentioned in #29648 

> > @glowwormX抱歉反馈晚了。我差点忘了这条 PR，因为已经有一段时间没有更新了……！ 我会检查一下@younesbelkada尽快发表评论
> 
> Thank you very much for your reply, I would also like to know if there is a solution for the potential issue you mentioned in #29648

well... i have no idea about this. 
i guess it's due to the nature of `torch.utils.checkpoint`. 
however, most cases where activation checkpointing needs to be selective CPU offloaded will be LLM finetuning, in which case you don't need to get the intermediate outputs, so i didn't try to fix this anymore.