Accelerate.unwrap should fully unwrap the model.

[ggoggam]

System Info

- `Accelerate` version: 0.28.0
- Platform: Linux-5.15.0-124-generic-x86_64-with-glibc2.35
- Python version: 3.10.12
- Numpy version: 1.26.3
- PyTorch version (GPU?): 2.5.1+cu118 (True)
- PyTorch XPU available: False
- PyTorch NPU available: False
- System RAM: 250.55 GB
- GPU type: NVIDIA T4
- `Accelerate` default config:
        Not found

Information

• The official example scripts
• 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)
• My own task or dataset (give details below)

Reproduction

When the model is wrapped in both distributed wrapper (e.g. DistributedDataParallel or DeepSpeedEngine) and compiled module (e.g. OptimizedModule), calling Accelerator.unwrap_model should return a fully unwrapped model (torch.nn.Module). Instead, it currently returns the distributed wrapper.

import accelerate

# some model
model = ...
optimizer = ...

accelerator = accelerate.Accelerator(
    deepspeed_plugin=accelerate.utils.DeepSpeedPlugin(
        **{"zero_stage": 2, "train_micro_batch_size_per_gpu": 4}
    ),
    dynamo_backend="INDUCTOR"
)
model, optimizer = accelerator.prepare(model, optimizer)

# unwrap
unwrapped = accelerator.unwrap_model(model)

# assert
# currently raises error
assert model == unwrapped

Expected behavior

Accelerator.unwrap_model should return fully unwrapped model, i.e. torch.nn.Module.