L-electron-Rare/mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

JACCL update (#3094)

Browse Source
This commit is contained in:
Angelos Katharopoulos
2026-02-05 15:16:07 -08:00
committed by GitHub
parent 99ca62c4d3
commit ceea571490
11 changed files with 2455 additions and 1256 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mlx/_distributed_utils/common.py
+89 -49
View File
@@ -4,9 +4,9 @@ import argparse
import ipaddress
import json
import sys
from dataclasses import dataclass
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional
from typing import Optional, Union
@dataclass
@@ -14,7 +14,93 @@ class Host:
rank: int
ssh_hostname: str
ips: list[str]
rdma: list[Optional[str]]
rdma: list[Optional[Union[str, list[str]]]]
@dataclass
class Hostfile:
hosts: list[Host]
backend: str = ""
envs: list[str] = field(default_factory=list)
def to_json(self):
return {
"backend": self.backend,
"envs": self.envs,
"hosts": [
{"ssh": h.ssh_hostname, "ips": h.ips, "rdma": h.rdma}
for h in self.hosts
],
}
@classmethod
def from_file(cls, hostfile):
"""Parse the json hostfile that contains both the hostnames to ssh into and
the ips to communicate over when using the ring backend. It can also
contain the backend to be used and environment variables to set when
launching a distributed job.
Example:
{
"backend": "jaccl",
"envs": [
"MLX_METAL_FAST_SYNCH=1"
],
"hosts": [
{"ssh": "hostname1", "ips": ["123.123.123.1"], "rdma": [null, "rdma_en2", "rdma_en3"]},
{"ssh": "hostname2", "ips": ["123.123.123.2"], "rdma": ["rdma_en2", null, "rdma_en3"]},
...
{"ssh": "hostnameN", "ips": ["123.123.123.N"], "rdma": ["rdma_en2", "rdma_en3", null]},
]
}
Args:
hostfile (str): The path to the json file containing the host
information
"""
hostfile = Path(hostfile)
if not hostfile.exists():
raise ValueError(f"Hostfile {str(hostfile)} doesn't exist")
try:
data = json.load(open(hostfile))
backend = ""
envs = []
hosts = []
if isinstance(data, dict):
backend = data["backend"]
envs = data["envs"]
hosts = data["hosts"]
elif isinstance(data, list):
hosts = data
hosts = [
Host(i, h["ssh"], h.get("ips", []), h.get("rdma", []))
for i, h in enumerate(hosts)
]
return cls(hosts, backend, envs)
except Exception as e:
raise ValueError(
f"Failed to parse hostfile {str(hostfile)} ({str(e)})"
) from e
@classmethod
def from_list(cls, hostlist, repeats=1):
hosts = []
for i, h in enumerate(hostlist.split(",")):
if h == "":
raise ValueError("Hostname cannot be empty")
try:
ipaddress.ip_address(h)
ips = [h]
except ValueError:
ips = []
for i in range(repeats):
hosts.append(Host(i, h, ips, []))
return cls(hosts)
class OptionalBoolAction(argparse.Action):
@@ -47,49 +133,3 @@ def log_warning(*args, **kwargs):
def log_error(*args, **kwargs):
kwargs["file"] = sys.stderr
print("\033[31m[ERROR]", *args, "\033[0m", **kwargs)
def parse_hostlist(parser, hostlist, repeats):
hosts = []
for i, h in enumerate(hostlist.split(",")):
if h == "":
raise ValueError("Hostname cannot be empty")
try:
ipaddress.ip_address(h)
ips = [h]
except ValueError:
ips = []
for i in range(repeats):
hosts.append(Host(i, h, ips, []))
return hosts
def parse_hostfile(parser, hostfile):
"""Parse the json hostfile that contains both the hostnames to ssh into and
the ips to communicate over when using the ring backend.
Example:
[
{"ssh": "hostname1", "ips": ["123.123.123.1"], "rdma": [null, "rdma_en2", "rdma_en3"]},
{"ssh": "hostname2", "ips": ["123.123.123.2"], "rdma": ["rdma_en2", null, "rdma_en3"]},
...
{"ssh": "hostnameN", "ips": ["123.123.123.N"], "rdma": ["rdma_en2", "rdma_en3", null]},
]
Args:
hostfile (str): The path to the json file containing the host
information
"""
hostfile = Path(hostfile)
if not hostfile.exists():
parser.error(f"Hostfile {str(hostfile)} doesn't exist")
try:
hosts = []
with open(hostfile) as f:
for i, h in enumerate(json.load(f)):
hosts.append(Host(i, h["ssh"], h.get("ips", []), h.get("rdma", [])))
return hosts
except Exception as e:
parser.error(f"Failed to parse hostfile {str(hostfile)} ({str(e)})")
python/mlx/_distributed_utils/config.py
+113 -65
View File
@@ -14,12 +14,11 @@ import mlx.core as mx
from .common import (
Host,
Hostfile,
OptionalBoolAction,
log,
log_error,
log_warning,
parse_hostfile,
parse_hostlist,
)
@@ -70,9 +69,20 @@ def add_ips(hosts, verbose=False):
log_warning("Could not extract ip for", h.ssh_hostname)
def check_rdma(hosts, verbose=False):
def save_hostfile(args, hostfile):
if args.output_hostfile:
with open(args.output_hostfile, "w") as f:
json.dump(hostfile.to_json(), f, indent=4)
else:
print("Hostfile")
print("========")
print(json.dumps(hostfile.to_json(), indent=4))
def check_rdma(hosts, verbose=False, strict=True):
# Check whether the hosts are capable of RDMA over thunderbolt
warn = False
log_f = log_warning if not strict else log_error
failed = False
for h in hosts:
log(verbose, "Checking that", h.ssh_hostname, "supports RDMA")
rdma_devs = (
@@ -82,19 +92,20 @@ def check_rdma(hosts, verbose=False):
)
rdma_devs = [d for d in rdma_devs if d.startswith("rdma_")]
if not rdma_devs:
log_warning(h.ssh_hostname, "does not seem to have RDMA enabled")
warn = True
log_f(h.ssh_hostname, "does not seem to have RDMA enabled")
failed = True
if warn:
log_warning()
log_warning(
"Some of the hosts don't have RDMA enabled or they don't support RDMA."
)
log_warning()
log_warning(
"See https://ml-explore.github.io/mlx/build/html/usage/distributed.html"
)
log_warning("for instructions on how to enable RDMA.")
if failed:
log_f()
log_f("Some of the hosts don't have RDMA enabled or they don't support RDMA.")
log_f()
log_f("See https://ml-explore.github.io/mlx/build/html/usage/distributed.html")
log_f("for instructions on how to enable RDMA.")
if failed and strict:
sys.exit(1)
return not failed
def can_auto_setup(hosts, sshinfo, auto_setup=False):
@@ -340,6 +351,20 @@ def check_valid_mesh(hosts, connectivity, strict=True):
return True
def check_valid_ring(hosts, rings, strict=True):
has_ring = len(rings) > 0 and len(rings[0][0]) == len(hosts)
if strict and not has_ring:
log_error("Could not find a full ring.")
log_error()
log_error("Try passing --dot to visualize the connectivity")
if len(rings) > 0:
log_error("Rings found:")
for r in rings:
log_error(f" - {','.join(hosts[i].ssh_hostname for i in r)}")
sys.exit(1)
return has_ring
def check_ssh_connections(hosts):
results = [None] * len(hosts)
@@ -408,52 +433,38 @@ def prepare_ethernet_hostfile(args, hosts):
log(args.verbose, f"Preparing an ethernet hostfile")
add_ips(hosts, args.verbose)
hostfile = []
for h in hosts:
hostfile.append(dict(ssh=h.ssh_hostname, ips=h.ips))
hostfile = Hostfile(
[Host(i, h.ssh_hostname, h.ips, []) for i, h in enumerate(hosts)], "", args.env
)
if args.output_hostfile:
with open(args.output_hostfile, "w") as f:
json.dump(hostfile, f, indent=4)
else:
print("Hostfile")
print("========")
print(json.dumps(hostfile, indent=4))
save_hostfile(args, hostfile)
def configure_ring(args, hosts, ips, ring, sshinfo):
log(args.verbose, "Prepare a ring hostfile")
ring, count = ring
hostfile = []
ring_hosts = []
for i, node in enumerate(ring):
h = hosts[node]
peer = ring[i - 1]
hostfile.append(
{
"ssh": h.ssh_hostname,
"ips": [ips.ips[node, peer][c][1] for c in range(count)],
"rdma": [],
}
ring_hosts.append(
Host(
i, h.ssh_hostname, [ips.ips[node, peer][c][1] for c in range(count)], []
)
)
hostfile = Hostfile(ring_hosts, "ring", args.env)
has_sudo = can_auto_setup(hosts, sshinfo, args.auto_setup)
ips.setup(verbose=args.verbose, auto_setup=args.auto_setup and has_sudo)
if args.output_hostfile:
with open(args.output_hostfile, "w") as f:
json.dump(hostfile, f, indent=4)
else:
print("Hostfile")
print("========")
print(json.dumps(hostfile, indent=4))
save_hostfile(args, hostfile)
def configure_jaccl(args, hosts, ips, sshinfo):
log(args.verbose, "Prepare a jaccl hostfile")
check_rdma(hosts, args.verbose)
add_ips(hosts, args.verbose)
hostfile = []
jaccl_hosts = []
for i, h in enumerate(hosts):
rdma = []
for j in range(len(hosts)):
@@ -461,18 +472,42 @@ def configure_jaccl(args, hosts, ips, sshinfo):
rdma.append(None)
else:
rdma.append(f"rdma_{ips.ips[i, j][0][0]}")
hostfile.append({"ssh": h.ssh_hostname, "ips": h.ips, "rdma": rdma})
jaccl_hosts.append(Host(i, h.ssh_hostname, h.ips, rdma))
hostfile = Hostfile(jaccl_hosts, "jaccl", args.env)
has_sudo = can_auto_setup(hosts, sshinfo, args.auto_setup)
ips.setup(verbose=args.verbose, auto_setup=args.auto_setup and has_sudo)
if args.output_hostfile:
with open(args.output_hostfile, "w") as f:
json.dump(hostfile, f, indent=4)
else:
print("Hostfile")
print("========")
print(json.dumps(hostfile, indent=4))
save_hostfile(args, hostfile)
def configure_jaccl_ring(args, hosts, ips, ring, sshinfo):
log(args.verbose, "Prepare a jaccl-ring hostfile")
add_ips(hosts, args.verbose)
jaccl_hosts = []
num_nodes = len(hosts)
ring, count = ring
for i, node in enumerate(ring):
h = hosts[node]
peer_left = ring[i - 1]
peer_right = ring[(i + 1) % num_nodes]
rdmas = []
for j in range(len(hosts)):
if j not in (peer_left, peer_right):
rdmas.append(None)
else:
rdma = []
for c in range(count):
rdma.append(f"rdma_{ips.ips[i, j][c][0]}")
rdmas.append(rdma[0] if count == 1 else rdma)
jaccl_hosts.append(Host(i, h.ssh_hostname, h.ips, rdmas))
hostfile = Hostfile(jaccl_hosts, "jaccl-ring", args.env)
has_sudo = can_auto_setup(hosts, sshinfo, args.auto_setup)
ips.setup(verbose=args.verbose, auto_setup=args.auto_setup and has_sudo)
save_hostfile(args, hostfile)
def prepare_tb_hostfile(args, hosts, sshinfo):
@@ -489,37 +524,44 @@ def prepare_tb_hostfile(args, hosts, sshinfo):
if args.backend is None:
rings = extract_rings(connectivity)
has_mesh = check_valid_mesh(hosts, connectivity, False)
has_ring = len(rings) > 0 and len(rings[0][0]) == len(hosts)
has_ring = check_valid_ring(hosts, rings, False)
has_rdma = check_rdma(hosts, args.verbose, False)
if not has_ring and not has_mesh:
log_error("Neither thunderbolt mesh nor ring found.")
log_error("Perhaps run with --dot to generate a plot of the connectivity.")
sys.exit(1)
elif has_rdma and has_mesh:
configure_jaccl(args, hosts, ips, sshinfo)
elif has_rdma and has_ring:
configure_jaccl_ring(args, hosts, ips, rings[0], sshinfo)
elif has_ring:
configure_ring(args, hosts, ips, rings[0], sshinfo)
else:
configure_jaccl(args, hosts, ips, sshinfo)
log_error("RDMA is not available and ring is not found.")
log_error("Perhaps run with --dot to generate a plot of the connectivity.")
sys.exit(1)
elif args.backend == "ring":
rings = extract_rings(connectivity)
has_ring = len(rings) > 0 and len(rings[0][0]) == len(hosts)
if not has_ring:
log_error("Could not find a full ring.")
log_error()
log_error("Try passing --dot to visualize the connectivity")
if len(rings) > 0:
log_error("Rings found:")
for r in rings:
log_error(f" - {','.join(hosts[i].ssh_hostname for i in r)}")
sys.exit(1)
check_valid_ring(hosts, rings)
configure_ring(args, hosts, ips, rings[0], sshinfo)
elif args.backend == "jaccl":
check_valid_mesh(hosts, connectivity)
check_rdma(hosts, args.verbose)
configure_jaccl(args, hosts, ips, sshinfo)
elif args.backend == "jaccl-ring":
rings = extract_rings(connectivity)
check_valid_ring(hosts, rings)
check_rdma(hosts, args.verbose)
configure_jaccl_ring(args, hosts, ips, rings[0], sshinfo)
def main():
parser = argparse.ArgumentParser(
@@ -555,16 +597,22 @@ def main():
)
parser.add_argument(
"--backend",
choices=["ring", "jaccl"],
choices=["ring", "jaccl", "jaccl-ring"],
default=None,
help="Which distributed backend to configure",
)
parser.add_argument(
"--env",
action="append",
default=[],
help="Set environment variables for the jobs",
)
args = parser.parse_args()
if args.hostfile is not None:
hosts = parse_hostfile(parser, args.hostfile)
hosts = Hostfile.from_file(args.hostfile).hosts
else:
hosts = parse_hostlist(parser, args.hosts, 1)
hosts = Hostfile.from_list(args.hosts).hosts
# Check that we can ssh
log(
python/mlx/_distributed_utils/launch.py
+24 -12
View File
@@ -19,7 +19,7 @@ from subprocess import PIPE, Popen, run
import mlx.core as mx
from .common import log, log_warning, parse_hostfile, parse_hostlist, positive_number
from .common import Hostfile, log, log_warning, positive_number
class CommandProcess:
@@ -367,6 +367,7 @@ def launch_jaccl(parser, hosts, args, command):
if not hosts[0].ips:
raise ValueError("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)
have_nulls = all(h.rdma[i] is None for i, h in enumerate(hosts))
if not have_rdmas or not have_nulls:
@@ -376,6 +377,8 @@ def launch_jaccl(parser, hosts, args, command):
env = args.env
cwd = args.cwd
env.append(f"MLX_JACCL_COORDINATOR={coordinator}:{args.starting_port}")
if jaccl_ring:
env.append("MLX_JACCL_RING=1")
files = {"MLX_IBV_DEVICES": json.dumps([h.rdma for h in hosts])}
log(args.verbose, "Running", shlex.join(command))
@@ -474,8 +477,6 @@ def main():
parser.add_argument("--hostfile", help="The file containing the hosts")
parser.add_argument(
"--backend",
choices=["ring", "mpi", "nccl", "jaccl"],
default="nccl" if mx.cuda.is_available() else "ring",
help="Which distributed backend to launch",
)
parser.add_argument(
@@ -535,9 +536,16 @@ def main():
# Try to extract a list of hosts and corresponding ips
if args.hostfile is not None:
hosts = parse_hostfile(parser, args.hostfile)
hostfile = Hostfile.from_file(args.hostfile)
else:
hosts = parse_hostlist(parser, args.hosts, args.repeat_hosts)
hostfile = Hostfile.from_list(args.hosts, args.repeat_hosts)
# Extract extra arguments from the hostfile
if hostfile.backend != "" and args.backend is None:
args.backend = hostfile.backend
if args.backend is None:
args.backend = "nccl" if mx.cuda.is_available() else "ring"
args.env = hostfile.envs + args.env
# Check if the script is a file and convert it to a full path
if (script := Path(rest[0])).exists() and script.is_file():
@@ -549,10 +557,14 @@ def main():
# Launch
if args.backend == "ring":
launch_ring(parser, hosts, args, rest)
if args.backend == "mpi":
launch_mpi(parser, hosts, args, rest)
if args.backend == "nccl":
launch_nccl(parser, hosts, args, rest)
if args.backend == "jaccl":
launch_jaccl(parser, hosts, args, rest)
launch_ring(parser, hostfile.hosts, args, rest)
elif args.backend == "mpi":
launch_mpi(parser, hostfile.hosts, args, rest)
elif args.backend == "nccl":
launch_nccl(parser, hostfile.hosts, args, rest)
elif args.backend == "jaccl" or args.backend == "jaccl-ring":
launch_jaccl(parser, hostfile.hosts, args, rest)
else:
parser.error(
"The backend should be one of {'ring', 'mpi', 'nccl', 'jaccl', 'jaccl-ring'}"
)