Share model (#871)
This commit is contained in:
Angelos Katharopoulos
committed by
GitHub
GitHub
parent
bcf630614f
commit
c6d9d3c9f5
@@ -22,6 +22,7 @@ def main():
|
||||
"gptq",
|
||||
"server",
|
||||
"upload",
|
||||
"share",
|
||||
)
|
||||
subpackages = {
|
||||
"awq": "quant",
|
||||
|
||||
+290
@@ -0,0 +1,290 @@
|
||||
# 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)
|
||||
Reference in New Issue
Block a user