Fix compile_regions running the uncompiled module

[hjinnkim]

What does this PR do?

Under mixed precision, compile_regions returns a module that runs the uncompiled original: nothing is traced, and gradient checkpointing set after prepare() never reaches the executing object. Its __dict__ copy carries instance methods still bound to the source module; this PR re-binds them.

Problem

# src/accelerate/accelerator.py:1818-1829 — prepare_model, before compiling
if self.native_amp:
    model._original_forward = model.forward
    ...
    model.forward = MethodType(convert_outputs_to_fp32(model.forward.__func__), model)

# src/accelerate/utils/other.py:158-160 — compile_regions
elif has_repeated_blocks(module):
    new_module = module.__class__.__new__(module.__class__)
    new_module.__dict__.update(module.__dict__)  # copies `forward`, bound to `module`

The autocast wrapper is an instance attribute bound to model; the copy compile_regions builds afterwards (accelerator.py:2064) inherits the binding, so _call_impl reads self.forward from __dict__ and runs the original with its own uncompiled children.

Trigger: mixed_precision != "no" with use_regional_compilation on a model with repeated blocks. Unaffected: plain torch.compile and the in-place compile_regions_fsdp2 / compile_regions_deepspeed.

Reproduction

CPU only, single process.

import torch
import torch._dynamo
from torch import nn
from accelerate import Accelerator
from accelerate.utils import TorchDynamoPlugin


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

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


class Tiny(nn.Module):
    def __init__(self):
        super().__init__()
        self.blocks = nn.ModuleList([Block(), Block()])

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


original = Tiny()
accelerator = Accelerator(
    mixed_precision="bf16",
    dynamo_plugin=TorchDynamoPlugin(backend="inductor", use_regional_compilation=True),
)
twin = accelerator.prepare_model(original)
print("twin.blocks[0]:", type(twin.blocks[0]).__name__)
print("bound to ORIGINAL:", twin.__dict__["forward"].__self__ is original)
torch._dynamo.utils.counters.clear()
twin(torch.ones(1, 4))
print("dynamo frames:", dict(torch._dynamo.utils.counters["frames"]))

v1.14.0      twin.blocks[0]: OptimizedModule   bound to ORIGINAL: True    dynamo frames: {}
this branch  twin.blocks[0]: OptimizedModule   bound to ORIGINAL: False   dynamo frames: {'total': 2, 'ok': 2}

Fix

             new_module.__dict__.update(module.__dict__)
+            for name, value in list(new_module.__dict__.items()):
+                if hasattr(value, "__func__") and getattr(value, "__self__", None) is module:
+                    new_module.__dict__[name] = MethodType(value.__func__, new_module)
             new_module._modules = {}

• The __self__ is module guard limits the rewrite to entries left pointing at the source; MethodType is already imported, and the _compile_regions recursion covers nested copies.
• _original_forward is rebound too, so unwrap_model(keep_fp32_wrapper=False) restores the right binding.
• Alternative: install the autocast wrapper after compiling. Not taken — it reorders the FP8/TE wrapping.

Tests

tests/test_compile.py gains RegionalCompilationRebindTester: CPU-only, backend="eager", separate from the @skipped RegionalCompilationTester.

• test_instance_bound_methods_are_rebound — the returned module runs, sees post-compile state, blocks are OptimizedModule. Fails on main.
• test_no_instance_bound_methods_is_a_no_op — negative control; passes either way.

pytest tests/test_compile.py: 2 passed, 5 skipped.

Measured effect

40 steps, bf16 LoRA on a diffusion transformer, one B200, checkpointing and use_regional_compilation on; only accelerate differs.

	v1.14.0	this branch
peak VRAM	65,378 MiB	30,496 MiB
inductor artifacts	0	270
dynamo frames	{}	15

On v1.14.0 neither checkpointing nor compilation took effect.

Note for reviewers

• DDP: measured on CPU only; prepare_model wraps before compiling, leaving the binding on a nested copy.
• MS-AMP installs a plain function, not a bound method: out of scope.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[SunMarc]

Thanks for the fix !