fit completion rule for swiglu op (#63290)

zhiqiu · web-flow · commit a61d20a5028e · 2024-04-08T13:34:19.000+08:00
diff --git a/python/paddle/distributed/auto_parallel/static/completion.py b/python/paddle/distributed/auto_parallel/static/completion.py
@@ -183,6 +183,7 @@ def _can_apply_infer_spmd_rule(dist_op):
         "silu",
         "concat",
         "expand_as_v2",
+        "swiglu",
     ]
     parallel_ce = os.getenv("PARALLEL_CROSS_ENTROPY")
     if parallel_ce == "true":
diff --git a/python/paddle/distributed/auto_parallel/static/operators/dist_default.py b/python/paddle/distributed/auto_parallel/static/operators/dist_default.py
@@ -15,7 +15,7 @@
 
 from paddle.distributed.fleet.meta_optimizers.common import OP_ROLE_KEY, OpRole
 
-from ..completion import get_phi_spmd_rule
+from ..completion import contains_spmd_rule, get_phi_spmd_rule
 from ..cost import (
     _g_op_cost_factory,
     build_comp_costs_from_descs,
@@ -137,10 +137,15 @@ def update_dims_mapping(dist_op):
             )
 
         # step2: infer spmd
-        rule = get_phi_spmd_rule("default_")
-        # tensor order following order in PHI definition
-        fw_results = rule.infer_forward(input_specs, output_specs)
-        bw_results = rule.infer_backward(input_specs, output_specs)
+        if contains_spmd_rule(dist_op.serial_op.type):
+            rule = get_phi_spmd_rule(dist_op.serial_op.type)
+            fw_results = rule.infer_forward(*input_specs)
+            bw_results = rule.infer_backward(*input_specs, output_specs)
+        else:
+            rule = get_phi_spmd_rule('default_')
+            # tensor order following order in PHI definition
+            fw_results = rule.infer_forward(input_specs, output_specs)
+            bw_results = rule.infer_backward(input_specs, output_specs)
 
         # step3: update dist_attr
         # tensor order following order in PHI definition

Original file line number	Diff line number	Diff line change
`@@ -183,6 +183,7 @@ def _can_apply_infer_spmd_rule(dist_op):`
`183`	`183`	`"silu",`
`184`	`184`	`"concat",`
`185`	`185`	`"expand_as_v2",`
	`186`	`+ "swiglu",`
`186`	`187`	`]`
`187`	`188`	`parallel_ce = os.getenv("PARALLEL_CROSS_ENTROPY")`
`188`	`189`	`if parallel_ce == "true":`