diff --git a/mlx_lm/examples/pipeline_generate.py b/mlx_lm/examples/pipeline_generate.py
index 6c883de..5d60b9e 100644
--- a/mlx_lm/examples/pipeline_generate.py
+++ b/mlx_lm/examples/pipeline_generate.py
@@ -17,71 +17,11 @@ https://ml-explore.github.io/mlx/build/html/usage/distributed.html).
"""
import argparse
-import json
-import resource
-from pathlib import Path
import mlx.core as mx
-from huggingface_hub import snapshot_download
-from mlx.utils import tree_flatten
-
-from mlx_lm import load, stream_generate
-from mlx_lm.utils import load_model, load_tokenizer
-
-# Needed for 8 bit model
-resource.setrlimit(resource.RLIMIT_NOFILE, (2048, 4096))
-
-
-def download(repo: str, allow_patterns: list[str]) -> Path:
- return Path(
- snapshot_download(
- repo,
- allow_patterns=allow_patterns,
- )
- )
-
-
-def shard_and_load(repo):
- # Get model path with everything but weight safetensors
- model_path = download(
- args.model,
- allow_patterns=["*.json", "*.py", "tokenizer.model", "*.tiktoken", "*.txt"],
- )
-
- # Lazy load and shard model to figure out
- # which weights we need
- model, config = load_model(model_path, lazy=True, strict=False)
-
- group = mx.distributed.init()
- rank = group.rank()
- model.model.pipeline(group)
-
- # Figure out which files we need for the local shard
- with open(model_path / "model.safetensors.index.json", "r") as fid:
- weight_index = json.load(fid)["weight_map"]
-
- local_files = set()
- for k, _ in tree_flatten(model.parameters()):
- local_files.add(weight_index[k])
-
- # Download weights for local shard
- download(args.model, allow_patterns=local_files)
-
- # Load and shard the model, and load the weights
- tokenizer = load_tokenizer(
- model_path,
- {"trust_remote_code": True},
- eos_token_ids=config.get("eos_token_id", None),
- )
- model, _ = load_model(model_path, lazy=True, strict=False)
- model.model.pipeline(group)
- mx.eval(model.parameters())
-
- # Synchronize processes before generation to avoid timeout if downloading
- # model for the first time.
- mx.eval(mx.distributed.all_sum(mx.array(1.0), stream=mx.cpu))
- return model, tokenizer
+from mlx_lm import stream_generate
+from mlx_lm.utils import pipeline_load
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="LLM pipelined inference example")
@@ -112,7 +52,7 @@ if __name__ == "__main__":
if rank == 0:
print(*args, **kwargs)
- model, tokenizer = shard_and_load(args.model)
+ model, tokenizer = pipeline_load(args.model)
messages = [{"role": "user", "content": args.prompt}]
prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
diff --git a/mlx_lm/models/deepseek_v2.py b/mlx_lm/models/deepseek_v2.py
index c7d428f..a3c4948 100644
--- a/mlx_lm/models/deepseek_v2.py
+++ b/mlx_lm/models/deepseek_v2.py
@@ -8,6 +8,7 @@ import mlx.core as mx
import mlx.nn as nn
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .pipeline import PipelineMixin
from .switch_layers import SwitchGLU
@@ -355,7 +356,7 @@ class DeepseekV2DecoderLayer(nn.Module):
return out
-class DeepseekV2Model(nn.Module):
+class DeepseekV2Model(PipelineMixin, nn.Module):
def __init__(self, config: ModelArgs):
super().__init__()
self.vocab_size = config.vocab_size
@@ -364,32 +365,8 @@ class DeepseekV2Model(nn.Module):
DeepseekV2DecoderLayer(config, idx)
for idx in range(config.num_hidden_layers)
]
- self.start_idx = 0
- self.end_idx = len(self.layers)
- self.num_layers = self.end_idx
-
self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
- self.pipeline_rank = 0
- self.pipeline_size = 1
-
- def pipeline(self, group):
- # Split layers in reverse so rank=0 gets the last layers and
- # rank=pipeline_size-1 gets the first
- self.pipeline_rank = group.rank()
- self.pipeline_size = group.size()
- layers_per_rank = len(self.layers) // self.pipeline_size
- extra = len(self.layers) - layers_per_rank * self.pipeline_size
- if self.pipeline_rank < extra:
- layers_per_rank += 1
-
- self.start_idx = (self.pipeline_size - self.pipeline_rank - 1) * layers_per_rank
- self.end_idx = self.start_idx + layers_per_rank
- self.num_layers = layers_per_rank
- self.layers = self.layers[: self.end_idx]
- self.layers[: self.start_idx] = [None] * self.start_idx
- self.num_layers = len(self.layers) - self.start_idx
-
def __call__(
self,
x: mx.array,
@@ -401,15 +378,15 @@ class DeepseekV2Model(nn.Module):
pipeline_size = self.pipeline_size
if cache is None:
- cache = [None] * self.num_layers
+ cache = [None] * len(self.pipeline_layers)
mask = create_attention_mask(h, cache[0])
# Receive from the previous process in the pipeline
if pipeline_rank < pipeline_size - 1:
h = mx.distributed.recv_like(h, (pipeline_rank + 1))
- for i in range(self.num_layers):
- h = self.layers[self.start_idx + i](h, mask, cache[i])
+ for l, c in zip(self.pipeline_layers, cache):
+ h = l(h, mask, cache=c)
# Send to the next process in the pipeline
if pipeline_rank != 0:
@@ -454,4 +431,4 @@ class Model(nn.Module):
@property
def layers(self):
- return self.model.layers[self.model.start_idx : self.model.end_idx]
+ return self.model.pipeline_layers
diff --git a/mlx_lm/models/deepseek_v3.py b/mlx_lm/models/deepseek_v3.py
index 012e2b9..a4abb3c 100644
--- a/mlx_lm/models/deepseek_v3.py
+++ b/mlx_lm/models/deepseek_v3.py
@@ -9,6 +9,7 @@ import mlx.core as mx
import mlx.nn as nn
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .pipeline import PipelineMixin
from .switch_layers import SwitchGLU
@@ -389,7 +390,7 @@ class DeepseekV3DecoderLayer(nn.Module):
return h + r
-class DeepseekV3Model(nn.Module):
+class DeepseekV3Model(PipelineMixin, nn.Module):
def __init__(self, config: ModelArgs):
super().__init__()
self.vocab_size = config.vocab_size
@@ -398,28 +399,7 @@ class DeepseekV3Model(nn.Module):
DeepseekV3DecoderLayer(config, idx)
for idx in range(config.num_hidden_layers)
]
- self.start_idx = 0
- self.end_idx = len(self.layers)
- self.num_layers = self.end_idx
-
self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
- self.pipeline_rank = 0
- self.pipeline_size = 1
-
- def pipeline(self, group):
- # Split layers in reverse so rank=0 gets the last layers and
- # rank=pipeline_size-1 gets the first
- self.pipeline_rank = group.rank()
- self.pipeline_size = group.size()
- layers_per_rank = len(self.layers) // self.pipeline_size
- extra = len(self.layers) - layers_per_rank * self.pipeline_size
- if self.pipeline_rank < extra:
- layers_per_rank += 1
- self.start_idx = (self.pipeline_size - self.pipeline_rank - 1) * layers_per_rank
- self.end_idx = self.start_idx + layers_per_rank
- self.layers = self.layers[: self.end_idx]
- self.layers[: self.start_idx] = [None] * self.start_idx
- self.num_layers = len(self.layers) - self.start_idx
def __call__(
self,
@@ -432,16 +412,15 @@ class DeepseekV3Model(nn.Module):
pipeline_size = self.pipeline_size
if cache is None:
- cache = [None] * self.num_layers
+ cache = [None] * len(self.pipeline_layers)
mask = create_attention_mask(h, cache[0])
# Receive from the previous process in the pipeline
-
if pipeline_rank < pipeline_size - 1:
h = mx.distributed.recv_like(h, (pipeline_rank + 1))
- for i in range(self.num_layers):
- h = self.layers[self.start_idx + i](h, mask, cache[i])
+ for l, c in zip(self.pipeline_layers, cache):
+ h = l(h, mask, cache=c)
# Send to the next process in the pipeline
if pipeline_rank != 0:
@@ -521,7 +500,7 @@ class Model(nn.Module):
@property
def layers(self):
- return self.model.layers[self.model.start_idx : self.model.end_idx]
+ return self.model.pipeline_layers
@property
def cast_predicate(self):
diff --git a/mlx_lm/models/glm4_moe.py b/mlx_lm/models/glm4_moe.py
index e03f181..07856f9 100644
--- a/mlx_lm/models/glm4_moe.py
+++ b/mlx_lm/models/glm4_moe.py
@@ -9,6 +9,7 @@ import mlx.core as mx
import mlx.nn as nn
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .pipeline import PipelineMixin
from .switch_layers import SwitchGLU
@@ -243,7 +244,7 @@ class DecoderLayer(nn.Module):
return h + r
-class LanguageModel(nn.Module):
+class LanguageModel(PipelineMixin, nn.Module):
def __init__(self, config: ModelArgs):
super().__init__()
self.vocab_size = config.vocab_size
@@ -264,13 +265,28 @@ class LanguageModel(nn.Module):
) -> mx.array:
h = self.embed_tokens(x)
- if cache is None:
- cache = [None] * self.num_layers
+ pipeline_rank = self.pipeline_rank
+ pipeline_size = self.pipeline_size
+ if cache is None:
+ cache = [None] * len(self.pipeline_layers)
mask = create_attention_mask(h, cache[0])
- for i in range(self.num_layers):
- h = self.layers[self.start_idx + i](h, mask, cache[i])
+ # Receive from the previous process in the pipeline
+ if pipeline_rank < pipeline_size - 1:
+ h = mx.distributed.recv_like(h, (pipeline_rank + 1))
+
+ for l, c in zip(self.pipeline_layers, cache):
+ h = l(h, mask, cache=c)
+
+ # Send to the next process in the pipeline
+ if pipeline_rank != 0:
+ h = mx.distributed.send(h, (pipeline_rank - 1) % pipeline_size)
+ if cache[-1] is not None:
+ cache[-1].keys = mx.depends(cache[-1].keys, h)
+
+ # Broadcast h while keeping it in the graph
+ h = mx.distributed.all_gather(h)[: h.shape[0]]
return self.norm(h)
@@ -315,7 +331,7 @@ class Model(nn.Module):
@property
def layers(self):
- return self.model.layers
+ return self.model.pipeline_layers
@property
def cast_predicate(self):
diff --git a/mlx_lm/models/pipeline.py b/mlx_lm/models/pipeline.py
new file mode 100644
index 0000000..3a056c6
--- /dev/null
+++ b/mlx_lm/models/pipeline.py
@@ -0,0 +1,31 @@
+# Copyright Ā© 2025 Apple Inc.
+
+import mlx.core as mx
+
+
+class PipelineMixin:
+ def __init__(self):
+ super().__init__()
+ self.pipeline_rank = 0
+ self.pipeline_size = 1
+ self.start_idx = 0
+ self.end_idx = None
+
+ @property
+ def pipeline_layers(self):
+ return self.layers[self.start_idx : self.end_idx]
+
+ def pipeline(self, group):
+ # Split layers in reverse so rank=0 gets the last layers and
+ # rank=pipeline_size-1 gets the first
+ self.pipeline_rank = group.rank()
+ self.pipeline_size = group.size()
+ layers_per_rank = len(self.layers) // self.pipeline_size
+ extra = len(self.layers) - layers_per_rank * self.pipeline_size
+ if self.pipeline_rank < extra:
+ layers_per_rank += 1
+ self.start_idx = (self.pipeline_size - self.pipeline_rank - 1) * layers_per_rank
+ self.end_idx = self.start_idx + layers_per_rank
+ self.layers = self.layers[: self.end_idx]
+ # Keep the layer numbers the same for model loading
+ self.layers[: self.start_idx] = [None] * self.start_idx
diff --git a/mlx_lm/quant/dwq.py b/mlx_lm/quant/dwq.py
index 223a95f..63c89d5 100644
--- a/mlx_lm/quant/dwq.py
+++ b/mlx_lm/quant/dwq.py
@@ -4,6 +4,7 @@ import argparse
import copy
import time
import types
+from pathlib import Path
import mlx.core as mx
import mlx.nn as nn
@@ -18,19 +19,62 @@ from mlx_lm.tuner.trainer import grad_checkpoint, iterate_batches
from mlx_lm.tuner.utils import print_trainable_parameters
from mlx_lm.utils import (
load,
+ load_tokenizer,
+ pipeline_load,
quantize_model,
save,
)
+def compute_dwq_targets(
+ model,
+ save_dir,
+ train_data,
+ valid_data,
+ batch_size,
+ max_seq_length,
+ seed,
+):
+ rank = mx.distributed.init().rank()
+
+ def _compute_targets(data, path, split):
+
+ if rank == 0:
+ path = path / split
+ path.mkdir(parents=True, exist_ok=True)
+ for i, (batch, _) in (
+ pbar := tqdm(
+ enumerate(iterate_batches(data, batch_size, max_seq_length, seed=seed)),
+ total=len(data) // batch_size,
+ desc=f"Computing targets for {split}",
+ disable=rank != 0,
+ )
+ ):
+ batch = batch[:, :-1]
+ logits = model(batch)
+ # Hack to make the last op pre-eval on the CPU to avoid even timeout
+ logits = mx.stop_gradient(logits, stream=mx.cpu)
+ mx.eval(logits)
+ if rank == 0:
+ idx = mx.argpartition(logits, kth=-1024, axis=-1)[..., -1024:]
+ logits = mx.take_along_axis(logits, idx, axis=-1)
+
+ file = path / f"{i:010d}.safetensors"
+ mx.save_safetensors(file, {"logits": logits, "indices": idx})
+
+ _compute_targets(valid_data, save_dir, "valid")
+ _compute_targets(train_data, save_dir, "train")
+
+
def dwq_quantize(
model,
- q_model,
+ target_fn,
opt,
train_data,
valid_data,
- batch_size: int = 2,
- max_seq_length: int = 2048,
+ batch_size,
+ max_seq_length,
+ seed,
dtype: mx.Dtype = mx.bfloat16,
gradient_checkpoint: bool = False,
temperature: float = 2.0,
@@ -52,18 +96,21 @@ def dwq_quantize(
):
m.unfreeze(keys=["scales", "biases"], recurse=False)
- q_model.train()
- q_model.apply_to_modules(unfreeze)
- print_trainable_parameters(q_model)
+ model.train()
+ model.apply_to_modules(unfreeze)
+ print_trainable_parameters(model)
if gradient_checkpoint:
- grad_checkpoint(q_model.layers[0])
+ grad_checkpoint(model.layers[0])
scale = 1 / temperature
def loss_fn(params, x, targets, lengths):
- q_model.update(tree_map(lambda x: x.astype(dtype), params))
- logits = q_model(x)
+ model.update(tree_map(lambda x: x.astype(dtype), params))
+ logits = model(x)
+ if isinstance(targets, tuple):
+ targets, ids = targets
+ logits = mx.take_along_axis(logits, ids, axis=-1)
losses = kl_div_loss(scale * logits, scale * targets)
mask = mx.arange(1, 1 + targets.shape[1]) < lengths[:, 1:]
ntoks = mask.sum()
@@ -81,14 +128,16 @@ def dwq_quantize(
def validate(params, it):
v_loss = 0.0
v_tokens = 0
- for batch, lengths in tqdm(
- iterate_batches(valid_data, batch_size, max_seq_length),
+ for i, (batch, lengths) in tqdm(
+ enumerate(
+ iterate_batches(valid_data, batch_size, max_seq_length, seed=seed)
+ ),
total=len(valid_data) // batch_size,
desc="Computing validation loss",
leave=False,
):
batch = batch[:, :-1]
- targets = model(batch)
+ targets = target_fn(batch, i, split="valid")
mx.eval(targets)
loss, ntoks = loss_fn(params, batch, targets, lengths)
mx.eval(loss, ntoks)
@@ -103,7 +152,7 @@ def dwq_quantize(
# Accumulate learned weights in higher precision
params = tree_map(
lambda x: x.astype(mx.float32),
- q_model.trainable_parameters(),
+ model.trainable_parameters(),
)
total_loss = 0.0
@@ -117,12 +166,14 @@ def dwq_quantize(
for it, (batch, lengths) in (
pbar := tqdm(
- enumerate(iterate_batches(train_data, batch_size, max_seq_length)),
+ enumerate(
+ iterate_batches(train_data, batch_size, max_seq_length, seed=seed)
+ ),
total=len(train_data) // batch_size,
)
):
batch = batch[:, :-1]
- targets = model(batch)
+ targets = target_fn(batch, it, split="train")
mx.eval(targets)
loss, ntoks, params = step(batch, targets, lengths, params)
mx.eval(loss, params)
@@ -155,7 +206,7 @@ def dwq_quantize(
" Model quality will likely be degraded.\nāāā"
)
- q_model.update(tree_map(lambda x: x.astype(dtype), params))
+ model.update(tree_map(lambda x: x.astype(dtype), params))
def load_data(
@@ -196,10 +247,12 @@ def main():
help="A model to distill from for DWQ. If `quantized-model` is not"
" given the student model will be this model quantized according"
" to `bits` and `group-size`.",
+ type=str,
required=True,
)
parser.add_argument(
"--quantized-model",
+ type=str,
default=None,
help="An already quantized model (the student model) to improve with DWQ.",
)
@@ -236,27 +289,78 @@ def main():
action="store_true",
help="Use gradient checkpointing to reduce memory use.",
)
+ parser.add_argument(
+ "--target-dir", type=str, default=None, help="Directory to save/load targets."
+ )
+ parser.add_argument(
+ "--targets-only", action="store_true", help="Compute the targets and exit."
+ )
+ parser.add_argument(
+ "--pipeline",
+ action="store_true",
+ help="Use pipeline parallel instead of data parallel.",
+ )
+
args = parser.parse_args()
group = mx.distributed.init()
num_samples = args.num_samples
- if num_samples % group.size() > 0:
+ if not args.pipeline and num_samples % group.size() > 0:
num_samples += group.size() - num_samples % group.size()
np.random.seed(args.seed)
mx.random.seed(args.seed)
- model, tokenizer, config = load(
- args.model,
- lazy=True,
- return_config=True,
- )
+ if args.target_dir is not None:
+ target_dir = Path(args.target_dir)
+ has_targets = target_dir.exists()
+ else:
+ has_targets = False
+ target_dir = None
+
+ tokenizer = load_tokenizer(args.model)
train_data, valid_data = load_data(
tokenizer, args.data_path, args.num_samples, args.max_seq_length
)
+ # Load the base model if we need it
+ if not has_targets or args.quantized_model is None:
+ if args.pipeline and group.size() > 1:
+ model, _, config = pipeline_load(args.model, return_config=True)
+ else:
+ model, _, config = load(args.model, return_config=True, lazy=True)
+ else:
+ model = None
+
+ # Pre-compute the targets
+ if not has_targets and target_dir is not None:
+ compute_dwq_targets(
+ model,
+ target_dir,
+ train_data,
+ valid_data,
+ batch_size=args.batch_size,
+ max_seq_length=args.max_seq_length,
+ seed=args.seed,
+ )
+ has_targets = True
+
+ if args.targets_only:
+ exit(0)
+
+ if has_targets:
+
+ def target_fn(_, idx, split):
+ targets = mx.load(target_dir / split / f"{idx:010d}.safetensors")
+ return targets["logits"], targets["indices"]
+
+ else:
+
+ def target_fn(batch, idx, split):
+ return model(batch)
+
if args.quantized_model is not None:
q_model, tokenizer, config = load(
args.quantized_model,
@@ -274,19 +378,24 @@ def main():
bits=args.bits,
)
+ # Delete the base model if it's not needed
+ if has_targets and model is not None:
+ del model
+
if mx.metal.is_available():
max_rec_size = mx.metal.device_info()["max_recommended_working_set_size"]
mx.set_wired_limit(max_rec_size)
opt = optimizers.Adam(learning_rate=args.learning_rate, bias_correction=True)
dwq_quantize(
- model,
q_model,
+ target_fn,
opt,
train_data,
valid_data,
batch_size=args.batch_size,
max_seq_length=args.max_seq_length,
+ seed=args.seed,
gradient_checkpoint=args.grad_checkpoint,
)
save(
diff --git a/mlx_lm/tokenizer_utils.py b/mlx_lm/tokenizer_utils.py
index c612c4b..af191b2 100644
--- a/mlx_lm/tokenizer_utils.py
+++ b/mlx_lm/tokenizer_utils.py
@@ -423,7 +423,7 @@ def _is_bpe_decoder(decoder):
return isinstance(decoder, dict) and decoder.get("type", None) == "ByteLevel"
-def load_tokenizer(
+def load(
model_path,
tokenizer_config_extra: Optional[Dict[str, Any]] = None,
return_tokenizer=True,
@@ -438,6 +438,7 @@ def load_tokenizer(
detokenizer_class = NaiveStreamingDetokenizer
tokenizer_file = model_path / "tokenizer.json"
+
if tokenizer_file.exists():
with open(tokenizer_file, "r", encoding="utf-8") as fid:
try:
diff --git a/mlx_lm/tuner/trainer.py b/mlx_lm/tuner/trainer.py
index a56301f..7b18c7f 100644
--- a/mlx_lm/tuner/trainer.py
+++ b/mlx_lm/tuner/trainer.py
@@ -92,7 +92,9 @@ def iterate_batches(
dataset,
batch_size,
max_seq_length,
- train=False,
+ loop=False,
+ seed=None,
+ comm_group=None,
):
# Sort by length:
if isinstance(dataset, CacheDataset):
@@ -108,8 +110,12 @@ def iterate_batches(
# If running in distributed mode (N machines) then each one should skip N-1
# samples
- offset = mx.distributed.init().rank()
- step = mx.distributed.init().size()
+ if comm_group is not None:
+ offset = comm_group.rank()
+ step = comm_group.size()
+ else:
+ offset = 0
+ step = 1
if batch_size % step != 0:
raise ValueError("The batch size must be divisible by the number of workers")
@@ -118,7 +124,8 @@ def iterate_batches(
idx[i + offset : i + offset + batch_size : step]
for i in range(0, len(idx) - batch_size + 1, batch_size)
]
-
+ if seed:
+ np.random.seed(seed)
while True:
indices = np.random.permutation(len(batch_idx))
for i in indices:
@@ -151,7 +158,7 @@ def iterate_batches(
batch = mx.array(batch_arr)
yield batch, mx.array(list(zip(offsets, lengths)))
- if not train:
+ if not loop:
break
@@ -177,6 +184,7 @@ def evaluate(
dataset=dataset,
batch_size=batch_size,
max_seq_length=max_seq_length,
+ comm_group=mx.distributed.init(),
),
),
desc="Calculating loss...",
@@ -254,7 +262,8 @@ def train(
dataset=train_dataset,
batch_size=args.batch_size,
max_seq_length=args.max_seq_length,
- train=True,
+ loop=True,
+ comm_group=world,
),
):
tic = time.perf_counter()
diff --git a/mlx_lm/utils.py b/mlx_lm/utils.py
index cac80cb..7e4a641 100644
--- a/mlx_lm/utils.py
+++ b/mlx_lm/utils.py
@@ -7,6 +7,7 @@ import inspect
import json
import logging
import os
+import resource
import shutil
from pathlib import Path
from textwrap import dedent
@@ -14,6 +15,7 @@ from typing import (
Any,
Callable,
Dict,
+ List,
Optional,
Tuple,
Type,
@@ -31,11 +33,14 @@ if os.getenv("MLXLM_USE_MODELSCOPE", "False").lower() == "true":
else:
from huggingface_hub import snapshot_download
+# For large models with lots of files
+resource.setrlimit(resource.RLIMIT_NOFILE, (2048, 4096))
+
from mlx.utils import tree_flatten, tree_map, tree_reduce, tree_unflatten
-from transformers import PreTrainedTokenizer
# Local imports
-from .tokenizer_utils import TokenizerWrapper, load_tokenizer
+from .tokenizer_utils import TokenizerWrapper
+from .tokenizer_utils import load as _load_tokenizer
# Constants
MODEL_REMAPPING = {
@@ -94,7 +99,11 @@ def compute_bits_per_weight(model):
return model_bytes * 8 / model_params
-def _download(path_or_hf_repo: str, revision: Optional[str] = None) -> Path:
+def _download(
+ path_or_hf_repo: str,
+ revision: Optional[str] = None,
+ allow_patterns: List[str] = None,
+) -> Path:
"""
Ensures the model is available locally. If the path does not exist locally,
it is downloaded from the Hugging Face Hub.
@@ -109,21 +118,22 @@ def _download(path_or_hf_repo: str, revision: Optional[str] = None) -> Path:
model_path = Path(path_or_hf_repo)
if not model_path.exists():
+ allow_patterns = allow_patterns or [
+ "*.json",
+ "model*.safetensors",
+ "*.py",
+ "tokenizer.model",
+ "*.tiktoken",
+ "tiktoken.model",
+ "*.txt",
+ "*.jsonl",
+ "*.jinja",
+ ]
model_path = Path(
snapshot_download(
path_or_hf_repo,
revision=revision,
- allow_patterns=[
- "*.json",
- "model*.safetensors",
- "*.py",
- "tokenizer.model",
- "*.tiktoken",
- "tiktoken.model",
- "*.txt",
- "*.jsonl",
- "*.jinja",
- ],
+ allow_patterns=allow_patterns,
)
)
@@ -244,6 +254,28 @@ def load_adapters(model: nn.Module, adapter_path: str) -> nn.Module:
return _load_adapters(model, adapter_path)
+def load_tokenizer(model_path, tokenizer_config_extra=None, eos_token_ids=None):
+ """Load a huggingface tokenizer and try to infer the type of streaming
+ detokenizer to use.
+ """
+ model_path = _download(
+ model_path,
+ allow_patterns=[
+ "*.json",
+ "*.py",
+ "tokenizer.model",
+ "*.tiktoken",
+ "tiktoken.model",
+ "*.txt",
+ "*.jsonl",
+ "*.jinja",
+ ],
+ )
+ return _load_tokenizer(
+ model_path, tokenizer_config_extra, eos_token_ids=eos_token_ids
+ )
+
+
def load(
path_or_hf_repo: str,
tokenizer_config: Optional[Dict[str, Any]] = None,
@@ -296,6 +328,62 @@ def load(
return model, tokenizer
+def pipeline_load(repo, return_config=False):
+ # Get model path with everything but weight safetensors
+ model_path = _download(
+ repo,
+ allow_patterns=[
+ "*.json",
+ "*.py",
+ "tokenizer.model",
+ "*.tiktoken",
+ "tiktoken.model",
+ "*.txt",
+ "*.jsonl",
+ "*.jinja",
+ ],
+ )
+
+ # Lazy load and shard model to figure out which weights we need
+ model, config = load_model(model_path, lazy=True, strict=False)
+
+ group = mx.distributed.init()
+ rank = group.rank()
+ model.model.pipeline(group)
+
+ # Figure out which files we need for the local shard
+ with open(model_path / "model.safetensors.index.json", "r") as fid:
+ weight_index = json.load(fid)["weight_map"]
+
+ local_files = set()
+ for k, _ in tree_flatten(model.parameters()):
+ if file_name := weight_index.get(k, None) is None:
+ raise ValueError(
+ "Pipeline loading is only supported for MLX converted models."
+ )
+ local_files.add(weight_index[k])
+
+ # Download weights for local shard
+ _download(repo, allow_patterns=local_files)
+
+ # Load and shard the model, and load the weights
+ tokenizer = load_tokenizer(
+ model_path,
+ {"trust_remote_code": True},
+ eos_token_ids=config.get("eos_token_id", None),
+ )
+ model, _ = load_model(model_path, lazy=True, strict=False)
+ model.model.pipeline(group)
+ mx.eval(model.parameters())
+
+ # Synchronize processes to avoid timeout
+ mx.eval(mx.distributed.all_sum(mx.array(1.0), stream=mx.cpu))
+ if return_config:
+ return model, tokenizer, config
+ else:
+ return model, tokenizer
+
+
def make_shards(weights: dict, max_file_size_gb: int = MAX_FILE_SIZE_GB) -> list:
"""
Splits the weights into smaller shards.
diff --git a/tests/test_finetune.py b/tests/test_finetune.py
index d54af90..3635ae3 100644
--- a/tests/test_finetune.py
+++ b/tests/test_finetune.py
@@ -5,7 +5,7 @@ import sys
import unittest
from contextlib import contextmanager
from io import StringIO
-from unittest.mock import MagicMock
+from unittest.mock import ANY, MagicMock
import mlx.core as mx
import mlx.nn as nn
@@ -405,6 +405,7 @@ class TestScheduleConfig(unittest.TestCase):
dataset=mock_dataset,
batch_size=2,
max_seq_length=2048,
+ comm_group=ANY,
)
self.assertEqual(mock_default_loss.call_count, 2)
@@ -441,6 +442,7 @@ class TestScheduleConfig(unittest.TestCase):
dataset=mock_dataset,
batch_size=2,
max_seq_length=2048,
+ comm_group=ANY,
)
self.assertEqual(mock_default_loss.call_count, 3)
diff --git a/tests/test_tokenizers.py b/tests/test_tokenizers.py
index 5cb12b1..c68eefb 100644
--- a/tests/test_tokenizers.py
+++ b/tests/test_tokenizers.py
@@ -9,27 +9,12 @@ from mlx_lm.tokenizer_utils import (
BPEStreamingDetokenizer,
NaiveStreamingDetokenizer,
SPMStreamingDetokenizer,
- load_tokenizer,
)
+from mlx_lm.utils import load_tokenizer
class TestTokenizers(unittest.TestCase):
- def download_tokenizer(self, repo):
- path = Path(
- snapshot_download(
- repo_id=repo,
- allow_patterns=[
- "tokenizer.json",
- "tokenizer_config.json",
- "special_tokens_map.json",
- "tokenizer.model",
- "chat_template.jinja",
- ],
- )
- )
- return load_tokenizer(path)
-
def check_tokenizer(self, tokenizer):
def check(tokens):
expected_text = tokenizer.decode(tokens)
@@ -77,19 +62,19 @@ class TestTokenizers(unittest.TestCase):
]
for tokenizer_repo, expected_detokenizer in tokenizer_repos:
with self.subTest(tokenizer=tokenizer_repo):
- tokenizer = self.download_tokenizer(tokenizer_repo)
+ tokenizer = load_tokenizer(tokenizer_repo)
tokenizer.decode([0, 1, 2])
self.assertTrue(isinstance(tokenizer.detokenizer, expected_detokenizer))
self.check_tokenizer(tokenizer)
# Try one with a naive detokenizer
- tokenizer = self.download_tokenizer("mlx-community/Llama-3.2-1B-Instruct-4bit")
+ tokenizer = load_tokenizer("mlx-community/Llama-3.2-1B-Instruct-4bit")
tokenizer._detokenizer = NaiveStreamingDetokenizer(tokenizer)
self.check_tokenizer(tokenizer)
def test_special_tokens(self):
tokenizer_repo = "mlx-community/DeepSeek-Coder-V2-Lite-Instruct-4bit-mlx"
- tokenizer = self.download_tokenizer(tokenizer_repo)
+ tokenizer = load_tokenizer(tokenizer_repo)
detokenizer = tokenizer.detokenizer
detokenizer.reset()
@@ -100,18 +85,18 @@ class TestTokenizers(unittest.TestCase):
def test_tool_calling(self):
tokenizer_repo = "mlx-community/Qwen3-4B-4bit"
- tokenizer = self.download_tokenizer(tokenizer_repo)
+ tokenizer = load_tokenizer(tokenizer_repo)
self.assertTrue(tokenizer.has_tool_calling)
self.assertEqual(tokenizer.tool_call_start, "")
self.assertEqual(tokenizer.tool_call_end, "")
tokenizer_repo = "mlx-community/Llama-3.2-1B-Instruct-4bit"
- tokenizer = self.download_tokenizer(tokenizer_repo)
+ tokenizer = load_tokenizer(tokenizer_repo)
self.assertFalse(tokenizer.has_tool_calling)
def test_thinking(self):
tokenizer_repo = "mlx-community/Qwen3-4B-4bit"
- tokenizer = self.download_tokenizer(tokenizer_repo)
+ tokenizer = load_tokenizer(tokenizer_repo)
self.assertTrue(tokenizer.has_thinking)
self.assertEqual(tokenizer.think_start, "")
self.assertEqual(tokenizer.think_end, "")
diff --git a/tests/test_tuner_trainer.py b/tests/test_tuner_trainer.py
index be1c866..09fe570 100644
--- a/tests/test_tuner_trainer.py
+++ b/tests/test_tuner_trainer.py
@@ -19,47 +19,35 @@ class MockDistributedGroup:
return self._size
-class MockDistributed:
- def __init__(self):
- self.rank = 0
- self.size = 1
-
- def init(self):
- return MockDistributedGroup(self.rank, self.size)
-
-
class TestTunerTrainer(unittest.TestCase):
def test_iterate_batches_ddp(self):
- olddist = mx.distributed
- try:
- mx.distributed = MockDistributed()
+ group = MockDistributedGroup(0, 1)
- def run(rank, size, batch):
- mx.distributed.rank = rank
- mx.distributed.size = size
+ def run(rank, size, batch):
+ group._rank = rank
+ group._size = size
- data = mx.arange(128).reshape(-1, 1).tolist()
- data = [(d, 0) for d in data]
+ data = mx.arange(128).reshape(-1, 1).tolist()
+ data = [(d, 0) for d in data]
- samples = set()
- for i, (b, l) in enumerate(iterate_batches(data, batch, 1)):
- samples.add(tuple(mx.flatten(b).tolist()))
+ samples = set()
+ for i, (b, l) in enumerate(
+ iterate_batches(data, batch, 1, comm_group=group)
+ ):
+ samples.add(tuple(mx.flatten(b).tolist()))
- ref_batches = mx.arange(128).reshape(-1, batch).tolist()
- for b in ref_batches:
- self.assertTrue(tuple(b[rank::size]) in samples)
+ ref_batches = mx.arange(128).reshape(-1, batch).tolist()
+ for b in ref_batches:
+ self.assertTrue(tuple(b[rank::size]) in samples)
- run(0, 1, 4)
- run(0, 1, 8)
- run(0, 2, 8)
- run(1, 2, 8)
- run(0, 4, 8)
- run(1, 4, 8)
- run(2, 4, 8)
- run(3, 4, 8)
-
- finally:
- mx.distributed = olddist
+ run(0, 1, 4)
+ run(0, 1, 8)
+ run(0, 2, 8)
+ run(1, 2, 8)
+ run(0, 4, 8)
+ run(1, 4, 8)
+ run(2, 4, 8)
+ run(3, 4, 8)
if __name__ == "__main__":