Cpu memory graph break

[cehongwang]

Description

Please include a summary of the change and which issue is fixed. Please also include relevant motivation and context. List any dependencies that are required for this change.

Fixes # (issue)

Type of change

Please delete options that are not relevant and/or add your own.

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• This change requires a documentation update

Checklist:

• My code follows the style guidelines of this project (You can use the linters)
• I have performed a self-review of my own code
• I have commented my code, particularly in hard-to-understand areas and hacks
• I have made corresponding changes to the documentation
• I have added tests to verify my fix or my feature
• New and existing unit tests pass locally with my changes
• I have added the relevant labels to my PR in so that relevant reviewers are notified

[github-actions]

There are some changes that do not conform to Python style guidelines:

--- /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/_compiler.py	2025-11-04 20:05:23.825034+00:00
+++ /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/_compiler.py	2025-11-04 20:05:55.253944+00:00
@@ -876,15 +876,14 @@
    # This is done to release CPU memory.
    for attr in dir(gm):
        if attr.startswith("_frozen_param"):
            delattr(gm, attr)

-
-
    from torch_tensorrt.dynamo.conversion._ConverterRegistry import DYNAMO_CONVERTERS
+
    DYNAMO_CONVERTERS.disallowed_targets = set()
-    
+
    for name, _ in partitioned_module.named_children():
        submodule = getattr(partitioned_module, name)
        # filter on the GraphModule
        if not isinstance(submodule, torch.fx.graph_module.GraphModule):
            continue

[narendasan]

Do you have a test case or something to demonstrate this feature?

[narendasan]

1. We should think about using this tech for refit vs non refit
2. Make refit apis work across graph breaks

[narendasan]

Improve usability by automating nn.Module -> atomic fx graph

[narendasan]

Its looking good, just add a quick example in the examples folder and list it under contributor documentation for now

[narendasan]

When will this fail? Do we need to check this or can FX just guarantee this for us?

[cehongwang]

This should never fail. I put this here just in case something changes in torch and we don't know

[narendasan]

Shoudnt this conditionally run based on if the user provided a budget at all?

[narendasan]

/ should we make this opt in while testing?

[narendasan]

We should also have a test with no fall back as well as a test which tests for making sure that the atomic subgraph system works

[narendasan]

Also one that tests registering a new atomic subgraph and then verifying that it has the desired effect

[cehongwang]

Also one that tests registering a new atomic subgraph and then verifying that it has the desired effect
Can this prove the atomic subgraph system works?