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

DWQ for very large models (#536)

Browse Source 
* pipeline parallel mixin

* Refactor pipeline parallel, add optional target saving to DWQ

* preserve batch order

* Fixes

* fix glm4 pipeline

* event timeout hack

* use full targets for regular training
This commit is contained in:
Awni Hannun
2025-11-07 06:43:40 -08:00
committed by GitHub
parent 3833c205c1
commit 6c1a459314
12 changed files with 351 additions and 226 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mlx_lm/examples/pipeline_generate.py
+3 -63
View File
@@ -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)
mlx_lm/models/deepseek_v2.py
+6 -29
View File
@@ -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
mlx_lm/models/deepseek_v3.py
+6 -27
View File
@@ -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):
mlx_lm/models/glm4_moe.py
+22 -6
View File
@@ -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):
mlx_lm/models/pipeline.py
+31
View File
@@ -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
mlx_lm/quant/dwq.py
+132 -23
View File
@@ -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(
mlx_lm/tokenizer_utils.py
+2 -1
View File
@@ -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:
mlx_lm/tuner/trainer.py
+15 -6
View File
@@ -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()
mlx_lm/utils.py
+102 -14
View File
@@ -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.
tests/test_finetune.py
+3 -1
View File
@@ -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)
tests/test_tokenizers.py
+7 -22
View File
@@ -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, "<tool_call>")
self.assertEqual(tokenizer.tool_call_end, "</tool_call>")
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, "<think>")
self.assertEqual(tokenizer.think_end, "</think>")
tests/test_tuner_trainer.py
+22 -34
View File
@@ -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__":