From 2a6abe2ff5c643ee853f2e0ef632b1e80d06918a Mon Sep 17 00:00:00 2001 From: Long Yixing Date: Wed, 8 Jul 2026 18:25:39 +0800 Subject: [PATCH] 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 Co-authored-by: danielhanchen Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../backend/core/inference/mlx_inference.py | 171 ++++++--- studio/backend/core/inference/orchestrator.py | 106 +++++- studio/backend/core/inference/worker.py | 123 ++++++- studio/backend/routes/inference.py | 46 +++ .../tests/test_mlx_inference_backend.py | 125 +++++++ unsloth_cli/_inference.py | 217 ++++++++++-- unsloth_cli/commands/chat.py | 178 +++++++--- unsloth_cli/commands/inference.py | 47 ++- unsloth_cli/tests/test_inference_chat.py | 330 +++++++++++++++++- 9 files changed, 1188 insertions(+), 155 deletions(-) diff --git a/studio/backend/core/inference/mlx_inference.py b/studio/backend/core/inference/mlx_inference.py index 62d268e15f..d84baa278d 100644 --- a/studio/backend/core/inference/mlx_inference.py +++ b/studio/backend/core/inference/mlx_inference.py @@ -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): diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index cf5d24c367..4fa0d3ed26 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -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( diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index d4b102e422..05dee39283 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -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) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index ce755ae1fb..d9901a5b2e 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -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) diff --git a/studio/backend/tests/test_mlx_inference_backend.py b/studio/backend/tests/test_mlx_inference_backend.py index ac4088fb25..55a3198a6b 100644 --- a/studio/backend/tests/test_mlx_inference_backend.py +++ b/studio/backend/tests/test_mlx_inference_backend.py @@ -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 diff --git a/unsloth_cli/_inference.py b/unsloth_cli/_inference.py index e8a3414d73..c3b188710e 100644 --- a/unsloth_cli/_inference.py +++ b/unsloth_cli/_inference.py @@ -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_BLOCK = re.compile(rf"{re.escape(_THINK_OPEN)}.*?", 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) diff --git a/unsloth_cli/commands/chat.py b/unsloth_cli/commands/chat.py index d3bfbbf96b..bc4a72f36c 100644 --- a/unsloth_cli/commands/chat.py +++ b/unsloth_cli/commands/chat.py @@ -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") diff --git a/unsloth_cli/commands/inference.py b/unsloth_cli/commands/inference.py index 1401fa9e9e..84a126163e 100644 --- a/unsloth_cli/commands/inference.py +++ b/unsloth_cli/commands/inference.py @@ -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() diff --git a/unsloth_cli/tests/test_inference_chat.py b/unsloth_cli/tests/test_inference_chat.py index 013b7c5a81..56633408fb 100644 --- a/unsloth_cli/tests/test_inference_chat.py +++ b/unsloth_cli/tests/test_inference_chat.py @@ -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 = "reasoninganswer" 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