[#4674][bugfix] AutoDeploy Fix memory leak in fuse_moe (#7844)

Delete the unstacked weights immediately to save GPU memory, cleanup occurs automatically after the transformation, but for large models we'll run out of memory during the transformation itself.

Signed-off-by: Gal Hubara Agam <[email protected]>

Author: galagam

tensorrt_llm/_torch/auto_deploy/config/default.yaml

@@ -107,7 +107,7 @@ transforms:
     backend: trtllm
   fuse_moe:
     stage: post_load_fusion
-    enabled: false # TODO: https://github.com/NVIDIA/TensorRT-LLM/issues/4674 this is causing OOMs
+    enabled: true
   fuse_allreduce_residual_rmsnorm:
     stage: post_load_fusion
   fuse_collectives:

tensorrt_llm/_torch/auto_deploy/transform/library/fused_moe.py

@@ -57,6 +57,12 @@ def _insert_fused_moe_ops(gm: GraphModule) -> int:
         node.replace_all_uses_with(new_node)
         graph.erase_node(node)
 
+        # Delete the unstacked weights immediately to save GPU memory
+        # This will happen automatically after the graph is canonicalized, but for large models we'll run out of memory
+        # during the transformation itself.
+        gm.graph.eliminate_dead_code()
+        gm.delete_all_unused_submodules()
+
     return fused_key_counter
 
 

tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_moe_fusion.py

@@ -368,3 +368,51 @@ def test_moe_fusion():
         num_param_nodes_fused < num_param_nodes
     ), f"""number of parameter nodes after fusion {num_param_nodes_fused} <
         number of parameter nodes before fusion {num_param_nodes}"""
+
+
+def test_fuse_moe_cleanup():
+    # Ensure deterministic allocations and a clean slate
+    torch.manual_seed(1234)
+    torch.cuda.manual_seed(1234)
+    torch.cuda.empty_cache()
+
+    device = "cuda"
+    dtype = torch.bfloat16
+
+    # Build model and export to GraphModule (pre-fusion)
+    model = MoEOpModel().to(device=device, dtype=dtype)
+    x = model.get_input(device=device, dtype=dtype)
+    gm = torch_export_to_gm(model, args=(x,), clone=True)
+
+    # Count parameters and measure memory before fusion
+    num_param_nodes_before = len(list(gm.named_parameters()))
+    torch.cuda.synchronize()
+    torch.cuda.empty_cache()
+    mem_before = torch.cuda.memory_allocated()
+
+    # Apply MoE fusion which should stack weights and clean up unstacked params
+    # We need to ensure the cleanup is done as part of the transformation to avoid OOM during the transformation itself.
+    gm_transformed = InferenceOptimizer(
+        None,
+        {
+            "fuse_moe": {
+                "stage": "post_load_fusion",
+                "run_graph_cleanup": False,  # verify cleanup is done as part of the transformation
+                "run_shape_prop": False,  # shape_prop can also trigger cleanup
+            },
+        },
+    )(None, gm)
+
+    # Ensure that parameter count decreased after fusion (unstacked params cleaned)
+    num_param_nodes_after = len(list(gm_transformed.named_parameters()))
+    assert num_param_nodes_after < num_param_nodes_before, (
+        f"Expected fewer parameters after fusion: before={num_param_nodes_before}, after={num_param_nodes_after}"
+    )
+
+    # Memory should not increase after fusion/cleanup
+    torch.cuda.synchronize()
+    torch.cuda.empty_cache()
+    mem_after = torch.cuda.memory_allocated()
+    assert mem_after <= mem_before, (
+        f"CUDA memory increased after fusion: before={mem_before} after={mem_after}"
+    )