enable the test

pytorch · Aug 16, 2024 · 6e23e8a · 6e23e8a
1 parent d8a9421
commit 6e23e8a
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 3 deletions.
diff --git a/test/run_tests.sh b/test/run_tests.sh
@@ -302,7 +302,7 @@ function run_mp_op_tests {
   # TODO(JackCaoG): enable this
   # run_test "$CDIR/dynamo/test_traceable_collectives.py"
   run_test "$CDIR/test_fsdp_auto_wrap.py"
-  # run_torchrun "$CDIR/test_mp_early_exit.py"
+  run_torchrun "$CDIR/test_mp_early_exit.py"
   run_pt_xla_debug "$CDIR/debug_tool/test_mp_pt_xla_debug.py"
   run_test "$CDIR/torch_distributed/test_torch_distributed_all_gather_xla_backend.py"
   run_test "$CDIR/torch_distributed/test_torch_distributed_all_reduce_xla_backend.py"

diff --git a/torch_xla/csrc/cross_replica_reduces.cpp b/torch_xla/csrc/cross_replica_reduces.cpp
@@ -115,7 +115,6 @@ std::shared_ptr<torch::lazy::Value> CreateToken(
 
 at::Tensor all_reduce(const at::Tensor& self, std::string reduceOp,
                       std::string /*group_name*/) {
-  std::cout << "trigger all_reduce lower" << std::endl;
   TORCH_LAZY_FN_COUNTER_TIMED_TRACING("xla::");
   auto self_tensor = bridge::GetXlaTensor(self);
   // TODO(alanwaketan): Use group_name to generate groups. Currently we just
@@ -260,7 +259,6 @@ AllGatherResult BuildAllGather(xla::XlaOp input, xla::XlaOp token, int64_t dim,
 at::Tensor all_gather_into_tensor(const at::Tensor& self, int64_t group_size,
                                   std::string group_name) {
   TORCH_LAZY_FN_COUNTER("xla::");
-  std::cout << "trigger all_gather_into_tensor lower" << std::endl;
   auto self_tensor = bridge::GetXlaTensor(self);
   std::vector<int64_t> all_groups(group_size);
   std::iota(all_groups.begin(), all_groups.end(), 0);