fix: nvbugs/5231298: pytorch allreduce issue (#3673)

VALLIS-NERIA · web-flow · commit 5bf8fdca6590 · 2025-04-18T15:50:40.000+08:00
Signed-off-by: Xiwen Yu &lt;13230610+VALLIS-NERIA@users.noreply.github.com&gt;
diff --git a/tensorrt_llm/plugin/plugin.py b/tensorrt_llm/plugin/plugin.py
@@ -704,7 +704,8 @@ def set_workspace_tensor(self,
         )
 
     @staticmethod
-    def max_workspace_size_auto(tp_size: int, support_deterministic) -> int:
+    def max_workspace_size_auto(tp_size: int,
+                                support_deterministic=True) -> int:
         if force_all_reduce_deterministic() and support_deterministic:
             workspace_size = os.getenv("FORCE_ALLREDUCE_KERNEL_WORKSPACE_SIZE",
                                        "1000000000")

Original file line number	Diff line number	Diff line change
`@@ -704,7 +704,8 @@ def set_workspace_tensor(self,`
`704`	`704`	`)`
`705`	`705`
`706`	`706`	`@staticmethod`
`707`		`- def max_workspace_size_auto(tp_size: int, support_deterministic) -> int:`
	`707`	`+ def max_workspace_size_auto(tp_size: int,`
	`708`	`+ support_deterministic=True) -> int:`
`708`	`709`	`if force_all_reduce_deterministic() and support_deterministic:`
`709`	`710`	`workspace_size = os.getenv("FORCE_ALLREDUCE_KERNEL_WORKSPACE_SIZE",`
`710`	`711`	`"1000000000")`