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