mirror of
https://github.com/unslothai/unsloth.git
synced 2026-07-09 15:58:41 +00:00
feat(cli): support MLX distributed inference (#6845)
* feat(cli): detect MLX distributed launch context * feat(mlx): wire distributed inference backend * feat(cli): broadcast MLX distributed chat turns * fix(cli): wait indefinitely for distributed chat turns * fix(cli): report MLX distributed load errors cleanly * fix(mlx): route distributed vlm through loader * fix(cli): detect inline MLX host JSON * fix(studio): harden distributed object sharing * fix(studio): select JACCL distributed backend * fix(cli): abort distributed error paths * Distinguish real stream errors from model text via GenStreamError in distributed CLI * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fail loud when MLX distributed init returns a singleton group The worker only reaches this block when distributed was explicitly requested. A singleton (size 1) group means the launch failed to form a real group (MLX built without distributed support, or an invalid launch env/hostfile); silently continuing leaves nonzero ranks looping forever on share_distributed_object. Raise instead so the surrounding handler returns a clear load error. * Tighten MLX distributed inference comments --------- Co-authored-by: Daniel Han <danielhanchen@gmail.com> Co-authored-by: danielhanchen <unslothai@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Long Yixing
GitHub
parent
38dacb8a1f
commit
2a6abe2ff5
9 changed files with 1188 additions and 155 deletions
|
|
@ -5,6 +5,7 @@ Drop-in replacement for InferenceBackend — same interface, uses mlx-lm/mlx-vlm
|
|||
instead of torch/transformers for model loading and generation.
|
||||
"""
|
||||
|
||||
import os
|
||||
import threading
|
||||
from typing import Optional, Generator
|
||||
from core.inference.runtime_context import runtime_context_length
|
||||
|
|
@ -41,6 +42,48 @@ def _build_generation_stats(prompt_n, prompt_tps, gen_n, gen_tps):
|
|||
}
|
||||
|
||||
|
||||
def _mlx_distributed_rank_size(group = None):
|
||||
"""Return ``(rank, world_size)`` for an optional MLX distributed group."""
|
||||
if group is None:
|
||||
return 0, 1
|
||||
rank = int(group.rank())
|
||||
world_size = int(group.size())
|
||||
if world_size < 1:
|
||||
raise ValueError(f"Invalid MLX distributed world_size={world_size}.")
|
||||
if rank < 0 or rank >= world_size:
|
||||
raise ValueError(f"Invalid MLX distributed rank={rank} for world_size={world_size}.")
|
||||
return rank, world_size
|
||||
|
||||
|
||||
def _mlx_distributed_backend_from_env():
|
||||
if os.environ.get("MLX_JACCL_COORDINATOR") and os.environ.get("MLX_IBV_DEVICES"):
|
||||
return "jaccl"
|
||||
return None
|
||||
|
||||
|
||||
def _init_mlx_distributed():
|
||||
"""Initialize MLX distributed state, falling back to singleton metadata."""
|
||||
import mlx.core as mx
|
||||
|
||||
group = None
|
||||
rank = 0
|
||||
world_size = 1
|
||||
distributed = getattr(mx, "distributed", None)
|
||||
init = getattr(distributed, "init", None) if distributed is not None else None
|
||||
if callable(init):
|
||||
backend = _mlx_distributed_backend_from_env()
|
||||
if backend is None:
|
||||
group = init()
|
||||
else:
|
||||
try:
|
||||
group = init(backend = backend)
|
||||
except TypeError:
|
||||
group = init()
|
||||
if group is not None:
|
||||
rank, world_size = _mlx_distributed_rank_size(group)
|
||||
return group, rank, world_size
|
||||
|
||||
|
||||
def _make_mlx_presence_penalty_processor(penalty: float):
|
||||
"""Presence penalty as an mlx_lm/mlx_vlm logits processor, matching the safetensors path.
|
||||
|
||||
|
|
@ -52,7 +95,7 @@ def _make_mlx_presence_penalty_processor(penalty: float):
|
|||
|
||||
def _processor(tokens, logits):
|
||||
if state["prompt_len"] is None:
|
||||
# First call = prompt only; latch its length.
|
||||
# First call is prompt-only; latch its length.
|
||||
state["prompt_len"] = int(tokens.shape[0])
|
||||
return logits
|
||||
generated = tokens[state["prompt_len"] :]
|
||||
|
|
@ -61,22 +104,17 @@ def _make_mlx_presence_penalty_processor(penalty: float):
|
|||
import mlx.core as mx
|
||||
|
||||
vocab = logits.shape[-1]
|
||||
# Bound generated ids to the valid range [0, vocab) before they index
|
||||
# logits. MLX does no bounds checking and out-of-bounds indexing is
|
||||
# documented undefined behavior (crash / memory corruption), unlike the
|
||||
# torch path's harmless negative wrap -- so this bound is load-bearing
|
||||
# here and matches the torch filter seen[(seen >= 0) & (seen < vocab)].
|
||||
# MLX has no boolean-mask filtering (data-dependent output shape is
|
||||
# unsupported), so instead of compacting the id list we route every
|
||||
# out-of-range or negative id to a scratch slot at index ``vocab`` that
|
||||
# is dropped before the subtract. That scratch slot can never collide
|
||||
# with a real token, so real ids (including id 0) are penalized exactly
|
||||
# once and stray ids are ignored.
|
||||
# Bound ids to [0, vocab) before indexing logits: MLX does no bounds
|
||||
# checking and out-of-bounds indexing is undefined behavior (crash /
|
||||
# corruption), unlike torch's harmless negative wrap. MLX also lacks
|
||||
# boolean-mask filtering, so out-of-range/negative ids route to a
|
||||
# scratch slot at index vocab (dropped before the subtract) that never
|
||||
# collides with a real token: real ids (including 0) are penalized
|
||||
# once, strays ignored.
|
||||
valid = (generated >= 0) & (generated < vocab)
|
||||
safe = mx.where(valid, generated, vocab).astype(mx.int32)
|
||||
# Scatter-assign a scalar penalty into a (vocab + 1)-wide mask: duplicate
|
||||
# ids are idempotent, so presence applies once per distinct token; the
|
||||
# scratch column is discarded and the full-width subtract stays on-device.
|
||||
# Scatter penalty into a (vocab + 1)-wide mask: duplicate ids are
|
||||
# idempotent (presence applies once per token); scratch column dropped.
|
||||
mask = mx.zeros((vocab + 1,), dtype = logits.dtype)
|
||||
mask[safe] = penalty
|
||||
logits = logits - mask[:vocab]
|
||||
|
|
@ -93,7 +131,7 @@ class MLXInferenceBackend:
|
|||
self.loaded_local_models = []
|
||||
self.device = "mlx"
|
||||
self._generation_lock = threading.Lock()
|
||||
# usage/timings of the latest generation; shipped on gen_done.
|
||||
# usage/timings of the latest generation, shipped on gen_done.
|
||||
self.last_generation_stats = None
|
||||
|
||||
self._model = None
|
||||
|
|
@ -101,6 +139,9 @@ class MLXInferenceBackend:
|
|||
self._processor = None
|
||||
self._is_vlm = False
|
||||
self._config = {}
|
||||
self._distributed_group = None
|
||||
self._distributed_rank = 0
|
||||
self._distributed_world_size = 1
|
||||
|
||||
# Recorded for unload to release pinned memory back to the OS.
|
||||
self._memory_limits_applied = {}
|
||||
|
|
@ -145,19 +186,26 @@ class MLXInferenceBackend:
|
|||
trust_remote_code = False,
|
||||
gpu_ids = None,
|
||||
dtype = None,
|
||||
parallel_mode = None,
|
||||
distributed_group = None,
|
||||
) -> bool:
|
||||
import mlx.core as mx
|
||||
|
||||
# Keep the token so the native-template fallback can fetch a
|
||||
# gated model's repo template later during generation.
|
||||
# Keep the token so the native-template fallback can fetch a gated
|
||||
# model's repo template during generation.
|
||||
self._hf_token = hf_token
|
||||
model_name = config.identifier if hasattr(config, "identifier") else str(config)
|
||||
is_vision = getattr(config, "is_vision", False)
|
||||
distributed_rank, distributed_size = _mlx_distributed_rank_size(distributed_group)
|
||||
is_distributed = distributed_group is not None and distributed_size > 1
|
||||
self._distributed_group = distributed_group
|
||||
self._distributed_rank = distributed_rank
|
||||
self._distributed_world_size = distributed_size
|
||||
|
||||
# GGUF guard. GGUF models are served by llama-server in the parent
|
||||
# process, not mlx-lm here. Reaching this with is_gguf=True means the
|
||||
# route's first detection flaked (transient HF Hub) but the subprocess
|
||||
# re-detected GGUF; raise loudly instead of a cryptic mlx_lm error.
|
||||
# GGUF guard: GGUF is served by llama-server in the parent process,
|
||||
# not mlx-lm. Reaching here with is_gguf=True means the route's
|
||||
# detection flaked but the subprocess re-detected GGUF; raise loudly
|
||||
# instead of a cryptic mlx_lm error.
|
||||
if getattr(config, "is_gguf", False):
|
||||
raise RuntimeError(
|
||||
f"MLXInferenceBackend cannot load GGUF model '{model_name}': "
|
||||
|
|
@ -176,11 +224,26 @@ class MLXInferenceBackend:
|
|||
is_lora = getattr(config, "is_lora", False)
|
||||
|
||||
logger.info(
|
||||
"Loading %s via %s (is_lora=%s)",
|
||||
"Loading %s via %s (is_lora=%s, distributed=%s, rank=%s/%s, mode=%s)",
|
||||
model_name,
|
||||
"mlx-vlm" if is_vision else "mlx-lm",
|
||||
is_lora,
|
||||
is_distributed,
|
||||
distributed_rank,
|
||||
distributed_size,
|
||||
parallel_mode,
|
||||
)
|
||||
if is_distributed and parallel_mode not in ("pipeline", "tensor"):
|
||||
raise ValueError(
|
||||
"Unsloth: distributed MLX inference requires parallel_mode='pipeline' "
|
||||
"or parallel_mode='tensor'."
|
||||
)
|
||||
if is_distributed and is_lora:
|
||||
raise ValueError(
|
||||
"Unsloth: distributed MLX inference for LoRA adapter repos "
|
||||
"is not supported yet. Merge/export the adapter into an MLX model "
|
||||
"before distributed inference."
|
||||
)
|
||||
|
||||
try:
|
||||
from unsloth_zoo.mlx.loader import FastMLXModel
|
||||
|
|
@ -190,14 +253,23 @@ class MLXInferenceBackend:
|
|||
"(unsloth_zoo.mlx.loader). Reinstall via install.sh on Apple Silicon."
|
||||
) from e
|
||||
|
||||
load_kwargs = {
|
||||
"max_seq_length": max_seq_length,
|
||||
"dtype": dtype,
|
||||
"load_in_4bit": load_in_4bit,
|
||||
"token": hf_token,
|
||||
"trust_remote_code": trust_remote_code,
|
||||
"text_only": False if is_vision else True,
|
||||
}
|
||||
if is_distributed:
|
||||
if parallel_mode == "pipeline":
|
||||
load_kwargs["pipeline_group"] = distributed_group
|
||||
else:
|
||||
load_kwargs["tensor_group"] = distributed_group
|
||||
|
||||
model, tokenizer_or_processor = FastMLXModel.from_pretrained(
|
||||
model_name,
|
||||
max_seq_length = max_seq_length,
|
||||
dtype = dtype,
|
||||
load_in_4bit = load_in_4bit,
|
||||
token = hf_token,
|
||||
trust_remote_code = trust_remote_code,
|
||||
text_only = False if is_vision else True,
|
||||
**load_kwargs,
|
||||
)
|
||||
|
||||
if is_vision:
|
||||
|
|
@ -217,8 +289,7 @@ class MLXInferenceBackend:
|
|||
self.models[model_name] = {
|
||||
# Per-model token for the native-template fallback (matches transformers).
|
||||
"hf_token": hf_token,
|
||||
# Per-model consent for the native-template reload: re-use the exact
|
||||
# trust_remote_code this model was loaded with (matches transformers).
|
||||
# Per-model trust_remote_code reused by the native-template reload (matches transformers).
|
||||
"trust_remote_code": trust_remote_code,
|
||||
"model": self._model,
|
||||
"tokenizer": self._tokenizer,
|
||||
|
|
@ -234,8 +305,7 @@ class MLXInferenceBackend:
|
|||
"has_audio_input": False,
|
||||
"context_length": runtime_context_length(self._model, max_seq_length),
|
||||
}
|
||||
# Capture chat_template_info so the worker IPC reply ships it back and
|
||||
# the route layer classifies capabilities like the other paths.
|
||||
# Capture chat_template_info for the worker IPC reply and route capability classification.
|
||||
self._populate_chat_template_info(model_name)
|
||||
|
||||
logger.info("Model %s loaded successfully", model_name)
|
||||
|
|
@ -293,6 +363,9 @@ class MLXInferenceBackend:
|
|||
self._model = None
|
||||
self._tokenizer = None
|
||||
self._processor = None
|
||||
self._distributed_group = None
|
||||
self._distributed_rank = 0
|
||||
self._distributed_world_size = 1
|
||||
if self.active_model_name == model_name:
|
||||
self.active_model_name = None
|
||||
gc.collect()
|
||||
|
|
@ -320,8 +393,7 @@ class MLXInferenceBackend:
|
|||
max_new_tokens = 256,
|
||||
repetition_penalty = 1.0,
|
||||
cancel_event = None,
|
||||
# Reasoning / tool kwargs forwarded by the route + worker; rendered via
|
||||
# apply_chat_template_for_generation like the transformers path.
|
||||
# Reasoning / tool kwargs, rendered via apply_chat_template_for_generation (transformers parity).
|
||||
tools = None,
|
||||
enable_thinking = None,
|
||||
reasoning_effort = None,
|
||||
|
|
@ -334,7 +406,6 @@ class MLXInferenceBackend:
|
|||
# Reset so a failed run cannot surface stale stats.
|
||||
self.last_generation_stats = None
|
||||
|
||||
# Build messages with system prompt
|
||||
full_messages = []
|
||||
if system_prompt:
|
||||
full_messages.append({"role": "system", "content": system_prompt})
|
||||
|
|
@ -351,7 +422,6 @@ class MLXInferenceBackend:
|
|||
{"type": "text", "text": content},
|
||||
]
|
||||
elif isinstance(content, list):
|
||||
# Prepend image if not already present
|
||||
has_image = any(
|
||||
p.get("type") == "image" for p in content if isinstance(p, dict)
|
||||
)
|
||||
|
|
@ -429,11 +499,11 @@ class MLXInferenceBackend:
|
|||
if prompt is None:
|
||||
raise RuntimeError("apply_chat_template returned None — tokenizer may be incompatible")
|
||||
|
||||
# Same parity fix as the transformers backend: if the template dropped the
|
||||
# requested tools, fall back to the native template so MLX text models keep
|
||||
# advertising them. ``self._tokenizer`` is this entry's model_info tokenizer,
|
||||
# so probe and native render share a renderer. (The VLM path renders via the
|
||||
# processor for image tokens and is intentionally not wired here.)
|
||||
# Parity with the transformers backend: if the template dropped the
|
||||
# requested tools, fall back to the native template so MLX text models
|
||||
# keep advertising them. self._tokenizer is this entry's tokenizer, so
|
||||
# probe and native render share a renderer. (VLM renders via the
|
||||
# processor for image tokens and is not wired here.)
|
||||
model_info = self.models.get(self.active_model_name, {})
|
||||
prompt = render_with_native_template_fallback(
|
||||
formatted_prompt = prompt,
|
||||
|
|
@ -455,7 +525,7 @@ class MLXInferenceBackend:
|
|||
min_p = float(min_p or 0.0),
|
||||
min_tokens_to_keep = 1,
|
||||
)
|
||||
# Repetition and/or presence penalty processors (parity with the GGUF/safetensors paths).
|
||||
# Repetition and/or presence penalty processors (GGUF/safetensors parity).
|
||||
logits_processors = []
|
||||
if repetition_penalty is not None and float(repetition_penalty) not in (
|
||||
0.0,
|
||||
|
|
@ -496,7 +566,6 @@ class MLXInferenceBackend:
|
|||
):
|
||||
final_response = response
|
||||
token_ids.append(response.token)
|
||||
# Decode full sequence with skip_special_tokens
|
||||
cumulative = self._tokenizer.decode(
|
||||
token_ids,
|
||||
skip_special_tokens = True,
|
||||
|
|
@ -544,8 +613,7 @@ class MLXInferenceBackend:
|
|||
)
|
||||
|
||||
# Pick the chat-template-aware caller: processors with their own
|
||||
# apply_chat_template + chat_template (e.g. Qwen2.5-VL) use it
|
||||
# directly; else fall back to the nested tokenizer.
|
||||
# apply_chat_template + chat_template (e.g. Qwen2.5-VL), else the nested tokenizer.
|
||||
chat_target = self._processor
|
||||
if (
|
||||
getattr(self._processor, "apply_chat_template", None) is None
|
||||
|
|
@ -572,10 +640,9 @@ class MLXInferenceBackend:
|
|||
len(prompt),
|
||||
image is not None,
|
||||
)
|
||||
# mlx_vlm.stream_generate forwards **kwargs into generate_step, which
|
||||
# builds the sampler + logits_processors internally.
|
||||
# GOTCHA: generate_step expects ``temperature=`` (long form); ``temp=``
|
||||
# silently falls into **kwargs and is ignored, stuck at greedy 0.0.
|
||||
# stream_generate forwards **kwargs into generate_step (builds the
|
||||
# sampler + logits_processors internally). GOTCHA: generate_step expects
|
||||
# temperature= (long form); temp= is silently ignored, stuck at greedy 0.0.
|
||||
vlm_kwargs = dict(
|
||||
max_tokens = max_new_tokens,
|
||||
temperature = temperature,
|
||||
|
|
@ -589,7 +656,7 @@ class MLXInferenceBackend:
|
|||
)
|
||||
if presence_penalty:
|
||||
# Presence needs a custom processor: pass the full list (repetition +
|
||||
# presence) instead of the repetition_penalty shortcut so both apply once.
|
||||
# presence) instead of the repetition_penalty shortcut so both apply.
|
||||
from mlx_lm.sample_utils import make_logits_processors
|
||||
|
||||
_vlm_processors = []
|
||||
|
|
@ -634,7 +701,7 @@ class MLXInferenceBackend:
|
|||
cancel_event = None,
|
||||
**gen_kwargs,
|
||||
) -> Generator[str, None, None]:
|
||||
# MLX LoRA adapter toggling not yet supported — generate normally
|
||||
# MLX LoRA adapter toggling not yet supported; generate normally
|
||||
yield from self.generate_chat_response(cancel_event = cancel_event, **gen_kwargs)
|
||||
|
||||
def reset_generation_state(self):
|
||||
|
|
|
|||
|
|
@ -50,6 +50,18 @@ _DISPATCH_DRAIN_TIMEOUT = 5.0
|
|||
_UNLOAD_GEN_LOCK_TIMEOUT = 15.0
|
||||
|
||||
|
||||
class GenStreamError(str):
|
||||
"""A stream chunk carrying a real backend/generation error, not model text.
|
||||
|
||||
Subclasses str so existing display/logging consumers are unaffected, while
|
||||
callers that must abort a distributed run on error (raise_on_streamed_error)
|
||||
can distinguish a real error from model output whose visible text starts with
|
||||
"Error:" by checking isinstance(chunk, GenStreamError).
|
||||
"""
|
||||
|
||||
__slots__ = ()
|
||||
|
||||
|
||||
class InferenceOrchestrator:
|
||||
"""
|
||||
Inference backend orchestrator — subprocess-based.
|
||||
|
|
@ -482,13 +494,13 @@ class InferenceOrchestrator:
|
|||
initial_resp_queue = self._resp_queue
|
||||
while True:
|
||||
if self._proc is not initial_proc or self._resp_queue is not initial_resp_queue:
|
||||
yield f"Error: {self._subprocess_crash_message(crash_context)}"
|
||||
yield GenStreamError(f"Error: {self._subprocess_crash_message(crash_context)}")
|
||||
return
|
||||
resp = read_one(read_timeout)
|
||||
if resp is None:
|
||||
# Check subprocess health
|
||||
if not self._ensure_subprocess_alive():
|
||||
yield f"Error: {self._subprocess_crash_message(crash_context)}"
|
||||
yield GenStreamError(f"Error: {self._subprocess_crash_message(crash_context)}")
|
||||
return
|
||||
continue
|
||||
|
||||
|
|
@ -498,7 +510,7 @@ class InferenceOrchestrator:
|
|||
# Subprocess-level error (no request_id); request-scoped failures
|
||||
# arrive as gen_error below.
|
||||
if rtype == "error" and not resp.get("request_id"):
|
||||
yield f"Error: {resp.get('error', 'Unknown error')}"
|
||||
yield GenStreamError(f"Error: {resp.get('error', 'Unknown error')}")
|
||||
return
|
||||
|
||||
if rtype == "token":
|
||||
|
|
@ -513,7 +525,7 @@ class InferenceOrchestrator:
|
|||
stats_holder["stats"] = resp.get("stats")
|
||||
return
|
||||
elif rtype == "gen_error":
|
||||
yield f"Error: {resp.get('error', 'Unknown error')}"
|
||||
yield GenStreamError(f"Error: {resp.get('error', 'Unknown error')}")
|
||||
return
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
|
|
@ -640,11 +652,11 @@ class InferenceOrchestrator:
|
|||
GPU work stays serialized; this only avoids orchestrator lock contention.
|
||||
"""
|
||||
if not self._ensure_subprocess_alive():
|
||||
yield "Error: Inference subprocess is not running"
|
||||
yield GenStreamError("Error: Inference subprocess is not running")
|
||||
return
|
||||
|
||||
if not self.active_model_name:
|
||||
yield "Error: No active model"
|
||||
yield GenStreamError("Error: No active model")
|
||||
return
|
||||
# Latch the target model so the recheck below can detect a switch that completed
|
||||
# between _start_dispatcher and mailbox registration (mirrors the locked path's
|
||||
|
|
@ -655,7 +667,7 @@ class InferenceOrchestrator:
|
|||
# so without this early-out a compare request would enqueue a generate on the
|
||||
# outgoing model and delay the switch.
|
||||
if self._unload_pending:
|
||||
yield "Error: model is being unloaded"
|
||||
yield GenStreamError("Error: model is being unloaded")
|
||||
return
|
||||
|
||||
# Ensure the dispatcher runs. _start_dispatcher serializes concurrent starters under
|
||||
|
|
@ -727,7 +739,7 @@ class InferenceOrchestrator:
|
|||
# _stop_dispatcher joins the dispatcher, which itself takes that lock.
|
||||
if orphaned_dispatcher:
|
||||
self._stop_dispatcher()
|
||||
yield "Error: model is being unloaded"
|
||||
yield GenStreamError("Error: model is being unloaded")
|
||||
return
|
||||
|
||||
try:
|
||||
|
|
@ -735,7 +747,7 @@ class InferenceOrchestrator:
|
|||
except RuntimeError as exc:
|
||||
with self._mailbox_lock:
|
||||
self._mailboxes.pop(request_id, None)
|
||||
yield f"Error: {exc}"
|
||||
yield GenStreamError(f"Error: {exc}")
|
||||
return
|
||||
|
||||
def read_mailbox(timeout):
|
||||
|
|
@ -813,6 +825,59 @@ class InferenceOrchestrator:
|
|||
self._stop_dispatcher()
|
||||
return True
|
||||
|
||||
def share_distributed_object(
|
||||
self,
|
||||
obj,
|
||||
timeout: Optional[float] = 300.0,
|
||||
):
|
||||
"""Share a small object through the worker's MLX distributed group."""
|
||||
if not self._ensure_subprocess_alive():
|
||||
raise RuntimeError("Inference subprocess is not running")
|
||||
|
||||
self._wait_dispatcher_idle()
|
||||
with self._mailbox_lock:
|
||||
if self._mailboxes:
|
||||
raise RuntimeError(
|
||||
"Cannot share distributed objects while compare requests are active"
|
||||
)
|
||||
request_id = str(uuid.uuid4())
|
||||
cmd = {
|
||||
"type": "share_object",
|
||||
"request_id": request_id,
|
||||
"object": obj,
|
||||
}
|
||||
|
||||
with self._gen_lock:
|
||||
self._send_cmd(cmd)
|
||||
deadline = None if timeout is None else time.monotonic() + timeout
|
||||
while deadline is None or time.monotonic() < deadline:
|
||||
remaining = 1.0 if deadline is None else max(0.1, deadline - time.monotonic())
|
||||
resp = self._read_resp(timeout = min(remaining, 1.0))
|
||||
if resp is None:
|
||||
if not self._ensure_subprocess_alive():
|
||||
raise RuntimeError(self._subprocess_crash_message("sharing chat turn"))
|
||||
continue
|
||||
|
||||
rtype = resp.get("type", "")
|
||||
rid = resp.get("request_id")
|
||||
if rid and rid != request_id:
|
||||
logger.debug(
|
||||
"Skipping response for request_id=%s while sharing request_id=%s",
|
||||
rid,
|
||||
request_id,
|
||||
)
|
||||
continue
|
||||
if rtype == "shared":
|
||||
return resp.get("object")
|
||||
if rtype == "share_error":
|
||||
raise RuntimeError(resp.get("error", "Failed to share object"))
|
||||
if rtype == "error":
|
||||
raise RuntimeError(resp.get("error", "Subprocess error"))
|
||||
if rtype == "status":
|
||||
continue
|
||||
|
||||
raise RuntimeError("Timeout waiting for distributed object share")
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Public API — same interface as InferenceBackend
|
||||
# ------------------------------------------------------------------
|
||||
|
|
@ -828,6 +893,8 @@ class InferenceOrchestrator:
|
|||
approved_remote_code_fingerprint: Optional[str] = None,
|
||||
gpu_ids: Optional[list[int]] = None,
|
||||
subject: Optional[str] = None,
|
||||
tensor_parallel: bool = False,
|
||||
mlx_distributed: bool = False,
|
||||
) -> bool:
|
||||
"""Load a model for inference.
|
||||
|
||||
|
|
@ -853,6 +920,11 @@ class InferenceOrchestrator:
|
|||
"approved_remote_code_fingerprint": approved_remote_code_fingerprint,
|
||||
"subject": subject,
|
||||
"gpu_ids": gpu_ids,
|
||||
"tensor_parallel": bool(tensor_parallel),
|
||||
"mlx_distributed": bool(mlx_distributed),
|
||||
"mlx_parallel_mode": ("tensor" if tensor_parallel else "pipeline")
|
||||
if mlx_distributed
|
||||
else None,
|
||||
}
|
||||
resolved_gpu_ids, gpu_selection = prepare_gpu_selection(
|
||||
gpu_ids,
|
||||
|
|
@ -1338,11 +1410,11 @@ class InferenceOrchestrator:
|
|||
readers don't consume each other's tokens off the shared resp_queue.
|
||||
"""
|
||||
if not self._ensure_subprocess_alive():
|
||||
yield "Error: Inference subprocess is not running"
|
||||
yield GenStreamError("Error: Inference subprocess is not running")
|
||||
return
|
||||
|
||||
if not self.active_model_name:
|
||||
yield "Error: No active model"
|
||||
yield GenStreamError("Error: No active model")
|
||||
return
|
||||
expected_model = self.active_model_name
|
||||
|
||||
|
|
@ -1359,7 +1431,7 @@ class InferenceOrchestrator:
|
|||
# so we never generate on the wrong one.
|
||||
if self._unload_pending or self.active_model_name != expected_model:
|
||||
# Won the lock handoff during a switch; don't start on the outgoing model.
|
||||
yield "Error: model is being unloaded"
|
||||
yield GenStreamError("Error: model is being unloaded")
|
||||
return
|
||||
request_id = str(uuid.uuid4())
|
||||
image_b64 = self._pil_to_base64(image) if image is not None else None
|
||||
|
|
@ -1385,7 +1457,7 @@ class InferenceOrchestrator:
|
|||
try:
|
||||
self._send_cmd(cmd)
|
||||
except RuntimeError as exc:
|
||||
yield f"Error: {exc}"
|
||||
yield GenStreamError(f"Error: {exc}")
|
||||
return
|
||||
|
||||
yield from self._consume_token_stream(
|
||||
|
|
@ -1544,10 +1616,10 @@ class InferenceOrchestrator:
|
|||
) -> Generator[str, None, None]:
|
||||
"""Shared inner logic for audio input generation (Whisper + ASR)."""
|
||||
if not self._ensure_subprocess_alive():
|
||||
yield "Error: Inference subprocess is not running"
|
||||
yield GenStreamError("Error: Inference subprocess is not running")
|
||||
return
|
||||
if not self.active_model_name:
|
||||
yield "Error: No active model"
|
||||
yield GenStreamError("Error: No active model")
|
||||
return
|
||||
expected_model = self.active_model_name
|
||||
|
||||
|
|
@ -1556,7 +1628,7 @@ class InferenceOrchestrator:
|
|||
# cleared or swapped the model while we waited.
|
||||
if self._unload_pending or self.active_model_name != expected_model:
|
||||
# Won the lock handoff during a switch; don't start on the outgoing model.
|
||||
yield "Error: model is being unloaded"
|
||||
yield GenStreamError("Error: model is being unloaded")
|
||||
return
|
||||
request_id = str(uuid.uuid4())
|
||||
|
||||
|
|
@ -1583,7 +1655,7 @@ class InferenceOrchestrator:
|
|||
try:
|
||||
self._send_cmd(cmd)
|
||||
except RuntimeError as exc:
|
||||
yield f"Error: {exc}"
|
||||
yield GenStreamError(f"Error: {exc}")
|
||||
return
|
||||
|
||||
yield from self._consume_token_stream(
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ mp.Queue, and exits on shutdown or unload. Pattern follows core/training/worker.
|
|||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import json
|
||||
from loggers import get_logger
|
||||
import os
|
||||
import queue as _queue
|
||||
|
|
@ -26,6 +27,9 @@ from typing import Any
|
|||
logger = get_logger(__name__)
|
||||
from utils.hardware import apply_gpu_ids
|
||||
|
||||
_SHARE_OBJECT_MAX_BYTES = 1 << 20
|
||||
_SHARE_OBJECT_ERROR_SIZE = -1
|
||||
|
||||
# studio/backend root, prepended to sys.path so the spawned subprocess can
|
||||
# import the utils/core packages.
|
||||
_BACKEND_PATH = str(Path(__file__).resolve().parent.parent.parent)
|
||||
|
|
@ -75,6 +79,17 @@ def _send_response(resp_queue: Any, response: dict) -> None:
|
|||
logger.error("Failed to send response: %s", exc)
|
||||
|
||||
|
||||
def _encode_share_object(obj: Any) -> bytes:
|
||||
data = json.dumps(obj, separators = (",", ":"), ensure_ascii = False).encode("utf-8")
|
||||
if len(data) > _SHARE_OBJECT_MAX_BYTES:
|
||||
raise ValueError("Distributed object share payload is too large")
|
||||
return data
|
||||
|
||||
|
||||
def _decode_share_object(data: Any) -> Any:
|
||||
return json.loads(bytes(data.tolist()).decode("utf-8"))
|
||||
|
||||
|
||||
def _clean_token(value: str | None) -> str | None:
|
||||
"""Normalize an HF token: blank or whitespace-only becomes None."""
|
||||
return value if value and value.strip() else None
|
||||
|
|
@ -329,14 +344,18 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
|
|||
xet_disabled = os.environ.get("HF_HUB_DISABLE_XET") == "1",
|
||||
)
|
||||
try:
|
||||
success = backend.load_model(
|
||||
config = mc,
|
||||
max_seq_length = config.get("max_seq_length", 2048),
|
||||
load_in_4bit = load_in_4bit,
|
||||
hf_token = hf_token,
|
||||
trust_remote_code = trust_remote_code,
|
||||
gpu_ids = config.get("resolved_gpu_ids"),
|
||||
)
|
||||
load_kwargs = {
|
||||
"config": mc,
|
||||
"max_seq_length": config.get("max_seq_length", 2048),
|
||||
"load_in_4bit": load_in_4bit,
|
||||
"hf_token": hf_token,
|
||||
"trust_remote_code": trust_remote_code,
|
||||
"gpu_ids": config.get("resolved_gpu_ids"),
|
||||
}
|
||||
if getattr(backend, "device", None) == "mlx":
|
||||
load_kwargs["parallel_mode"] = config.get("mlx_parallel_mode")
|
||||
load_kwargs["distributed_group"] = config.get("_mlx_distributed_group")
|
||||
success = backend.load_model(**load_kwargs)
|
||||
finally:
|
||||
heartbeat_stop.set()
|
||||
|
||||
|
|
@ -521,6 +540,67 @@ def _handle_generate(backend, cmd: dict, resp_queue: Any, cancel_event) -> None:
|
|||
)
|
||||
|
||||
|
||||
def _handle_share_object(backend, cmd: dict, resp_queue: Any) -> None:
|
||||
"""Share a small Python object across MLX distributed ranks."""
|
||||
request_id = cmd.get("request_id", "")
|
||||
group = getattr(backend, "_distributed_group", None)
|
||||
rank = int(getattr(backend, "_distributed_rank", 0) or 0)
|
||||
world_size = int(getattr(backend, "_distributed_world_size", 1) or 1)
|
||||
obj = cmd.get("object")
|
||||
|
||||
try:
|
||||
if group is None or world_size <= 1:
|
||||
shared = obj
|
||||
else:
|
||||
import mlx.core as mx
|
||||
if rank == 0:
|
||||
if obj is None:
|
||||
mx.eval(mx.distributed.all_sum(mx.array(0), group = group))
|
||||
shared = None
|
||||
else:
|
||||
try:
|
||||
data = mx.array(_encode_share_object(obj), dtype = mx.uint8)
|
||||
except Exception:
|
||||
mx.eval(
|
||||
mx.distributed.all_sum(
|
||||
mx.array(_SHARE_OBJECT_ERROR_SIZE),
|
||||
group = group,
|
||||
)
|
||||
)
|
||||
raise
|
||||
mx.eval(mx.distributed.all_sum(mx.array(data.size), group = group))
|
||||
mx.eval(mx.distributed.all_sum(data, group = group))
|
||||
shared = obj
|
||||
else:
|
||||
size = int(mx.distributed.all_sum(mx.array(0), group = group).item())
|
||||
if size == _SHARE_OBJECT_ERROR_SIZE:
|
||||
raise RuntimeError("Failed to share distributed object")
|
||||
if size == 0:
|
||||
shared = None
|
||||
else:
|
||||
data = mx.zeros(size, dtype = mx.uint8)
|
||||
data = mx.distributed.all_sum(data, group = group)
|
||||
shared = _decode_share_object(data)
|
||||
_send_response(
|
||||
resp_queue,
|
||||
{
|
||||
"type": "shared",
|
||||
"request_id": request_id,
|
||||
"object": shared,
|
||||
},
|
||||
)
|
||||
except Exception as exc:
|
||||
_send_response(
|
||||
resp_queue,
|
||||
{
|
||||
"type": "share_error",
|
||||
"request_id": request_id,
|
||||
"error": str(exc),
|
||||
"stack": traceback.format_exc(limit = 20),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _handle_generate_audio(backend, cmd: dict, resp_queue: Any) -> None:
|
||||
"""Handle TTS audio generation — returns WAV bytes + sample_rate."""
|
||||
request_id = cmd.get("request_id", "")
|
||||
|
|
@ -720,9 +800,29 @@ def run_inference_process(
|
|||
exc,
|
||||
)
|
||||
try:
|
||||
from core.inference.mlx_inference import MLXInferenceBackend
|
||||
from core.inference.mlx_inference import MLXInferenceBackend, _init_mlx_distributed
|
||||
|
||||
backend = MLXInferenceBackend()
|
||||
if config.get("mlx_distributed"):
|
||||
group, rank, size = _init_mlx_distributed()
|
||||
config["_mlx_distributed_group"] = group
|
||||
if size <= 1:
|
||||
# A singleton group (MLX built without distributed support,
|
||||
# or an invalid launch env/hostfile) would leave nonzero ranks
|
||||
# looping forever on share_distributed_object. Fail the load
|
||||
# instead of silently continuing without sharding.
|
||||
raise RuntimeError(
|
||||
"MLX distributed launch requested but initialized a singleton "
|
||||
"group (size 1). Ensure the installed MLX has distributed "
|
||||
"support and the launch environment/hostfile is valid, or run "
|
||||
"without distributed."
|
||||
)
|
||||
logger.info(
|
||||
"MLX distributed initialized in worker: rank=%s size=%s mode=%s",
|
||||
rank,
|
||||
size,
|
||||
config.get("mlx_parallel_mode"),
|
||||
)
|
||||
_send_response(
|
||||
resp_queue,
|
||||
{"type": "status", "message": "Loading model..."},
|
||||
|
|
@ -764,6 +864,8 @@ def run_inference_process(
|
|||
if _drain_skip_generate(cmd, resp_queue, drain_event):
|
||||
continue
|
||||
_handle_generate(backend, cmd, resp_queue, cancel_event)
|
||||
elif cmd_type == "share_object":
|
||||
_handle_share_object(backend, cmd, resp_queue)
|
||||
elif cmd_type == "load":
|
||||
if backend.active_model_name:
|
||||
backend.unload_model(backend.active_model_name)
|
||||
|
|
@ -977,6 +1079,9 @@ def run_inference_process(
|
|||
continue
|
||||
_handle_generate(backend, cmd, resp_queue, cancel_event)
|
||||
|
||||
elif cmd_type == "share_object":
|
||||
_handle_share_object(backend, cmd, resp_queue)
|
||||
|
||||
elif cmd_type == "load":
|
||||
if backend.active_model_name:
|
||||
backend.unload_model(backend.active_model_name)
|
||||
|
|
|
|||
|
|
@ -40,6 +40,43 @@ def _positive_int_or_none(value: Any) -> Optional[int]:
|
|||
return value_int if value_int > 0 else None
|
||||
|
||||
|
||||
def _nonnegative_int_or_none(value: Any) -> Optional[int]:
|
||||
if isinstance(value, bool):
|
||||
return None
|
||||
try:
|
||||
value_int = int(value)
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
return value_int if value_int >= 0 else None
|
||||
|
||||
|
||||
_MLX_MPI_DISTRIBUTED_ENV_PAIRS = (
|
||||
("OMPI_COMM_WORLD_RANK", "OMPI_COMM_WORLD_SIZE"),
|
||||
("PMI_RANK", "PMI_SIZE"),
|
||||
("PMIX_RANK", "PMIX_SIZE"),
|
||||
("MPI_RANK", "MPI_WORLD_SIZE"),
|
||||
("MV2_COMM_WORLD_RANK", "MV2_COMM_WORLD_SIZE"),
|
||||
)
|
||||
|
||||
|
||||
def _mlx_distributed_launch_detected() -> bool:
|
||||
if _nonnegative_int_or_none(os.environ.get("MLX_RANK")) is not None:
|
||||
world_size = _positive_int_or_none(os.environ.get("MLX_WORLD_SIZE"))
|
||||
if world_size is not None and world_size > 1:
|
||||
return True
|
||||
return bool(
|
||||
os.environ.get("MLX_HOSTFILE")
|
||||
or os.environ.get("MLX_IBV_DEVICES")
|
||||
or os.environ.get("MLX_JACCL_COORDINATOR")
|
||||
or (os.environ.get("NCCL_HOST_IP") and os.environ.get("NCCL_PORT"))
|
||||
)
|
||||
return any(
|
||||
_nonnegative_int_or_none(os.environ.get(rank_env)) is not None
|
||||
and (_positive_int_or_none(os.environ.get(size_env)) or 0) > 1
|
||||
for rank_env, size_env in _MLX_MPI_DISTRIBUTED_ENV_PAIRS
|
||||
)
|
||||
|
||||
|
||||
def _install_httpcore_asyncgen_silencer() -> None:
|
||||
"""Silence benign httpx/httpcore asyncgen GC noise on Python 3.13.
|
||||
|
||||
|
|
@ -3426,6 +3463,15 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
|
|||
status_code = 400,
|
||||
detail = "gpu_ids is not supported for GGUF models yet.",
|
||||
)
|
||||
if not config.is_gguf and _mlx_distributed_launch_detected():
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = (
|
||||
"Studio does not support distributed MLX inference under "
|
||||
"mlx.launch. Use `mlx.launch ... unsloth chat` or run Studio "
|
||||
"without the distributed launcher."
|
||||
),
|
||||
)
|
||||
|
||||
# Effective quantization (LoRA can flip 4-bit -> 16-bit); guard + load reuse it.
|
||||
effective_load_in_4bit = _effective_load_in_4bit(config, request.load_in_4bit)
|
||||
|
|
|
|||
|
|
@ -4,6 +4,8 @@ import sys
|
|||
import types
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
class _DummyMetal:
|
||||
@staticmethod
|
||||
|
|
@ -185,6 +187,129 @@ def test_mlx_inference_vlm_lora_uses_unsloth_loader_without_native_adapter_rewri
|
|||
assert isinstance(backend._tokenizer, _DummyTokenizer)
|
||||
|
||||
|
||||
def test_mlx_inference_distributed_vlm_forwards_group_to_fast_mlx(monkeypatch):
|
||||
_install_fake_mlx(monkeypatch)
|
||||
calls = []
|
||||
_install_fake_fast_mlx(monkeypatch, calls)
|
||||
from core.inference.mlx_inference import MLXInferenceBackend
|
||||
|
||||
group = SimpleNamespace(size = lambda: 2, rank = lambda: 0)
|
||||
config = SimpleNamespace(identifier = "fake/vlm", is_vision = True, is_lora = False)
|
||||
for mode, group_key in (("tensor", "tensor_group"), ("pipeline", "pipeline_group")):
|
||||
calls.clear()
|
||||
assert MLXInferenceBackend().load_model(config, parallel_mode = mode, distributed_group = group)
|
||||
_, kwargs = calls.pop()
|
||||
assert kwargs["text_only"] is False and kwargs[group_key] is group
|
||||
|
||||
calls.clear()
|
||||
singleton = SimpleNamespace(size = lambda: 1, rank = lambda: 0)
|
||||
assert MLXInferenceBackend().load_model(
|
||||
config, parallel_mode = "tensor", distributed_group = singleton
|
||||
)
|
||||
assert not {"tensor_group", "pipeline_group"} & set(calls.pop()[1])
|
||||
|
||||
config = SimpleNamespace(identifier = "fake/adapter", is_vision = False, is_lora = True)
|
||||
with pytest.raises(ValueError, match = "LoRA adapter repos"):
|
||||
MLXInferenceBackend().load_model(config, parallel_mode = "tensor", distributed_group = group)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("accepts_backend", (True, False))
|
||||
def test_mlx_distributed_init_selects_jaccl_backend(monkeypatch, accepts_backend):
|
||||
_install_fake_mlx(monkeypatch)
|
||||
from core.inference.mlx_inference import _init_mlx_distributed
|
||||
|
||||
group = SimpleNamespace(rank = lambda: 1, size = lambda: 2)
|
||||
calls = []
|
||||
|
||||
def _init(**kwargs):
|
||||
calls.append(kwargs)
|
||||
if kwargs and not accepts_backend:
|
||||
raise TypeError("backend keyword unsupported")
|
||||
return group
|
||||
|
||||
sys.modules["mlx.core"].distributed = SimpleNamespace(init = _init)
|
||||
monkeypatch.setenv("MLX_JACCL_COORDINATOR", "127.0.0.1:12345")
|
||||
monkeypatch.setenv("MLX_IBV_DEVICES", "/tmp/devices.json")
|
||||
|
||||
assert _init_mlx_distributed() == (group, 1, 2)
|
||||
assert calls == ([{"backend": "jaccl"}] if accepts_backend else [{"backend": "jaccl"}, {}])
|
||||
|
||||
|
||||
def test_worker_share_object_receives_distributed_payload(monkeypatch):
|
||||
from core.inference import worker
|
||||
|
||||
shared_obj = {"type": "turn", "text": "hi"}
|
||||
payload = worker._encode_share_object(shared_obj)
|
||||
|
||||
def _array(value):
|
||||
val = value.item() if hasattr(value, "item") else value
|
||||
return SimpleNamespace(
|
||||
item = lambda: val,
|
||||
tolist = lambda: list(val) if hasattr(val, "__iter__") else [val],
|
||||
)
|
||||
|
||||
mlx_pkg = types.ModuleType("mlx")
|
||||
mlx_core = types.ModuleType("mlx.core")
|
||||
mlx_core.uint8 = "uint8"
|
||||
mlx_core.array = _array
|
||||
mlx_core.zeros = lambda *_a, **_k: _array([])
|
||||
|
||||
def _all_sum(value, group = None):
|
||||
value = value.item() if hasattr(value, "item") else value
|
||||
return _array(len(payload)) if value == 0 else _array(payload)
|
||||
|
||||
mlx_core.distributed = SimpleNamespace(all_sum = _all_sum)
|
||||
mlx_pkg.core = mlx_core
|
||||
monkeypatch.setitem(sys.modules, "mlx", mlx_pkg)
|
||||
monkeypatch.setitem(sys.modules, "mlx.core", mlx_core)
|
||||
|
||||
responses = []
|
||||
worker._handle_share_object(
|
||||
SimpleNamespace(
|
||||
_distributed_group = object(),
|
||||
_distributed_rank = 1,
|
||||
_distributed_world_size = 2,
|
||||
),
|
||||
{"type": "share_object", "request_id": "rid", "object": None},
|
||||
SimpleNamespace(put = responses.append),
|
||||
)
|
||||
|
||||
response = responses[0]
|
||||
assert response["object"] == shared_obj
|
||||
|
||||
|
||||
def test_worker_share_object_oversize_notifies_peers(monkeypatch):
|
||||
from core.inference import worker
|
||||
|
||||
calls = []
|
||||
|
||||
mlx_pkg = types.ModuleType("mlx")
|
||||
mlx_core = types.ModuleType("mlx.core")
|
||||
mlx_core.array = lambda value, **_kwargs: SimpleNamespace(item = lambda: value)
|
||||
mlx_core.eval = lambda value: value
|
||||
mlx_core.distributed = SimpleNamespace(
|
||||
all_sum = lambda value, group = None: calls.append(value.item()) or value
|
||||
)
|
||||
mlx_pkg.core = mlx_core
|
||||
monkeypatch.setitem(sys.modules, "mlx", mlx_pkg)
|
||||
monkeypatch.setitem(sys.modules, "mlx.core", mlx_core)
|
||||
monkeypatch.setattr(worker, "_SHARE_OBJECT_MAX_BYTES", 8)
|
||||
|
||||
responses = []
|
||||
worker._handle_share_object(
|
||||
SimpleNamespace(
|
||||
_distributed_group = object(),
|
||||
_distributed_rank = 0,
|
||||
_distributed_world_size = 2,
|
||||
),
|
||||
{"type": "share_object", "request_id": "rid", "object": {"text": "too long"}},
|
||||
SimpleNamespace(put = responses.append),
|
||||
)
|
||||
|
||||
assert calls == [worker._SHARE_OBJECT_ERROR_SIZE]
|
||||
assert responses[0]["type"] == "share_error"
|
||||
|
||||
|
||||
# Regression: generate_chat_response must accept the four template kwargs
|
||||
# (tools / enable_thinking / reasoning_effort / preserve_thinking) so the route
|
||||
# layer can forward UI toggles. The old signature raised
|
||||
|
|
|
|||
|
|
@ -4,9 +4,11 @@
|
|||
"""Model loading and streaming shared by `inference` and `chat`."""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
from contextlib import contextmanager, redirect_stderr, redirect_stdout
|
||||
from pathlib import Path
|
||||
from typing import List, Optional
|
||||
|
||||
|
|
@ -14,10 +16,18 @@ import typer
|
|||
|
||||
_THINK_OPEN = "<think>"
|
||||
_THINK_BLOCK = re.compile(rf"{re.escape(_THINK_OPEN)}.*?</think>", re.DOTALL)
|
||||
_STREAMED_ERROR_PREFIX = "Error: "
|
||||
|
||||
# Cloudflare (in front of remote Studio proxies like RunPod) 403s the default
|
||||
# "Python-urllib/X.Y" User-Agent as a bot; send a real one on every request.
|
||||
_USER_AGENT = "unsloth-cli"
|
||||
_MPI_ENV_PAIRS = (
|
||||
("OMPI_COMM_WORLD_RANK", "OMPI_COMM_WORLD_SIZE"),
|
||||
("PMI_RANK", "PMI_SIZE"),
|
||||
("PMIX_RANK", "PMIX_SIZE"),
|
||||
("MPI_RANK", "MPI_WORLD_SIZE"),
|
||||
("MV2_COMM_WORLD_RANK", "MV2_COMM_WORLD_SIZE"),
|
||||
)
|
||||
|
||||
# Built lazily; urllib stays function-local to match this module.
|
||||
_no_redirect_opener = None
|
||||
|
|
@ -61,6 +71,108 @@ def configure_quiet_logging() -> None:
|
|||
os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
|
||||
|
||||
|
||||
def _parse_nonnegative_int(value: Optional[str]) -> Optional[int]:
|
||||
if value is None:
|
||||
return None
|
||||
try:
|
||||
parsed = int(value)
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
return parsed if parsed >= 0 else None
|
||||
|
||||
|
||||
def _first_mpi_env_pair() -> tuple[Optional[int], Optional[int]]:
|
||||
for rank_name, size_name in _MPI_ENV_PAIRS:
|
||||
rank = _parse_nonnegative_int(os.environ.get(rank_name))
|
||||
world_size = _parse_nonnegative_int(os.environ.get(size_name))
|
||||
if rank is not None and world_size is not None and world_size > 1 and rank < world_size:
|
||||
return rank, world_size
|
||||
return None, None
|
||||
|
||||
|
||||
def _json_rank_count_from_env(name: str) -> Optional[int]:
|
||||
value = os.environ.get(name)
|
||||
if not value:
|
||||
return None
|
||||
try:
|
||||
if value.lstrip().startswith(("[", "{")):
|
||||
data = json.loads(value)
|
||||
else:
|
||||
with open(value, "r") as f:
|
||||
data = json.load(f)
|
||||
except (OSError, json.JSONDecodeError):
|
||||
return None
|
||||
if isinstance(data, list):
|
||||
return len(data)
|
||||
if isinstance(data, dict) and isinstance(data.get("hosts"), list):
|
||||
return len(data["hosts"])
|
||||
return None
|
||||
|
||||
|
||||
def mlx_distributed_info() -> tuple[bool, int, Optional[int]]:
|
||||
"""Return launch-context metadata without initializing MLX distributed."""
|
||||
rank = _parse_nonnegative_int(os.environ.get("MLX_RANK"))
|
||||
world_size = _parse_nonnegative_int(os.environ.get("MLX_WORLD_SIZE"))
|
||||
if rank is not None:
|
||||
if (
|
||||
world_size is not None
|
||||
and world_size > 1
|
||||
and rank < world_size
|
||||
and os.environ.get("NCCL_HOST_IP")
|
||||
and os.environ.get("NCCL_PORT")
|
||||
):
|
||||
return True, rank, world_size
|
||||
inferred_size = _json_rank_count_from_env("MLX_HOSTFILE")
|
||||
if inferred_size is not None and inferred_size > 1 and rank < inferred_size:
|
||||
return True, rank, inferred_size
|
||||
inferred_size = _json_rank_count_from_env("MLX_IBV_DEVICES")
|
||||
if (
|
||||
inferred_size is not None
|
||||
and inferred_size > 1
|
||||
and rank < inferred_size
|
||||
and os.environ.get("MLX_JACCL_COORDINATOR")
|
||||
):
|
||||
return True, rank, inferred_size
|
||||
return False, 0, None
|
||||
|
||||
mpi_rank, mpi_world_size = _first_mpi_env_pair()
|
||||
return mpi_rank is not None, mpi_rank or 0, mpi_world_size
|
||||
|
||||
|
||||
def mlx_distributed_uses_mpi() -> bool:
|
||||
"""Whether the current distributed context was launched through MPI."""
|
||||
return (
|
||||
_parse_nonnegative_int(os.environ.get("MLX_RANK")) is None
|
||||
and _first_mpi_env_pair()[0] is not None
|
||||
)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def quiet_if_nonzero_mlx_rank():
|
||||
"""Silence parent and child-process stdout/stderr on nonzero ranks."""
|
||||
if mlx_distributed_info()[1] == 0:
|
||||
yield
|
||||
return
|
||||
|
||||
sys.stdout.flush()
|
||||
sys.stderr.flush()
|
||||
saved_stdout_fd = os.dup(1)
|
||||
saved_stderr_fd = os.dup(2)
|
||||
with open(os.devnull, "w") as devnull:
|
||||
try:
|
||||
os.dup2(devnull.fileno(), 1)
|
||||
os.dup2(devnull.fileno(), 2)
|
||||
with redirect_stdout(devnull), redirect_stderr(devnull):
|
||||
yield
|
||||
finally:
|
||||
sys.stdout.flush()
|
||||
sys.stderr.flush()
|
||||
os.dup2(saved_stdout_fd, 1)
|
||||
os.dup2(saved_stderr_fd, 2)
|
||||
os.close(saved_stdout_fd)
|
||||
os.close(saved_stderr_fd)
|
||||
|
||||
|
||||
def visible_text(text: str, show_thinking: bool) -> str:
|
||||
if show_thinking:
|
||||
return text
|
||||
|
|
@ -120,6 +232,21 @@ def collect_stream(stream, show_thinking: bool) -> str:
|
|||
return visible_text(raw, show_thinking)
|
||||
|
||||
|
||||
def raise_on_streamed_error(stream):
|
||||
# Match real backend errors by type (GenStreamError), not the "Error:" text
|
||||
# prefix, so a completion whose text opens with "Error:" is not misread as a
|
||||
# failure that aborts a distributed run.
|
||||
try:
|
||||
ensure_studio_backend_path()
|
||||
from core.inference.orchestrator import GenStreamError
|
||||
except Exception:
|
||||
GenStreamError = None
|
||||
for chunk in stream:
|
||||
if GenStreamError is not None and isinstance(chunk, GenStreamError):
|
||||
raise RuntimeError(str(chunk)[len(_STREAMED_ERROR_PREFIX) :].strip() or "Unknown error")
|
||||
yield chunk
|
||||
|
||||
|
||||
def render_columns(
|
||||
left_label: str,
|
||||
left_text: str,
|
||||
|
|
@ -200,6 +327,19 @@ class ChatBackend:
|
|||
except Exception:
|
||||
pass
|
||||
|
||||
def share_distributed_object(
|
||||
self,
|
||||
obj,
|
||||
*,
|
||||
timeout = 300.0,
|
||||
):
|
||||
if self._kind != "unsloth" or not hasattr(self._backend, "share_distributed_object"):
|
||||
raise RuntimeError(
|
||||
"Distributed MLX chat requires the Unsloth MLX backend; "
|
||||
f"backend '{self._kind}' cannot broadcast chat turns."
|
||||
)
|
||||
return self._backend.share_distributed_object(obj, timeout = timeout)
|
||||
|
||||
|
||||
def resolve_model_config(model: str, *, hf_token: Optional[str]):
|
||||
ensure_studio_backend_path()
|
||||
|
|
@ -293,36 +433,59 @@ def load_chat_backend(
|
|||
fresh_backend uses a private orchestrator so a second model (compare's
|
||||
base column) can run alongside the main one.
|
||||
"""
|
||||
if model_config is None:
|
||||
model_config = resolve_model_config(model, hf_token = hf_token)
|
||||
with quiet_if_nonzero_mlx_rank():
|
||||
is_mlx_distributed, rank, _world_size = mlx_distributed_info()
|
||||
if model_config is None:
|
||||
model_config = resolve_model_config(model, hf_token = hf_token)
|
||||
|
||||
typer.echo(f"Loading {model}", err = True)
|
||||
if is_mlx_distributed and model_config.is_gguf:
|
||||
if rank == 0:
|
||||
typer.echo(
|
||||
"Distributed MLX inference does not support GGUF/llama.cpp models. "
|
||||
"Use a non-GGUF MLX model under mlx.launch, or run GGUF without "
|
||||
"mlx.launch.",
|
||||
err = True,
|
||||
)
|
||||
raise typer.Exit(code = 1)
|
||||
|
||||
if model_config.is_gguf:
|
||||
return _load_gguf_backend(
|
||||
model_config,
|
||||
hf_token = hf_token,
|
||||
max_seq_length = max_seq_length,
|
||||
tensor_parallel = tensor_parallel,
|
||||
llama_extra_args = llama_extra_args,
|
||||
)
|
||||
if rank == 0:
|
||||
typer.echo(f"Loading {model}", err = True)
|
||||
|
||||
if fresh_backend:
|
||||
ensure_studio_backend_path()
|
||||
from core.inference import InferenceOrchestrator
|
||||
backend = InferenceOrchestrator()
|
||||
else:
|
||||
ensure_studio_backend_path()
|
||||
from core.inference import get_inference_backend
|
||||
backend = get_inference_backend()
|
||||
if not backend.load_model(
|
||||
config = model_config,
|
||||
max_seq_length = max_seq_length,
|
||||
load_in_4bit = load_in_4bit,
|
||||
hf_token = hf_token,
|
||||
):
|
||||
typer.echo("Model load failed", err = True)
|
||||
raise typer.Exit(code = 1)
|
||||
if model_config.is_gguf:
|
||||
return _load_gguf_backend(
|
||||
model_config,
|
||||
hf_token = hf_token,
|
||||
max_seq_length = max_seq_length,
|
||||
tensor_parallel = tensor_parallel,
|
||||
llama_extra_args = llama_extra_args,
|
||||
)
|
||||
|
||||
if fresh_backend:
|
||||
ensure_studio_backend_path()
|
||||
from core.inference import InferenceOrchestrator
|
||||
backend = InferenceOrchestrator()
|
||||
else:
|
||||
ensure_studio_backend_path()
|
||||
from core.inference import get_inference_backend
|
||||
backend = get_inference_backend()
|
||||
try:
|
||||
loaded = backend.load_model(
|
||||
config = model_config,
|
||||
max_seq_length = max_seq_length,
|
||||
load_in_4bit = load_in_4bit,
|
||||
hf_token = hf_token,
|
||||
tensor_parallel = tensor_parallel,
|
||||
mlx_distributed = is_mlx_distributed,
|
||||
)
|
||||
except Exception as exc:
|
||||
if not is_mlx_distributed:
|
||||
raise
|
||||
if rank == 0:
|
||||
typer.echo(str(exc) or "Model load failed", err = True)
|
||||
raise typer.Exit(code = 1)
|
||||
if not loaded:
|
||||
typer.echo("Model load failed", err = True)
|
||||
raise typer.Exit(code = 1)
|
||||
return ChatBackend("unsloth", backend)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import sys
|
||||
from typing import List, Optional
|
||||
|
||||
import typer
|
||||
|
|
@ -12,6 +13,10 @@ from unsloth_cli._inference import (
|
|||
connect_studio_server,
|
||||
ensure_studio_backend_path,
|
||||
load_chat_backend,
|
||||
mlx_distributed_info,
|
||||
mlx_distributed_uses_mpi,
|
||||
quiet_if_nonzero_mlx_rank,
|
||||
raise_on_streamed_error,
|
||||
render_columns,
|
||||
resolve_model_config,
|
||||
stream_markdown,
|
||||
|
|
@ -107,6 +112,20 @@ def _compare_needs_second_model() -> bool:
|
|||
return False
|
||||
|
||||
|
||||
def _drain_available_stdin() -> None:
|
||||
"""Drain already-buffered launcher stdin on nonzero distributed ranks."""
|
||||
try:
|
||||
import os
|
||||
from select import select
|
||||
|
||||
fd = sys.stdin.fileno()
|
||||
while select([fd], [], [], 0)[0]:
|
||||
if not os.read(fd, 8192):
|
||||
break
|
||||
except Exception:
|
||||
return
|
||||
|
||||
|
||||
def _pick_trained_model(console) -> str:
|
||||
ensure_studio_backend_path()
|
||||
from utils.models import scan_trained_models
|
||||
|
|
@ -158,7 +177,8 @@ def chat(
|
|||
"--tensor-parallel/--no-tensor-parallel",
|
||||
help = (
|
||||
"Split a GGUF across GPUs by tensor (--split-mode tensor) instead "
|
||||
"of by layer. Ignored for non-GGUF models."
|
||||
"of by layer. Under non-MPI mlx.launch, select MLX tensor "
|
||||
"parallel mode instead of pipeline mode."
|
||||
),
|
||||
),
|
||||
llama_extra_args: Optional[List[str]] = typer.Option(
|
||||
|
|
@ -195,15 +215,43 @@ def chat(
|
|||
|
||||
console = Console()
|
||||
err = Console(stderr = True)
|
||||
is_mlx_distributed, rank, _world_size = mlx_distributed_info()
|
||||
should_print = rank == 0
|
||||
|
||||
if is_mlx_distributed and mlx_distributed_uses_mpi():
|
||||
if should_print:
|
||||
err.print(
|
||||
"Distributed `unsloth chat` with MPI needs rank-0 prompt broadcast, "
|
||||
"which is not enabled yet. Use a non-MPI MLX launcher backend "
|
||||
"such as ring/JACCL for now.",
|
||||
style = "red",
|
||||
markup = False,
|
||||
)
|
||||
raise typer.Exit(code = 1)
|
||||
|
||||
if model is None:
|
||||
if is_mlx_distributed:
|
||||
if should_print:
|
||||
err.print(
|
||||
"Distributed `unsloth chat` requires an explicit model id or path.",
|
||||
style = "red",
|
||||
markup = False,
|
||||
)
|
||||
raise typer.Exit(code = 1)
|
||||
model = _pick_trained_model(console)
|
||||
|
||||
# Resolve first so --compare can be rejected before the slow load.
|
||||
model_config = resolve_model_config(model, hf_token = hf_token)
|
||||
with quiet_if_nonzero_mlx_rank():
|
||||
model_config = resolve_model_config(model, hf_token = hf_token)
|
||||
compare_blocked = _compare_blocked_reason(model_config)
|
||||
if is_mlx_distributed:
|
||||
compare_blocked = (
|
||||
"distributed MLX chat does not support compare mode yet because it "
|
||||
"would need a second distributed worker group on the same ranks"
|
||||
)
|
||||
if compare and compare_blocked:
|
||||
err.print(f"--compare unavailable: {compare_blocked}", style = "red", markup = False)
|
||||
if should_print:
|
||||
err.print(f"--compare unavailable: {compare_blocked}", style = "red", markup = False)
|
||||
raise typer.Exit(code = 1)
|
||||
|
||||
load_opts = dict(
|
||||
|
|
@ -215,9 +263,11 @@ def chat(
|
|||
)
|
||||
|
||||
# Prefer a running Studio server: instant starts, model shared with the UI.
|
||||
chat_backend = None if no_server else connect_studio_server(model, **load_opts)
|
||||
chat_backend = (
|
||||
None if (no_server or is_mlx_distributed) else connect_studio_server(model, **load_opts)
|
||||
)
|
||||
server_mode = chat_backend is not None
|
||||
if server_mode:
|
||||
if server_mode and should_print:
|
||||
console.print(
|
||||
"(Studio server connected — model stays warm after /exit)",
|
||||
style = "bright_black",
|
||||
|
|
@ -242,23 +292,26 @@ def chat(
|
|||
return True
|
||||
base_id = model_config.base_model
|
||||
if not base_id:
|
||||
console.print(
|
||||
"(compare unavailable: this adapter doesn't record its base model)",
|
||||
style = "yellow",
|
||||
)
|
||||
if should_print:
|
||||
console.print(
|
||||
"(compare unavailable: this adapter doesn't record its base model)",
|
||||
style = "yellow",
|
||||
)
|
||||
return False
|
||||
console.print(
|
||||
f"(loading base model {base_id} for compare — keeps two models in memory)",
|
||||
style = "bright_black",
|
||||
markup = False,
|
||||
)
|
||||
if should_print:
|
||||
console.print(
|
||||
f"(loading base model {base_id} for compare — keeps two models in memory)",
|
||||
style = "bright_black",
|
||||
markup = False,
|
||||
)
|
||||
try:
|
||||
# Use the same precision as the tuned model for fair comparison
|
||||
base_load_opts = dict(load_opts) # Copy original options
|
||||
base_load_opts["load_in_4bit"] = _get_base_load_in_4bit(model_config)
|
||||
base_backend = load_chat_backend(base_id, fresh_backend = True, **base_load_opts)
|
||||
except Exception as exc:
|
||||
err.print(f"(base model load failed: {exc})", style = "red", markup = False)
|
||||
if should_print:
|
||||
err.print(f"(base model load failed: {exc})", style = "red", markup = False)
|
||||
return False
|
||||
return True
|
||||
|
||||
|
|
@ -267,7 +320,7 @@ def chat(
|
|||
|
||||
def generate(backend = None, use_adapter = None):
|
||||
# Reads messages and show_thinking live, so /reset and /think apply.
|
||||
return (backend or chat_backend).stream(
|
||||
stream = (backend or chat_backend).stream(
|
||||
messages,
|
||||
system_prompt = system_prompt,
|
||||
temperature = temperature,
|
||||
|
|
@ -278,22 +331,48 @@ def chat(
|
|||
enable_thinking = show_thinking,
|
||||
use_adapter = use_adapter,
|
||||
)
|
||||
return raise_on_streamed_error(stream) if is_mlx_distributed else stream
|
||||
|
||||
console.print()
|
||||
console.print(f"Chatting with {name}", style = "bold green", markup = False)
|
||||
console.print(_HELP, style = "bright_black")
|
||||
if should_print:
|
||||
console.print()
|
||||
console.print(f"Chatting with {name}", style = "bold green", markup = False)
|
||||
console.print(_HELP, style = "bright_black")
|
||||
|
||||
# legacy_windows: pre-VT consoles print raw ANSI as ←[1;36m garbage.
|
||||
you_prompt = _you_prompt(console.is_terminal and not console.legacy_windows)
|
||||
you_prompt = (
|
||||
_you_prompt(console.is_terminal and not console.legacy_windows) if should_print else ""
|
||||
)
|
||||
assistant_label = "[bold magenta]Assistant:[/bold magenta]"
|
||||
|
||||
try:
|
||||
while True:
|
||||
try:
|
||||
user = input(you_prompt).strip()
|
||||
except (EOFError, KeyboardInterrupt):
|
||||
console.print()
|
||||
break
|
||||
if should_print:
|
||||
try:
|
||||
user = input(you_prompt).strip()
|
||||
except (EOFError, KeyboardInterrupt):
|
||||
if should_print:
|
||||
console.print()
|
||||
user = "/exit"
|
||||
turn = {"type": "turn", "text": user}
|
||||
else:
|
||||
turn = None
|
||||
|
||||
if is_mlx_distributed:
|
||||
try:
|
||||
turn = chat_backend.share_distributed_object(turn, timeout = None)
|
||||
if not should_print:
|
||||
_drain_available_stdin()
|
||||
except Exception as exc:
|
||||
if should_print:
|
||||
err.print(
|
||||
f"\n(error sharing chat turn: {exc})",
|
||||
style = "red",
|
||||
markup = False,
|
||||
)
|
||||
raise typer.Exit(code = 1)
|
||||
if not turn:
|
||||
continue
|
||||
user = str(turn.get("text", "")).strip()
|
||||
|
||||
if not user:
|
||||
continue
|
||||
|
|
@ -301,55 +380,69 @@ def chat(
|
|||
break
|
||||
if user == "/reset":
|
||||
messages = []
|
||||
console.print("(history cleared)", style = "bright_black")
|
||||
if should_print:
|
||||
console.print("(history cleared)", style = "bright_black")
|
||||
continue
|
||||
if user == "/think":
|
||||
show_thinking = not show_thinking
|
||||
state = "on" if show_thinking else "off"
|
||||
console.print(f"(thinking {state})", style = "bright_black")
|
||||
if should_print:
|
||||
state = "on" if show_thinking else "off"
|
||||
console.print(f"(thinking {state})", style = "bright_black")
|
||||
continue
|
||||
if user == "/compare":
|
||||
if compare_blocked:
|
||||
console.print(f"(compare unavailable: {compare_blocked})", style = "yellow")
|
||||
if should_print:
|
||||
console.print(f"(compare unavailable: {compare_blocked})", style = "yellow")
|
||||
continue
|
||||
if not compare_mode and dual_compare and not load_base_for_compare():
|
||||
continue
|
||||
compare_mode = not compare_mode
|
||||
state = "on" if compare_mode else "off"
|
||||
console.print(f"(compare {state})", style = "bright_black")
|
||||
if should_print:
|
||||
state = "on" if compare_mode else "off"
|
||||
console.print(f"(compare {state})", style = "bright_black")
|
||||
continue
|
||||
if user in ("/help", "/?"):
|
||||
console.print(_HELP, style = "bright_black")
|
||||
if should_print:
|
||||
console.print(_HELP, style = "bright_black")
|
||||
continue
|
||||
|
||||
messages.append({"role": "user", "content": user})
|
||||
|
||||
try:
|
||||
if compare_mode:
|
||||
console.print("(comparing base vs tuned…)", style = "bright_black")
|
||||
if should_print:
|
||||
console.print("(comparing base vs tuned…)", style = "bright_black")
|
||||
if dual_compare:
|
||||
base_text = collect_stream(generate(backend = base_backend), show_thinking)
|
||||
tuned_text = collect_stream(generate(), show_thinking)
|
||||
else:
|
||||
base_text = collect_stream(generate(use_adapter = False), show_thinking)
|
||||
tuned_text = collect_stream(generate(use_adapter = True), show_thinking)
|
||||
console.print()
|
||||
render_columns(
|
||||
"base", base_text, f"{name} (tuned)", tuned_text, console = console
|
||||
)
|
||||
if should_print:
|
||||
console.print()
|
||||
render_columns(
|
||||
"base", base_text, f"{name} (tuned)", tuned_text, console = console
|
||||
)
|
||||
# History continues as the tuned model; base is just the reference.
|
||||
answer = tuned_text
|
||||
else:
|
||||
console.print(assistant_label)
|
||||
answer = stream_markdown(generate(), show_thinking, console = console)
|
||||
if should_print:
|
||||
console.print(assistant_label)
|
||||
answer = stream_markdown(generate(), show_thinking, console = console)
|
||||
else:
|
||||
answer = collect_stream(generate(), show_thinking)
|
||||
except KeyboardInterrupt:
|
||||
# Ctrl-C aborts this answer only; drop the unanswered turn.
|
||||
console.print("\n(interrupted)", style = "bright_black")
|
||||
if should_print:
|
||||
console.print("\n(interrupted)", style = "bright_black")
|
||||
messages.pop()
|
||||
continue
|
||||
except Exception as exc:
|
||||
err.print(f"\n(error: {exc})", style = "red", markup = False)
|
||||
if should_print:
|
||||
err.print(f"\n(error: {exc})", style = "red", markup = False)
|
||||
messages.pop()
|
||||
if is_mlx_distributed:
|
||||
raise typer.Exit(code = 1)
|
||||
continue
|
||||
|
||||
messages.append(
|
||||
|
|
@ -359,4 +452,5 @@ def chat(
|
|||
chat_backend.close()
|
||||
if base_backend is not None:
|
||||
base_backend.close()
|
||||
err.print("\nBye.", style = "bright_black")
|
||||
if should_print:
|
||||
err.print("\nBye.", style = "bright_black")
|
||||
|
|
|
|||
|
|
@ -6,9 +6,13 @@ from typing import List, Optional
|
|||
import typer
|
||||
|
||||
from unsloth_cli._inference import (
|
||||
collect_stream,
|
||||
configure_quiet_logging,
|
||||
connect_studio_server,
|
||||
load_chat_backend,
|
||||
mlx_distributed_info,
|
||||
mlx_distributed_uses_mpi,
|
||||
raise_on_streamed_error,
|
||||
stream_to_stdout,
|
||||
)
|
||||
|
||||
|
|
@ -36,7 +40,8 @@ def inference(
|
|||
"--tensor-parallel/--no-tensor-parallel",
|
||||
help = (
|
||||
"Split a GGUF across GPUs by tensor (--split-mode tensor) instead "
|
||||
"of by layer. Ignored for non-GGUF models."
|
||||
"of by layer. Under non-MPI mlx.launch, select MLX tensor "
|
||||
"parallel mode instead of pipeline mode."
|
||||
),
|
||||
),
|
||||
llama_extra_args: Optional[List[str]] = typer.Option(
|
||||
|
|
@ -69,8 +74,20 @@ def inference(
|
|||
if not verbose:
|
||||
configure_quiet_logging()
|
||||
|
||||
# A running Studio server keeps the model warm between runs, which is
|
||||
# exactly what a one-shot command wants.
|
||||
is_mlx_distributed, rank, _world_size = mlx_distributed_info()
|
||||
if is_mlx_distributed and mlx_distributed_uses_mpi():
|
||||
if rank == 0:
|
||||
typer.echo(
|
||||
"Distributed `unsloth inference` with MPI is not supported by "
|
||||
"the current subprocess backend. Use a non-MPI MLX launcher "
|
||||
"backend such as ring/JACCL for now.",
|
||||
err = True,
|
||||
)
|
||||
raise typer.Exit(code = 1)
|
||||
|
||||
# A running Studio server keeps the model warm between runs. Under
|
||||
# mlx.launch, every rank must enter the local MLX path instead of rank 0
|
||||
# alone talking to a server.
|
||||
load_opts = dict(
|
||||
hf_token = hf_token,
|
||||
max_seq_length = max_seq_length,
|
||||
|
|
@ -78,7 +95,9 @@ def inference(
|
|||
tensor_parallel = tensor_parallel,
|
||||
llama_extra_args = llama_extra_args,
|
||||
)
|
||||
chat_backend = None if no_server else connect_studio_server(model, **load_opts)
|
||||
chat_backend = (
|
||||
None if (no_server or is_mlx_distributed) else connect_studio_server(model, **load_opts)
|
||||
)
|
||||
if chat_backend is None:
|
||||
chat_backend = load_chat_backend(model, **load_opts)
|
||||
try:
|
||||
|
|
@ -92,7 +111,23 @@ def inference(
|
|||
repetition_penalty = repetition_penalty,
|
||||
enable_thinking = think,
|
||||
)
|
||||
typer.echo("Assistant:")
|
||||
stream_to_stdout(stream, show_thinking = think)
|
||||
if is_mlx_distributed:
|
||||
stream = raise_on_streamed_error(stream)
|
||||
if rank == 0:
|
||||
typer.echo("Assistant:")
|
||||
try:
|
||||
stream_to_stdout(stream, show_thinking = think)
|
||||
except RuntimeError as exc:
|
||||
if not is_mlx_distributed:
|
||||
raise
|
||||
typer.echo(f"Error: {exc}", err = True)
|
||||
raise typer.Exit(code = 1)
|
||||
else:
|
||||
try:
|
||||
collect_stream(stream, show_thinking = think)
|
||||
except RuntimeError:
|
||||
if not is_mlx_distributed:
|
||||
raise
|
||||
raise typer.Exit(code = 1)
|
||||
finally:
|
||||
chat_backend.close()
|
||||
|
|
|
|||
|
|
@ -26,6 +26,8 @@ from unsloth_cli._inference import (
|
|||
ChatBackend,
|
||||
HttpChatBackend,
|
||||
collect_stream,
|
||||
mlx_distributed_info,
|
||||
mlx_distributed_uses_mpi,
|
||||
render_columns,
|
||||
visible_text,
|
||||
)
|
||||
|
|
@ -39,6 +41,16 @@ class _FakeConfig:
|
|||
path = None
|
||||
|
||||
|
||||
_EXPECTED_MPI_ENV_PAIRS = [
|
||||
("OMPI_COMM_WORLD_RANK", "OMPI_COMM_WORLD_SIZE"),
|
||||
("PMI_RANK", "PMI_SIZE"),
|
||||
("PMIX_RANK", "PMIX_SIZE"),
|
||||
("MPI_RANK", "MPI_WORLD_SIZE"),
|
||||
("MV2_COMM_WORLD_RANK", "MV2_COMM_WORLD_SIZE"),
|
||||
]
|
||||
_IGNORED_DISTRIBUTED_ENV_PAIRS = [("SLURM_PROCID", "SLURM_NTASKS")]
|
||||
|
||||
|
||||
def _chat_app():
|
||||
cli = typer.Typer()
|
||||
cli.command()(chatmod.chat)
|
||||
|
|
@ -53,6 +65,39 @@ def _inference_app():
|
|||
return cli
|
||||
|
||||
|
||||
def _clear_mlx_distributed_env(monkeypatch):
|
||||
for name in (
|
||||
"MLX_RANK",
|
||||
"MLX_HOSTFILE",
|
||||
"MLX_WORLD_SIZE",
|
||||
"MLX_IBV_DEVICES",
|
||||
"MLX_JACCL_COORDINATOR",
|
||||
"NCCL_HOST_IP",
|
||||
"NCCL_PORT",
|
||||
*(rank for rank, _size in _EXPECTED_MPI_ENV_PAIRS + _IGNORED_DISTRIBUTED_ENV_PAIRS),
|
||||
*(size for _rank, size in _EXPECTED_MPI_ENV_PAIRS + _IGNORED_DISTRIBUTED_ENV_PAIRS),
|
||||
):
|
||||
monkeypatch.delenv(name, raising = False)
|
||||
|
||||
|
||||
def _set_mlx_nccl_env(
|
||||
monkeypatch,
|
||||
*,
|
||||
rank: str = "0",
|
||||
size: str = "2",
|
||||
):
|
||||
monkeypatch.setenv("MLX_RANK", rank)
|
||||
monkeypatch.setenv("MLX_WORLD_SIZE", size)
|
||||
monkeypatch.setenv("NCCL_HOST_IP", "127.0.0.1")
|
||||
monkeypatch.setenv("NCCL_PORT", "12345")
|
||||
|
||||
|
||||
@pytest.fixture(autouse = True)
|
||||
def _isolate_mlx_distributed_env(monkeypatch):
|
||||
_clear_mlx_distributed_env(monkeypatch)
|
||||
monkeypatch.delenv("HF_TOKEN", raising = False)
|
||||
|
||||
|
||||
def test_visible_text_passthrough_when_shown():
|
||||
text = "<think>reasoning</think>answer"
|
||||
assert visible_text(text, show_thinking = True) == text
|
||||
|
|
@ -101,6 +146,46 @@ def test_inference_exposes_gguf_runtime_options():
|
|||
assert "--llama-extra-arg" in (getattr(extra, "param_decls", None) or [])
|
||||
|
||||
|
||||
def test_mlx_distributed_info_reads_launch_env(monkeypatch, tmp_path):
|
||||
_clear_mlx_distributed_env(monkeypatch)
|
||||
assert mlx_distributed_info() == (False, 0, None)
|
||||
assert mlx_distributed_uses_mpi() is False
|
||||
|
||||
monkeypatch.setenv("MLX_RANK", "1")
|
||||
monkeypatch.setenv("MLX_WORLD_SIZE", "2")
|
||||
assert mlx_distributed_info() == (False, 0, None)
|
||||
monkeypatch.setenv("NCCL_HOST_IP", "127.0.0.1")
|
||||
monkeypatch.setenv("NCCL_PORT", "12345")
|
||||
assert mlx_distributed_info() == (True, 1, 2)
|
||||
assert mlx_distributed_uses_mpi() is False
|
||||
|
||||
_clear_mlx_distributed_env(monkeypatch)
|
||||
ring_hostfile = tmp_path / "ring.json"
|
||||
ring_hostfile.write_text('[["127.0.0.1:5000"], ["127.0.0.1:5001"]]\n')
|
||||
monkeypatch.setenv("MLX_RANK", "0")
|
||||
monkeypatch.setenv("MLX_HOSTFILE", str(ring_hostfile))
|
||||
assert mlx_distributed_info() == (True, 0, 2)
|
||||
assert mlx_distributed_uses_mpi() is False
|
||||
|
||||
_clear_mlx_distributed_env(monkeypatch)
|
||||
monkeypatch.setenv("MLX_RANK", "1")
|
||||
monkeypatch.setenv("MLX_IBV_DEVICES", '[["node-a"], ["node-b"]]')
|
||||
monkeypatch.setenv("MLX_JACCL_COORDINATOR", "node-a:12345")
|
||||
assert mlx_distributed_info() == (True, 1, 2)
|
||||
assert mlx_distributed_uses_mpi() is False
|
||||
|
||||
_clear_mlx_distributed_env(monkeypatch)
|
||||
monkeypatch.setenv("OMPI_COMM_WORLD_RANK", "1")
|
||||
monkeypatch.setenv("OMPI_COMM_WORLD_SIZE", "2")
|
||||
assert mlx_distributed_info() == (True, 1, 2)
|
||||
assert mlx_distributed_uses_mpi() is True
|
||||
|
||||
_clear_mlx_distributed_env(monkeypatch)
|
||||
monkeypatch.setenv("MLX_RANK", "bad")
|
||||
monkeypatch.setenv("MLX_WORLD_SIZE", "-3")
|
||||
assert mlx_distributed_info() == (False, 0, None)
|
||||
|
||||
|
||||
def test_chat_command_is_registered_with_options():
|
||||
params = inspect.signature(chatmod.chat).parameters
|
||||
assert "model" in params
|
||||
|
|
@ -751,7 +836,6 @@ def test_chat_server_mode_compare_loads_base_locally(monkeypatch):
|
|||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "(compare on)" in result.output
|
||||
# Only the base model loaded locally, on its own private backend.
|
||||
assert base_loads == [("fake/base", True)]
|
||||
assert streamed == ["base", "tuned"]
|
||||
assert set(closed) == {"http", "base"}
|
||||
|
|
@ -785,6 +869,248 @@ def test_chat_compare_on_mlx_loads_base_model_side_by_side(monkeypatch):
|
|||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert loads == [("tuned-run", False), ("fake/base", True)]
|
||||
# Both models answered the turn, via plain generation (no adapter toggle).
|
||||
assert ("base", None) in streamed and ("tuned", None) in streamed
|
||||
assert set(closed) == {"tuned", "base"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("chunk_kind", "expected_exit"),
|
||||
[
|
||||
("answer", 0),
|
||||
("model_text_error", 0),
|
||||
("real_error", 1),
|
||||
],
|
||||
)
|
||||
def test_inference_under_mlx_launch_handles_stream(monkeypatch, chunk_kind, expected_exit):
|
||||
from unsloth_cli.commands import inference as infermod
|
||||
from unsloth_cli._inference import ensure_studio_backend_path
|
||||
|
||||
ensure_studio_backend_path()
|
||||
from core.inference.orchestrator import GenStreamError
|
||||
|
||||
if chunk_kind == "answer":
|
||||
chunks = ["answer"]
|
||||
elif chunk_kind == "model_text_error":
|
||||
# Model output whose visible text starts with "Error:" must not abort.
|
||||
chunks = ["Error: printed by the model, not a backend failure"]
|
||||
else:
|
||||
chunks = [GenStreamError("Error: generation failed")]
|
||||
|
||||
loads, closed = [], []
|
||||
|
||||
class _FakeBackend:
|
||||
def stream(self, messages, **kwargs):
|
||||
return iter(chunks)
|
||||
|
||||
def close(self):
|
||||
closed.append(True)
|
||||
|
||||
_set_mlx_nccl_env(monkeypatch, rank = "0")
|
||||
monkeypatch.setattr(
|
||||
infermod,
|
||||
"connect_studio_server",
|
||||
lambda *_a, **_k: (_ for _ in ()).throw(AssertionError("server disabled")),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
infermod,
|
||||
"load_chat_backend",
|
||||
lambda model, **kwargs: (loads.append((model, kwargs)), _FakeBackend())[1],
|
||||
)
|
||||
|
||||
result = CliRunner().invoke(
|
||||
_inference_app(),
|
||||
["fake-model", "hello", "--tensor-parallel"],
|
||||
)
|
||||
|
||||
assert result.exit_code == expected_exit, result.output
|
||||
assert loads[0][1]["tensor_parallel"] is True
|
||||
if chunk_kind == "real_error":
|
||||
assert "generation failed" in result.output
|
||||
|
||||
|
||||
def test_chat_under_mlx_launch_nonzero_rank_drains_stdin(monkeypatch):
|
||||
drains, closed = [], []
|
||||
turns = iter(
|
||||
[
|
||||
{"type": "turn", "text": "hi"},
|
||||
{"type": "turn", "text": "/exit"},
|
||||
]
|
||||
)
|
||||
|
||||
class _FakeChatBackend:
|
||||
def share_distributed_object(
|
||||
self,
|
||||
obj,
|
||||
*,
|
||||
timeout = 300.0,
|
||||
):
|
||||
assert obj is None
|
||||
return next(turns)
|
||||
|
||||
def stream(self, messages, **kwargs):
|
||||
return iter(["hidden"])
|
||||
|
||||
def close(self):
|
||||
closed.append(True)
|
||||
|
||||
_set_mlx_nccl_env(monkeypatch, rank = "1")
|
||||
monkeypatch.setattr(chatmod, "resolve_model_config", lambda *a, **k: _FakeConfig())
|
||||
monkeypatch.setattr(
|
||||
chatmod,
|
||||
"connect_studio_server",
|
||||
lambda *_a, **_k: (_ for _ in ()).throw(AssertionError("server disabled")),
|
||||
)
|
||||
monkeypatch.setattr(chatmod, "load_chat_backend", lambda *a, **k: _FakeChatBackend())
|
||||
monkeypatch.setattr(chatmod, "_compare_needs_second_model", lambda: False)
|
||||
monkeypatch.setattr(chatmod, "_drain_available_stdin", lambda: drains.append(True))
|
||||
|
||||
result = CliRunner().invoke(_chat_app(), ["fake-model"], input = "hi\n/exit\n")
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "Chatting with" not in result.output
|
||||
assert drains == [True, True]
|
||||
assert closed == [True]
|
||||
|
||||
|
||||
def test_chat_under_mlx_launch_rank0_bypasses_studio_and_prints(monkeypatch):
|
||||
loads, shares, closed = [], [], []
|
||||
|
||||
class _FakeChatBackend:
|
||||
def share_distributed_object(
|
||||
self,
|
||||
obj,
|
||||
*,
|
||||
timeout = 300.0,
|
||||
):
|
||||
shares.append((obj, timeout))
|
||||
return obj
|
||||
|
||||
def stream(self, messages, **kwargs):
|
||||
return iter(["hello"])
|
||||
|
||||
def close(self):
|
||||
closed.append(True)
|
||||
|
||||
_set_mlx_nccl_env(monkeypatch, rank = "0")
|
||||
monkeypatch.setattr(chatmod, "resolve_model_config", lambda *a, **k: _FakeConfig())
|
||||
monkeypatch.setattr(
|
||||
chatmod,
|
||||
"connect_studio_server",
|
||||
lambda *_a, **_k: (_ for _ in ()).throw(AssertionError("server disabled")),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
chatmod,
|
||||
"load_chat_backend",
|
||||
lambda model, **kwargs: (loads.append((model, kwargs)), _FakeChatBackend())[1],
|
||||
)
|
||||
monkeypatch.setattr(chatmod, "_compare_needs_second_model", lambda: False)
|
||||
|
||||
result = CliRunner().invoke(
|
||||
_chat_app(),
|
||||
["fake-model", "--tensor-parallel"],
|
||||
input = "hi\n/exit\n",
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "Chatting with fake-model" in result.output
|
||||
assert "hello" in result.output
|
||||
assert loads and loads[0][0] == "fake-model"
|
||||
assert loads[0][1]["tensor_parallel"] is True
|
||||
assert shares == [
|
||||
({"type": "turn", "text": "hi"}, None),
|
||||
({"type": "turn", "text": "/exit"}, None),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("stream_error", "expected_exit"),
|
||||
[("exception", 1), ("chunk", 1), ("model_text", 0)],
|
||||
)
|
||||
def test_chat_under_mlx_launch_exits_on_generation_error(monkeypatch, stream_error, expected_exit):
|
||||
from unsloth_cli._inference import ensure_studio_backend_path
|
||||
|
||||
ensure_studio_backend_path()
|
||||
from core.inference.orchestrator import GenStreamError
|
||||
|
||||
closed = []
|
||||
|
||||
class _FakeChatBackend:
|
||||
def share_distributed_object(
|
||||
self,
|
||||
obj,
|
||||
*,
|
||||
timeout = 300.0,
|
||||
):
|
||||
return obj
|
||||
|
||||
def stream(self, messages, **kwargs):
|
||||
if stream_error == "exception":
|
||||
raise RuntimeError("generation failed")
|
||||
if stream_error == "model_text":
|
||||
# Plain model text starting with "Error:" must not abort the run.
|
||||
return iter(["Error: printed by the model"])
|
||||
return iter([GenStreamError("Error: generation failed")])
|
||||
|
||||
def close(self):
|
||||
closed.append(True)
|
||||
|
||||
_set_mlx_nccl_env(monkeypatch, rank = "0")
|
||||
monkeypatch.setattr(chatmod, "resolve_model_config", lambda *a, **k: _FakeConfig())
|
||||
monkeypatch.setattr(
|
||||
chatmod,
|
||||
"connect_studio_server",
|
||||
lambda *_a, **_k: (_ for _ in ()).throw(AssertionError("server disabled")),
|
||||
)
|
||||
monkeypatch.setattr(chatmod, "load_chat_backend", lambda *a, **k: _FakeChatBackend())
|
||||
monkeypatch.setattr(chatmod, "_compare_needs_second_model", lambda: False)
|
||||
|
||||
result = CliRunner().invoke(_chat_app(), ["fake-model"], input = "hi\n/exit\n")
|
||||
|
||||
assert result.exit_code == expected_exit
|
||||
if expected_exit:
|
||||
assert "generation failed" in result.output
|
||||
assert closed == [True]
|
||||
|
||||
|
||||
def test_load_chat_backend_forwards_mlx_distributed_options(monkeypatch):
|
||||
import unsloth_cli._inference as inference
|
||||
|
||||
calls = []
|
||||
|
||||
class _FakeBackend:
|
||||
def load_model(self, **kwargs):
|
||||
calls.append(kwargs)
|
||||
return True
|
||||
|
||||
class _FakeModelConfig:
|
||||
is_gguf = False
|
||||
|
||||
@classmethod
|
||||
def from_identifier(cls, **_kwargs):
|
||||
return cls()
|
||||
|
||||
fake_backend = _FakeBackend()
|
||||
fake_inference = types.ModuleType("core.inference")
|
||||
fake_inference.get_inference_backend = lambda: fake_backend
|
||||
fake_utils = types.ModuleType("utils")
|
||||
fake_utils.__path__ = []
|
||||
fake_models = types.ModuleType("utils.models")
|
||||
fake_models.ModelConfig = _FakeModelConfig
|
||||
|
||||
_set_mlx_nccl_env(monkeypatch, rank = "0")
|
||||
monkeypatch.setitem(sys.modules, "core", types.ModuleType("core"))
|
||||
monkeypatch.setitem(sys.modules, "core.inference", fake_inference)
|
||||
monkeypatch.setitem(sys.modules, "utils", fake_utils)
|
||||
monkeypatch.setitem(sys.modules, "utils.models", fake_models)
|
||||
monkeypatch.setattr(inference, "ensure_studio_backend_path", lambda: None)
|
||||
|
||||
inference.load_chat_backend(
|
||||
"fake-model",
|
||||
hf_token = None,
|
||||
max_seq_length = 2048,
|
||||
load_in_4bit = True,
|
||||
tensor_parallel = True,
|
||||
)
|
||||
|
||||
assert calls[0]["tensor_parallel"] is True
|
||||
assert calls[0]["mlx_distributed"] is True
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue