Repository: huggingface/accelerate
Pull request: fix torch compile with FSDP
# What does this PR do?
1. Fixes https://github.com/huggingface/accelerate/issues/1889

Problem statement:
ValueError: Expected _orig_mod to NOT be FullyShardedDataParallel if using an `auto_wrap_policy`
### System Info

```Shell
- `Accelerate` version: 0.22.0
- Platform: Linux-5.4.0-121-generic-x86_64-with-glibc2.31
- Python version: 3.10.11
- Numpy version: 1.24.3
- PyTorch version (GPU?): 2.0.1 (True)
- PyTorch XPU available: False
- PyTorch NPU available: False
- System RAM: 125.74 GB
- GPU type: NVIDIA GeForce RTX 3090
- `Accelerate` default config:
        - compute_environment: LOCAL_MACHINE
        - distributed_type: FSDP
        - mixed_precision: fp16        
        - use_cpu: False        
        - debug: True
        - num_processes: 4
        - machine_rank: 0
        - num_machines: 1
        - rdzv_backend: static
        - same_network: True
        - main_training_function: main        
        - fsdp_config: {'fsdp_auto_wrap_policy': 'TRANSFORMER_BASED_WRAP',
           'fsdp_backward_prefetch_policy': 'BACKWARD_PRE',
           'fsdp_forward_prefetch': False,
           'fsdp_offload_params': False,
           'fsdp_sharding_strategy': 1,
           'fsdp_state_dict_type': 'SHARDED_STATE_DICT', 
           'fsdp_sync_module_states': True, 
           'fsdp_transformer_layer_cls_to_wrap': 'LlamaDecoderLayer',
           'fsdp_use_orig_params': False}
        - downcast_bf16: no
        - tpu_use_cluster: False
        - tpu_use_sudo: False
        - tpu_env: []
        - dynamo_config: {'dynamo_backend': 'INDUCTOR'}
```


### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] One of the scripts in the examples/ folder of Accelerate or an officially supported `no_trainer` script in the `examples` folder of the `transformers` repo (such as `run_no_trainer_glue.py`)
- [X] My own task or dataset (give details below)

### Reproduction

I am trying to fine tune LLama-2 using FSDP and accelerate, but when i run `accelerate launch finetune_base_model.py` i get `ValueError: Expected _orig_mod to NOT be FullyShardedDataParallel if using an auto_wrap_policy`. Below is the code that should reproduce the issue

```python
from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer
import torch

model_name = "meta-llama/Llama-2-7b-hf"

model = AutoModelForCausalLM.from_pretrained(
        model_name, torch_dtype=torch.float16)

dataset = {'input_ids': [1,23,1,31,41]}

tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token

training_args = TrainingArguments(
    output_dir="./finetuned_model",
    overwrite_output_dir=True,
    max_steps=500,
    per_device_train_batch_size=1,
    gradient_accumulation_steps=1,
    learning_rate=1e-5,
    lr_scheduler_type="linear",
    dataloader_num_workers=20,
    warmup_steps=100,
    save_steps=250,
    save_total_limit=1,
    report_to="wandb",
    run_name="finetune",
    remove_unused_columns=False,
    push_to_hub=True,
    hub_model_id="Llama-2",
    hub_strategy="every_save",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    tokenizer=tokenizer
)

trainer.train()
```

### Expected behavior

Model to be distributed across the gpus and training to begin

Hints from issue discussion:
cc @pacman100 
Hello, I'm able to run it if I remove the torch compile config `dynamo_config: {'dynamo_backend': 'INDUCTOR'}`. torch compile seems to be causing issues with FSDP. Could you raise this issue with the PyTorch team as it doesn't pertain to the integration.
Hello, I looked into this, and the above PR should fix this.

1. fsdp_config.json:
```
{
    "backward_prefetch": "backward_pre",
    "forward_prefetch": "True",
    "limit_all_gathers": "False",
    "use_orig_params": "True",
    "sync_module_states": "True",
    "activation_checkpointing": "False",
    "transformer_layer_cls_to_wrap": "BertLayer"
}
```
2. running `run_glue.py` example from transformers
```
cd transformers
export TASK_NAME=mrpc
export CUDA_VISIBLE_DEVICES="0,1"
torchrun --nnodes 1 --nproc-per-node 2 ./examples/pytorch/text-classification/run_glue.py   --model_name_or_path bert-base-cased   --task_name $TASK_NAME   --do_train   --do_eval   --max_seq_length 128   --per_device_train_batch_size 16   --learning_rate 5e-5   --num_train_epochs 3   --output_dir /tmp/$TASK_NAME/ --overwrite_output_dir --fsdp "full_shard auto_wrap" --bf16 --fsdp_config fsdp_config.json --torch_compile --pad_to_max_length --dataloader_drop_last
```

3. output logs:
```
...
18 [rank0]:[2023-09-01 12:01:34,636] torch.utils._sympy.value_ranges: [WARNING] unhandled ValueRange op load_seed
...

40 [rank0]:[2023-09-01 12:01:45,330] torch.utils._sympy.value_ranges: [WARNING] unhandled ValueRange op rand
41  13%|███████████▍                                                                             | 44/342 [00:48<00:56,  5.27it/s]
```

4. Versions:
`accelerate env`
```
- `Accelerate` version: 0.23.0.dev0
- Platform: Linux-5.4.0-125-generic-x86_64-with-glibc2.31
- Python version: 3.10.11
- Numpy version: 1.25.0
- PyTorch version (GPU?): 2.1.0.dev20230725+cu118 (True)
- PyTorch XPU available: False
- PyTorch NPU available: False
- GPU type: NVIDIA A100-SXM4-80GB
- `Accelerate` default config:
	Not found
```

`transformers-cli env`
```
- `transformers` version: 4.33.0.dev0
- Platform: Linux-5.4.0-125-generic-x86_64-with-glibc2.31
- Python version: 3.10.11
- Huggingface_hub version: 0.16.4
- Safetensors version: 0.3.1
- Accelerate version: 0.23.0.dev0
- Accelerate config: 	not found
- PyTorch version (GPU?): 2.1.0.dev20230725+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?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>
```

Triton version: `2.1.0`

Commit message: addressing comments

Patch to review:
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 701dfb2ba30..f9f9b9fa2ab 100755
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -96,6 +96,7 @@
     wait_for_everyone,
 )
 from .utils.constants import FSDP_PYTORCH_VERSION
+from .utils.other import is_compiled_module
 
 
 if is_deepspeed_available():
@@ -1215,7 +1216,11 @@ def prepare(self, *args, device_placement=None):
             for obj in args:
                 if isinstance(obj, torch.nn.Module):
                     model_count += 1
-                    is_type_fsdp = type(obj) == FSDP
+                    # if the model is compiled using PyTorch 2.0,
+                    # check that the wrapped model is FSDP or not;
+                    # else check if it is FSDP or not;
+                    unwrapped_model = obj._orig_mod if is_compiled_module(obj) else obj
+                    is_type_fsdp = isinstance(obj, FSDP) or isinstance(unwrapped_model, FSDP)
                 if isinstance(obj, torch.optim.Optimizer):
                     optimizer_present = True
             if model_count > 1 and optimizer_present:
@@ -1371,7 +1376,7 @@ def prepare_model(self, model: torch.nn.Module, device_placement: bool = None, e
         elif device_placement and not self.verify_device_map(model):
             model = model.to(self.device)
 
-        if self.native_amp and self.distributed_type != DistributedType.FSDP:
+        if self.native_amp:
             model._original_forward = model.forward
             model_forward_func = model.forward.__func__ if hasattr(model.forward, "__func__") else model.forward
             autocast_context = get_mixed_precision_context_manager(self.native_amp, self.autocast_handler)
@@ -1423,7 +1428,12 @@ def prepare_model(self, model: torch.nn.Module, device_placement: bool = None, e
 
                 # Check if the model is already a FSDP model due to `Manual Wrapping` and if so,
                 # don't wrap it again
-                if type(model) != FSDP:
+                # In case the model is already compiled using PyTorch 2.0 and the wrapped model in it
+                # is a FSDP model, don't wrap it again
+                unwrapped_model = model._orig_mod if is_compiled_module(model) else model
+                is_type_fsdp = isinstance(model, FSDP) or isinstance(unwrapped_model, FSDP)
+
+                if not is_type_fsdp:
                     self.state.fsdp_plugin.set_auto_wrap_policy(model)
                     fsdp_plugin = self.state.fsdp_plugin
                     kwargs = {
@@ -1456,15 +1466,17 @@ def prepare_model(self, model: torch.nn.Module, device_placement: bool = None, e
                             ),
                             auto_wrap_policy=fsdp_plugin.auto_wrap_policy,
                         )
-
+                # if the previous and current models are same, delete the previous one
+                if len(self._models) > 1 and (self._models[-2] is self._models[-1]):
+                    del self._models[-2]
                 self._models[-1] = model
             elif self.distributed_type == DistributedType.MULTI_CPU:
                 kwargs = self.ddp_handler.to_kwargs() if self.ddp_handler is not None else {}
                 model = torch.nn.parallel.DistributedDataParallel(model, **kwargs)
             elif self.distributed_type == DistributedType.TPU and self.state.fork_launched:
                 model = xmp.MpModelWrapper(model).to(self.device)
-        # torch.compile should be called last.
-        if self.state.dynamo_plugin.backend != DynamoBackend.NO:
+        # torch.compile should be called last and only if the model isn't already compiled.
+        if self.state.dynamo_plugin.backend != DynamoBackend.NO and not is_compiled_module(model):
             if not is_torch_version(">=", "2.0"):
                 raise ValueError("Using `torch.compile` requires PyTorch 2.0 or higher.")
             model = torch.compile(model, **self.state.dynamo_plugin.to_kwargs())


Code review comments:
diff_hunk: @@ -1215,7 +1216,11 @@ def prepare(self, *args, device_placement=None):
             for obj in args:
                 if isinstance(obj, torch.nn.Module):
                     model_count += 1
-                    is_type_fsdp = type(obj) == FSDP
+                    # if the model is compiled using PyTorch 2.0,
+                    # check that the wrapped model is FSDP or not;
+                    # else check if it is FSDP or not;
+                    unwrapped_model = obj._orig_mod if is_compiled_module(obj) else obj
+                    is_type_fsdp = isinstance(obj, FSDP) or isinstance(unwrapped_model, FSDP)
path: src/accelerate/accelerator.py
text: @user1:
Minor nit, but note that if the obj is not compiled, the `isinstance` check is performed on the object twice. It doesn't really matter for performance, it could just be confusing for readers of the code.

@author:
Done, resolved it in latest commit.