From 65f5e50ecbe5b6ff7553ec13e997eb0d5fe783ab Mon Sep 17 00:00:00 2001 From: Jason Lu Date: Thu, 2 Jul 2026 11:27:05 -0700 Subject: [PATCH] feat(placement): bond parallel RDMA links in JACCL device matrix The JACCL device matrix builder (get_mlx_jaccl_devices_matrix) previously selected a single RDMA interface per device pair and discarded any additional parallel links via an early break. This left extra Thunderbolt (and on Linux, InfiniBand) links unused even when the underlying JACCL ring backend is built to stripe across up to RING_MAX_CONNS (=4) wires per direction. This change collects every RDMA interface per pair and emits the matrix cells as lists of interface names (list[list[list[str] | None]]), which the JACCL backend's JSON parser already accepts (string or array of strings per cell). Link counts are normalised down to the minimum observed across all pairs, so the backend's ring validator (which requires uniform connection counts on every ring edge) never rejects an asymmetric matrix. Bonding is effective for ring/tensor-parallel workloads; the mesh fallback path only consumes the first interface per pair (a JACCL limitation). - Widen MlxJacclInstance.jaccl_devices type to list[list[list[str] | None]] - Collect all RDMAConnection edges per pair instead of breaking on the first - Normalise to uniform per-pair count for ring validator compatibility - Update existing tensor-parallel matrix test assertions to list cells - Add test_jaccl_devices_matrix_bonds_parallel_rdma_links - Add test_jaccl_devices_matrix_normalises_asymmetric_link_counts --- src/exo/master/placement_utils.py | 49 +++++- src/exo/master/tests/test_placement.py | 188 ++++++++++++++++++++++- src/exo/shared/types/worker/instances.py | 5 +- 3 files changed, 230 insertions(+), 12 deletions(-) diff --git a/src/exo/master/placement_utils.py b/src/exo/master/placement_utils.py index 0375e97e01..25f1c2f59e 100644 --- a/src/exo/master/placement_utils.py +++ b/src/exo/master/placement_utils.py @@ -296,15 +296,27 @@ def get_shard_assignments( def get_mlx_jaccl_devices_matrix( selected_cycle: list[NodeId], cycle_digraph: Topology, -) -> list[list[str | None]]: +) -> list[list[list[str] | None]]: """Build connectivity matrix mapping device i to device j via RDMA interface names. - The matrix element [i][j] contains the interface name on device i that connects - to device j, or None if no connection exists or no interface name is found. - Diagonal elements are always None. + The matrix element [i][j] contains the list of interface names on device i that + connect to device j, or None if no connection exists. Diagonal elements are always + None. + + Multiple parallel RDMA links between the same pair are returned as a list so the + JACCL backend can bond them. The backend's ring path stripes across up to + RING_MAX_CONNS (=4) wires per direction; extra links beyond that are dropped here + to stay within the backend's limit. Note the mesh path only consumes the first + interface per pair, so bonding is effective for ring/tensor-parallel workloads. + + The JACCL ring validator requires every ring edge to carry the same number of + connections, so the counts are normalised down to the minimum across all pairs to + avoid the backend rejecting an asymmetric matrix. """ num_nodes = len(selected_cycle) - matrix: list[list[str | None]] = [ + # First pass: collect every RDMA interface per pair (order is stable: insertion + # order from get_all_connections_between). + interfaces_per_pair: list[list[list[str] | None]] = [ [None for _ in range(num_nodes)] for _ in range(num_nodes) ] @@ -313,14 +325,35 @@ def get_mlx_jaccl_devices_matrix( if i == j: continue + names: list[str] = [] for conn in cycle_digraph.get_all_connections_between(node_i, node_j): if isinstance(conn, RDMAConnection): - matrix[i][j] = conn.source_rdma_iface - break - else: + names.append(conn.source_rdma_iface) + + if not names: raise ValueError( "Current jaccl backend requires all-to-all RDMA connections" ) + interfaces_per_pair[i][j] = names + + # Second pass: normalise to a uniform connection count across all pairs. The + # backend's ring validator rejects matrices where ring edges have differing + # numbers of connections, and bonding only helps when it's consistent, so trim + # every pair down to the smallest count observed anywhere in the matrix. + non_null_counts = [ + len(names) for row in interfaces_per_pair for names in row if names is not None + ] + target_count = min(non_null_counts) if non_null_counts else 0 + + matrix: list[list[list[str] | None]] = [ + [None for _ in range(num_nodes)] for _ in range(num_nodes) + ] + for i in range(num_nodes): + for j in range(num_nodes): + if i == j: + continue + assert interfaces_per_pair[i][j] is not None + matrix[i][j] = interfaces_per_pair[i][j][:target_count] # type: ignore[slice] return matrix diff --git a/src/exo/master/tests/test_placement.py b/src/exo/master/tests/test_placement.py index b891187266..b802f38c54 100644 --- a/src/exo/master/tests/test_placement.py +++ b/src/exo/master/tests/test_placement.py @@ -505,9 +505,9 @@ def test_tensor_rdma_backend_connectivity_matrix( idx_b = node_to_idx[node_b] idx_c = node_to_idx[node_c] - assert matrix[idx_a][idx_b] == "rdma_en3" - assert matrix[idx_b][idx_c] == "rdma_en4" - assert matrix[idx_c][idx_a] == "rdma_en5" + assert matrix[idx_a][idx_b] == ["rdma_en3"] + assert matrix[idx_b][idx_c] == ["rdma_en4"] + assert matrix[idx_c][idx_a] == ["rdma_en5"] # Verify coordinators are set for all nodes assert len(instance.jaccl_coordinators) == 3 @@ -1056,3 +1056,185 @@ def test_mlx_jaccl_rejects_cuda_only_cycle(model_card: ModelCard): node_backends, node_rdma_ctl=node_rdma_ctl, ) + + +def _build_two_node_bonded_rdma_topology( + links_per_pair: int = 2, +) -> tuple[Topology, NodeId, NodeId, dict[NodeId, NodeNetworkInfo]]: + """Two-node topology with multiple parallel RDMA links between the pair. + + Each direction gets ``links_per_pair`` RDMA edges using distinct interfaces + (rdma_en3, rdma_en4, ...). This mirrors two Macs connected by more than one + Thunderbolt cable. + """ + topology = Topology() + node_a = NodeId() + node_b = NodeId() + + ethernet_interface = NetworkInterfaceInfo(name="en0", ip_address="10.0.0.1") + ethernet_conn = SocketConnection( + sink_multiaddr=Multiaddr(address="/ip4/10.0.0.1/tcp/8000") + ) + node_network = { + node_a: NodeNetworkInfo(interfaces=[ethernet_interface]), + node_b: NodeNetworkInfo(interfaces=[ethernet_interface]), + } + + for n in (node_a, node_b): + topology.add_node(n) + + # Parallel RDMA links in both directions, each on its own interface. + for offset in range(links_per_pair): + iface = 3 + offset # rdma_en3, rdma_en4, ... + topology.add_connection( + Connection(source=node_a, sink=node_b, edge=create_rdma_connection(iface)) + ) + topology.add_connection( + Connection(source=node_b, sink=node_a, edge=create_rdma_connection(iface)) + ) + + topology.add_connection(Connection(source=node_a, sink=node_b, edge=ethernet_conn)) + topology.add_connection(Connection(source=node_b, sink=node_a, edge=ethernet_conn)) + + return topology, node_a, node_b, node_network + + +def test_jaccl_devices_matrix_bonds_parallel_rdma_links(model_card: ModelCard): + """Multiple Thunderbolt links between a pair should be emitted as a list so the + JACCL ring backend can stripe across them, rather than picking one and discarding + the rest.""" + # arrange + model_card = model_card.model_copy( + update={"n_layers": 12, "storage_size": Memory.from_bytes(1500)} + ) + topology, node_a, node_b, node_network = _build_two_node_bonded_rdma_topology( + links_per_pair=2 + ) + node_memory = { + node_a: create_node_memory(1000), + node_b: create_node_memory(1000), + } + node_rdma_ctl = { + node_a: NodeRdmaCtlStatus(enabled=True), + node_b: NodeRdmaCtlStatus(enabled=True), + } + + cic = PlaceInstance( + sharding=Sharding.Tensor, + instance_meta=InstanceMeta.MlxJaccl, + command_id=CommandId(), + model_card=model_card, + min_nodes=1, + ) + + placements = place_instance( + cic, + topology, + {}, + node_memory, + node_network, + _metal_only(node_memory), + node_rdma_ctl=node_rdma_ctl, + ) + + assert len(placements) == 1 + instance = next(iter(placements.values())) + assert isinstance(instance, MlxJacclInstance) + + matrix = instance.jaccl_devices + assert len(matrix) == 2 + assert matrix[0][0] is None + assert matrix[1][1] is None + # Both parallel interfaces survive in the cell. Order is not guaranteed (the + # topology graph returns edges in reverse insertion order), so compare as sets. + assert set(matrix[0][1]) == {"rdma_en3", "rdma_en4"} and len(matrix[0][1]) == 2 + assert set(matrix[1][0]) == {"rdma_en3", "rdma_en4"} and len(matrix[1][0]) == 2 + + +def test_jaccl_devices_matrix_normalises_asymmetric_link_counts(model_card: ModelCard): + """When one pair has more parallel links than another, the matrix is normalised + down to the minimum count so the JACCL ring validator (which requires every ring + edge to carry the same number of connections) accepts it.""" + # arrange: three nodes where the A<->B pair has 2 links but B<->C and A<->C have 1 + model_card = model_card.model_copy( + update={"n_layers": 12, "storage_size": Memory.from_bytes(1500)} + ) + topology = Topology() + node_a, node_b, node_c = NodeId(), NodeId(), NodeId() + + ethernet_interface = NetworkInterfaceInfo(name="en0", ip_address="10.0.0.1") + ethernet_conn = SocketConnection( + sink_multiaddr=Multiaddr(address="/ip4/10.0.0.1/tcp/8000") + ) + node_network = { + n: NodeNetworkInfo(interfaces=[ethernet_interface]) + for n in (node_a, node_b, node_c) + } + for n in (node_a, node_b, node_c): + topology.add_node(n) + + # A<->B gets two parallel links; the other two pairs get one each. + for iface in (3, 4): + topology.add_connection( + Connection(source=node_a, sink=node_b, edge=create_rdma_connection(iface)) + ) + topology.add_connection( + Connection(source=node_b, sink=node_a, edge=create_rdma_connection(iface)) + ) + for src, sink, iface in [ + (node_b, node_c, 5), + (node_c, node_b, 5), + (node_a, node_c, 6), + (node_c, node_a, 6), + ]: + topology.add_connection( + Connection(source=src, sink=sink, edge=create_rdma_connection(iface)) + ) + for src, sink in [ + (node_a, node_b), + (node_b, node_a), + (node_b, node_c), + (node_c, node_b), + (node_a, node_c), + (node_c, node_a), + ]: + topology.add_connection(Connection(source=src, sink=sink, edge=ethernet_conn)) + + node_memory = {n: create_node_memory(1000) for n in (node_a, node_b, node_c)} + node_rdma_ctl = { + n: NodeRdmaCtlStatus(enabled=True) for n in (node_a, node_b, node_c) + } + + cic = PlaceInstance( + sharding=Sharding.Tensor, + instance_meta=InstanceMeta.MlxJaccl, + command_id=CommandId(), + model_card=model_card, + min_nodes=1, + ) + + placements = place_instance( + cic, + topology, + {}, + node_memory, + node_network, + _metal_only(node_memory), + node_rdma_ctl=node_rdma_ctl, + ) + + instance = next(iter(placements.values())) + assert isinstance(instance, MlxJacclInstance) + matrix = instance.jaccl_devices + + # place_instance selects the smallest sufficient cycle, so the matrix may be + # 2x2 (one pair) rather than the full 3x3. Either way, every off-diagonal cell + # must be normalised to a uniform length (the ring validator requires this). + n = len(matrix) + off_diag_lengths = {len(matrix[i][j]) for i in range(n) for j in range(n) if i != j} + assert len(off_diag_lengths) == 1, ( + f"ring edges must carry a uniform connection count, got {off_diag_lengths}" + ) + # The only single-link pairs are count 1; the A<->B pair (2 physical links) must + # have been trimmed down to that minimum, not left at 2. + assert off_diag_lengths == {1} diff --git a/src/exo/shared/types/worker/instances.py b/src/exo/shared/types/worker/instances.py index 16233f3f05..4181352566 100644 --- a/src/exo/shared/types/worker/instances.py +++ b/src/exo/shared/types/worker/instances.py @@ -31,7 +31,10 @@ class MlxRingInstance(BaseInstance): class MlxJacclInstance(BaseInstance): - jaccl_devices: list[list[str | None]] + # Each cell is the list of RDMA interface names from this device to another + # (multiple when parallel Thunderbolt links are bonded), or None on the + # diagonal. The JACCL backend accepts either a single name or a list per cell. + jaccl_devices: list[list[list[str] | None]] jaccl_coordinators: dict[NodeId, str]