huggingface / huggingface/accelerate

compile_regions runs the uncompiled model when the root module has an accelerate hook

Open
#4,253 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
9.9k
Forks
1.5k
Avg merge
5d 2h
Merged PRs (30d)
27

Description

### System Info

```Shell
- `Accelerate` version: 1.16.0.dev0
- Platform: Linux-5.15.133.1-microsoft-standard-WSL2-x86_64-with-glibc2.35
- `accelerate` bash location: Not found
- Python version: 3.13.15
- Numpy version: 2.5.2
- PyTorch version: 2.11.0+cu130
- PyTorch accelerator: N/A
- System RAM: 15.47 GB
- `Accelerate` default config:
Not found
```

Source checkout at `f13f7c13b64c10b6eb7e2d73171ea1b94c748701`.

### 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

When the root module carries an accelerate hook, the model returned by `compile_regions` never calls its compiled blocks. Nothing errors, the forward just runs the original uncompiled blocks, so regional compilation silently does nothing. The script counts calls to the `OptimizedModule` wrappers that `compile_regions` creates:

```python
import torch
from torch import nn
from accelerate import cpu_offload
from accelerate.hooks import ModelHook, add_hook_to_module
from accelerate.utils import compile_regions

class Block(nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(4, 4)

def forward(self, x):
return self.linear(x)

class Model(nn.Module):
def __init__(self):
super().__init__()
self.blocks = nn.ModuleList([Block() for _ in range(3)])

def forward(self, x):
for block in self.blocks:
x = block(x)
return x

def count_compiled_block_calls(model):
compiled = compile_regions(model, backend="eager")
calls = []
for block in compiled.blocks: # the OptimizedModule wrappers created by compile_regions
block.register_forward_pre_hook(lambda module, args: calls.append(module))
compiled(torch.randn(2, 4))
return type(compiled.blocks[0]).__name__, len(calls)

hooked = Model()
add_hook_to_module(hooked, ModelHook())

print("plain model: ", count_compiled_block_calls(Model()))
print("add_hook_to_module: ", count_compiled_block_calls(hooked))
print("cpu_offload: ", count_compiled_block_calls(cpu_offload(Model(), execution_device=torch.device("cpu"))))
```

```
plain model: ('OptimizedModule', 3)
add_hook_to_module: ('OptimizedModule', 0)
cpu_offload: ('OptimizedModule', 0)
```

The same happens for `dispatch_model` whenever it attaches hooks: with `device_map={"blocks.0": "cpu", "blocks.1": "disk", "blocks.2": "cpu"}` the root gets an `AlignDevicesHook` and the count is 0, while `device_map={"": "cpu"}` (no hooks) gives 3. It also goes through `Accelerator.prepare`: with `TorchDynamoPlugin(backend="eager", use_regional_compilation=True)`, preparing the hooked model gives 0 calls both with `mixed_precision="no"` and with `mixed_precision="bf16"`, and a plain model gives 3 in both cases.

Why: for a module that contains repeated blocks, `_compile_regions` builds a shallow copy (`module.__class__.__new__` followed by `__dict__.update`) and then replaces the children. #4188 rebinds instance attributes that are bound methods of the original module, but `add_hook_to_module` sets

```python
module.forward = functools.update_wrapper(functools.partial(new_forward, module), old_forward)
```

A `functools.partial` has no `__func__`, so the copy keeps a `forward` that calls `new_forward` with the original module. That runs `module._old_forward` on the original, whose children are still the uncompiled blocks.

Under mixed precision there is a second layer. For a hooked model `model.forward` has no `__func__`, so `prepare_model` takes the first branch and sets `model.forward = convert_outputs_to_fp32(autocast_context(model_forward_func))`, a plain function that closes over the same partial. Rebinding partials in `_compile_regions` would therefore fix `compile_regions`, `cpu_offload` and `dispatch_model` models and `prepare` without mixed precision, but not `prepare` with `fp16`/`bf16`.

### Expected behavior

The model returned by `compile_regions` (and by `Accelerator.prepare` with regional compilation) runs its compiled blocks whether or not the root module has a hook, i.e. 3 calls for all three cases in the script.

I can see two parts to a fix: rebinding a `functools.partial` whose first argument is the original module when `_compile_regions` copies the instance dict, and making the mixed precision wrapper in `prepare_model` rebindable for hooked models. A quick local patch that rebinds such partials brings the count back to 3 for `compile_regions` on the hooked, `cpu_offload` and `dispatch_model` models and for `prepare` with `mixed_precision="no"`, while `prepare` with `bf16` stays at 0. I am happy to open a PR, but wanted to check first which way you would prefer for the mixed precision part.

Contributor guide

Open the contributing guide

Research direction

Start by tracing _compile_regions and prepare_model, focusing on how copied instance attributes and mixed-precision wrappers retain the hooked module. Reproduce the three-block call counts for add_hook_to_module, cpu_offload, dispatch_model, and Accelerator.prepare, then add regression coverage. Done means every hooked and mixed-precision case invokes all three compiled blocks.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
52/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.