diff --git a/src/mindtorch_v2/distributed/__init__.py b/src/mindtorch_v2/distributed/__init__.py index ed8411055..4469b2902 100644 --- a/src/mindtorch_v2/distributed/__init__.py +++ b/src/mindtorch_v2/distributed/__init__.py @@ -628,7 +628,7 @@ def all_to_all(output_tensor_list, input_tensor_list, group=None, return work -def _validate_all_to_all_single_splits(pg, input_split_sizes, output_split_sizes): +def _validate_all_to_all_single_splits(pg, input_split_sizes, output_split_sizes, input_numel, output_numel): if len(input_split_sizes) != pg.size(): raise ValueError( f"input_split_sizes length {len(input_split_sizes)} must equal world_size {pg.size()}" @@ -640,6 +640,17 @@ def _validate_all_to_all_single_splits(pg, input_split_sizes, output_split_sizes if any(int(s) < 0 for s in input_split_sizes + output_split_sizes): raise ValueError("all_to_all_single split sizes must be non-negative") + input_split_sum = sum(int(s) for s in input_split_sizes) + output_split_sum = sum(int(s) for s in output_split_sizes) + if input_split_sum != int(input_numel): + raise ValueError( + f"all_to_all_single input numel {int(input_numel)} must equal sum(input_split_sizes) {input_split_sum}" + ) + if output_split_sum != int(output_numel): + raise ValueError( + f"all_to_all_single output numel {int(output_numel)} must equal sum(output_split_sizes) {output_split_sum}" + ) + def _validate_hccl_all_to_all_single_pairwise(pg, input_split_sizes, output_split_sizes): if pg not in _pg_map: @@ -699,7 +710,7 @@ def all_to_all_single(output, input, output_split_sizes=None, chunk_size = output.numel() // world_size output_split_sizes = [chunk_size] * world_size - _validate_all_to_all_single_splits(pg, input_split_sizes, output_split_sizes) + _validate_all_to_all_single_splits(pg, input_split_sizes, output_split_sizes, input.numel(), output.numel()) if isinstance(pg, ProcessGroupHCCL): _validate_hccl_all_to_all_single_pairwise(pg, input_split_sizes, output_split_sizes) diff --git a/tests/mindtorch_v2/test_hccl_all_to_all_single_split_numel_validation_multicard.py b/tests/mindtorch_v2/test_hccl_all_to_all_single_split_numel_validation_multicard.py new file mode 100644 index 000000000..bd669ea3f --- /dev/null +++ b/tests/mindtorch_v2/test_hccl_all_to_all_single_split_numel_validation_multicard.py @@ -0,0 +1,150 @@ +"""HCCL all_to_all_single split/numel validation on 2/4/8 NPUs.""" + +import os +import subprocess +import sys +import time + +import pytest + + +SCRIPT = r''' +import os, sys, time +src_dir = os.environ.get("MINDTORCH_V2_SRC") +if src_dir: + sys.path.insert(0, src_dir) + +import mindtorch_v2 as torch +import mindtorch_v2.distributed as dist + +rank = int(os.environ["RANK"]) +world_size = int(os.environ["WORLD_SIZE"]) + +mode = os.environ["CASE_MODE"] + +device = torch.Device(f"npu:{rank}") +time.sleep(0.05 * rank) +dist.init_process_group("hccl", device_id=device) + +base_in = [1 if i == rank else 2 for i in range(world_size)] +base_out = [1 if j == rank else 2 for j in range(world_size)] + +if mode == "input_sum_mismatch": + input_split = list(base_in) + output_split = list(base_out) + # Make split sum larger than input numel. + input_split[rank] += 1 + inp_numel = sum(base_in) + out_numel = sum(output_split) +elif mode == "output_sum_mismatch": + input_split = list(base_in) + output_split = list(base_out) + # Make split sum larger than output numel. + output_split[rank] += 1 + inp_numel = sum(input_split) + out_numel = sum(base_out) +else: + raise RuntimeError(f"unexpected mode: {mode}") + +inp = torch.zeros(inp_numel, device=device) +out = torch.zeros(out_numel, device=device) + +try: + dist.all_to_all_single( + out, + inp, + output_split_sizes=output_split, + input_split_sizes=input_split, + async_op=True, + ) +except ValueError as exc: + msg = str(exc) + assert "numel" in msg and "split" in msg, msg +else: + raise AssertionError("expected ValueError for split sum and tensor numel mismatch") + +dist.destroy_process_group() +print(f"[rank {rank}] HCCL split/numel validation {mode} {world_size}card PASS") +''' + + +def _run_once(world_size, master_port, mode): + env = os.environ.copy() + env["MASTER_ADDR"] = "127.0.0.1" + env["MASTER_PORT"] = str(master_port) + env["WORLD_SIZE"] = str(world_size) + env["CASE_MODE"] = mode + src_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "src")) + env["MINDTORCH_V2_SRC"] = src_dir + env["PYTHONPATH"] = src_dir + \ + (":" + env["PYTHONPATH"] if "PYTHONPATH" in env else "") + + worker_file = f"/tmp/_hccl_all_to_all_single_split_numel_validation_{world_size}card.py" + with open(worker_file, "w") as f: + f.write(SCRIPT) + + failed = [] + outputs = [] + procs = [] + + for r in range(world_size): + p = subprocess.Popen( + [sys.executable, worker_file], + env={**env, "RANK": str(r)}, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + ) + procs.append(p) + + timeout = 420 if world_size <= 4 else 900 + for r, p in enumerate(procs): + try: + out, _ = p.communicate(timeout=timeout) + txt = out.decode("utf-8", errors="replace") + except subprocess.TimeoutExpired: + p.kill() + out, _ = p.communicate() + txt = "TIMEOUT\n" + out.decode("utf-8", errors="replace") + outputs.append(txt) + if p.returncode != 0: + failed.append(r) + + return failed, outputs + + +def _run_case(world_size, master_port, mode): + retries = 3 + for attempt in range(1, retries + 1): + failed, outputs = _run_once(world_size, master_port, mode) + if not failed: + return + + joined = "\n".join(outputs) + transient = "resource unavailable" in joined + if transient and attempt < retries: + print( + f"HCCL transient init failure on {world_size} cards ({mode}), " + f"retry {attempt}/{retries}" + ) + time.sleep(5) + continue + + for r, txt in enumerate(outputs): + print(f"=== RANK {r} ===") + print(txt) + raise AssertionError( + f"HCCL split/numel validation {mode} {world_size}card failed on ranks: {failed}" + ) + + +@pytest.mark.parametrize( + "world_size,master_port", + [ + (2, 29716), + (4, 29726), + (8, 29736), + ], +) +@pytest.mark.parametrize("mode", ["input_sum_mismatch", "output_sum_mismatch"]) +def test_hccl_all_to_all_single_split_numel_validation_multicard(world_size, master_port, mode): + _run_case(world_size, master_port, mode)