[Feature] Add Muse Glimmer model support (#34262)
Co-authored-by: sglang-bot <232288953+sglang-bot@users.noreply.github.com> Co-authored-by: Brayden Zhong <brayden.zhong@radixark.ai> Co-authored-by: Jimmy Shong <69131491+Jiminator@users.noreply.github.com> Co-authored-by: hnyls2002 <lsyincs@gmail.com> Co-authored-by: Alex Nails <alex.nails@radixark.ai> Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com>
This commit is contained in:
sglang-bot
co-authored by
sglang-bot
Brayden Zhong
Jimmy Shong
hnyls2002
Alex Nails
Liangsheng Yin
sglang-bot
Brayden Zhong
Jimmy Shong
hnyls2002
Alex Nails
Liangsheng Yin
parent
9c1517df4a
commit
fde9ad2531
@@ -0,0 +1,94 @@
|
||||
import unittest
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.kits.eval_accuracy_kit import GSM8KMixin
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=900, stage="nightly", runner_config="1-gpu-large")
|
||||
|
||||
TARGET_MODEL = "meta-models/Muse-Glimmer-30B"
|
||||
DRAFT_MODEL = "meta-models/Muse-Glimmer-30B-assistant"
|
||||
|
||||
|
||||
class TestMuseGlimmerDflashAssistantGSM8K(CustomTestCase, GSM8KMixin):
|
||||
"""GSM8K + DFlash accept-length regression test for the native
|
||||
MuseGlimmerAssistantModel draft (``meta-models/Muse-Glimmer-30B-assistant``).
|
||||
|
||||
This checkpoint loads through sglang's own native ``models/dflash.py`` /
|
||||
``configs/muse_glimmer.py::MuseGlimmerAssistantConfig`` with no extra wheel
|
||||
dependency. Integrating it surfaced two bugs that fail *silently*, never raise, and do not move
|
||||
GSM8K accuracy at temperature 0 (speculative decoding always falls back
|
||||
to the target's own correct token on a draft miss, so a broken draft
|
||||
still produces exactly the target's answers -- just slower):
|
||||
|
||||
1. The vendor's weight names (``encoder.fc.weight`` /
|
||||
``encoder.output_norm_enc.weight``) didn't match what
|
||||
``DFlashDraftModel.load_weights`` expected (``fc.weight`` /
|
||||
``hidden_norm.weight``), so those two tensors silently stayed at
|
||||
random init.
|
||||
2. The vendor's ``target_layer_ids`` are in the HF "output of layer k"
|
||||
convention; ``models/muse_glimmer.py::set_dflash_layers_to_capture``
|
||||
uses ids as-is (Muse Glimmer's own draft configs carry llama.cpp's
|
||||
layer-*input* convention), so every captured layer was off by one.
|
||||
|
||||
Both together collapsed real (non-simulated) accept_length to ~1.00 at
|
||||
``--speculative-dflash-block-size 5`` -- effectively no speculation, only
|
||||
the mandatory bonus token -- with GSM8K accuracy unaffected throughout.
|
||||
``gsm8k_accept_length_thres`` is therefore the actual regression guard
|
||||
here; the accuracy threshold alone would not have caught this.
|
||||
|
||||
Measured after both fixes, same block size, same target: median
|
||||
accept_length 3.12 (this draft) vs 3.09 (our own GGUF-converted draft,
|
||||
the previous ground truth) across batch sizes 1/4/8 at 512in/256out --
|
||||
statistically indistinguishable. 2.5 leaves headroom below that for
|
||||
workload variance while sitting well above the ~1.0 broken-state value.
|
||||
"""
|
||||
|
||||
model = TARGET_MODEL
|
||||
gsm8k_backend = (
|
||||
"sgl_eval" # chat completions API, not /generate or raw /completions
|
||||
)
|
||||
gsm8k_score_threshold = 0.85
|
||||
gsm8k_num_examples = 200
|
||||
gsm8k_accept_length_thres = 2.5
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=[
|
||||
"--reasoning-parser",
|
||||
"muse",
|
||||
"--tool-call-parser",
|
||||
"muse",
|
||||
"--language-model-only",
|
||||
"--speculative-algorithm",
|
||||
"DFLASH",
|
||||
"--speculative-draft-model-path",
|
||||
DRAFT_MODEL,
|
||||
"--speculative-draft-load-format",
|
||||
"auto",
|
||||
"--speculative-dflash-block-size",
|
||||
"5",
|
||||
"--mem-fraction-static",
|
||||
"0.85",
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
if hasattr(cls, "process") and cls.process:
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -311,6 +311,80 @@ class ReasoningRequestForwardingTestCase(unittest.TestCase):
|
||||
self.assertFalse(parser_cls.call_args.kwargs["force_reasoning"])
|
||||
|
||||
|
||||
class SkipSpecialTokensForwardingTestCase(CustomTestCase):
|
||||
"""The skip_special_tokens override from _process_messages must reach the
|
||||
engine sampling params; muse's channel markers die in detok otherwise."""
|
||||
|
||||
def _create_responses_sampling_params(self, serving):
|
||||
serving.default_chat_template_kwargs = None
|
||||
rendered = MessageProcessingResult(
|
||||
prompt="prompt",
|
||||
prompt_ids=[1, 2, 3],
|
||||
image_data=None,
|
||||
audio_data=None,
|
||||
video_data=None,
|
||||
modalities=[],
|
||||
stop=[],
|
||||
)
|
||||
captured = {}
|
||||
|
||||
async def fake_generate(
|
||||
request_id,
|
||||
request_prompt,
|
||||
adapted_request,
|
||||
sampling_params,
|
||||
context,
|
||||
**kwargs,
|
||||
):
|
||||
captured["sampling_params"] = sampling_params
|
||||
context.append_output(
|
||||
{
|
||||
"text": "done",
|
||||
"meta_info": {
|
||||
"prompt_tokens": 3,
|
||||
"completion_tokens": 1,
|
||||
"cached_tokens": 0,
|
||||
},
|
||||
}
|
||||
)
|
||||
yield context
|
||||
|
||||
serving._generate_with_builtin_tools = fake_generate
|
||||
request = ResponsesRequest(
|
||||
model="x",
|
||||
input="answer",
|
||||
request_id="resp_skip_special",
|
||||
store=False,
|
||||
)
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
serving, "_apply_conversation_template", return_value=rendered
|
||||
),
|
||||
patch(
|
||||
"sglang.srt.entrypoints.openai.serving_responses.ReasoningParser"
|
||||
) as parser_cls,
|
||||
):
|
||||
parser_cls.return_value.parse_non_stream.return_value = (None, "done")
|
||||
response = asyncio.run(serving.create_responses(request))
|
||||
|
||||
self.assertEqual(response.status, "completed")
|
||||
return captured["sampling_params"]
|
||||
|
||||
def test_marker_preserving_parser_disables_skip_special_tokens(self):
|
||||
serving = make_serving()
|
||||
serving.reasoning_parser = "muse"
|
||||
params = self._create_responses_sampling_params(serving)
|
||||
self.assertFalse(params["skip_special_tokens"])
|
||||
|
||||
def test_default_parser_keeps_skip_special_tokens(self):
|
||||
serving = make_serving()
|
||||
params = self._create_responses_sampling_params(serving)
|
||||
# The chat request's True is a synthesized default (ResponsesRequest has
|
||||
# no such field), so leave it unset for --preferred-sampling-params.
|
||||
self.assertNotIn("skip_special_tokens", params)
|
||||
|
||||
|
||||
class InputItemNormalizationTestCase(CustomTestCase):
|
||||
def test_function_call_becomes_assistant_tool_call(self):
|
||||
normalized = OpenAIServingResponses._normalize_response_message_for_chat(
|
||||
|
||||
@@ -167,6 +167,7 @@ class NonHarmonyStreamTestCase(CustomTestCase):
|
||||
parser_cls.return_value.parse_stream_chunk.side_effect = (
|
||||
fake_parse_stream_chunk
|
||||
)
|
||||
parser_cls.return_value.parse_stream_end.return_value = ("", [])
|
||||
fixture = StreamFixture(serving, request)
|
||||
events = fixture.run(chunks)
|
||||
|
||||
@@ -178,6 +179,35 @@ class NonHarmonyStreamTestCase(CustomTestCase):
|
||||
self.assertEqual(output[1]["name"], "get_weather")
|
||||
self.assertEqual(output[2]["content"][0]["text"], "It's sunny.")
|
||||
|
||||
def test_reasoning_parser_flushed_at_stream_end(self):
|
||||
"""Bug regression: the stream loop never drained text the reasoning
|
||||
parser held back as a possible marker prefix, so a response whose text
|
||||
genuinely ends with e.g. "<|e" lost that tail on /v1/responses (chat
|
||||
flushes via parse_stream_end; responses did not)."""
|
||||
serving = make_serving()
|
||||
serving.reasoning_parser = "muse"
|
||||
serving.tool_call_parser = None
|
||||
|
||||
request = ResponsesRequest(model="x", input="hi", stream=True, store=False)
|
||||
text = (
|
||||
" to=self<|message|>think<|eom|>"
|
||||
"<|start|>assistant to=user<|message|>Answer<|e"
|
||||
)
|
||||
fixture = StreamFixture(serving, request)
|
||||
events = fixture.run(
|
||||
[
|
||||
engine_chunk(text[:30], 4),
|
||||
engine_chunk(text, 9, finish=True),
|
||||
]
|
||||
)
|
||||
|
||||
streamed = "".join(
|
||||
p["delta"]
|
||||
for ev, p in zip(event_types(events), event_payloads(events))
|
||||
if ev == "response.output_text.delta"
|
||||
)
|
||||
self.assertEqual(streamed, "Answer<|e")
|
||||
|
||||
|
||||
class MultiToolCallStreamingOrderTestCase(CustomTestCase):
|
||||
"""The wire order of message / function_call items across tool-call deltas."""
|
||||
|
||||
@@ -0,0 +1,432 @@
|
||||
"""Unit tests for the Muse Glimmer ATEM tool-call detector — no server, no model loading.
|
||||
|
||||
The expectations here are pinned to the checkpoint's own ``response_template``
|
||||
(``MUSE_GLIMMER_RESPONSE_SCHEMA`` in ``tokenizer_config.json``) and to the vendor's
|
||||
reference parser, with particular attention to channel scoping: an
|
||||
``<atem:invoke>`` that only appears inside a reasoning block or a final answer
|
||||
must never become a real tool call.
|
||||
"""
|
||||
|
||||
import json
|
||||
|
||||
from sglang.srt.entrypoints.openai.protocol import Function, Tool
|
||||
from sglang.srt.function_call.function_call_parser import FunctionCallParser
|
||||
from sglang.srt.function_call.muse_glimmer_detector import MuseGlimmerDetector
|
||||
from sglang.srt.parser.reasoning_parser import ReasoningParser
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(1.0, "base-a-test-cpu")
|
||||
|
||||
DOUBLED = "get_weather.get_weather"
|
||||
|
||||
|
||||
def atem(name: str, **params: str) -> str:
|
||||
body = "".join(
|
||||
f'<atem:parameter name="{k}">{v}</atem:parameter>\n' for k, v in params.items()
|
||||
)
|
||||
return (
|
||||
f'<atem:function_calls>\n<atem:invoke name="{name}">\n{body}'
|
||||
f"</atem:invoke>\n</atem:function_calls>"
|
||||
)
|
||||
|
||||
|
||||
class TestMuseGlimmerDetector(CustomTestCase):
|
||||
def setUp(self):
|
||||
self.tools = [
|
||||
Tool(
|
||||
type="function",
|
||||
function=Function(
|
||||
name="get_weather",
|
||||
description="Get weather",
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {"city": {"type": "string"}},
|
||||
"required": ["city"],
|
||||
},
|
||||
),
|
||||
),
|
||||
]
|
||||
|
||||
# ---- helpers ----------------------------------------------------------
|
||||
|
||||
def parse(self, text):
|
||||
"""Non-streaming parse -> (normal_text, [(name, args), ...])."""
|
||||
result = MuseGlimmerDetector().detect_and_parse(text, self.tools)
|
||||
return result.normal_text, [
|
||||
(c.name, json.loads(c.parameters)) for c in result.calls if c.name
|
||||
]
|
||||
|
||||
def parse_streaming(self, text, chunk_size):
|
||||
detector = MuseGlimmerDetector()
|
||||
normal, calls = [], []
|
||||
for i in range(0, len(text), chunk_size):
|
||||
result = detector.parse_streaming_increment(
|
||||
text[i : i + chunk_size], self.tools
|
||||
)
|
||||
normal.append(result.normal_text)
|
||||
calls.extend(
|
||||
(c.name, json.loads(c.parameters)) for c in result.calls if c.name
|
||||
)
|
||||
return "".join(normal), calls
|
||||
|
||||
def assert_streaming_matches(self, text):
|
||||
"""Streaming must agree with one-shot parsing at every chunk boundary."""
|
||||
expected = self.parse(text)
|
||||
for chunk_size in (1, 2, 3, 5, 7, 13, 29, 100):
|
||||
self.assertEqual(
|
||||
self.parse_streaming(text, chunk_size),
|
||||
expected,
|
||||
f"streaming diverged at chunk_size={chunk_size}",
|
||||
)
|
||||
|
||||
# ---- tool extraction --------------------------------------------------
|
||||
|
||||
def test_single_tool_call(self):
|
||||
text = (
|
||||
f" to=self<|message|>Need weather.<|eom|>"
|
||||
f"<|start|>assistant to={DOUBLED}<|message|>{atem(DOUBLED, city='Paris')}"
|
||||
)
|
||||
normal, calls = self.parse(text)
|
||||
self.assertEqual(calls, [("get_weather", {"city": "Paris"})])
|
||||
self.assertEqual(normal, "Need weather.")
|
||||
self.assert_streaming_matches(text)
|
||||
|
||||
def test_parallel_tool_calls(self):
|
||||
text = (
|
||||
f" to=self<|message|>Two cities.<|eom|>"
|
||||
f"<|start|>assistant to={DOUBLED}<|message|>"
|
||||
f"{atem(DOUBLED, city='Paris')}<|eom|>"
|
||||
f"<|start|>assistant to={DOUBLED}<|message|>{atem(DOUBLED, city='Tokyo')}"
|
||||
)
|
||||
_, calls = self.parse(text)
|
||||
self.assertEqual(
|
||||
calls,
|
||||
[("get_weather", {"city": "Paris"}), ("get_weather", {"city": "Tokyo"})],
|
||||
)
|
||||
self.assert_streaming_matches(text)
|
||||
|
||||
def test_tool_call_then_final_answer(self):
|
||||
text = (
|
||||
f" to=self<|message|>r<|eom|>"
|
||||
f"<|start|>assistant to={DOUBLED}<|message|>"
|
||||
f"{atem(DOUBLED, city='Paris')}<|eom|>"
|
||||
f"<|start|>assistant to=user<|message|>It is sunny."
|
||||
)
|
||||
normal, calls = self.parse(text)
|
||||
self.assertEqual(calls, [("get_weather", {"city": "Paris"})])
|
||||
self.assertIn("It is sunny.", normal)
|
||||
self.assert_streaming_matches(text)
|
||||
|
||||
def test_namespaced_name_passes_through(self):
|
||||
tools = [
|
||||
Tool(
|
||||
type="function",
|
||||
function=Function(
|
||||
name="weather.get",
|
||||
description="d",
|
||||
parameters={"type": "object", "properties": {}},
|
||||
),
|
||||
)
|
||||
]
|
||||
text = (
|
||||
f" to=self<|message|>r<|eom|><|start|>assistant to=weather.get<|message|>"
|
||||
f"{atem('weather.get', city='Paris')}"
|
||||
)
|
||||
result = MuseGlimmerDetector().detect_and_parse(text, tools)
|
||||
self.assertEqual([c.name for c in result.calls], ["weather.get"])
|
||||
|
||||
def test_parameter_value_typing(self):
|
||||
"""``allow_non_json: True`` — JSON literals decode, bare strings do not."""
|
||||
invoke = (
|
||||
'<atem:function_calls>\n<atem:invoke name="get_weather">\n'
|
||||
'<atem:parameter name="s">hello world</atem:parameter>\n'
|
||||
'<atem:parameter name="i">42</atem:parameter>\n'
|
||||
'<atem:parameter name="b">true</atem:parameter>\n'
|
||||
'<atem:parameter name="n">null</atem:parameter>\n'
|
||||
'<atem:parameter name="o">{"a": 1}</atem:parameter>\n'
|
||||
'<atem:parameter name="l">[1, 2]</atem:parameter>\n'
|
||||
"</atem:invoke>\n</atem:function_calls>"
|
||||
)
|
||||
text = f"<|start|>assistant to=get_weather<|message|>{invoke}"
|
||||
_, calls = self.parse(text)
|
||||
self.assertEqual(
|
||||
calls[0][1],
|
||||
{
|
||||
"s": "hello world",
|
||||
"i": 42,
|
||||
"b": True,
|
||||
"n": None,
|
||||
"o": {"a": 1},
|
||||
"l": [1, 2],
|
||||
},
|
||||
)
|
||||
|
||||
def test_multiline_parameter_value(self):
|
||||
value = 'line1\nline2\n"quoted"\n'
|
||||
text = (
|
||||
f"<|start|>assistant to=get_weather<|message|>"
|
||||
f'<atem:function_calls>\n<atem:invoke name="get_weather">\n'
|
||||
f'<atem:parameter name="code">{value}</atem:parameter>\n'
|
||||
f"</atem:invoke>\n</atem:function_calls>"
|
||||
)
|
||||
_, calls = self.parse(text)
|
||||
self.assertEqual(calls[0][1], {"code": value})
|
||||
self.assert_streaming_matches(text)
|
||||
|
||||
# ---- channel scoping (safety) -----------------------------------------
|
||||
|
||||
def test_invoke_inside_reasoning_is_not_a_call(self):
|
||||
text = (
|
||||
f" to=self<|message|>Maybe I call {atem(DOUBLED, city='X')} — no.<|eom|>"
|
||||
f"<|start|>assistant to=user<|message|>I will not call it."
|
||||
)
|
||||
_, calls = self.parse(text)
|
||||
self.assertEqual(calls, [])
|
||||
self.assert_streaming_matches(text)
|
||||
|
||||
def test_invoke_inside_final_answer_is_not_a_call(self):
|
||||
text = (
|
||||
f" to=self<|message|>r<|eom|><|start|>assistant to=user<|message|>"
|
||||
f"You would write:\n{atem(DOUBLED, city='X')}"
|
||||
)
|
||||
_, calls = self.parse(text)
|
||||
self.assertEqual(calls, [])
|
||||
self.assert_streaming_matches(text)
|
||||
|
||||
def test_invoke_inside_truncated_reasoning_is_not_a_call(self):
|
||||
"""Generation cut mid-CoT leaves no closing ``<|eom|>`` to anchor on."""
|
||||
text = f" to=self<|message|>I could call {atem(DOUBLED, city='X')} but"
|
||||
_, calls = self.parse(text)
|
||||
self.assertEqual(calls, [])
|
||||
self.assert_streaming_matches(text)
|
||||
|
||||
def test_truncated_tool_channel_drops_partial_invoke(self):
|
||||
"""A token cap mid-invoke must not fabricate a call from partial
|
||||
arguments, and ATEM scaffolding must not leak into content."""
|
||||
text = (
|
||||
f" to=self<|message|>Need weather.<|eom|>"
|
||||
f"<|start|>assistant to={DOUBLED}<|message|>"
|
||||
f'<atem:function_calls>\n<atem:invoke name="{DOUBLED}">\n'
|
||||
f'<atem:parameter name="city">Par'
|
||||
)
|
||||
normal, calls = self.parse(text)
|
||||
self.assertEqual(calls, [])
|
||||
self.assertEqual(normal, "Need weather.")
|
||||
self.assert_streaming_matches(text)
|
||||
|
||||
def test_prose_opening_with_to_equals_does_not_stall(self):
|
||||
"""A bare ``to=...`` header opens the stream, so prose that happens to
|
||||
start the same way is ambiguous. Mis-reading it as a header parks the
|
||||
parser waiting for a ``<|message|>`` that never arrives and strands the
|
||||
whole response in the buffer."""
|
||||
for text in ("to=x is the syntax.", "to= is an assignment", "to=a<b is false"):
|
||||
detector = MuseGlimmerDetector()
|
||||
streamed = "".join(
|
||||
detector.parse_streaming_increment(ch, self.tools).normal_text
|
||||
for ch in text
|
||||
)
|
||||
self.assertEqual(streamed, text)
|
||||
self.assertEqual(detector._buffer, "", "text stranded in the buffer")
|
||||
|
||||
def test_unframed_atem_is_content_not_a_call(self):
|
||||
"""Deliberately stricter than the vendor; see ``_is_tool_channel``."""
|
||||
text = atem(DOUBLED, city="Paris")
|
||||
normal, calls = self.parse(text)
|
||||
self.assertEqual(calls, [])
|
||||
self.assertEqual(normal, text)
|
||||
|
||||
# ---- integration with the reasoning parser ----------------------------
|
||||
|
||||
def test_pipeline_with_reasoning_parser(self):
|
||||
"""The real serving order: reasoning parser first, then this detector."""
|
||||
raw = (
|
||||
f" to=self<|message|>Need weather.<|eom|>"
|
||||
f"<|start|>assistant to={DOUBLED}<|message|>"
|
||||
f"{atem(DOUBLED, city='Paris')}<|eom|>"
|
||||
f"<|start|>assistant to=user<|message|>It is sunny in Paris."
|
||||
)
|
||||
reasoning, remainder = ReasoningParser("muse").parse_non_stream(raw)
|
||||
content, calls = FunctionCallParser(self.tools, "muse").parse_non_stream(
|
||||
remainder
|
||||
)
|
||||
self.assertEqual(reasoning, "Need weather.")
|
||||
self.assertEqual(content, "It is sunny in Paris.")
|
||||
self.assertEqual(
|
||||
[(c.name, json.loads(c.parameters)) for c in calls],
|
||||
[("get_weather", {"city": "Paris"})],
|
||||
)
|
||||
|
||||
def test_pipeline_quoted_invoke_stays_content(self):
|
||||
"""A quoted ATEM block must survive as text, not become a call."""
|
||||
raw = (
|
||||
f" to=self<|message|>r<|eom|><|start|>assistant to=user<|message|>"
|
||||
f"Example:\n{atem(DOUBLED, city='X')}"
|
||||
)
|
||||
_, remainder = ReasoningParser("muse").parse_non_stream(raw)
|
||||
content, calls = FunctionCallParser(self.tools, "muse").parse_non_stream(
|
||||
remainder
|
||||
)
|
||||
self.assertEqual(calls, [])
|
||||
self.assertIn("<atem:invoke", content)
|
||||
|
||||
def pipeline_stream(self, raw, chunk_size, tool_call_parser_active=True):
|
||||
"""The real streaming order: reasoning deltas feed the tool parser."""
|
||||
rp = ReasoningParser("muse", tool_call_parser_active=tool_call_parser_active)
|
||||
fp = FunctionCallParser(self.tools, "muse")
|
||||
reasoning_parts, content_parts, calls = [], [], []
|
||||
chunks = [raw[i : i + chunk_size] for i in range(0, len(raw), chunk_size)]
|
||||
for i, chunk in enumerate(chunks):
|
||||
reasoning, normal = rp.parse_stream_chunk(chunk)
|
||||
if i == len(chunks) - 1:
|
||||
end_reasoning, end_normal = rp.parse_stream_end()
|
||||
reasoning = (reasoning or "") + (end_reasoning or "")
|
||||
normal = (normal or "") + (end_normal or "")
|
||||
if reasoning:
|
||||
reasoning_parts.append(reasoning)
|
||||
if normal:
|
||||
content, chunk_calls = fp.parse_stream_chunk(normal)
|
||||
content_parts.append(content)
|
||||
calls.extend(
|
||||
(c.name, json.loads(c.parameters)) for c in chunk_calls if c.name
|
||||
)
|
||||
end_content, end_calls = fp.parse_stream_end()
|
||||
content_parts.append(end_content)
|
||||
calls.extend((c.name, json.loads(c.parameters)) for c in end_calls if c.name)
|
||||
return "".join(reasoning_parts), "".join(content_parts), calls
|
||||
|
||||
def test_pipeline_quoted_header_in_answer_is_not_a_call(self):
|
||||
"""The answer keeps its channel framing on the way to this detector, so
|
||||
a quoted header inside it stays inside the ``to=user`` body. Goes red if
|
||||
the reasoning parser unwraps the answer before the tool parser runs, or
|
||||
if ``detect_and_parse`` regains a scan-ahead ``<|start|>`` search."""
|
||||
for quoted_at in ("after prose:\n", ""):
|
||||
raw = (
|
||||
f" to=self<|message|>r<|eom|><|start|>assistant to=user<|message|>"
|
||||
f"{quoted_at}<|start|>assistant to={DOUBLED}<|message|>"
|
||||
f"{atem(DOUBLED, city='X')}"
|
||||
)
|
||||
_, remainder = ReasoningParser(
|
||||
"muse", tool_call_parser_active=True
|
||||
).parse_non_stream(raw)
|
||||
content, calls = FunctionCallParser(self.tools, "muse").parse_non_stream(
|
||||
remainder
|
||||
)
|
||||
self.assertEqual(
|
||||
calls, [], f"quoted header parsed as a call ({quoted_at!r})"
|
||||
)
|
||||
self.assertIn("<atem:invoke", content)
|
||||
for chunk_size in (1, 7, 100):
|
||||
_, s_content, s_calls = self.pipeline_stream(raw, chunk_size)
|
||||
self.assertEqual(s_calls, [])
|
||||
self.assertIn("<atem:invoke", s_content)
|
||||
|
||||
def test_pipeline_preamble_before_tool_call_streams(self):
|
||||
"""A real ``to=user`` message may precede the tool channel; its
|
||||
terminator is what re-arms the header state. Goes red if the hand-off
|
||||
drops the ``to=user`` terminator again."""
|
||||
raw = (
|
||||
f" to=self<|message|>r<|eom|><|start|>assistant to=user<|message|>"
|
||||
f"Let me check.<|eom|><|start|>assistant to={DOUBLED}<|message|>"
|
||||
f"{atem(DOUBLED, city='Paris')}"
|
||||
)
|
||||
want_calls = [("get_weather", {"city": "Paris"})]
|
||||
_, remainder = ReasoningParser(
|
||||
"muse", tool_call_parser_active=True
|
||||
).parse_non_stream(raw)
|
||||
content, calls = FunctionCallParser(self.tools, "muse").parse_non_stream(
|
||||
remainder
|
||||
)
|
||||
self.assertEqual(
|
||||
[(c.name, json.loads(c.parameters)) for c in calls], want_calls
|
||||
)
|
||||
self.assertIn("Let me check.", content)
|
||||
for chunk_size in (1, 7, 100):
|
||||
_, s_content, s_calls = self.pipeline_stream(raw, chunk_size)
|
||||
self.assertEqual(s_calls, want_calls, f"chunk_size={chunk_size}")
|
||||
self.assertIn("Let me check.", s_content)
|
||||
|
||||
def test_pipeline_plain_answer_stays_clean_without_tool_parse(self):
|
||||
"""Serving skips the tool detector when ``has_tool_call()`` is false, so
|
||||
a turn with no ATEM block must come out of the reasoning parser already
|
||||
unwrapped. Goes red if framing is preserved unconditionally."""
|
||||
raw = (
|
||||
" to=self<|message|>r<|eom|>"
|
||||
"<|start|>assistant to=user<|message|>Hello there."
|
||||
)
|
||||
reasoning, remainder = ReasoningParser(
|
||||
"muse", tool_call_parser_active=True
|
||||
).parse_non_stream(raw)
|
||||
self.assertFalse(
|
||||
FunctionCallParser(self.tools, "muse").has_tool_call(remainder)
|
||||
)
|
||||
self.assertEqual(reasoning, "r")
|
||||
self.assertEqual(remainder, "Hello there.")
|
||||
|
||||
def test_reasoning_finish_flushes_unframed_text(self):
|
||||
"""An unframed turn never emits ``<|message|>``, so nothing leaves the
|
||||
buffer until the end-of-stream flush. Goes red if the Muse Glimmer reasoning
|
||||
detector loses its ``finish()`` override."""
|
||||
rp = ReasoningParser("muse")
|
||||
_, streamed = rp.parse_stream_chunk("Just plain text.")
|
||||
_, flushed = rp.parse_stream_end()
|
||||
self.assertEqual((streamed or "") + (flushed or ""), "Just plain text.")
|
||||
|
||||
def test_interleaved_reasoning_blocks_join_with_newline(self):
|
||||
"""A turn may reason, call a tool, then reason again; the reference
|
||||
schema joins the blocks with a newline. Goes red if the reasoning
|
||||
detector concatenates the bodies directly, gluing two thoughts into
|
||||
one word ("...thoughtsecond...")."""
|
||||
raw = (
|
||||
f" to=self<|message|>first thought<|eom|>"
|
||||
f"<|start|>assistant to={DOUBLED}<|message|>"
|
||||
f"{atem(DOUBLED, city='Paris')}<|eom|>"
|
||||
f"<|start|>assistant to=self<|message|>second thought<|eom|>"
|
||||
f"<|start|>assistant to=user<|message|>done"
|
||||
)
|
||||
reasoning, _ = ReasoningParser(
|
||||
"muse", tool_call_parser_active=True
|
||||
).parse_non_stream(raw)
|
||||
self.assertEqual(reasoning, "first thought\nsecond thought")
|
||||
for chunk_size in (1, 7, 100):
|
||||
s_reasoning, s_content, s_calls = self.pipeline_stream(raw, chunk_size)
|
||||
self.assertEqual(s_reasoning, "first thought\nsecond thought")
|
||||
self.assertEqual(s_calls, [("get_weather", {"city": "Paris"})])
|
||||
self.assertIn("done", s_content)
|
||||
|
||||
def test_stream_end_flushes_partial_marker(self):
|
||||
"""An answer ending in a marker prefix (``<|st``) is held back while
|
||||
streaming in case it grows into ``<|start|>``; the stream's end proves
|
||||
it never will. Goes red if the tool parser loses its stream-end flush
|
||||
(``parse_stream_end`` / detector ``finish``)."""
|
||||
raw = (
|
||||
" to=self<|message|>r<|eom|>"
|
||||
"<|start|>assistant to=user<|message|>answer<|st"
|
||||
)
|
||||
for chunk_size in (1, 7, 100):
|
||||
_, content, calls = self.pipeline_stream(raw, chunk_size)
|
||||
self.assertEqual(calls, [])
|
||||
self.assertEqual(content, "answer<|st", f"chunk_size={chunk_size}")
|
||||
|
||||
def test_whitespace_before_header_is_tolerated(self):
|
||||
"""The model may put whitespace between ``<|eom|>`` and the next
|
||||
``<|start|>``; it travels with the header text. Goes red if the header
|
||||
state requires ``<|start|>`` at exactly the first byte again."""
|
||||
text = (
|
||||
f" to=self<|message|>r<|eom|>\n<|start|>assistant to={DOUBLED}"
|
||||
f"<|message|>{atem(DOUBLED, city='Paris')}"
|
||||
)
|
||||
_, calls = self.parse(text)
|
||||
self.assertEqual(calls, [("get_weather", {"city": "Paris"})])
|
||||
self.assert_streaming_matches(text)
|
||||
|
||||
def test_registered_in_parser_enum(self):
|
||||
self.assertIs(
|
||||
FunctionCallParser.ToolCallParserEnum["muse"], MuseGlimmerDetector
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import unittest
|
||||
|
||||
unittest.main()
|
||||
@@ -0,0 +1,186 @@
|
||||
"""Unit tests for the MLX remote-code gate.
|
||||
|
||||
mlx-lm executes ``config.json``'s ``model_file`` unconditionally at load
|
||||
time, so SGLang refuses such checkpoints before any checkpoint Python runs
|
||||
unless the server was started with ``--trust-remote-code``. The refusal
|
||||
tests prove non-execution with a sentinel ``model_file`` whose import would
|
||||
leave an observable marker.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import json
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
from sglang.test.ci.ci_register import register_cpu_ci, register_mlx_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
||||
register_mlx_ci(est_time=5, suite="stage-a-unit-test-mlx")
|
||||
|
||||
_HAS_MLX = (
|
||||
importlib.util.find_spec("mlx") is not None
|
||||
and importlib.util.find_spec("mlx_lm") is not None
|
||||
)
|
||||
|
||||
from sglang.srt.hardware_backend.mlx.remote_code_gate import ( # noqa: E402
|
||||
RemoteCodeGateError,
|
||||
ensure_remote_code_allowed,
|
||||
)
|
||||
|
||||
_SENTINEL = "GATE FAILED: checkpoint python executed"
|
||||
|
||||
|
||||
def _make_checkpoint(tmp: Path, config: dict, *, with_sentinel: bool = True) -> Path:
|
||||
(tmp / "config.json").write_text(json.dumps(config))
|
||||
if with_sentinel:
|
||||
# Importing this file would create marker.txt — the refusal tests
|
||||
# assert it never appears.
|
||||
(tmp / "evil.py").write_text(
|
||||
"from pathlib import Path\n"
|
||||
f"Path(__file__).parent.joinpath('marker.txt').write_text({_SENTINEL!r})\n"
|
||||
)
|
||||
return tmp
|
||||
|
||||
|
||||
class TestEnsureRemoteCodeAllowed(CustomTestCase):
|
||||
def setUp(self):
|
||||
self._tmp = tempfile.TemporaryDirectory()
|
||||
self.dir = Path(self._tmp.name)
|
||||
|
||||
def tearDown(self):
|
||||
self._tmp.cleanup()
|
||||
|
||||
def _assert_sentinel_not_executed(self):
|
||||
self.assertFalse(
|
||||
(self.dir / "marker.txt").exists(),
|
||||
"checkpoint python executed despite gate refusal",
|
||||
)
|
||||
|
||||
def test_refuses_model_file_without_trust(self):
|
||||
_make_checkpoint(
|
||||
self.dir, {"model_type": "muse_glimmer", "model_file": "evil.py"}
|
||||
)
|
||||
with self.assertRaisesRegex(RemoteCodeGateError, "--trust-remote-code"):
|
||||
ensure_remote_code_allowed(self.dir, trust_remote_code=False)
|
||||
self._assert_sentinel_not_executed()
|
||||
|
||||
def test_allows_model_file_with_trust(self):
|
||||
_make_checkpoint(
|
||||
self.dir, {"model_type": "muse_glimmer", "model_file": "evil.py"}
|
||||
)
|
||||
ensure_remote_code_allowed(self.dir, trust_remote_code=True)
|
||||
# The gate itself never imports the file either way.
|
||||
self._assert_sentinel_not_executed()
|
||||
|
||||
def test_builtin_checkpoint_passes_without_trust(self):
|
||||
_make_checkpoint(self.dir, {"model_type": "qwen3"}, with_sentinel=False)
|
||||
ensure_remote_code_allowed(self.dir, trust_remote_code=False)
|
||||
|
||||
def test_missing_config_rejected(self):
|
||||
with self.assertRaisesRegex(RemoteCodeGateError, "no config.json"):
|
||||
ensure_remote_code_allowed(self.dir, trust_remote_code=True)
|
||||
|
||||
def test_malformed_config_rejected(self):
|
||||
(self.dir / "config.json").write_text("{not json")
|
||||
with self.assertRaisesRegex(RemoteCodeGateError, "not valid JSON"):
|
||||
ensure_remote_code_allowed(self.dir, trust_remote_code=True)
|
||||
|
||||
def test_non_object_config_rejected(self):
|
||||
(self.dir / "config.json").write_text('["a", "b"]')
|
||||
with self.assertRaisesRegex(RemoteCodeGateError, "JSON object"):
|
||||
ensure_remote_code_allowed(self.dir, trust_remote_code=True)
|
||||
|
||||
def test_missing_model_file_target_rejected(self):
|
||||
_make_checkpoint(
|
||||
self.dir,
|
||||
{"model_type": "muse_glimmer", "model_file": "nope.py"},
|
||||
with_sentinel=False,
|
||||
)
|
||||
with self.assertRaisesRegex(RemoteCodeGateError, "does not exist"):
|
||||
ensure_remote_code_allowed(self.dir, trust_remote_code=True)
|
||||
|
||||
def test_absolute_model_file_rejected(self):
|
||||
_make_checkpoint(
|
||||
self.dir,
|
||||
{"model_type": "muse_glimmer", "model_file": "/etc/anything.py"},
|
||||
with_sentinel=False,
|
||||
)
|
||||
with self.assertRaisesRegex(RemoteCodeGateError, "relative path"):
|
||||
ensure_remote_code_allowed(self.dir, trust_remote_code=True)
|
||||
|
||||
def test_traversal_model_file_rejected(self):
|
||||
_make_checkpoint(
|
||||
self.dir,
|
||||
{"model_type": "muse_glimmer", "model_file": "../outside.py"},
|
||||
with_sentinel=False,
|
||||
)
|
||||
with self.assertRaisesRegex(RemoteCodeGateError, "relative path"):
|
||||
ensure_remote_code_allowed(self.dir, trust_remote_code=True)
|
||||
|
||||
def test_non_string_model_file_rejected(self):
|
||||
_make_checkpoint(
|
||||
self.dir,
|
||||
{"model_type": "muse_glimmer", "model_file": 42},
|
||||
with_sentinel=False,
|
||||
)
|
||||
with self.assertRaisesRegex(RemoteCodeGateError, "non-string"):
|
||||
ensure_remote_code_allowed(self.dir, trust_remote_code=True)
|
||||
|
||||
|
||||
@unittest.skipUnless(_HAS_MLX, "requires mlx + mlx_lm")
|
||||
class TestModelRunnerGateWiring(CustomTestCase):
|
||||
"""The runner must gate BEFORE calling mlx_lm's loader, on the same
|
||||
resolved directory it then loads from."""
|
||||
|
||||
class _StopInit(Exception):
|
||||
pass
|
||||
|
||||
def setUp(self):
|
||||
self._tmp = tempfile.TemporaryDirectory()
|
||||
self.dir = Path(self._tmp.name)
|
||||
|
||||
def tearDown(self):
|
||||
self._tmp.cleanup()
|
||||
|
||||
def test_refusal_precedes_loader_call(self):
|
||||
from sglang.srt.hardware_backend.mlx.model_runner import MlxModelRunner
|
||||
|
||||
_make_checkpoint(
|
||||
self.dir, {"model_type": "muse_glimmer", "model_file": "evil.py"}
|
||||
)
|
||||
with patch(
|
||||
"sglang.srt.hardware_backend.mlx.model_runner.mlx_lm_load"
|
||||
) as loader:
|
||||
with self.assertRaisesRegex(RemoteCodeGateError, "--trust-remote-code"):
|
||||
MlxModelRunner(model_path=str(self.dir), trust_remote_code=False)
|
||||
loader.assert_not_called()
|
||||
self.assertFalse((self.dir / "marker.txt").exists())
|
||||
|
||||
def test_trusted_load_uses_resolved_directory(self):
|
||||
from sglang.srt.hardware_backend.mlx.model_runner import MlxModelRunner
|
||||
|
||||
_make_checkpoint(
|
||||
self.dir, {"model_type": "muse_glimmer", "model_file": "evil.py"}
|
||||
)
|
||||
with patch(
|
||||
"sglang.srt.hardware_backend.mlx.model_runner.mlx_lm_load",
|
||||
side_effect=self._StopInit,
|
||||
) as loader:
|
||||
with self.assertRaises(self._StopInit):
|
||||
MlxModelRunner(model_path=str(self.dir), trust_remote_code=True)
|
||||
loader.assert_called_once()
|
||||
called_path = loader.call_args.args[0]
|
||||
self.assertEqual(
|
||||
Path(called_path).resolve(),
|
||||
self.dir.resolve(),
|
||||
"loader must receive the same directory the gate inspected",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main(verbosity=2)
|
||||
@@ -0,0 +1,278 @@
|
||||
"""Unit tests for the Muse Glimmer MLX model file's load path — no weights, no server.
|
||||
|
||||
End-to-end coverage needs the private packaged artifact, so the checkpoint-schema
|
||||
logic is pinned here with tiny synthetic weights instead:
|
||||
|
||||
1. ``flatten_rc_config`` — RC nested-config translation, including the two
|
||||
convention conversions (qk_scale_factor gains sqrt(head_dim); NoPE layers
|
||||
come from zeros in ``layer_rope_theta``).
|
||||
2. ``ModelArgs`` validation — derived ``no_rope_layers``/``layer_types``,
|
||||
rejection of inconsistent or malformed lists, format-version check.
|
||||
3. ``sanitize`` — all three accepted weight layouts (raw HF, RC multimodal,
|
||||
packaged) plus rejection of incomplete or mislabeled checkpoints. The
|
||||
positional RC norm renames and the per-head q/gate interleave are verified
|
||||
numerically, since getting either silently wrong still yields a model that
|
||||
runs but computes garbage.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import unittest
|
||||
|
||||
from sglang.test.ci.ci_register import register_cpu_ci, register_mlx_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=6, suite="base-a-test-cpu")
|
||||
register_mlx_ci(est_time=6, suite="stage-a-unit-test-mlx")
|
||||
|
||||
_HAS_MLX = (
|
||||
importlib.util.find_spec("mlx") is not None
|
||||
and importlib.util.find_spec("mlx_lm") is not None
|
||||
)
|
||||
_SKIP_REASON = "requires mlx + mlx_lm"
|
||||
|
||||
if _HAS_MLX:
|
||||
import mlx.core as mx
|
||||
|
||||
from sglang.srt.hardware_backend.mlx.models.muse_glimmer_mlx import (
|
||||
Model,
|
||||
ModelArgs,
|
||||
flatten_rc_config,
|
||||
)
|
||||
|
||||
# Tiny architecture: 4 layers so the derived NoPE pattern (last layer NoPE
|
||||
# with every_n_layers_nope=4) exercises both layer types.
|
||||
_TINY = dict(
|
||||
hidden_size=8,
|
||||
num_hidden_layers=4,
|
||||
num_attention_heads=2,
|
||||
num_key_value_heads=1,
|
||||
head_dim=4,
|
||||
intermediate_size=16,
|
||||
vocab_size=32,
|
||||
every_n_layers_nope=4,
|
||||
sliding_window=4,
|
||||
max_position_embeddings=64,
|
||||
)
|
||||
|
||||
_RC_TEXT_CONFIG = dict(
|
||||
hidden_size=8,
|
||||
num_hidden_layers=4,
|
||||
num_attention_heads=2,
|
||||
num_key_value_heads=1,
|
||||
head_dim=4,
|
||||
intermediate_size=16,
|
||||
vocab_size=32,
|
||||
rms_norm_eps=1e-5,
|
||||
post_norm_eps=1e-8,
|
||||
max_position_embeddings=64,
|
||||
qk_scale_factor=0.5,
|
||||
output_multiplier=0.2,
|
||||
final_logit_softcapping=20.0,
|
||||
sliding_window=4,
|
||||
layer_rope_theta=[500000.0, 500000.0, 500000.0, 0],
|
||||
)
|
||||
|
||||
|
||||
def _raw_weights(args):
|
||||
"""A complete raw HF export with deterministic values."""
|
||||
mx.random.seed(0)
|
||||
H, D, hid = args.num_attention_heads, args.head_dim, args.hidden_size
|
||||
kv = args.num_key_value_heads * D
|
||||
weights = {
|
||||
"model.embed_tokens.weight": mx.random.normal((args.vocab_size, hid)),
|
||||
"model.norm.weight": mx.random.normal((hid,)),
|
||||
"lm_head.weight": mx.random.normal((args.vocab_size, hid)),
|
||||
}
|
||||
for i in range(args.num_hidden_layers):
|
||||
p = f"model.layers.{i}."
|
||||
weights.update(
|
||||
{
|
||||
p + "self_attn.q_proj.weight": mx.random.normal((H * D, hid)),
|
||||
p + "self_attn.k_proj.weight": mx.random.normal((kv, hid)),
|
||||
p + "self_attn.v_proj.weight": mx.random.normal((kv, hid)),
|
||||
p + "self_attn.o_proj.weight": mx.random.normal((hid, H * D)),
|
||||
p + "self_attn.output_gate_proj.weight": mx.random.normal((H * D, hid)),
|
||||
p + "input_layernorm.weight": mx.full((hid,), 0.10),
|
||||
p + "post_attn_norm.weight": mx.full((hid,), 0.20),
|
||||
p + "post_attention_layernorm.weight": mx.full((hid,), 0.30),
|
||||
p + "post_ffn_norm.weight": mx.full((hid,), 0.40),
|
||||
p
|
||||
+ "mlp.gate_proj.weight": mx.random.normal(
|
||||
(args.intermediate_size, hid)
|
||||
),
|
||||
p
|
||||
+ "mlp.up_proj.weight": mx.random.normal((args.intermediate_size, hid)),
|
||||
p
|
||||
+ "mlp.down_proj.weight": mx.random.normal(
|
||||
(hid, args.intermediate_size)
|
||||
),
|
||||
}
|
||||
)
|
||||
return weights
|
||||
|
||||
|
||||
def _rc_weights(args):
|
||||
"""The same export in the RC multimodal layout (nested prefix, RC names,
|
||||
a vision tower to be dropped, no output-gate fusion)."""
|
||||
raw = _raw_weights(args)
|
||||
rc = {}
|
||||
renames = {
|
||||
"self_attn.output_gate_proj.weight": "self_attn.gate_proj.weight",
|
||||
"post_attn_norm.weight": "post_attention_layernorm.weight",
|
||||
"post_attention_layernorm.weight": "pre_feedforward_layernorm.weight",
|
||||
"post_ffn_norm.weight": "post_feedforward_layernorm.weight",
|
||||
}
|
||||
for name, w in raw.items():
|
||||
if not name.startswith("model."):
|
||||
rc[name] = w # lm_head stays top-level in the RC layout too
|
||||
continue
|
||||
rest = name[len("model.") :]
|
||||
for raw_suffix, rc_suffix in renames.items():
|
||||
if rest.endswith(raw_suffix):
|
||||
rest = rest[: -len(raw_suffix)] + rc_suffix
|
||||
break
|
||||
rc["model.language_model." + rest] = w
|
||||
rc["model.vision_tower.patch_embed.weight"] = mx.zeros((4, 4))
|
||||
return rc
|
||||
|
||||
|
||||
@unittest.skipUnless(_HAS_MLX, _SKIP_REASON)
|
||||
class TestFlattenRcConfig(CustomTestCase):
|
||||
def test_field_mapping_and_conversions(self):
|
||||
flat = flatten_rc_config({"text_config": dict(_RC_TEXT_CONFIG)})
|
||||
self.assertEqual(flat["model_type"], "muse_glimmer")
|
||||
# RC qk_scale_factor is against SDPA's 1/sqrt(head_dim).
|
||||
self.assertAlmostEqual(flat["qk_scale_factor"], 0.5 * 4**0.5)
|
||||
self.assertEqual(flat["output_soft_cap_temp"], 20.0)
|
||||
# The current vendor export stores q/k in the NeoX rotary
|
||||
# layout; the RC path always reads that convention.
|
||||
self.assertIs(flat["rope_is_neox_style"], True)
|
||||
# normalize_tok_embeddings must stay at the ModelArgs default (True):
|
||||
# the 20260806 export ships the raw table (needs the runtime norm),
|
||||
# and on older baked-table exports the scaleless RMS norm is
|
||||
# idempotent, so always-on covers both generations.
|
||||
self.assertNotIn("normalize_tok_embeddings", flat)
|
||||
self.assertIs(ModelArgs.normalize_tok_embeddings, True)
|
||||
# Zeros in layer_rope_theta mark NoPE layers.
|
||||
self.assertEqual(flat["no_rope_layers"], [1, 1, 1, 0])
|
||||
|
||||
def test_non_silu_activation_rejected(self):
|
||||
cfg = dict(_RC_TEXT_CONFIG, hidden_activation="gelu")
|
||||
with self.assertRaisesRegex(ValueError, "silu"):
|
||||
flatten_rc_config({"text_config": cfg})
|
||||
|
||||
def test_model_args_from_dict_accepts_rc_schema(self):
|
||||
args = ModelArgs.from_dict({"text_config": dict(_RC_TEXT_CONFIG)})
|
||||
self.assertEqual(args.no_rope_layers, [1, 1, 1, 0])
|
||||
self.assertEqual(
|
||||
args.layer_types,
|
||||
["sliding_attention"] * 3 + ["full_attention"],
|
||||
)
|
||||
|
||||
|
||||
@unittest.skipUnless(_HAS_MLX, _SKIP_REASON)
|
||||
class TestModelArgsValidation(CustomTestCase):
|
||||
def test_derives_nope_and_layer_types(self):
|
||||
args = ModelArgs(**_TINY)
|
||||
self.assertEqual(args.no_rope_layers, [1, 1, 1, 0])
|
||||
self.assertEqual(
|
||||
args.layer_types,
|
||||
["sliding_attention"] * 3 + ["full_attention"],
|
||||
)
|
||||
|
||||
def test_layer_types_must_match_no_rope_layers(self):
|
||||
with self.assertRaisesRegex(ValueError, "disagrees"):
|
||||
ModelArgs(**_TINY, layer_types=["full_attention"] * 4)
|
||||
|
||||
def test_wrong_length_no_rope_layers_rejected(self):
|
||||
with self.assertRaisesRegex(ValueError, "entries"):
|
||||
ModelArgs(**_TINY, no_rope_layers=[1, 0])
|
||||
|
||||
def test_non_binary_no_rope_flags_rejected(self):
|
||||
with self.assertRaisesRegex(ValueError, "non-binary"):
|
||||
ModelArgs(**_TINY, no_rope_layers=[1, 1, 2, 0])
|
||||
|
||||
def test_unknown_format_version_rejected(self):
|
||||
with self.assertRaisesRegex(ValueError, "muse_glimmer_mlx_format"):
|
||||
ModelArgs(**_TINY, muse_glimmer_mlx_format=99)
|
||||
|
||||
|
||||
@unittest.skipUnless(_HAS_MLX, _SKIP_REASON)
|
||||
class TestSanitize(CustomTestCase):
|
||||
def _model(self, **overrides):
|
||||
return Model(ModelArgs(**dict(_TINY, **overrides)))
|
||||
|
||||
def test_raw_export_folds_norms_and_fuses_gate(self):
|
||||
model = self._model()
|
||||
raw = _raw_weights(model.args)
|
||||
out = model.sanitize(dict(raw))
|
||||
|
||||
# Offset norms gain +1.0; the final norm does not.
|
||||
norm = out["model.layers.0.input_layernorm.weight"]
|
||||
self.assertTrue(mx.allclose(norm, mx.full(norm.shape, 1.10)))
|
||||
self.assertTrue(mx.allclose(out["model.norm.weight"], raw["model.norm.weight"]))
|
||||
|
||||
# q/gate interleave is per-head [q_head; gate_head].
|
||||
H, D = model.args.num_attention_heads, model.args.head_dim
|
||||
hid = model.args.hidden_size
|
||||
fused = out["model.layers.0.self_attn.q_proj.weight"]
|
||||
self.assertEqual(tuple(fused.shape), (2 * H * D, hid))
|
||||
per_head = fused.reshape(H, 2 * D, hid)
|
||||
q = raw["model.layers.0.self_attn.q_proj.weight"].reshape(H, D, hid)
|
||||
g = raw["model.layers.0.self_attn.output_gate_proj.weight"].reshape(H, D, hid)
|
||||
self.assertTrue(mx.allclose(per_head[:, :D, :], q))
|
||||
self.assertTrue(mx.allclose(per_head[:, D:, :], g))
|
||||
self.assertNotIn("model.layers.0.self_attn.output_gate_proj.weight", out)
|
||||
|
||||
def test_sanitized_raw_weights_load_and_forward(self):
|
||||
model = self._model()
|
||||
out = model.sanitize(_raw_weights(model.args))
|
||||
model.load_weights(list(out.items()))
|
||||
logits = model(mx.array([[1, 2, 3]], dtype=mx.int32), cache=model.make_cache())
|
||||
self.assertEqual(tuple(logits.shape), (1, 3, model.args.vocab_size))
|
||||
self.assertTrue(bool(mx.all(mx.isfinite(logits))))
|
||||
|
||||
def test_rc_layout_positional_renames(self):
|
||||
model = self._model()
|
||||
out = model.sanitize(_rc_weights(model.args))
|
||||
# Distinct per-norm constants prove each RC name landed in its
|
||||
# positional slot (+1 folded): RC post_attention_layernorm ->
|
||||
# raw post_attn_norm (0.20), RC pre_feedforward_layernorm ->
|
||||
# raw post_attention_layernorm (0.30).
|
||||
for raw_name, value in (
|
||||
("input_layernorm", 1.10),
|
||||
("post_attn_norm", 1.20),
|
||||
("post_attention_layernorm", 1.30),
|
||||
("post_ffn_norm", 1.40),
|
||||
):
|
||||
w = out[f"model.layers.0.{raw_name}.weight"]
|
||||
self.assertTrue(
|
||||
mx.allclose(w, mx.full(w.shape, value)),
|
||||
f"{raw_name} expected {value}",
|
||||
)
|
||||
self.assertFalse(any("vision" in k for k in out))
|
||||
self.assertFalse(any("language_model" in k for k in out))
|
||||
|
||||
def test_packaged_artifact_passes_through(self):
|
||||
model = self._model(muse_glimmer_mlx_format=1)
|
||||
packaged = {"model.embed_tokens.weight": mx.zeros((32, 8))}
|
||||
self.assertIs(model.sanitize(packaged), packaged)
|
||||
|
||||
def test_packaged_marker_with_raw_keys_rejected(self):
|
||||
model = self._model(muse_glimmer_mlx_format=1)
|
||||
raw = _raw_weights(ModelArgs(**_TINY))
|
||||
with self.assertRaisesRegex(ValueError, "raw-checkpoint keys"):
|
||||
model.sanitize(raw)
|
||||
|
||||
def test_incomplete_raw_checkpoint_rejected(self):
|
||||
model = self._model()
|
||||
raw = _raw_weights(model.args)
|
||||
del raw["model.layers.0.mlp.up_proj.weight"]
|
||||
with self.assertRaisesRegex(ValueError, "missing"):
|
||||
model.sanitize(raw)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -75,7 +75,7 @@ class TestKVCacheQuantRegistry(CustomTestCase):
|
||||
from sglang.srt.runtime_context import get_context
|
||||
|
||||
runner = object.__new__(ModelRunner)
|
||||
runner.server_args = SimpleNamespace()
|
||||
runner.server_args = SimpleNamespace(speculative_draft_kv_cache_dtype=None)
|
||||
runner.draft_attention_backend = None
|
||||
# The runner reads the requested dtype off the model bag, so the double
|
||||
# publishes it rather than carrying it on a stand-in config.
|
||||
|
||||
@@ -89,6 +89,7 @@ class TestModelOverridableWhitelist(CustomTestCase):
|
||||
"decode_attention_backend",
|
||||
"flashinfer_allreduce_fusion_backend",
|
||||
"fp8_gemm_runner_backend",
|
||||
"fp4_gemm_runner_backend",
|
||||
"disable_custom_all_reduce",
|
||||
"enable_aiter_allreduce_fusion",
|
||||
"enable_symm_mem",
|
||||
|
||||
Reference in New Issue
Block a user