Preserve shared modules in regional compilation

[AaravMehta-07]

What does this PR do?

Preserve shared module registrations in compile_regions.

When a model contains repeated blocks, regional compilation rebuilds its module hierarchy. The rebuild iterates through named_children(), which deduplicates shared children. If a child is registered under two names, the second registration disappears from the compiled model and a forward pass that uses that name raises AttributeError. Separately, a block registered at multiple positions or under different parents receives different compiled wrappers, breaking module identity even though the original model shares the module.

This change traverses all registered children and caches the result for each original module within one compile_regions call. Aliased names, repeated block references, and shared descendants then reference the same rebuilt or compiled module. Registered None children are retained as well. The source hierarchy is not rewritten, and the existing regional compilation boundaries are unchanged.

No runtime dependencies are added. This does not change the in-place FSDP2 or DeepSpeed paths.

Reproduction

import torch
from accelerate.utils import compile_regions


class Shared(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.blocks = torch.nn.ModuleList([torch.nn.Linear(4, 4), torch.nn.Linear(4, 4)])
        self.blocks_alias = self.blocks

    def forward(self, x):
        return self.blocks_alias[0](x)


compile_regions(Shared(), backend="eager")(torch.ones(1, 4))
# Before: AttributeError: 'Shared' object has no attribute 'blocks_alias'

Validation

Windows, Python 3.11.9, PyTorch 2.13.0+cpu, pytest 8.4.2. Tests use the real repository implementation and torch.compile(..., backend="eager"), without checkpoint downloads.

• Three new regression tests fail on the unmodified base at 4ebc726849a2d7f03c16198847c22954e224dd87 and pass after this change.
• python -m pytest tests/test_compile.py -v: 5 passed, 5 skipped. The existing model/performance tests are explicitly skipped by their upstream decorator.
• python -m pytest tests/test_utils.py -k extract_model -v: 3 passed, 1 skipped, 48 deselected.
• The new tests check alias identity, unchanged source registrations, forward outputs, input gradients, and parameter gradients against an independently copied eager model.
• Repository-pinned Ruff 0.13.1: checks and formatting checks passed for both changed files. git diff --check passed.
• Independent agent review reran the shared-module, bound-method, and unwrapping tests together: 8 passed, 3 skipped, 51 deselected.

The full suite, GPU backends, Inductor code generation, distributed execution, and performance were not validated locally. The eager backend verifies the compilation wrapping and real forward/backward behavior, not hardware speedups.

Before submitting

• Read the contribution guidelines and searched existing issues and pull requests.
• Discussed or approved via an issue/forum. This independently reproduced bug has not been pre-approved.
• Added regression tests.
• Documentation: no new feature or public API is introduced; the documented return behavior is preserved.

AI disclosure: this patch, regression tests, and description were prepared by a Codex coding agent. The agent ran the reported validation. This disclosure does not represent the work as unaided human authorship or claim human review or maintainer approval.

[AaravMehta-07]

Hi maintainers, whenever you have a chance, I'd appreciate a review of this fix for preserving shared module aliases during regional compilation. Five focused CPU tests passed locally; five hardware-gated cases were skipped. Forward and gradient checks are included, without claiming hardware coverage.

Please let me know if you'd prefer a different approach or additional focused coverage. Thank you for your time.

AI-assisted follow-up.