[Distributed] Add send and recv helpers (#5719)

2024-06-23 17:42:28 -04:00
parent 6c916ac8a8
commit 5d4d90536f
6 changed files with 278 additions and 24 deletions
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -129,7 +129,7 @@ def init_test_distributed_environment(
    ensure_model_parallel_initialized(tp_size, pp_size)


-def multi_process_tensor_parallel(
+def multi_process_parallel(
    tp_size: int,
    pp_size: int,
    test_target,