[Fixbug] Prevent allreduce op from being fused (#304)

hidet-org · Jul 5, 2023 · 1c1cd11 · 1c1cd11
1 parent f937ad2
commit 1c1cd11
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/python/hidet/graph/transforms/fuse_operator.py b/python/hidet/graph/transforms/fuse_operator.py
@@ -15,6 +15,7 @@
 from hidet.graph.flow_graph import FlowGraph, Operator, Tensor
 from hidet.graph.ops.special import BarrierOp
 from hidet.graph.ops.transfer import TransferOp
+from hidet.graph.ops.distributed import AllReduceOp
 from hidet.graph.graph_utils.functors import analyze_usage
 from hidet.graph.transforms.base import GraphPass
 from hidet.utils.structure import DirectedGraph
@@ -24,7 +25,7 @@
 
 
 # the following operators are not fusible
-NOT_FUSIBLE = {BarrierOp, TransferOp}
+NOT_FUSIBLE = {BarrierOp, TransferOp, AllReduceOp}
 
 
 class FusibleGraph: