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

System Info

- `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
• 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 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:

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

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.