cannot use activation_checkpoint in torch native fsdp
Author: hanwen-sun
Pull request: no
Labels: name: bug

name: PyTorch FSDP

### System Info

- `transformers` version: 4.43.0.dev0
- Platform: Linux-5.10.112-005.ali5000.alios7.x86_64-x86_64-with-glibc2.31
- Python version: 3.10.13
- Huggingface_hub version: 0.23.3
- Safetensors version: 0.4.3
- Accelerate version: 0.30.1
- Accelerate config:    not found
- PyTorch version (GPU?): 2.2.2 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: <fill in>
- Using GPU in script?: <fill in>
- GPU type: NVIDIA A100-SXM4-80GB

### Who can help?

_No response_

### Information

- [X] The official example scripts
- [ ] 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

we use the run_clm.py script to train model.

the training shell:
```shell
set -ex

echo "Running a native torch job ..."

export USE_TORCH_XLA=0

[ -z "$RANK" ] && RANK=0
[ -z "$WORLD_SIZE" ] && WORLD_SIZE=1
[ -z "$MASTER_ADDR" ] && MASTER_ADDR=127.0.0.1
[ -z "$MASTER_PORT" ] && MASTER_PORT=9010

BS=1
SEQLEN=4096
NPROC_PER_NODE=2
PRECISION="bf16=true"
FSDP_CONFIG="llama3_fsdp_native.json"
JOB_NAME="LLAMA3_FSDP_NATIVE_GPU${NPROC_PER_NODE}_BS${BS}_SEQLEN${SEQLEN}_BF16_FA"


torchrun --nproc_per_node $NPROC_PER_NODE \
    --nnodes $WORLD_SIZE \
    --node_rank $RANK \
    --master_port $MASTER_PORT \
    --master_addr $MASTER_ADDR \
   your_path/transformers/examples/pytorch/language-modeling/run_clm.py \
    --num_train_epochs 2 \
    --dataset_name wikitext \
    --dataset_config_name wikitext-2-raw-v1 \
    --use_fast_tokenizer false \
    --per_device_train_batch_size $BS \
    --per_device_eval_batch_size $BS \
    --do_train \
    --output_dir /tmp/test-clm \
    --overwrite_output_dir \
    --config_name your_path/Meta-Llama-3-8B/ \
    --tokenizer_name your_path/Meta-Llama-3-8B/ \
    --trust_remote_code true \
    --cache_dir ./cache \
    --block_size $SEQLEN \
    --optim adamw_torch \
    --save_strategy no \
    --logging_strategy steps \
    --gradient_checkpointing no \
    --logging_steps 100 \
    --$PRECISION \
    --fsdp "auto_wrap" \
    --fsdp_config $FSDP_CONFIG 2>&1 | tee ./$JOB_NAME.log
```

the fsdp_config:
```
{
    "fsdp_transformer_layer_cls_to_wrap": [
        "LlamaDecoderLayer"
    ],
    "activation_checkpointing": true
}
```

And you can replace the llama model by any other model like gpt2

then run the shell.


### Expected behavior

you will meet the error like:
```
tensor at position 18:
saved metadata: {'shape': torch.Size([2, 4096, 8, 128]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=3)}
recomputed metadata: {'shape': torch.Size([2, 8192, 8, 128]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=3)}
raise CheckpointError(
torch.utils.checkpoint: Recomputed values for the following tensors have different metadata than during the forward pass.
```

Comments:
cc @muellerzr @SunMarc 

Hi, I have a question about the scripts. In the scripts, do_train is not explicitly set to true, but --do_train \ is included and seems to be empty. Does this mean that when I call trainer.train(), it will perform training and not evaluation?

https://github.com/huggingface/transformers/blob/e316c5214fe51de0bf8e824245bfd6225c9925aa/src/transformers/trainer.py#L1837

@hanwen-sun  @amyeroberts 

> do_train is not explicitly set to true, but --do_train \ is included and seems to be empty

`--do_train` is a flag such that if it's included in the command it's evaluated as `True`. This controls the value of `training_args.do_train` set when running the script.

Calling `trainer.train` will running training. 

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.

I recently ran into what appears to be the same issue.

If I run with `gradient_checkpointing: true`, I get a warning during training that reads: `activation_checkpointing: "true" # When using FSDP full shard, instead of using `gradient_checkpointing` in TrainingArguments, please use `activation_checkpointing` in `fsdp_config`. The former introduces a redundant AllGather operation in backward pass. Reference: https://github.com/huggingface/transformers/issues/30404`. 

But when I set `gradient_checkpointing: false` and `activation_checkpointing: true`, then I get the error message originally indicated in this issue: 

```
torch.utils.checkpoint.CheckpointError: torch.utils.checkpoint: Recomputed values for the following tensors have different metadata than during the forward pass
```

I'm running my training script with the following command: `torchrun --nnodes 1 --nproc_per_node 8 train.py`

```
gradient_checkpointing: false 
gradient_checkpointing_kwargs:
  use_reentrant: false
attn_implementation: "flash_attention_2"
packing: false
bf16: "auto"

fsdp: "full_shard auto_wrap offload"
fsdp_config:
  limit_all_gathers: true
  backward_prefetch: "backward_pre"
  forward_prefetch: "false"
  use_orig_params: "false"
  min_num_params: 0
  activation_checkpointing: "true"
```