ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add GGUF --tensor-parallel CLI option (#6561)

Browse source 
---------

Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
This commit is contained in:
Avaya Aggarwal 2026-06-27 00:00:10 +05:30 • committed by GitHub
commit cb274484a6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 494 additions and 24 deletions
Download patch file Download diff file Expand all files Collapse all files

3
studio/backend/models/inference.py
View file

@ -106,8 +106,7 @@ class LoadRequest(BaseModel):
"Extra arguments forwarded verbatim to llama-server for GGUF models. "
"One token per list entry, e.g. ['--top-k', '20', '--seed', '42']. "
"Studio-managed flags (model identity, port, context length, GPU placement, "
"auth, --flash-attn, --no-context-shift, --jinja) are rejected. Ignored for "
"non-GGUF models."
"auth, UI/server mode) are rejected. Ignored for non-GGUF models."
),
)

108
unsloth_cli/_inference.py
View file

@ -3,11 +3,12 @@
"""Model loading and streaming shared by `inference` and `chat`."""
import asyncio
import os
import re
import sys
from pathlib import Path
from typing import Optional
from typing import List, Optional
import typer
@ -211,28 +212,65 @@ def resolve_model_config(model: str, *, hf_token: Optional[str]):
return model_config
def _load_gguf_backend(model_config, *, hf_token, max_seq_length):
def _validate_llama_extra_args_or_exit(llama_extra_args: Optional[List[str]]) -> list[str]:
from core.inference.llama_server_args import validate_extra_args
try:
return validate_extra_args(llama_extra_args)
except ValueError as exc:
typer.echo(f"Error: {exc}", err = True)
raise typer.Exit(code = 1)
def _load_gguf_backend(
model_config,
*,
hf_token,
max_seq_length,
tensor_parallel: bool = False,
llama_extra_args: Optional[List[str]] = None,
):
ensure_studio_backend_path()
from core.inference.llama_cpp import LlamaCppBackend
from core.inference.tensor_fallback import load_with_tensor_fallback
llama_backend = LlamaCppBackend()
extra_args = _validate_llama_extra_args_or_exit(llama_extra_args)
common = dict(
hf_variant = model_config.gguf_variant,
model_identifier = model_config.identifier,
is_vision = model_config.is_vision,
n_ctx = max_seq_length,
)
if model_config.gguf_hf_repo:
loaded = llama_backend.load_model(
hf_repo = model_config.gguf_hf_repo, hf_token = hf_token, **common
async def _attempt_gguf_load(
requested_tensor_parallel: bool, attempt_extra_args: Optional[List[str]]
) -> bool:
attempt_common = dict(
common,
tensor_parallel = requested_tensor_parallel,
extra_args = attempt_extra_args,
)
else:
loaded = llama_backend.load_model(
if model_config.gguf_hf_repo:
return llama_backend.load_model(
hf_repo = model_config.gguf_hf_repo,
hf_token = hf_token,
**attempt_common,
)
return llama_backend.load_model(
gguf_path = model_config.gguf_file,
mmproj_path = model_config.gguf_mmproj_file,
mtp_draft_path = model_config.gguf_mtp_file,
**common,
**attempt_common,
)
loaded = asyncio.run(
load_with_tensor_fallback(
_attempt_gguf_load,
requested_tensor = tensor_parallel,
extra_args = extra_args,
label = model_config.identifier,
)
)
if not loaded:
typer.echo("Model load failed", err = True)
raise typer.Exit(code = 1)
@ -245,6 +283,8 @@ def load_chat_backend(
hf_token: Optional[str],
max_seq_length: int,
load_in_4bit: bool,
tensor_parallel: bool = False,
llama_extra_args: Optional[List[str]] = None,
model_config = None,
fresh_backend: bool = False,
):
@ -259,7 +299,13 @@ def load_chat_backend(
typer.echo(f"Loading {model}", err = True)
if model_config.is_gguf:
return _load_gguf_backend(model_config, hf_token = hf_token, max_seq_length = max_seq_length)
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()
@ -447,18 +493,31 @@ class HttpChatBackend:
# No redirects: this carries a bearer token (see urlopen_no_redirect).
return urlopen_no_redirect(request, timeout = timeout)
def ensure_loaded(self, model: str, *, hf_token, max_seq_length, load_in_4bit) -> None:
def ensure_loaded(
self,
model: str,
*,
hf_token,
max_seq_length,
load_in_4bit,
tensor_parallel: bool = False,
llama_extra_args: Optional[List[str]] = None,
) -> None:
typer.echo(f"Loading {model} on the Studio server", err = True)
payload = {
"model_path": model,
"hf_token": hf_token,
"max_seq_length": max_seq_length,
"load_in_4bit": load_in_4bit,
"tensor_parallel": tensor_parallel,
}
if llama_extra_args:
payload["llama_extra_args"] = llama_extra_args
try:
self._request(
"POST",
"/api/inference/load",
{
"model_path": model,
"hf_token": hf_token,
"max_seq_length": max_seq_length,
"load_in_4bit": load_in_4bit,
},
payload,
).close()
except Exception as exc:
typer.echo(f"Model load failed: {exc}", err = True)
@ -538,7 +597,15 @@ class HttpChatBackend:
pass
def connect_studio_server(model: str, *, hf_token, max_seq_length, load_in_4bit):
def connect_studio_server(
model: str,
*,
hf_token,
max_seq_length,
load_in_4bit,
tensor_parallel: bool = False,
llama_extra_args: Optional[List[str]] = None,
):
"""Backend on a running Studio server, or None (caller loads locally)."""
base_url = find_studio_server()
if not base_url:
@ -576,6 +643,11 @@ def connect_studio_server(model: str, *, hf_token, max_seq_length, load_in_4bit)
return _refuse("couldn't self-issue a Studio token (is Studio set up here?).")
backend = HttpChatBackend(base_url, token)
backend.ensure_loaded(
model, hf_token = hf_token, max_seq_length = max_seq_length, load_in_4bit = load_in_4bit
model,
hf_token = hf_token,
max_seq_length = max_seq_length,
load_in_4bit = load_in_4bit,
tensor_parallel = tensor_parallel,
llama_extra_args = llama_extra_args,
)
return backend

26
unsloth_cli/commands/chat.py
View file

@ -1,7 +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
from typing import Optional
from typing import List, Optional
import typer
from rich.console import Console
@ -153,6 +153,22 @@ def chat(
),
max_seq_length: int = typer.Option(4096, "--max-seq-length"),
load_in_4bit: bool = typer.Option(True, "--load-in-4bit/--no-load-in-4bit"),
tensor_parallel: bool = typer.Option(
False,
"--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."
),
),
llama_extra_args: Optional[List[str]] = typer.Option(
None,
"--llama-extra-arg",
help = (
"Extra llama-server arg for GGUF models. Repeat for multiple "
"tokens, e.g. --llama-extra-arg=--top-k --llama-extra-arg 20."
),
),
think: bool = typer.Option(
False,
"--think/--no-think",
@ -190,7 +206,13 @@ def chat(
err.print(f"--compare unavailable: {compare_blocked}", style = "red", markup = False)
raise typer.Exit(code = 1)
load_opts = dict(hf_token = hf_token, max_seq_length = max_seq_length, load_in_4bit = load_in_4bit)
load_opts = dict(
hf_token = hf_token,
max_seq_length = max_seq_length,
load_in_4bit = load_in_4bit,
tensor_parallel = tensor_parallel,
llama_extra_args = llama_extra_args,
)
# 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)

26
unsloth_cli/commands/inference.py
View file

@ -1,7 +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
from typing import Optional
from typing import List, Optional
import typer
@ -31,6 +31,22 @@ def inference(
),
max_seq_length: int = typer.Option(2048, "--max-seq-length"),
load_in_4bit: bool = typer.Option(True, "--load-in-4bit/--no-load-in-4bit"),
tensor_parallel: bool = typer.Option(
False,
"--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."
),
),
llama_extra_args: Optional[List[str]] = typer.Option(
None,
"--llama-extra-arg",
help = (
"Extra llama-server arg for GGUF models. Repeat for multiple "
"tokens, e.g. --llama-extra-arg=--top-k --llama-extra-arg 20."
),
),
think: bool = typer.Option(
False,
"--think/--no-think",
@ -55,7 +71,13 @@ def inference(
# A running Studio server keeps the model warm between runs, which is
# exactly what a one-shot command wants.
load_opts = dict(hf_token = hf_token, max_seq_length = max_seq_length, load_in_4bit = load_in_4bit)
load_opts = dict(
hf_token = hf_token,
max_seq_length = max_seq_length,
load_in_4bit = load_in_4bit,
tensor_parallel = tensor_parallel,
llama_extra_args = llama_extra_args,
)
chat_backend = None if no_server else connect_studio_server(model, **load_opts)
if chat_backend is None:
chat_backend = load_chat_backend(model, **load_opts)

355
unsloth_cli/tests/test_inference_chat.py
View file

@ -9,6 +9,7 @@ import inspect
import sys
import types
from pathlib import Path
from types import SimpleNamespace
_REPO_ROOT = Path(__file__).resolve().parents[2]
if str(_REPO_ROOT) not in sys.path:
@ -16,6 +17,7 @@ if str(_REPO_ROOT) not in sys.path:
import typer
import pytest
from rich.console import Console
from typer.testing import CliRunner
@ -43,6 +45,14 @@ def _chat_app():
return cli
def _inference_app():
from unsloth_cli.commands.inference import inference
cli = typer.Typer()
cli.command()(inference)
return cli
def test_visible_text_passthrough_when_shown():
text = "<think>reasoning</think>answer"
assert visible_text(text, show_thinking = True) == text
@ -81,6 +91,16 @@ def test_inference_think_defaults_off():
assert "--think/--no-think" in (getattr(opt, "param_decls", None) or [])
def test_inference_exposes_gguf_runtime_options():
from unsloth_cli.commands.inference import inference
tensor = _option(inference, "tensor_parallel")
assert "--tensor-parallel/--no-tensor-parallel" in (getattr(tensor, "param_decls", None) or [])
extra = _option(inference, "llama_extra_args")
assert "--llama-extra-arg" in (getattr(extra, "param_decls", None) or [])
def test_chat_command_is_registered_with_options():
params = inspect.signature(chatmod.chat).parameters
assert "model" in params
@ -94,6 +114,12 @@ def test_chat_command_is_registered_with_options():
verbose = _option(chatmod.chat, "verbose")
assert {"--verbose", "-v"} <= set(getattr(verbose, "param_decls", None) or [])
tensor = _option(chatmod.chat, "tensor_parallel")
assert "--tensor-parallel/--no-tensor-parallel" in (getattr(tensor, "param_decls", None) or [])
extra = _option(chatmod.chat, "llama_extra_args")
assert "--llama-extra-arg" in (getattr(extra, "param_decls", None) or [])
class _FakeBackend:
def __init__(self):
@ -324,6 +350,243 @@ def test_http_backend_streams_cumulative_text(monkeypatch):
assert out == ["He", "Hello"]
def test_http_backend_load_forwards_gguf_runtime_options(monkeypatch):
backend = HttpChatBackend("http://localhost:8888", "token")
requests = []
class _OK:
def close(self):
pass
def fake_request(
method,
path,
payload = None,
timeout = None,
):
requests.append((method, path, payload, timeout))
return _OK()
monkeypatch.setattr(backend, "_request", fake_request)
backend.ensure_loaded(
"org/model-GGUF",
hf_token = "hf_x",
max_seq_length = 8192,
load_in_4bit = False,
tensor_parallel = True,
llama_extra_args = ["--top-k", "20"],
)
assert requests == [
(
"POST",
"/api/inference/load",
{
"model_path": "org/model-GGUF",
"hf_token": "hf_x",
"max_seq_length": 8192,
"load_in_4bit": False,
"tensor_parallel": True,
"llama_extra_args": ["--top-k", "20"],
},
None,
)
]
def test_http_backend_load_sends_explicit_false_tensor_parallel(monkeypatch):
backend = HttpChatBackend("http://localhost:8888", "token")
requests = []
class _OK:
def close(self):
pass
monkeypatch.setattr(
backend,
"_request",
lambda method, path, payload = None, timeout = None: (
requests.append((method, path, payload, timeout)),
_OK(),
)[1],
)
backend.ensure_loaded(
"org/model-GGUF",
hf_token = None,
max_seq_length = 4096,
load_in_4bit = True,
tensor_parallel = False,
)
assert requests[0][2]["tensor_parallel"] is False
def test_load_gguf_backend_forwards_local_runtime_options(monkeypatch):
import unsloth_cli._inference as inference
calls = []
class _FakeLlamaCppBackend:
def load_model(self, **kwargs):
calls.append(kwargs)
return True
fake_llama_cpp = types.ModuleType("core.inference.llama_cpp")
fake_llama_cpp.LlamaCppBackend = _FakeLlamaCppBackend
fake_args = types.ModuleType("core.inference.llama_server_args")
fake_args.validate_extra_args = lambda args: list(args or [])
fake_tensor_fallback = types.ModuleType("core.inference.tensor_fallback")
async def _passthrough(
attempt_load,
*,
requested_tensor,
extra_args,
label = "",
cancelled = None,
):
return await attempt_load(requested_tensor, extra_args)
fake_tensor_fallback.load_with_tensor_fallback = _passthrough
monkeypatch.setitem(sys.modules, "core", types.ModuleType("core"))
monkeypatch.setitem(sys.modules, "core.inference", types.ModuleType("core.inference"))
monkeypatch.setitem(sys.modules, "core.inference.llama_cpp", fake_llama_cpp)
monkeypatch.setitem(sys.modules, "core.inference.llama_server_args", fake_args)
monkeypatch.setitem(sys.modules, "core.inference.tensor_fallback", fake_tensor_fallback)
monkeypatch.setattr(inference, "ensure_studio_backend_path", lambda: None)
config = SimpleNamespace(
gguf_variant = "Q4_K_M",
identifier = "org/model-GGUF",
is_vision = False,
gguf_hf_repo = "org/model-GGUF",
)
backend = inference._load_gguf_backend(
config,
hf_token = "hf_x",
max_seq_length = 8192,
tensor_parallel = True,
llama_extra_args = ["--top-k", "20"],
)
assert isinstance(backend, ChatBackend)
assert calls == [
{
"hf_repo": "org/model-GGUF",
"hf_token": "hf_x",
"hf_variant": "Q4_K_M",
"model_identifier": "org/model-GGUF",
"is_vision": False,
"n_ctx": 8192,
"tensor_parallel": True,
"extra_args": ["--top-k", "20"],
}
]
def test_load_gguf_backend_exits_cleanly_on_invalid_extra_args(monkeypatch):
import unsloth_cli._inference as inference
fake_llama_cpp = types.ModuleType("core.inference.llama_cpp")
fake_llama_cpp.LlamaCppBackend = object
fake_args = types.ModuleType("core.inference.llama_server_args")
def _raise(_args):
raise ValueError("llama-server flag '--model' is managed by Unsloth Studio")
fake_args.validate_extra_args = _raise
fake_tensor_fallback = types.ModuleType("core.inference.tensor_fallback")
fake_tensor_fallback.load_with_tensor_fallback = None
monkeypatch.setitem(sys.modules, "core", types.ModuleType("core"))
monkeypatch.setitem(sys.modules, "core.inference", types.ModuleType("core.inference"))
monkeypatch.setitem(sys.modules, "core.inference.llama_cpp", fake_llama_cpp)
monkeypatch.setitem(sys.modules, "core.inference.llama_server_args", fake_args)
monkeypatch.setitem(sys.modules, "core.inference.tensor_fallback", fake_tensor_fallback)
monkeypatch.setattr(inference, "ensure_studio_backend_path", lambda: None)
config = SimpleNamespace(
gguf_variant = "Q4_K_M",
identifier = "org/model-GGUF",
is_vision = False,
gguf_hf_repo = "org/model-GGUF",
)
with pytest.raises(typer.Exit) as excinfo:
inference._load_gguf_backend(
config,
hf_token = "hf_x",
max_seq_length = 8192,
llama_extra_args = ["--model"],
)
assert excinfo.value.exit_code == 1
def test_load_gguf_backend_uses_tensor_fallback(monkeypatch):
import unsloth_cli._inference as inference
calls = []
fallback_calls = []
class _FakeLlamaCppBackend:
def load_model(self, **kwargs):
calls.append(kwargs)
return kwargs["tensor_parallel"] is False
fake_llama_cpp = types.ModuleType("core.inference.llama_cpp")
fake_llama_cpp.LlamaCppBackend = _FakeLlamaCppBackend
fake_args = types.ModuleType("core.inference.llama_server_args")
fake_args.validate_extra_args = lambda args: list(args or [])
fake_tensor_fallback = types.ModuleType("core.inference.tensor_fallback")
async def _fallback(
attempt_load,
*,
requested_tensor,
extra_args,
label = "",
cancelled = None,
):
fallback_calls.append((requested_tensor, extra_args, label))
ok = await attempt_load(requested_tensor, extra_args)
if ok:
return True
return await attempt_load(False, ["--split-mode", "layer"])
fake_tensor_fallback.load_with_tensor_fallback = _fallback
monkeypatch.setitem(sys.modules, "core", types.ModuleType("core"))
monkeypatch.setitem(sys.modules, "core.inference", types.ModuleType("core.inference"))
monkeypatch.setitem(sys.modules, "core.inference.llama_cpp", fake_llama_cpp)
monkeypatch.setitem(sys.modules, "core.inference.llama_server_args", fake_args)
monkeypatch.setitem(sys.modules, "core.inference.tensor_fallback", fake_tensor_fallback)
monkeypatch.setattr(inference, "ensure_studio_backend_path", lambda: None)
config = SimpleNamespace(
gguf_variant = "Q4_K_M",
identifier = "org/model-GGUF",
is_vision = False,
gguf_hf_repo = "org/model-GGUF",
)
backend = inference._load_gguf_backend(
config,
hf_token = "hf_x",
max_seq_length = 8192,
tensor_parallel = True,
)
assert isinstance(backend, ChatBackend)
assert fallback_calls == [(True, [], "org/model-GGUF")]
assert [call["tensor_parallel"] for call in calls] == [True, False]
assert calls[1]["extra_args"] == ["--split-mode", "layer"]
def test_http_backend_merges_emoji_split_across_deltas(monkeypatch):
backend = HttpChatBackend("http://localhost:8888", "token")
response = _FakeSSEResponse(
@ -365,6 +628,98 @@ def test_chat_prefers_running_studio_server(monkeypatch):
assert closed == ["http"]
def test_chat_forwards_gguf_runtime_options_to_loader(monkeypatch):
loads = []
class _FakeHttpBackend:
def close(self):
pass
monkeypatch.setattr(chatmod, "resolve_model_config", lambda *a, **k: _FakeConfig())
monkeypatch.setattr(
chatmod,
"connect_studio_server",
lambda model, **kwargs: (loads.append((model, kwargs)), _FakeHttpBackend())[1],
)
monkeypatch.setattr(chatmod, "load_chat_backend", lambda *a, **k: None)
monkeypatch.setattr(chatmod, "_compare_needs_second_model", lambda: False)
result = CliRunner().invoke(
_chat_app(),
[
"fake-model",
"--tensor-parallel",
"--llama-extra-arg=--top-k",
"--llama-extra-arg",
"20",
],
input = "/exit\n",
)
assert result.exit_code == 0, result.output
assert loads == [
(
"fake-model",
{
"hf_token": None,
"max_seq_length": 4096,
"load_in_4bit": True,
"tensor_parallel": True,
"llama_extra_args": ["--top-k", "20"],
},
)
]
def test_inference_forwards_gguf_runtime_options_to_loader(monkeypatch):
from unsloth_cli.commands import inference as infermod
loads, streams, closed = [], [], []
class _FakeBackend:
def stream(self, messages, **kwargs):
streams.append((messages, kwargs))
return iter(["answer"])
def close(self):
closed.append(True)
monkeypatch.setattr(
infermod,
"connect_studio_server",
lambda model, **kwargs: (loads.append((model, kwargs)), _FakeBackend())[1],
)
monkeypatch.setattr(infermod, "load_chat_backend", lambda *a, **k: None)
result = CliRunner().invoke(
_inference_app(),
[
"fake-model",
"hello",
"--tensor-parallel",
"--llama-extra-arg=--top-k",
"--llama-extra-arg",
"20",
],
)
assert result.exit_code == 0, result.output
assert loads == [
(
"fake-model",
{
"hf_token": None,
"max_seq_length": 2048,
"load_in_4bit": True,
"tensor_parallel": True,
"llama_extra_args": ["--top-k", "20"],
},
)
]
assert streams[0][0] == [{"role": "user", "content": "hello"}]
assert closed == [True]
def test_chat_server_mode_compare_loads_base_locally(monkeypatch):
streamed, closed, base_loads = [], [], []