ThunderFX: Modify GraphModule in-place (#2399)

Co-authored-by: Masato Shinokawa <shino16@users.noreply.github.com>
Co-authored-by: Masato Shinokawa <mshinokawa@viking-prod-235.ipp2u1.colossus.nvidia.com>
Co-authored-by: Masaki Kozuki <mkozuki@nvidia.com>
Co-authored-by: Yan Wang <kiya00wy@gmail.com>
Co-authored-by: Thomas Viehmann <tv.code@beamnet.de>

Author: 5 people

thunder/dynamo/compiler.py

@@ -106,7 +106,7 @@ def __init__(self, **thunder_options):
         self._torch_compile = torch.compile
 
     def __call__(self, gm: torch.fx.GraphModule, sample_args: list[torch.SymInt, torch.Tensor]):
-        gm = remove_empty_autocast(gm)
+        remove_empty_autocast(gm)
 
         # Dynamo uses lazy generation of the underlying Python code, so we need to
         # force recompilation of the GraphModule before passing it to Thunder.

thunder/dynamo/report.py

@@ -1094,8 +1094,9 @@ def foo(x):
     from thunder.dynamo.splitter import _splitter
     from thunder import jit
 
+    gm = copy.deepcopy(report.graph)
     # Splits the FX graph module using Thunder splitter
-    gm = remove_empty_autocast(report.graph)
+    remove_empty_autocast(gm)
     # Dynamo uses lazy generation of the underlying Python code, so we need to
     # force recompilation of the GraphModule before passing it to Thunder.
     recompile_graph(gm)

thunder/dynamo/splitter.py

@@ -144,17 +144,20 @@ def callback(node) -> int:
     gm.recompile()
 
     # `split_module` iterates over nodes and determines the partition to place them based on the callback.
-    original_split_gm: torch.fx.GraphModule = split_module(
+    split_gm: torch.fx.GraphModule = split_module(
         gm, root_m=None, split_callback=callback, keep_original_order=True, keep_original_node_name=True
     )
 
     # Workaround for the Torch bug https://github.com/pytorch/pytorch/pull/139275
-    for submodule in original_split_gm.children():
+    for submodule in split_gm.children():
         if not submodule.graph.find_nodes(op="output"):
             submodule.graph.output(())
-    if not original_split_gm.graph.find_nodes(op="output"):
-        original_split_gm.graph.output(())
-    split_gm = copy.deepcopy(original_split_gm)
+    if not split_gm.graph.find_nodes(op="output"):
+        split_gm.graph.output(())
+
+    # If split_gm contains Parameters or Tensors then deepcopy would also create their copies.
+    # TODO: Eliminate deepcopy
+    original_split_gm = copy.deepcopy(split_gm)
 
     def is_thunder_supported_partition(node: torch.fx.Node) -> bool:
         return node.name.startswith("submod") and int(node.name.replace("submod_", "")) in supported_partitions

thunder/dynamo/utils.py

@@ -697,13 +697,10 @@ def remove_empty_autocast(graph_module: torch.fx.GraphModule) -> torch.fx.GraphM
         graph_module: Graph module to which this pass is applied.
 
     """
-
-    empty_autocast_removed_graph_module = copy.deepcopy(graph_module)
-
     # Dummy init node.
     prev_node = torch.fx.node.Node(graph_module.graph, "start_node", "call_function", lambda: None, None, None)
     nodes_to_erase = []
-    for node in empty_autocast_removed_graph_module.graph.nodes:
+    for node in graph_module.graph.nodes:
         # As _enter_autocast and _exit_autocast functions map the regions created by context manager,
         # previous `_enter_autocast` will always correspond with current `_exit_autocast`.
         if (
@@ -721,9 +718,9 @@ def remove_empty_autocast(graph_module: torch.fx.GraphModule) -> torch.fx.GraphM
 
     # Erase the marked nodes.
     for node in nodes_to_erase:
-        empty_autocast_removed_graph_module.graph.erase_node(node)
+        graph_module.graph.erase_node(node)
 
-    return empty_autocast_removed_graph_module
+    return graph_module
 
 
 def arg_like_tensor(arg: torch.Tensor | ExampleInputMetaData):

thunder/tests/test_dynamo.py

@@ -474,7 +474,6 @@ def func(x):
     dtypes=(dtypes.float32,),
     executors=(DynamoThunderExecutor,),
     decorators=(
-        pytest.mark.skip(reason="https://github.com/Lightning-AI/lightning-thunder/issues/1821"),
         pytest.mark.parametrize(
             "optim",
             (