[Distributed][HCCL] Validate all_to_all_single split sums against tensor numel

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)

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)