From c4a1d47b1ad893d30972a5d84442849ac8205a4f Mon Sep 17 00:00:00 2001 From: Erwin Zhang Date: Sat, 15 Aug 2026 19:44:04 -0400 Subject: [PATCH 1/2] Reject a jaccl hostfile that does not describe the links the backend uses launch_jaccl checked the shape of the rdma matrix, its dimensions and its null diagonal, but never whether the pairs the backend is about to communicate over are actually in it. A hostfile whose rank order places unconnected nodes beside each other passes both checks. What follows is the worst failure available. The group does not fail to form and it does not run slowly: every rank ends up alone in a group of size one, reports success, and the job computes N separate wrong answers at full speed with nothing on stderr. Verified on four M4 Pro minis wired as a full mesh, where permuting two hosts in a working hostfile, rows and columns together so the file stays internally consistent, produced four independent groups of one and four 'membership OK' lines. This is the same class as #4210, where a group of size 1 counted as a successful init. missing_jaccl_links reports the pairs that are used but absent: for a ring each rank's two neighbours, for a mesh every pair. With two ranks the previous and next neighbour are the same node, which is harmless because the pair is simply examined twice. The error names the hosts and says that for jaccl-ring the host order is the ring order, since that is the mistake that produces it. Checked against the live mesh: the misordered hostfile is now rejected, and valid mesh and valid ring hostfiles still run at 20.87 and 16.33 GB/s. --- python/mlx/_distributed_utils/launch.py | 38 +++++++++++ python/tests/test_distributed_launch.py | 85 +++++++++++++++++++++++++ 2 files changed, 123 insertions(+) create mode 100644 python/tests/test_distributed_launch.py diff --git a/python/mlx/_distributed_utils/launch.py b/python/mlx/_distributed_utils/launch.py index 4771c1fb5b..7dcf94cd0e 100644 --- a/python/mlx/_distributed_utils/launch.py +++ b/python/mlx/_distributed_utils/launch.py @@ -370,6 +370,29 @@ def launch_nccl(parser, hosts, args, command): ) +def missing_jaccl_links(hosts, jaccl_ring): + """Pairs the backend will use that the hostfile has no device for. + + The shape checks in `launch_jaccl` cover the size of the matrix and its null + diagonal, but not whether the links about to be used are in it. A hostfile + whose rank order puts unconnected nodes next to each other passes both, and + then fails in the worst way available: every rank ends up alone in a group + of size one, reports success, and the job computes N separate wrong answers + at full speed with nothing on stderr. + + A ring only uses each rank's two neighbours; a mesh uses every pair. With + two ranks the previous and next neighbour are the same node, which is + harmless here because the same pair is simply examined twice. + """ + n = len(hosts) + return [ + (i, j) + for i, h in enumerate(hosts) + for j in (((i - 1) % n, (i + 1) % n) if jaccl_ring else range(n)) + if i != j and h.rdma[j] is None + ] + + def launch_jaccl(parser, hosts, args, command): if not hosts[0].ips: parser.error("Rank 0 should have an IP reachable from all other ranks") @@ -380,6 +403,21 @@ def launch_jaccl(parser, hosts, args, command): if not have_rdmas or not have_nulls: parser.error("Malformed hostfile for jaccl backend") + missing = missing_jaccl_links(hosts, jaccl_ring) + if missing: + shape = "ring neighbours" if jaccl_ring else "a full mesh" + detail = ", ".join( + f"{hosts[i].ssh_hostname} to {hosts[j].ssh_hostname}" + for i, j in missing[:3] + ) + more = f" and {len(missing) - 3} more" if len(missing) > 3 else "" + parser.error( + f"The hostfile does not describe {shape}. No RDMA device is listed " + f"for {detail}{more}. Regenerate it with mlx.distributed_config; " + "note that for jaccl-ring the order of the hosts is the order of " + "the ring." + ) + coordinator = hosts[0].ips[0] env = args.env cwd = args.cwd diff --git a/python/tests/test_distributed_launch.py b/python/tests/test_distributed_launch.py new file mode 100644 index 0000000000..e9aee9ce56 --- /dev/null +++ b/python/tests/test_distributed_launch.py @@ -0,0 +1,85 @@ +# Copyright © 2026 Apple Inc. + +import unittest + +from mlx._distributed_utils.common import Host +from mlx._distributed_utils.launch import missing_jaccl_links + + +def hosts_from(matrix): + """Hosts whose rdma matrix is `matrix`, named a, b, c ... by position.""" + return [ + Host(i, chr(ord("a") + i), [f"10.0.0.{i + 1}"], row) + for i, row in enumerate(matrix) + ] + + +# Four nodes in a ring: each reaches its two neighbours and not the one +# opposite. Valid for jaccl-ring, and not a mesh. +RING4 = [ + [None, "rdma_en2", None, "rdma_en3"], + ["rdma_en2", None, "rdma_en3", None], + [None, "rdma_en2", None, "rdma_en3"], + ["rdma_en3", None, "rdma_en2", None], +] + + +class TestMissingJacclLinks(unittest.TestCase): + """A hostfile has to describe the links the backend is going to use. + + A matrix of the right shape whose rank order puts unconnected nodes beside + each other is not merely slow: it leaves every rank alone in a group of + size one, where "rank 0 of 1" is internally consistent and completely + wrong. Catching it needs a check against the topology, not the shape. + """ + + def test_mesh_reports_a_missing_pair(self): + # b and c share no device, so this is not the full mesh jaccl needs. + matrix = [ + [None, "rdma_en2", "rdma_en3"], + ["rdma_en2", None, None], + ["rdma_en3", None, None], + ] + self.assertEqual( + missing_jaccl_links(hosts_from(matrix), False), [(1, 2), (2, 1)] + ) + + def test_ring_reports_unconnected_neighbours(self): + # A triangle reordered so ranks 0 and 1 have no cable between them. + # The shape is right and the diagonal is null, which is exactly what + # let this through before. + matrix = [ + [None, None, "rdma_en3"], + [None, None, "rdma_en2"], + ["rdma_en3", "rdma_en2", None], + ] + self.assertEqual( + missing_jaccl_links(hosts_from(matrix), True), [(0, 1), (1, 0)] + ) + + def test_ring_accepts_a_ring_that_is_not_a_mesh(self): + self.assertEqual(missing_jaccl_links(hosts_from(RING4), True), []) + + def test_mesh_rejects_that_same_ring(self): + # Every rank is missing the node opposite it: four ranks, one gap each. + self.assertEqual(len(missing_jaccl_links(hosts_from(RING4), False)), 4) + + def test_two_ranks(self): + # The previous and next neighbour are the same node here, so the check + # must not trip over asking about it twice. + matrix = [[None, "rdma_en2"], ["rdma_en2", None]] + self.assertEqual(missing_jaccl_links(hosts_from(matrix), True), []) + self.assertEqual(missing_jaccl_links(hosts_from(matrix), False), []) + + def test_full_mesh_passes_both(self): + matrix = [ + [None, "rdma_en2", "rdma_en3"], + ["rdma_en2", None, "rdma_en4"], + ["rdma_en3", "rdma_en4", None], + ] + self.assertEqual(missing_jaccl_links(hosts_from(matrix), False), []) + self.assertEqual(missing_jaccl_links(hosts_from(matrix), True), []) + + +if __name__ == "__main__": + unittest.main() From 7b0499a68662e40a27b21a38433d8d804a8ae4fc Mon Sep 17 00:00:00 2001 From: Cheng Date: Sat, 22 Aug 2026 17:25:14 +0900 Subject: [PATCH 2/2] nit --- python/mlx/_distributed_utils/launch.py | 58 ++++++----------- python/tests/test_distributed_launch.py | 85 ------------------------- 2 files changed, 21 insertions(+), 122 deletions(-) delete mode 100644 python/tests/test_distributed_launch.py diff --git a/python/mlx/_distributed_utils/launch.py b/python/mlx/_distributed_utils/launch.py index 7dcf94cd0e..0e661e358f 100644 --- a/python/mlx/_distributed_utils/launch.py +++ b/python/mlx/_distributed_utils/launch.py @@ -370,53 +370,37 @@ def launch_nccl(parser, hosts, args, command): ) -def missing_jaccl_links(hosts, jaccl_ring): - """Pairs the backend will use that the hostfile has no device for. - - The shape checks in `launch_jaccl` cover the size of the matrix and its null - diagonal, but not whether the links about to be used are in it. A hostfile - whose rank order puts unconnected nodes next to each other passes both, and - then fails in the worst way available: every rank ends up alone in a group - of size one, reports success, and the job computes N separate wrong answers - at full speed with nothing on stderr. - - A ring only uses each rank's two neighbours; a mesh uses every pair. With - two ranks the previous and next neighbour are the same node, which is - harmless here because the same pair is simply examined twice. - """ - n = len(hosts) - return [ - (i, j) - for i, h in enumerate(hosts) - for j in (((i - 1) % n, (i + 1) % n) if jaccl_ring else range(n)) - if i != j and h.rdma[j] is None - ] - - def launch_jaccl(parser, hosts, args, command): if not hosts[0].ips: parser.error("Rank 0 should have an IP reachable from all other ranks") jaccl_ring = args.backend == "jaccl-ring" have_rdmas = all(len(h.rdma) == len(hosts) for h in hosts) + if not have_rdmas: + parser.error( + "The hostfile is malformed: number of RDMA devices does not match hosts" + ) have_nulls = all(h.rdma[i] is None for i, h in enumerate(hosts)) - if not have_rdmas or not have_nulls: - parser.error("Malformed hostfile for jaccl backend") + if not have_nulls: + parser.error("The hostfile is malformed: RDMA device of self should be null") - missing = missing_jaccl_links(hosts, jaccl_ring) - if missing: - shape = "ring neighbours" if jaccl_ring else "a full mesh" - detail = ", ".join( + # Find pairs that miss rmda in hostfile. + n = len(hosts) + missing_rdma = [ + (i, j) + for i, h in enumerate(hosts) + for j in (((i - 1) % n, (i + 1) % n) if jaccl_ring else range(n)) + if i != j and h.rdma[j] is None + ] + + if missing_rdma: + pairs = ", ".join( f"{hosts[i].ssh_hostname} to {hosts[j].ssh_hostname}" - for i, j in missing[:3] - ) - more = f" and {len(missing) - 3} more" if len(missing) > 3 else "" - parser.error( - f"The hostfile does not describe {shape}. No RDMA device is listed " - f"for {detail}{more}. Regenerate it with mlx.distributed_config; " - "note that for jaccl-ring the order of the hosts is the order of " - "the ring." + for i, j in missing_rdma[:3] ) + if len(missing_rdma) > 3: + pairs += f" and {len(missing_rdma) - 3} more" + parser.error(f"The hostfile is malformed: no RDMA device is listed for {pairs}") coordinator = hosts[0].ips[0] env = args.env diff --git a/python/tests/test_distributed_launch.py b/python/tests/test_distributed_launch.py deleted file mode 100644 index e9aee9ce56..0000000000 --- a/python/tests/test_distributed_launch.py +++ /dev/null @@ -1,85 +0,0 @@ -# Copyright © 2026 Apple Inc. - -import unittest - -from mlx._distributed_utils.common import Host -from mlx._distributed_utils.launch import missing_jaccl_links - - -def hosts_from(matrix): - """Hosts whose rdma matrix is `matrix`, named a, b, c ... by position.""" - return [ - Host(i, chr(ord("a") + i), [f"10.0.0.{i + 1}"], row) - for i, row in enumerate(matrix) - ] - - -# Four nodes in a ring: each reaches its two neighbours and not the one -# opposite. Valid for jaccl-ring, and not a mesh. -RING4 = [ - [None, "rdma_en2", None, "rdma_en3"], - ["rdma_en2", None, "rdma_en3", None], - [None, "rdma_en2", None, "rdma_en3"], - ["rdma_en3", None, "rdma_en2", None], -] - - -class TestMissingJacclLinks(unittest.TestCase): - """A hostfile has to describe the links the backend is going to use. - - A matrix of the right shape whose rank order puts unconnected nodes beside - each other is not merely slow: it leaves every rank alone in a group of - size one, where "rank 0 of 1" is internally consistent and completely - wrong. Catching it needs a check against the topology, not the shape. - """ - - def test_mesh_reports_a_missing_pair(self): - # b and c share no device, so this is not the full mesh jaccl needs. - matrix = [ - [None, "rdma_en2", "rdma_en3"], - ["rdma_en2", None, None], - ["rdma_en3", None, None], - ] - self.assertEqual( - missing_jaccl_links(hosts_from(matrix), False), [(1, 2), (2, 1)] - ) - - def test_ring_reports_unconnected_neighbours(self): - # A triangle reordered so ranks 0 and 1 have no cable between them. - # The shape is right and the diagonal is null, which is exactly what - # let this through before. - matrix = [ - [None, None, "rdma_en3"], - [None, None, "rdma_en2"], - ["rdma_en3", "rdma_en2", None], - ] - self.assertEqual( - missing_jaccl_links(hosts_from(matrix), True), [(0, 1), (1, 0)] - ) - - def test_ring_accepts_a_ring_that_is_not_a_mesh(self): - self.assertEqual(missing_jaccl_links(hosts_from(RING4), True), []) - - def test_mesh_rejects_that_same_ring(self): - # Every rank is missing the node opposite it: four ranks, one gap each. - self.assertEqual(len(missing_jaccl_links(hosts_from(RING4), False)), 4) - - def test_two_ranks(self): - # The previous and next neighbour are the same node here, so the check - # must not trip over asking about it twice. - matrix = [[None, "rdma_en2"], ["rdma_en2", None]] - self.assertEqual(missing_jaccl_links(hosts_from(matrix), True), []) - self.assertEqual(missing_jaccl_links(hosts_from(matrix), False), []) - - def test_full_mesh_passes_both(self): - matrix = [ - [None, "rdma_en2", "rdma_en3"], - ["rdma_en2", None, "rdma_en4"], - ["rdma_en3", "rdma_en4", None], - ] - self.assertEqual(missing_jaccl_links(hosts_from(matrix), False), []) - self.assertEqual(missing_jaccl_links(hosts_from(matrix), True), []) - - -if __name__ == "__main__": - unittest.main()