Skip to content

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

Merged
lanluo-nvidia merged 1 commit into
pytorch:mainfrom
shoumikhin:trt-nested-buffer-lookup
Aug 18, 2026
Merged

fix: lift mutated buffers owned by a submodule#4472
lanluo-nvidia merged 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") # Falsemodule.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)
exceptAttributeError:
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.

@github-actionsgithub-actionsBot 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
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.
@shoumikhin
shoumikhinforce-pushed the trt-nested-buffer-lookup branch from 5408a18 to f4060baCompareAugust 18, 2026 18:45
@lanluo-nvidialanluo-nvidia added this to the v2.14.0 milestone Aug 18, 2026

@lanluo-nvidialanluo-nvidia left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lgtm

@lanluo-nvidia
lanluo-nvidia merged commit 75e7132 into pytorch:mainAug 18, 2026
35 of 38 checks passed
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cla signedcomponent: api [Python]Issues re: Python APIcomponent: coreIssues re: The core compilercomponent: dynamoIssues relating to the `torch.compile` or `torch._dynamo.export` pathscomponent: loweringIssues re: The lowering / preprocessing passescomponent: testsIssues re: Tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants

@shoumikhin@lanluo-nvidia