Angelos Katharopoulos
291 lines
8.6 KiB
Python
291 lines
8.6 KiB
Python
# Copyright © 2026 Apple Inc.
|
|
|
|
import argparse
|
|
import os
|
|
import pickle
|
|
import sys
|
|
import time
|
|
from dataclasses import dataclass
|
|
from functools import partial, total_ordering
|
|
from pathlib import Path
|
|
from tempfile import TemporaryDirectory
|
|
from typing import Literal, Optional
|
|
|
|
import mlx.core as mx
|
|
from huggingface_hub.errors import LocalEntryNotFoundError
|
|
from mlx._distributed_utils.common import Hostfile
|
|
from mlx._distributed_utils.launch import launch_jaccl, launch_ring
|
|
from tqdm import tqdm
|
|
|
|
from .utils import hf_repo_to_path
|
|
|
|
CHUNK_SIZE = 100 * 1024 * 1024
|
|
|
|
|
|
@total_ordering
|
|
@dataclass
|
|
class DirectoryEntry:
|
|
entry_type: Literal["directory", "symlink", "file"]
|
|
path: str
|
|
dst: Optional[str]
|
|
|
|
def __lt__(self, other):
|
|
order_type = dict(directory=0, symlink=1, file=2)
|
|
o1 = order_type[self.entry_type]
|
|
o2 = order_type[other.entry_type]
|
|
return o1 < o2 or (o1 == o2 and self.path < other.path)
|
|
|
|
def __eq__(self, other):
|
|
return (
|
|
self.entry_type == other.entry_type
|
|
and self.path == other.path
|
|
and self.dst == other.dst
|
|
)
|
|
|
|
@classmethod
|
|
def from_path(cls, root, path):
|
|
entry_type = {
|
|
(True, False): "directory",
|
|
(False, True): "symlink",
|
|
(False, False): "file",
|
|
}[path.is_dir(), path.is_symlink()]
|
|
dst = path.readlink() if path.is_symlink() else None
|
|
|
|
return cls(entry_type, str(path.relative_to(root)), str(dst))
|
|
|
|
|
|
def error(*args, **kwargs):
|
|
kwargs["file"] = sys.stderr
|
|
print("\033[31m[ERROR]", *args, "\033[0m", **kwargs)
|
|
|
|
|
|
def launch(args):
|
|
if args.hostfile is None:
|
|
raise ValueError("No hostfile provided")
|
|
|
|
hostfile = Hostfile.from_file(args.hostfile)
|
|
if hostfile.backend == "":
|
|
raise ValueError("Backend needs to be defined in the hostfile.")
|
|
if len(hostfile.hosts) == 1:
|
|
raise ValueError("More than one node needs to be in the hostfile")
|
|
|
|
launch_args = argparse.Namespace(
|
|
backend=hostfile.backend,
|
|
cwd=str(Path.cwd()),
|
|
env=hostfile.envs,
|
|
verbose=False,
|
|
python=None,
|
|
starting_port=32323,
|
|
connections_per_ip=1,
|
|
)
|
|
cmd = [
|
|
sys.executable,
|
|
"-m",
|
|
"mlx_lm",
|
|
"share",
|
|
]
|
|
if args.path is not None:
|
|
cmd += ["--path", args.path]
|
|
if args.model is not None:
|
|
cmd += ["--model", args.model]
|
|
if args.tmpdir is not None:
|
|
cmd += ["--tmpdir", args.tmpdir]
|
|
if args.dst is not None:
|
|
cmd += ["--dst", args.dst]
|
|
|
|
if hostfile.backend == "ring":
|
|
launch_ring(None, hostfile.hosts, launch_args, cmd)
|
|
elif hostfile.backend == "jaccl" or hostfile.backend == "jaccl-ring":
|
|
launch_jaccl(None, hostfile.hosts, launch_args, cmd)
|
|
else:
|
|
raise ValueError("Only ring, jaccl and jaccl-ring backends are supported.")
|
|
|
|
|
|
def get_files(path):
|
|
if not path.is_dir():
|
|
return path.parent, [DirectoryEntry.from_path(path.parent, path)]
|
|
|
|
files = [DirectoryEntry.from_path(path, f) for f in path.rglob("*")]
|
|
return path, sorted(files)
|
|
|
|
|
|
def format_bw(x):
|
|
if x >= 1e9:
|
|
return f"{x / 1e9:.2} GB/s"
|
|
if x >= 1e6:
|
|
return f"{x / 1e6:.2} MB/s"
|
|
if x >= 1e3:
|
|
return f"{x / 1e3:.2} KB/s"
|
|
return f"{x:.2} B/s"
|
|
|
|
|
|
def share_file(path, file, src, group=None):
|
|
group = group or mx.distributed.init()
|
|
all_sum = partial(mx.distributed.all_sum, group=group)
|
|
total_size = 0
|
|
start_time = time.time()
|
|
|
|
if group.rank() == src:
|
|
with open(path / file, "rb") as f:
|
|
f.seek(0, 2)
|
|
total_size = f.tell()
|
|
f.seek(0)
|
|
|
|
pbar = tqdm(
|
|
total=total_size,
|
|
unit="B",
|
|
unit_scale=True,
|
|
desc=file,
|
|
position=1,
|
|
leave=False,
|
|
)
|
|
while True:
|
|
data = f.read(CHUNK_SIZE)
|
|
if not data:
|
|
mx.eval(all_sum(0))
|
|
break
|
|
|
|
mx.eval(all_sum(len(data)))
|
|
mx.async_eval(all_sum(data))
|
|
pbar.update(len(data))
|
|
pbar.close()
|
|
|
|
else:
|
|
with open(path / file, "wb") as f:
|
|
data = None
|
|
chunk_size = all_sum(0).item()
|
|
if chunk_size > 0:
|
|
data = all_sum(mx.zeros(chunk_size, dtype=mx.uint8))
|
|
mx.eval(data)
|
|
|
|
while chunk_size > 0:
|
|
next_data = None
|
|
chunk_size = all_sum(0).item()
|
|
if chunk_size > 0:
|
|
next_data = all_sum(mx.zeros(chunk_size, dtype=mx.uint8))
|
|
mx.async_eval(next_data)
|
|
|
|
f.write(bytes(data))
|
|
data = next_data
|
|
|
|
return total_size, time.time() - start_time
|
|
|
|
|
|
def share_files(path, files, src, group=None):
|
|
group = group or mx.distributed.init()
|
|
all_sum = partial(mx.distributed.all_sum, group=group)
|
|
|
|
if group.rank() == src:
|
|
# Share the list first
|
|
file_list = pickle.dumps(files)
|
|
mx.eval(all_sum(len(file_list)))
|
|
mx.eval(all_sum(file_list))
|
|
|
|
else:
|
|
# Get the list first
|
|
file_list_size = all_sum(0).item()
|
|
data = all_sum(mx.zeros(file_list_size, dtype=mx.uint8))
|
|
files = pickle.loads(bytes(data))
|
|
|
|
# Make the directories and symlinks
|
|
for file in files:
|
|
if file.entry_type == "directory":
|
|
(path / file.path).mkdir()
|
|
elif file.entry_type == "symlink":
|
|
(path / file.path).symlink_to(file.dst)
|
|
|
|
# Everybody shares the files
|
|
total_size = 0
|
|
total_time = 1e-6
|
|
pbar = tqdm(total=len(files), desc="Files", position=0, disable=group.rank() != src)
|
|
for file in files:
|
|
if file.entry_type == "file":
|
|
s, t = share_file(path, file.path, src, group)
|
|
total_size += s
|
|
total_time += t
|
|
pbar.update(1)
|
|
pbar.set_postfix(speed=format_bw(total_size / total_time))
|
|
pbar.close()
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="Distribute a model to other nodes using MLX distributed."
|
|
)
|
|
parser.add_argument("--path", type=str, help="Path to a file or folder to share.")
|
|
parser.add_argument(
|
|
"--model", type=str, help="The path to a local model or Hugging Face repo"
|
|
)
|
|
parser.add_argument(
|
|
"--hostfile",
|
|
type=str,
|
|
help="The file containing the hosts and connection information",
|
|
)
|
|
parser.add_argument(
|
|
"--dst",
|
|
type=str,
|
|
help="The destination path in other nodes (defaults to --path or --model)",
|
|
)
|
|
parser.add_argument(
|
|
"--tmpdir",
|
|
type=str,
|
|
help="Intermediate temporary directory to ensure successfull transfer",
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
if args.path is args.model is None:
|
|
parser.error("One of --path or --model must be provided")
|
|
|
|
mx.set_default_device(mx.cpu)
|
|
world = mx.distributed.init()
|
|
|
|
if world.size() == 1:
|
|
launch(args)
|
|
return
|
|
|
|
# Check if any node has the data
|
|
path = None
|
|
files = []
|
|
if args.path is not None and (path := Path(args.path)).exists():
|
|
path, files = get_files(path)
|
|
elif args.model is not None:
|
|
try:
|
|
path = hf_repo_to_path(args.model)
|
|
if path.parent.name != "snapshots":
|
|
raise ValueError(
|
|
f"The model repository appears to be corrupted, it resolved to {str(path)}"
|
|
)
|
|
path, files = get_files(path.parent.parent)
|
|
except Exception as e:
|
|
pass
|
|
has_file = mx.distributed.all_gather(len(files) > 0)
|
|
src = has_file.argmax().item()
|
|
has_file = has_file.any().item()
|
|
|
|
if not has_file:
|
|
error("The --path needs to exist in at least one node.")
|
|
error("If it is a remote repository download it first with `hf download`")
|
|
sys.exit(1)
|
|
|
|
# Share the path that is resolved
|
|
if args.dst is None:
|
|
if world.rank() == src:
|
|
data = str(path).encode("utf-8")
|
|
mx.eval(mx.distributed.all_sum(len(data)))
|
|
mx.eval(mx.distributed.all_sum(data))
|
|
else:
|
|
data_size = mx.distributed.all_sum(0).item()
|
|
data = mx.distributed.all_sum(mx.zeros(data_size, dtype=mx.uint8))
|
|
path = Path(bytes(data).decode("utf-8"))
|
|
elif world.rank() != src:
|
|
path = Path(args.dst)
|
|
|
|
with TemporaryDirectory(dir=args.tmpdir) as tmp:
|
|
if world.rank() == src:
|
|
share_files(path, files, src, world)
|
|
else:
|
|
share_files(Path(tmp), files, src, world)
|
|
path.mkdir(parents=True, exist_ok=True)
|
|
os.rename(tmp, path)
|