Fix regional compilation not executing under mixed precision

[hjinnkim]

Summary

• dynamo_use_regional_compilation: true under mixed_precision other than no builds a compiled
  model that never executes; the uncompiled one runs instead.
• State set on the prepared model afterwards lands on that unused object, so
  gradient_checkpointing: true changes nothing.
• This re-binds the mixed-precision forward onto the module prepare() returned. It affects every
  model family; the FSDP2 and DeepSpeed prepare paths are not.

Root cause

Accelerator.prepare_model binds the autocast wrapper as an instance attribute on its argument
(accelerator.py:1818-1829), then compiles (:2064). compile_regions
(utils/other.py:106-175) builds a twin with new_module.__dict__.update(module.__dict__),
replacing only _modules, so the copied forward stays bound to the pre-compile module, which
the twin re-enters.

Trigger: mixed_precision != "no" with dynamo_use_regional_compilation: true. Unaffected: FSDP2
and DeepSpeed, whose compile_regions_* helpers (utils/other.py:178-225) compile in place and
build no twin.

Changes

• helpers/training/wrappers.py: rebind_prepared_forward re-binds forward and
  _original_forward in place on the module prepare returned, only where they still point at
  the original; DDP, FSDP, torch.compile, mixed_precision: no and no-dynamo are untouched.
• trainer.py:4577 and :4640, the two sites that store a prepared model.
  prepare_model(evaluation_mode=True) builds the same twin; without the rebind, the .eval()
  below would leave the executing module at training=True.
• The defect is accelerate's; this helper is a no-op once accelerate re-binds the twin itself.

Validation

• .venv/bin/python -m unittest -v -f tests.test_wrappers_rebind — 9 tests driven by accelerate's
  own compile_regions and the neighbouring wrapper shapes, six of them no-ops.
• 40 steps, B200, bf16 krea2 LoRA, 1024 px. On 632d9595f checkpointing on and off both peak at
  65,378 MiB, 0 inductor artifacts; with the fix, checkpointing on peaks at 30,496 MiB, writes
  270, and runs 0.4667 s/it after warm-up against the same tree's 0.600 with compilation off.
• Compilation off: both trees 31,398 MiB, step_loss 0.0481.

Notes

• Not covered: DDP's nested wrapper (no instance forward, so a no-op) and the text-encoder
  prepare at trainer.py:4615.

[bghira]

@copilot is this a single gpu only issue or impacted ddp as well? regional compile definitely changes performance in testing on multigpu systems. this is confusing me

[hjinnkim]

DDP is affected the same way, one level down, and this PR's helper does not reach it; huggingface/accelerate#4188 does.

Why (accelerate 1.14.0 line numbers): prepare_model installs the autocast forward on the model (accelerator.py:1818-1829), wraps it in DDP (:2053), then compile_regions (:2064) copies the DDP wrapper and the inner model, and the inner twin's forward stays bound to the original inner model. The DDP wrapper has no instance forward, so rebind_prepared_forward from this PR is a no-op there (the Notes line).

Reproducer — 2 processes, no data or checkpoint (CPU box: gloo, GPUs: NCCL). bound to ORIGINAL checks which object the autocast forward accelerate installed is bound to; frames = torch._dynamo.utils.counters["frames"]; file count = a fresh TORCHINDUCTOR_CACHE_DIR:

import torch, 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


accelerator = Accelerator(
    cpu=not torch.cuda.is_available(),
    mixed_precision="bf16",
    dynamo_plugin=TorchDynamoPlugin(backend="inductor", use_regional_compilation=True),
)
original = Tiny()
prepared = accelerator.prepare(original)
inner = prepared.module  # DDP wrapper -> inner twin
torch._dynamo.utils.counters.clear()
prepared(torch.ones(1, 4, device=accelerator.device))
print(
    f"rank {accelerator.process_index} | {type(prepared).__name__} "
    f"| inner forward bound to ORIGINAL: {inner.__dict__['forward'].__self__ is original} "
    f"| blocks[0]={type(inner.blocks[0]).__name__} "
    f"| dynamo frames={dict(torch._dynamo.utils.counters['frames'])}"
)

$ pip install accelerate==1.14.0 && TORCHINDUCTOR_CACHE_DIR=/tmp/ind torchrun --nproc_per_node 2 repro.py
rank 0 | DistributedDataParallel | inner forward bound to ORIGINAL: True | blocks[0]=OptimizedModule | dynamo frames={}
rank 1 | DistributedDataParallel | inner forward bound to ORIGINAL: True | blocks[0]=OptimizedModule | dynamo frames={}
$ find /tmp/ind -type f | wc -l
0
$ pip install git+https://github.com/hjinnkim/accelerate@compile-regions-rebind-instance-methods   # accelerate#4188
$ torchrun --nproc_per_node 2 repro.py
rank 0 | DistributedDataParallel | inner forward bound to ORIGINAL: False | blocks[0]=OptimizedModule | dynamo frames={'total': 2, 'ok': 2}
rank 1 | DistributedDataParallel | inner forward bound to ORIGINAL: False | blocks[0]=OptimizedModule | dynamo frames={'total': 2, 'ok': 2}
$ find /tmp/ind -type f | wc -l
63

torch 2.14.0.dev; same result on CPU (gloo) and with backend="eager".

Trainer (train.py from each tree): accelerate launch --multi_gpu --num_processes 2 --mixed_precision bf16, 2 × B200, krea2 LoRA rank 64, batch 1 per GPU, 1024 px, 40 steps, gradient_checkpointing: true, dynamo_backend: inductor, dynamo_use_regional_compilation: true, identical config in all three. Peak = 1 Hz nvidia-smi maximum minus idle usage; artifacts = files under a fresh per-rank TORCHINDUCTOR_CACHE_DIR; frames = the counter above at exit. The third column installs the #4188 branch over 1.14.0.

	main	this PR	this PR + accelerate#4188
peak VRAM per GPU, above baseline	66,494 MiB	66,494 MiB	31,614 MiB
inductor artifacts per rank	0	0	270
dynamo frames, both ranks	{}	{}	19/19

Regional compile changing multi-GPU performance is consistent with FSDP2/DeepSpeed (compile_regions_fsdp2/_deepspeed compile in place, no twin) or mixed_precision: no — both unaffected.

If you want DDP covered here rather than waiting on accelerate, I can extend the helper to walk .module wrappers and add the 2-process test.

[bghira]

did an Accelerate update break this or something? because running tests with regional compile just a few months back (or maybe it's been a year, time flies) was materially changing vram consumption (often for the worse, which this would explain) and performance, but not as much as i'd expect. still, sometimes i'd see the ~30% gain in speed - this is without deepspeed or fsdp2 enabled.

[hjinnkim]

This bug has existed since accelerate v1.7.0, where regional compilation first shipped (huggingface/accelerate#3529): compile_regions has built the copy the same way ever since, and prepare_model has installed the instance-bound autocast forward before it in every release through 1.14.0, so any bf16 run with regional compilation on has hit it. What you saw fits that: VRAM getting worse is checkpointing landing on the copy that never runs, and the occasional ~30 % speedup would be a run where the copy does execute — mixed_precision: no, or a model without a repeated-block ModuleList. I am confirming that on krea2 with the trainer now (precision × compile mode × single/DDP) and will post the table with VRAM and step time.

PR status: the first commit (939505d2e) only rebound the object prepare() returns, which under DDP is the DistributedDataParallel wrapper — so it fixed single-process runs and was a no-op for DDP. The second commit (6520cd052) walks .module wrappers and rebinds the inner twin, and is what the branch carries now.

Verified, 2 × B200, bf16 krea2 LoRA, checkpointing and regional compilation on, 40 steps, accelerate 1.14.0:

	1 GPU, main	1 GPU, this PR	2 GPU DDP, main	2 GPU DDP, this PR
peak VRAM per GPU, above baseline	65,378 MiB	30,496 MiB	66,494 MiB	31,614 MiB
inductor artifacts per rank	0	270	0	270

[bghira]

did you single bf16 out because certain quant methods bypass the issue somehow? for what it's worth, i was trying compile with eg. sdnq and torchao to fuse their compute paths via triton kernels. but i won't push back on merging this one in, it's just that historically it's not been clear where we have to unwrap_model and when operations are done in-place. the fact that Accelerate treats deepspeed and fsdp2 in-place while DDP and single GPU get a twin model is really annoying of Accelerate to have done.

[hjinnkim]

bf16 is singled out because it is what installs the instance-bound forward, not because of quantization — a quantized base behaves exactly the same (measured below with int8-quanto): prepare_model wraps forward with autocast when native_amp is on (mixed_precision bf16, or fp16 off CPU; accelerator.py:1818-1829), and the twin's copy of that attribute is what goes stale. The bug has existed since regional compilation shipped in accelerate v1.7.0 (huggingface/accelerate#3529) and needs all four:

1. mixed_precision bf16 or fp16 — no installs nothing.
2. dynamo_backend set and dynamo_use_regional_compilation: true — plain torch.compile wraps the original.
3. A model with repeated blocks (has_repeated_blocks: a ModuleList of one class).
4. Single process or DDP — FSDP2 and DeepSpeed compile in place.

Measured with the trainer, upstream main, accelerate 1.14.0, B200; every row krea2 LoRA (rank 64), batch 1 per GPU, 1024 px, 40 steps, gradient checkpointing on; peak = nvidia-smi maximum minus idle, artifacts = inductor cache files. mixed_precision: no keeps fp32 base weights unless the base is quantized with base_model_default_dtype: bf16 (cmd_args.py:1109).

base	mixed_precision	compile	bug	1 GPU main: peak MiB · artifacts	1 GPU this PR	2 GPU DDP main: per GPU · per rank	2 GPU DDP this PR
bf16	bf16	regional	yes	65,378 · 0	30,496 · 270	66,494 / 66,494 · 0 / 0	31,614 / 31,614 · 270 / 270
bf16	bf16	plain torch.compile	no	31,348 · 284	—	35,908 / 35,854 · 284 / 284	—
bf16	bf16	none	no	31,398 · 0	—	32,512 / 32,512 · 0 / 0	—
fp32	no	regional	no	59,604 · 276	—	60,878 / 60,878 · 276 / 276	—
fp32	no	plain torch.compile	no	64,210 · 279	—	65,466 / 70,964 · 279 / 279	—
fp32	no	none	no	60,662 · 0	—	61,938 / 61,938 · 0 / 0	—
int8-quanto	bf16	regional	yes	63,142 · 0	28,142 · 308	64,298 / 64,298 · 0 / 0	29,118 / 29,118 · 308 / 308
int8-quanto	no (bf16 weights)	regional	no	28,142 · 308	—	29,118 / 29,118 · 308 / 308	—

Bug rows run the original: nothing compiled, VRAM at the checkpointing-off level; this branch (6520cd052) brings them to the healthy rows and is a no-op elsewhere. A quantized base with mixed_precision: no never hit this; those would be the ~30 % runs. int8-torchao and int8-sdnq fail inside inductor on my torch nightly once the twin executes (unrelated).

On the twin: agreed. accelerate#4188 makes compile_regions return a copy whose bound methods point at itself, so the twin path behaves like the in-place ones.