[SPMD] Add debug of SPMD for single/multi host (#5742)

pytorch · Jan 12, 2024 · 4d8c0c8 · 4d8c0c8
1 parent f7666a4
commit 4d8c0c8
Show file tree

Hide file tree

Showing 5 changed files with 991 additions and 4 deletions.
diff --git a/test/run_tests.sh b/test/run_tests.sh
@@ -203,6 +203,7 @@ function run_xla_op_tests3 {
   run_test "$CDIR/spmd/test_xla_sharding_hlo.py"
   run_test "$CDIR/spmd/test_xla_virtual_device.py"
   run_test "$CDIR/spmd/test_dynamo_spmd.py"
+  run_test "$CDIR/spmd/test_spmd_debugging.py"
   run_test "$CDIR/spmd/test_xla_distributed_checkpoint.py"
   run_test "$CDIR/spmd/test_xla_spmd_python_api_interaction.py"
   run_test "$CDIR/test_operations_hlo.py" "$@" --verbosity=$VERBOSITY