No effect of gradient_checkpointing when training llama-2
Author: getao
Pull request: no

### System Info

transformers == 4.35.2
pytorch == 2.1.1

### Who can help?

 @ArthurZucker 
Hello, I'm training Llama-2 with flash-attn 2 using torchrun. However, I found that using gradient_checkpointing doesn't help save GPU memory. The training speed doesn't reduce either.

I doubt there is something wrong with the gradient_checkpointing. Could you please help take a look at the issue?

I enable gradient_checkpointing in the training_args.

### Information

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

### Tasks

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

### Reproduction

You can continuously train the llama-2-7b-hf with any text data with batch=1, seq_len=1024, using flash-attn.

One enables gradient_checkpointing and the other doesn't enable.

To compare the difference.

### Expected behavior

Training speed reduces and GPU memory cost reduced.

Comments:
Could you share a minimal reproducer that shows that there is no improvements / how you are enabling it? 

> Could you share a minimal reproducer that shows that there is no improvements / how you are enabling it?

```

def main():
    parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
    model_args, data_args, training_args = parser.parse_args_into_dataclasses()

    print(model_args)
    print(data_args)
    print(training_args)

    data_prefix = data_args.data_path
    
    train_file = f"{data_prefix}.train.json"
    eval_file = f"{data_prefix}.eval.json"
    
    dataset = load_dataset("json", data_files={"train": train_file, "eval": eval_file})
    train_dataset = dataset["train"]
    eval_dataset = dataset["eval"]
    tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path)
    train_dataset = train_dataset.map(tokenize_function, batched=True, fn_kwargs={"tokenizer": tokenizer, "max_seq_length": data_args.max_seq_length)
    eval_dataset = eval_dataset.map(tokenize_function, batched=True, fn_kwargs={"tokenizer": tokenizer, "max_seq_length": data_args.max_seq_length)

    model_download_flag = False
    while model_download_flag is False: 
        try:
            model = AutoModelForCausalLM.from_pretrained(model_args.model_name_or_path, torch_dtype=torch.float16 if training_args.bf16 is False else torch.bfloat16, use_flash_attention_2=True, resume_download=True)
            model_download_flag = True
        except:
            print("Downloading incomplete... Retry")

    train_model(model, train_dataset, eval_dataset, training_args)

main()

```

Above is the training script for training with llama-2.

The command line is:

```
torchrun --nproc-per-node=4 baseline.py --adam_beta2 0.95 --adam_epsilon 1e-6 --num_train_epochs $epoch --per_device_train_batch_size $batch --per_device_eval_batch_size $batch --gradient_accumulation_steps 1  \
          --learning_rate $lr --seed $seed --data_seed $seed **--gradient_checkpointing** --logging_steps 10 --save_strategy 'no' --evaluation_strategy 'steps' --eval_steps $eval_steps --save_steps $save_steps --bf16 --max_seq_length $max_seq_len --output_dir $OUT_DIR --logging_dir $OUT_DIR/logs --data_path $INPUT_DATA | tee $OUT_DIR/train.log

```

Sorry, you're gonna have to give more details about what the `train_model` function does. If you don't use the official trainer / training script then maybe you are just not enabling it.         `model.enable_gradient_checkpointing()` should be called. 


> model.enable_gradient_checkpointing()

I used the official trainer:

```
def train_model(model, train_dataset, eval_dataset, training_args, data_collator=None):

    last_checkpoint = None
    if os.path.isdir(training_args.output_dir):
        last_checkpoint = get_last_checkpoint(training_args.output_dir)
        if last_checkpoint is not None and training_args.resume_from_checkpoint is None:
            print(
                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
            )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        data_collator=data_collator
    )

    checkpoint = None
    if training_args.resume_from_checkpoint is not None:
        checkpoint = training_args.resume_from_checkpoint
    elif last_checkpoint is not None:
        checkpoint = last_checkpoint

    print(f"Loaded from the checkpoint: {checkpoint}")

    train_result = trainer.train(resume_from_checkpoint=checkpoint)
    
    trainer.save_model()
    trainer.log_metrics("train", train_result.metrics)
    metrics = trainer.evaluate()
    trainer.log_metrics("eval", metrics)
    trainer.save_metrics("eval", metrics)
```

Or, do you mean that I still need to add ```model.enable_gradient_checkpointing()``` even if I use the default official Trainer?

Not necessarily no, the code here https://github.com/huggingface/transformers/blob/04e67503d804c5f51a667e2db8e817c1d744048b/src/transformers/trainer.py#L1650 should automatically enable it if the arg is properly passed. As it is one of the most important aspect w.r.t training with peft, this is heavily tested. 
Could you make sure that it is set by printing the `training_args`? 

> Not necessarily no, the code here
> 
> https://github.com/huggingface/transformers/blob/04e67503d804c5f51a667e2db8e817c1d744048b/src/transformers/trainer.py#L1650
> 
> should automatically enable it if the arg is properly passed. As it is one of the most important aspect w.r.t training with peft, this is heavily tested.
> Could you make sure that it is set by printing the `training_args`?

Sure, I checked this and found it was enabled by printing training_args:

gradient_checkpointing=True,
gradient_checkpointing_kwargs=None,

Also, I have observed some logs showing the gradient_checkpointing is enabled before the training progress bar appears:

`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...

Unfortunately, no GPU memory saving is observed.

I am probably missing something so pinging @younesbelkada here! 🤗 


hi @getao 
Thanks for the issue! I don't think training should be faster with GC - do you observe the same behavioru without FA?

> hi @getao Thanks for the issue! I don't think training should be faster with GC - do you observe the same behavioru without FA?

Hi @younesbelkada , I didn't say that GC helps training become faster. GC is thought to save memory but at a cost of slowing down training. What I said is that I found enabling GC didn't save memory and didn't slow down training, which is weird to me.

BTW, What is FA?

I see thanks for explaining, OK! 
Sorry I meant Flash Attention, I wonder if it is not a weird interaction between Flash Attention and GC

> I see thanks for explaining, OK! Sorry I meant Flash Attention, I wonder if it is not a weird interaction between Flash Attention and GC

Oh, I see. I'll take a look at it without Flash attention.

Regarding potential weird interaction between Flash Attention and GC -- The HF gradient checkpointing will cause inefficient redundant computation when using flash attention. TL;DR: it (computationally equivalent) computes a redundant FA forward.

The reason is the flash attention backward kernel will recompute the Q^tK for rematerialization, which is similar to the recomputation in the backward process of the original HF gradient checkpointing. 

If you're interested, we implemented a more efficient flash attention friendly gradient checkpointing in [FastCkpt](https://github.com/RulinShao/FastCkpt/tree/main) to mitigate this issue, where you just need to pip install fastckpt and import a monkey patch to accelerate your training by saving one flash attention forward in every layer compared with the HF gradient checkpointing.

> Regarding potential weird interaction between Flash Attention and GC -- The HF gradient checkpointing will cause inefficient redundant computation when using flash attention. TL;DR: it (computationally equivalent) computes a redundant FA forward.
> 
> The reason is the flash attention backward kernel will recompute the Q^tK for rematerialization, which is similar to the recomputation in the backward process of the original HF gradient checkpointing.
> 
> If you're interested, we implemented a more efficient flash attention friendly gradient checkpointing in [FastCkpt](https://github.com/RulinShao/FastCkpt/tree/main) to mitigate this issue, where you just need to pip install fastckpt and import a monkey patch to accelerate your training by saving one flash attention forward in every layer compared with the HF gradient checkpointing.

Thank you for your answer. It makes sense for the speed issue. However, does it make sense that the GPU memory is not saved?

Hello, just an update.

I find that --gradient_checkpointing is highly useful to save memory when it is used with deepspeed along with the deepspeed config file. 

However, if I use it in torchrun (without any specified config), it doesn't work to save memory. I wonder if there is anything that I missed for properly enabling gradient checkpointing.

Any help would be appreciated.

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.

@getao 
Hi. Did you solve issue?