[FIX] add sharding propagation after auto sharding (#136)

alpa-projects · Nov 8, 2022 · 24d5f1b · 24d5f1b
1 parent 2079e74
commit 24d5f1b
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/tensorflow/compiler/xla/service/spmd/alpa_compiler.cc b/tensorflow/compiler/xla/service/spmd/alpa_compiler.cc
@@ -148,6 +148,9 @@ Status RunAutoShardingPass(HloModule* hlo_module,
       spmd_simplify.AddPass<HloDCE>();
 
       spmd_pipeline.AddPass<AutoSharding>();
+      spmd_pipeline.AddPass<ShardingPropagation>(
+          /*is_spmd=*/true, /*propagate_metadata=*/false,
+          /*allow_spmd_sharding_propagation_to_output=*/true);
       spmd_pipeline.AddPass<SliceAutoShardedStages>();
     } else {
       spmd_pipeline.AddPass<CallInliner>();