Skip to content

fix: lift mutated buffers owned by a submodule - #4472

Open
shoumikhin wants to merge 1 commit into
pytorch:mainfrom
shoumikhin:trt-nested-buffer-lookup
Open

fix: lift mutated buffers owned by a submodule#4472
shoumikhin wants to merge 1 commit into
pytorch:mainfrom
shoumikhin:trt-nested-buffer-lookup

Conversation

@shoumikhin

Copy link
Copy Markdown
Contributor

The problem

A mutated buffer that belongs to a submodule is never lifted. The rewrite skips it with a
warning and carries on:

lift_mutated_buffers: get_attr target inner.cache not found on gm; skipping

The buffer is there. The lookup is what fails. A get_attr target is fully qualified, so a
buffer owned by a submodule arrives as inner.cache, and neither hasattr nor getattr
walks a dotted path:

hasattr(module, "inner.cache")     # False
module.get_buffer("inner.cache")   # the tensor

So only a buffer owned directly by the top-level module resolves. Any model that keeps state
inside its layers, which is the usual way to write one, is skipped.

The fix

Use get_buffer, which resolves through the submodules:

        try:
            buffer_tensor = gm.get_buffer(buffer_name)
        except AttributeError:
            logger.warning(
                "lift_mutated_buffers: get_attr target %s not found on gm; skipping",
                buffer_name,
            )
            continue

The warning is kept for a target that genuinely does not exist, so a real problem is still
reported. Nothing changes for a top-level buffer.

Testing

Added test_nested_buffer_lifted, which is the same shape as the existing
test_single_buffer_lifted but with the buffer one level down. It asserts the buffer is
lifted, the qualified name is preserved as inner.cache, the flattened placeholder is
buf_inner_cache, the rebuilt forward accepts it, and no trailing copy_ remains.

The test fails on the current code with exactly the warning above, and passes with the fix:

without the fix:  1 failed
with the fix:     8 passed

Worth noting why the existing tests do not catch this: all of them use a buffer owned by
the top-level module, so the name has no dots and getattr happens to work.

A mutated buffer that belongs to a submodule is never lifted, so the rewrite
silently does nothing for it:

  lift_mutated_buffers: get_attr target inner.cache not found on gm; skipping

A get_attr target is fully qualified, so such a buffer arrives as "inner.cache",
and neither hasattr nor getattr walks a dotted path:

  hasattr(module, "inner.cache")      False
  module.get_buffer("inner.cache")    the tensor

Only a buffer owned directly by the top-level module resolves today, while any
model that keeps state inside its layers is skipped with a warning.

Use get_buffer, which resolves through the submodules, and keep the warning for
a target that genuinely does not exist.
@meta-cla meta-cla Bot added the cla signed label Aug 9, 2026
@github-actions github-actions Bot added component: tests Issues re: Tests component: lowering Issues re: The lowering / preprocessing passes component: core Issues re: The core compiler component: api [Python] Issues re: Python API component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths labels Aug 9, 2026
@github-actions
github-actions Bot requested a review from zewenli98 August 9, 2026 00:05
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cla signed component: api [Python] Issues re: Python API component: core Issues re: The core compiler component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths component: lowering Issues re: The lowering / preprocessing passes component: tests Issues re: Tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant