diff --git a/torchtitan/experiments/rl/examples/search_r1/README.md b/torchtitan/experiments/rl/examples/search_r1/README.md index 373bd51651..26ac0ad041 100644 --- a/torchtitan/experiments/rl/examples/search_r1/README.md +++ b/torchtitan/experiments/rl/examples/search_r1/README.md @@ -53,6 +53,19 @@ python /local_dense_retriever/retrieval_server.py \ Override `message_env.search_url` / `message_env.topk` in the config if needed. +### 3. Checkpoint +Download the base checkpoint the config expects. `download_hf_assets.py` writes to a +subdirectory named after the repo, which is the path in `hf_assets_path`: + +```bash +python scripts/download_hf_assets.py \ + --repo_id meta-models/Muse-Glimmer-30B \ + --local_dir torchtitan/experiments/rl/example_checkpoint \ + --all +``` + +Swap `--repo_id` for the model your config selects (e.g. `Qwen/Qwen3-1.7B`). + ## Run ```bash diff --git a/torchtitan/experiments/rl/examples/search_r1/config_registry.py b/torchtitan/experiments/rl/examples/search_r1/config_registry.py index 1dfd11c695..84ad0428f7 100644 --- a/torchtitan/experiments/rl/examples/search_r1/config_registry.py +++ b/torchtitan/experiments/rl/examples/search_r1/config_registry.py @@ -27,6 +27,7 @@ ParallelismConfig, TrainingConfig, ) +from torchtitan.distributed.activation_checkpoint import FullAC from torchtitan.experiments.rl.actors.generator import ( SamplingConfig, VLLMCudagraphConfig, @@ -47,6 +48,10 @@ from torchtitan.experiments.rl.observability.metrics import MetricsProcessor from torchtitan.experiments.rl.renderer import RendererConfig from torchtitan.experiments.rl.rollout.advantage import AdvantageEstimator +from torchtitan.models.muse_glimmer import model_registry as muse_glimmer_model_registry +from torchtitan.models.muse_glimmer.state_dict_adapter import ( + MuseGlimmerStateDictAdapter, +) from torchtitan.models.qwen3 import model_registry @@ -245,3 +250,96 @@ def rl_grpo_qwen3_30b_a3b_deepep_search_r1_perf() -> Controller.Config: # from this scheduler limit, CUDA graph capture sizes, CP, and SP. config.generator.max_num_batched_tokens = 2048 # TODO: TBD return config + + +def rl_grpo_muse_glimmer_30b_search_r1() -> Controller.Config: + """GRPO/DAPO Search-R1 for Muse Glimmer 30B. + + 8 GPUs: 6 trainer (FSDP=3 x TP=2) + 2 generator (TP=2), with a dense retrieval + server on spare capacity. Requires a running retrieval server and the QA parquet + data; see ``README.md``. + + Two constraints are specific to this model: + + * **Generator TP <= 2.** Muse Glimmer has 2 KV heads, so attention cannot be + tensor-split further. Scale the trainer with FSDP rather than TP. + * **Full activation checkpointing is required.** Adam's m/v are allocated on the + *first* ``optimizer.step()``, so per-GPU memory jumps by roughly 8 bytes/param + between step 1 and step 2 (~37 GB/GPU here, sharded 6 ways). With the default + ``SelectiveAC`` that jump OOMs at step 2; ``FullAC`` frees the activation + headroom it needs. + + varlen attention is used for both roles so the trainer and the vLLM generator run + one ModelSpec. The state-dict adapter handles the HF checkpoint's Q/K RoPE layout + on load, and the renderer (registered below) handles Muse Glimmer's harmony chat + format and ATEM tool calls. + """ + # Muse Glimmer's renderer ships in torchtitan rather than the `renderers` library; + # registering makes RendererConfig(name="muse_glimmer") resolve it. + + model_spec = muse_glimmer_model_registry("30B", attn_backend="varlen") + model_spec = dataclasses.replace( + model_spec, state_dict_adapter=MuseGlimmerStateDictAdapter + ) + + return Controller.Config( + model_spec=model_spec, + hf_assets_path="torchtitan/experiments/rl/example_checkpoint/Muse-Glimmer-30B", + async_loop=AsyncLoopConfig( + num_training_steps=500, + num_prompts_per_train_step=8, + num_samples_per_prompt=8, + validation=ValidationConfig(num_samples=500), + ), + compile=CompileConfig(enable=False), + rollouter=SearchR1Rollouter.Config( + worker=SearchR1Worker.Config( + advantage=AdvantageEstimator.Config(should_std_normalize=True), + ), + ), + renderer=RendererConfig(name="muse_glimmer", enable_thinking=True), + metrics=MetricsProcessor.Config(enable_wandb=True), + trainer=PolicyTrainer.Config( + optimizer=default_adamw(lr=1e-6), + lr_scheduler=LRSchedulersContainer.Config( + warmup_steps=2, decay_type="linear", min_lr_factor=1.0 + ), + training=TrainingConfig( + num_tokens_per_microbatch_per_dp_rank=4096, + max_context_length=4096, + ), + ac_config=FullAC.Config(), + parallelism=ParallelismConfig( + data_parallel_shard_degree=3, + tensor_parallel_degree=2, + ), + checkpoint=CheckpointManager.Config( + enable=True, + initial_load_in_hf=True, # first run loads HF; restarts resume from DCP + interval=50, + last_save_model_only=False, + keep_latest_k=3, + ), + loss=ChunkedLossWrapper.Config( + num_chunks=8, + loss_fn=DAPOLoss.Config( + ratio_clip_low=0.2, + ratio_clip_high=0.28, + ), + ), + ), + generator=VLLMGenerator.Config( + model_dtype="bfloat16", + parallelism=InferenceParallelismConfig( + data_parallel_degree=1, + tensor_parallel_degree=2, # <= 2 KV heads + ), + cudagraph=VLLMCudagraphConfig(enable=False), + checkpoint=CheckpointManager.Config(enable=False), + sampling=SamplingConfig( + temperature=1.0, + top_p=1.0, + max_tokens=4096, + ), + ), + ) diff --git a/torchtitan/experiments/rl/models/muse_glimmer/__init__.py b/torchtitan/experiments/rl/models/muse_glimmer/__init__.py new file mode 100644 index 0000000000..2e41cd717f --- /dev/null +++ b/torchtitan/experiments/rl/models/muse_glimmer/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. diff --git a/torchtitan/experiments/rl/models/muse_glimmer/atem.py b/torchtitan/experiments/rl/models/muse_glimmer/atem.py new file mode 100644 index 0000000000..af2c6a86ae --- /dev/null +++ b/torchtitan/experiments/rl/models/muse_glimmer/atem.py @@ -0,0 +1,90 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Muse Glimmer ATEM tool-call parse/render (the novel core of the Muse Glimmer renderer). + +Muse Glimmer emits tool calls in Anthropic-style ATEM XML inside the harmony envelope: + + + + who wrote Blade Runner + + + +parse: model text -> [{"name", "arguments"}] (what env.step() needs) +render: a tool call -> ATEM text (what goes back into the prompt) +""" + +from __future__ import annotations + +import json +import re + +_FUNCTION_CALLS = re.compile( + r"(.*?)", re.DOTALL +) +_INVOKE = re.compile( + r'(?P.*?)', re.DOTALL +) +_PARAMETER = re.compile( + r'(?P.*?)', re.DOTALL +) + + +def parse_atem_tool_calls(text: str) -> list[dict]: + """Parse every ATEM tool call in `text` into `[{"name", "arguments"}]`. + + Values are JSON-decoded when possible (dicts/lists/numbers/bools), else kept + as the raw string. Supports multiple parallel invokes in one block. + """ + calls: list[dict] = [] + for block in _FUNCTION_CALLS.findall(text): + for invoke in _INVOKE.finditer(block): + arguments: dict = {} + for param in _PARAMETER.finditer(invoke.group("body")): + raw = param.group("value") + try: + arguments[param.group("key")] = json.loads(raw) + except (json.JSONDecodeError, ValueError): + arguments[param.group("key")] = raw + calls.append({"name": invoke.group("name"), "arguments": arguments}) + return calls + + +def render_atem_tool_call(name: str, arguments: dict) -> str: + """Render one tool call as ATEM text (matches Muse Glimmer's chat template).""" + lines = ["", f''] + for key, value in arguments.items(): + if isinstance(value, bool): + sval = "true" if value else "false" + elif value is None: + sval = "null" + elif isinstance(value, (dict, list)): + sval = json.dumps(value) + else: + sval = str(value) + lines.append(f'{sval}') + lines += ["", ""] + return "\n".join(lines) + + +if __name__ == "__main__": + # round-trip self-test (no deps, no GPU) + sample = ( + 'thinking...\n\n\n' + 'who wrote Blade Runner\n' + "\n" + ) + calls = parse_atem_tool_calls(sample) + assert calls == [ + {"name": "search", "arguments": {"query": "who wrote Blade Runner"}} + ], calls + # no tool call -> empty (this is how env.step() detects "final answer") + assert parse_atem_tool_calls("Philip K. Dick") == [] + # render -> parse round-trip + rendered = render_atem_tool_call("search", {"query": "x", "topk": 3}) + assert parse_atem_tool_calls(rendered)[0]["arguments"]["topk"] == 3 + print("muse_glimmer atem: all checks passed") diff --git a/torchtitan/experiments/rl/models/muse_glimmer/renderer.py b/torchtitan/experiments/rl/models/muse_glimmer/renderer.py new file mode 100644 index 0000000000..2be5350f70 --- /dev/null +++ b/torchtitan/experiments/rl/models/muse_glimmer/renderer.py @@ -0,0 +1,788 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Muse Glimmer renderer: chat messages <-> tokens for TorchTitan RL. + +RL needs two directions. ``render_ids`` turns messages + tool definitions into token +ids; ``parse_response`` turns generated tokens back into ``(content, +reasoning_content, tool_calls)`` so the rollout loop can tell "call a tool and +continue" apart from "final answer, score it". + +Muse Glimmer uses the harmony chat format -- an assistant turn is a sequence of +``to=<|message|>`` channels, where ``to=self`` is private reasoning and +other recipients carry user-visible content -- and expresses tool calls as ATEM XML +inside those channels (see ``atem.py``). + +The format is built natively in Python rather than by calling +``tokenizer.apply_chat_template``. That matches how the ``renderers`` library implements +every model-specific renderer (only ``DefaultRenderer`` wraps Jinja), and it buys two +things a template wrapper cannot: per-token loss attribution is exact rather than +recovered by diffing prefixes, and ``bridge_to_next_turn`` can extend a sampled +completion without re-rendering it. + +``_render_text`` is kept byte-exact against the published ``chat_template.jinja``; the +renderer test asserts equality with ``apply_chat_template`` across roles, tool shapes and +reasoning states. Treat that test as the spec -- if the template changes upstream, it +fails first. + +Implements the ``renderers.Renderer`` Protocol. ``register()`` installs it into the +``renderers`` library's public registry (``RENDERER_REGISTRY`` / ``_CONFIG_BY_NAME``), +which is that library's supported extension path -- no fork or upstream change needed. + +Every other TorchTitan model resolves to a renderer that lives in +PrimeIntellect-ai/renderers. This one ships here because Muse Glimmer is not in that +library yet. + +TODO: upstream this to PrimeIntellect-ai/renderers (renderer -> renderers/muse_glimmer.py, +atem.py -> a tool parser in renderers/parsers.py), then delete both files and +``register()``, leaving only the ``_RENDERER_BY_MODEL`` entry in +experiments/rl/renderer.py. + +It lives under ``experiments/rl`` rather than ``torchtitan/models/muse_glimmer`` because +RL is its only consumer and ``renderers`` is an RL-only optional dependency; keeping it +here leaves the core model package importable without it. +""" + +from __future__ import annotations + +import datetime +import json +import re +from typing import ClassVar, Literal, NamedTuple + +from renderers.base import ( + extract_message_tool_names, + ParsedResponse, + ParsedToolCall, + reject_assistant_in_extension, + RenderedTokens, + resolve_thinking_retention, + should_rerender_for_thinking_retention, + trim_to_turn_close, +) +from renderers.configs import BaseRendererConfig + +from .atem import parse_atem_tool_calls, render_atem_tool_call + +RENDERER_NAME = "muse_glimmer" + + +class MuseGlimmerRendererConfig(BaseRendererConfig): + """Muse Glimmer (harmony chat format + ATEM tool calls) renderer config.""" + + name: Literal["muse_glimmer"] = RENDERER_NAME + + # renderers validates in BaseRendererConfig.__pydantic_init_subclass__ that every + # non-base field is classified as either a chat-template kwarg or a renderer-internal + # knob; the two sets must be disjoint and together cover all of them. Declared + # unconditionally -- versions without the validator ignore these ClassVars, so this + # is compatible with both. The template fields mirror kwargs the published + # chat_template.jinja reads, which is what the library's parity matrix varies. + _template_fields: ClassVar[frozenset[str]] = frozenset( + {"reasoning_strength", "knowledge_cutoff", "current_date"} + ) + _internal_fields: ClassVar[frozenset[str]] = frozenset( + {"retain_reasoning_in_history", "answer_from_reasoning_fallback"} + ) + + reasoning_strength: str | None = None + """Sizes the reasoning budget, rendered as ``Reasoning strength: .`` + + ``None`` (default) uses the template's own default of ``"high"``. Set ``"low"`` for + agentic tasks with a tight token budget: at high strength the model can spend the + whole budget reasoning and get truncated before it emits a tool call or an answer, + which scores as a failed rollout. + """ + + knowledge_cutoff: str | None = None + """Knowledge-cutoff date in the default system prompt. ``None`` uses the template's + own default. Only rendered when the caller supplies no system message.""" + + current_date: str | None = None + """Pins ``Current date:`` in the default system prompt. + + ``None`` reproduces the template's behaviour of substituting today's date, which + makes the rendered prompt change from one day to the next. Pin it for runs that need + to be reproducible across days. + """ + + retain_reasoning_in_history: bool = True + """Whether prior assistant turns keep their ``reasoning_content`` when re-rendered. + + A ``to=self`` channel is emitted for any assistant message carrying + ``reasoning_content``, so history grows with every turn's reasoning. Harmony-style + models are often trained with prior analysis *dropped* from context (gpt-oss does + this via ``auto_drop_analysis``); set this False to match that convention and to + keep multi-turn prompts short. + + Defaults True to preserve the template's own behaviour -- flip it only if you have + checked what the model was trained against. + """ + + answer_from_reasoning_fallback: bool = False + """When the model produces only reasoning channels -- no user-facing content and no + tool call -- treat the last non-empty reasoning line as the answer. + + Off by default: it promotes private reasoning to user-visible content, which is + usually not what you want. Useful for outcome-scored RL, where a rollout with an + empty ``content`` is unscoreable and the answer is often the final reasoning line. + """ + + +# Muse Glimmer special tokens. The ids are checked against the tokenizer in __init__ +# rather than trusted, since they are baked into parse_response and the loss mask. +START_STR, MESSAGE_STR = "<|start|>", "<|message|>" +EOM_STR, EOT_STR = "<|eom|>", "<|eot|>" + +START_ID = 200022 # <|start|> begins a harmony message header +MESSAGE_ID = 200023 # <|message|> ends the header, body follows +EOT_ID = 200008 # <|eot|> end of turn +EOM_ID = 200007 # <|eom|> end of message +EOS_ID = 200001 # <|end_of_text|> + +_DEFAULT_REASONING_STRENGTH = "high" +_DEFAULT_KNOWLEDGE_CUTOFF = "2026-01-04" + +_FUNCTION_CALLS_BLOCK = re.compile( + r".*?", re.DOTALL +) +# One harmony channel: "to=<|message|>" up to <|eom|>/<|eot|>/next channel. +_CHANNEL = re.compile( + r"(?:to=(?P[^\s<|]+))?\s*<\|message\|>(?P.*?)" + r"(?=<\|eom\|>|<\|eot\|>|<\|end_of_text\|>|<\|start\|>|\Z)", + re.DOTALL, +) + +# The template normalises "reasoning effort" to "reasoning strength" in a caller-supplied +# system prompt. Jinja has no case-insensitive replace, so it spells out four casings; +# reproduce exactly those four, or renders diverge on any other casing. +_EFFORT_TO_STRENGTH = ( + ("Reasoning effort", "Reasoning strength"), + ("Reasoning Effort", "Reasoning Strength"), + ("reasoning effort", "reasoning strength"), + ("REASONING EFFORT", "REASONING STRENGTH"), +) + + +def _tojson(value) -> str: + """Jinja's ``tojson`` as transformers configures it: insertion order, raw unicode.""" + return json.dumps(value, ensure_ascii=False) + + +def _render_content(content) -> str: + """A message body: a plain string, or a list of typed multimodal parts.""" + if content is None: + return "" + if isinstance(content, str): + return content + out = [] + for part in content: + kind = part.get("type") + if kind == "image": + out.append("<|patch|>") + elif kind == "video": + out.append("<|video|>") + elif kind == "text": + out.append(part["text"]) + return "".join(out) + + +def _tool_fn(tool) -> dict: + """Tools arrive either OpenAI-nested (``{"function": {...}}``) or flat.""" + if isinstance(tool, dict) and tool.get("function") is not None: + return tool["function"] + return tool + + +def _tool_namespaces(tools) -> list[str]: + """Leading dotted segment of each tool name, first-seen order, deduplicated.""" + seen: list[str] = [] + for tool in tools: + namespace = _tool_fn(tool)["name"].split(".")[0] + if namespace not in seen: + seen.append(namespace) + return seen + + +def _render_tool_defs(tools) -> str: + """The tool-definition block injected into the system message.""" + parts = [ + "In this environment you have access to a set of tools you can use to answer " + "the user's question.\n\n", + 'You can invoke a function by writing a "" block like the ' + "following:\n", + '\n\n' + '$PARAMETER_VALUE\n' + "...\n\n\n\n", + "String and scalar parameters should be specified as is, while lists and " + "objects should use JSON format. Note that spaces for string values are not " + "stripped. The output is not expected to be valid XML and is parsed with " + "regular expressions.\n", + "Here are the functions available in JSONSchema format:\n", + "// Tool metadata\n", + ] + for namespace in _tool_namespaces(tools): + parts.append( + f'{{"name": {_tojson(namespace)}, "description": {_tojson("")}}}\n' + ) + parts.append("// Function schemas") + for tool in tools: + fn = _tool_fn(tool) + parts.append( + f'\n{{"name": {_tojson(fn["name"])}, ' + f'"description": {_tojson(fn.get("description"))}, ' + f'"parameters": {_tojson(fn.get("parameters"))}}}' + ) + parts.append( + "\n\nHere's an example of how to call a function in the tool set:\n" + "(If the tool namespace is not specified, invoke the function directly as " + "`example_function_name` rather than " + "`example_tool_name.example_function_name`)\n\n" + "to=example_tool_name.example_function_name\n\n" + "\n" + '\n' + 'value_1\n' + 'This is the value for the second ' + 'parameter\nthat can span\n"multiple" lines\n\n' + "\n" + ) + return "".join(parts) + + +def _render_system_meta(tools) -> str: + """The ``# Valid recipients:`` line closing every system message.""" + recipients = ['"self"'] + if tools: + recipients += [f'"{ns}.*"' for ns in _tool_namespaces(tools)] + recipients.append('"user"') + return "# Valid recipients: " + ", ".join(recipients) + "." + + +def _render_atem(tool_call) -> str: + """An assistant tool call as an ATEM block.""" + fn = tool_call["function"] + args = fn.get("arguments") + if not isinstance(args, dict): + raise ValueError( + "Muse Glimmer tool_call.function.arguments must be a dict, got " + f"{type(args).__name__}. A JSON string cannot be rendered." + ) + return render_atem_tool_call(fn["name"], args) + + +def _strip_reasoning_history(messages): + """Drop ``reasoning_content`` from assistant messages (see the config field).""" + out = [] + for m in messages: + if ( + isinstance(m, dict) + and m.get("role") == "assistant" + and m.get("reasoning_content") + ): + m = {k: v for k, v in m.items() if k != "reasoning_content"} + out.append(m) + return out + + +def _normalize_tool_calls(messages): + """Our ParsedToolCall is flat (.name/.arguments); rendering expects OpenAI-nested.""" + out = [] + for m in messages: + tcs = m.get("tool_calls") if isinstance(m, dict) else None + if not tcs: + out.append(m) + continue + norm = [] + for tc in tcs: + if isinstance(tc, dict) and "function" in tc: + norm.append(tc) + continue + name = tc.get("name") if isinstance(tc, dict) else getattr(tc, "name", None) + args = ( + tc.get("arguments") + if isinstance(tc, dict) + else getattr(tc, "arguments", None) + ) + norm.append( + {"type": "function", "function": {"name": name, "arguments": args}} + ) + out.append({**m, "tool_calls": norm}) + return out + + +def _parse_channels(text: str) -> list[tuple[str, str]]: + """Split a Muse Glimmer assistant completion into (recipient, body) harmony channels.""" + return [ + (m.group("rcpt") or "user", m.group("body").strip()) + for m in _CHANNEL.finditer(text) + ] + + +class _Piece(NamedTuple): + """One span of the render: either a literal special token or text to encode. + + ``sampled`` marks spans the model itself produced, which is what the trainer's loss + mask keys off. Template scaffolding (``<|start|>``, the role header, ``<|message|>``) + is never sampled, even inside an assistant turn. + + ``is_body`` marks the message body specifically, which the bridge reports as content + even on non-assistant roles that are never sampled. ``is_generation_prompt`` marks + the trailing ``<|start|>assistant``, which belongs to no message. + """ + + text: str + token_id: int | None + sampled: bool + msg_idx: int + is_body: bool = False + is_generation_prompt: bool = False + + +class MuseGlimmerRenderer: + def __init__(self, tokenizer, config: MuseGlimmerRendererConfig | None = None): + # (tokenizer, config) is the renderers-library constructor contract, so + # ``create_renderer`` can instantiate this from RENDERER_REGISTRY. + self._tok = tokenizer + self._config = config or MuseGlimmerRendererConfig() + # The controller reads renderer._tokenizer (e.g. for pad_id=eos_token_id). + self._tokenizer = tokenizer + self._bos = tokenizer.bos_token or "" + # BaseRendererConfig.thinking_retention is the library-wide knob every renderer + # is expected to honour in its bridge. Muse Glimmer's published chat template + # renders reasoning_content for every assistant turn unconditionally -- no + # query-boundary drop like gpt-oss's auto_drop_analysis or Qwen3's think-block + # stripping -- so "all" is the template-faithful implied policy. An explicit + # thinking_retention on the config overrides it. + self.effective_thinking_retention = resolve_thinking_retention( + self._config, + "all" if self._config.retain_reasoning_in_history else "tool_cycle", + ) + self._verify_special_ids() + + def _verify_special_ids(self) -> None: + """The special-token ids are baked into parse_response and the loss mask. + + A tokenizer that disagrees would mask the wrong spans and mis-split channels, + both of which corrupt training silently, so check rather than assume. + """ + for token, expected in ( + (START_STR, START_ID), + (MESSAGE_STR, MESSAGE_ID), + (EOM_STR, EOM_ID), + (EOT_STR, EOT_ID), + ): + actual = self._tok.convert_tokens_to_ids(token) + if actual != expected: + raise ValueError( + f"Tokenizer maps {token} to id {actual}, but the muse_glimmer " + f"renderer expects {expected}. This tokenizer is not compatible." + ) + + # ---------------------------------------------------------------- rendering + + def _prepare(self, messages): + """Apply the history policy, then the tool-call shape rendering expects.""" + if not self._config.retain_reasoning_in_history: + messages = _strip_reasoning_history(messages) + return _normalize_tool_calls(messages) + + def _reasoning_line(self) -> str: + strength = self._config.reasoning_strength or _DEFAULT_REASONING_STRENGTH + return f"Reasoning strength: {strength}." + + def _default_system_body(self, tools) -> str: + """System block synthesised when the caller supplies no system message.""" + cutoff = self._config.knowledge_cutoff or _DEFAULT_KNOWLEDGE_CUTOFF + date = self._config.current_date or datetime.datetime.now().strftime("%Y-%m-%d") + body = ( + "You are a helpful AI assistant." + f"\nKnowledge cutoff: {cutoff}." + f"\nCurrent date: {date}." + f"\n\n{self._reasoning_line()}" + ) + if tools: + body += "\n\n" + _render_tool_defs(tools) + return body + "\n\n" + _render_system_meta(tools) + + def _explicit_system_body(self, message, tools) -> str: + text = _render_content(message.get("content")) + for old, new in _EFFORT_TO_STRENGTH: + text = text.replace(old, new) + body = text + if "reasoning strength" not in text.lower(): + body += "\n\n" + self._reasoning_line() + if tools: + body += "\n\n" + _render_tool_defs(tools) + return body + "\n\n" + _render_system_meta(tools) + + @staticmethod + def _tool_name(message, messages) -> str: + """Tool messages name their tool directly, or via the call id that produced them.""" + name = message.get("name") + if name: + return name + call_id = message.get("tool_call_id") + resolved = call_id if call_id else "" + for m in messages: + for tc in m.get("tool_calls") or (): + if call_id is not None and tc.get("id") == call_id: + resolved = tc["function"]["name"] + return resolved + + def _build(self, messages, *, tools, add_generation_prompt) -> list[_Piece]: + """Render to spans. Concatenating ``.text`` reproduces the chat template exactly. + + One assistant message can expand into several harmony blocks -- a ``to=self`` + reasoning channel plus one per tool call -- which is why attribution is tracked + per span rather than per message. + """ + pieces: list[_Piece] = [] + + def emit( + text: str, + *, + token_id: int | None = None, + sampled=False, + idx=0, + is_body=False, + is_generation_prompt=False, + ): + if text: + pieces.append( + _Piece(text, token_id, sampled, idx, is_body, is_generation_prompt) + ) + + def block(header: str, body: str, end: str, *, idx: int, sampled: bool): + emit(START_STR, token_id=START_ID, idx=idx) + emit(header, idx=idx) + emit(MESSAGE_STR, token_id=MESSAGE_ID, idx=idx) + emit(body, sampled=sampled, idx=idx, is_body=True) + # The terminator counts as content only on assistant turns, where the model + # emits its own stop token; on history roles it is template scaffolding. + emit( + end, + token_id=EOT_ID if end == EOT_STR else EOM_ID, + sampled=sampled, + idx=idx, + is_body=sampled, + ) + + emit(self._bos, token_id=self._tok.bos_token_id) + + if not any(m.get("role") == "system" for m in messages): + block( + "system", + self._default_system_body(tools), + EOT_STR, + idx=0, + sampled=False, + ) + + for i, message in enumerate(messages): + role = message.get("role") + # The template picks the terminator by looking at the NEXT message's role: + # two consecutive same-role messages are joined with <|eom|>. + same_role_next = ( + i + 1 < len(messages) and messages[i + 1].get("role") == role + ) + end_token = EOM_STR if same_role_next else EOT_STR + + if role == "system": + block( + "system", + self._explicit_system_body(message, tools), + EOT_STR, + idx=i, + sampled=False, + ) + elif role == "user": + block( + "user", + _render_content(message.get("content")), + EOT_STR, + idx=i, + sampled=False, + ) + elif role == "tool": + name = self._tool_name(message, messages) + body = ( + f'\n' + f'{_render_content(message.get("content"))}\n' + ) + block(f"tool {name}", body, EOT_STR, idx=i, sampled=False) + elif role == "assistant": + if message.get("reasoning_content"): + block( + "assistant to=self", + message["reasoning_content"], + EOM_STR, + idx=i, + sampled=True, + ) + tool_calls = message.get("tool_calls") + if tool_calls: + for j, tc in enumerate(tool_calls): + last = j == len(tool_calls) - 1 + block( + f'assistant to={tc["function"]["name"]}', + _render_atem(tc), + end_token if last else EOM_STR, + idx=i, + sampled=True, + ) + else: + recipient = message.get("recipient") or "user" + end_turn = message.get("end_turn") + if end_turn is None: + end_turn = recipient == "user" + block( + f"assistant to={recipient}", + _render_content(message.get("content")), + EOT_STR if end_turn else EOM_STR, + idx=i, + sampled=True, + ) + + if add_generation_prompt: + idx = max(len(messages) - 1, 0) + emit(START_STR, token_id=START_ID, idx=idx, is_generation_prompt=True) + emit("assistant", idx=idx, is_generation_prompt=True) + + return pieces + + def _render_text(self, messages, *, tools=None, add_generation_prompt=False) -> str: + """The rendered prompt as text. Byte-exact against chat_template.jinja.""" + pieces = self._build( + self._prepare(messages), + tools=tools, + add_generation_prompt=add_generation_prompt, + ) + return "".join(p.text for p in pieces) + + def _encode(self, piece: _Piece) -> list[int]: + if piece.token_id is not None: + return [piece.token_id] + return self._tok.encode(piece.text, add_special_tokens=False) + + def get_stop_token_ids(self) -> list[int]: + # NOT <|eom|> (200007): it ends the *reasoning* channel, after which the model emits + # the tool call / final answer. Stopping at eom truncates before the answer. + return [EOT_ID, EOS_ID] + + def render_ids( + self, messages, *, tools=None, add_generation_prompt: bool = False + ) -> list[int]: + pieces = self._build( + self._prepare(messages), + tools=tools, + add_generation_prompt=add_generation_prompt, + ) + ids: list[int] = [] + for piece in pieces: + ids += self._encode(piece) + return ids + + def render( + self, messages, *, tools=None, add_generation_prompt: bool = False + ) -> RenderedTokens: + """Render with per-token attribution for the trainer loss mask. + + Attribution is exact: spans are emitted already labelled, so assistant bodies and + their terminators are marked sampled while the surrounding scaffolding is not. + Splits only ever fall on special-token boundaries, which are atomic in the + tokenizer, so encoding per span concatenates to the same ids as encoding the + whole string at once. + """ + pieces = self._build( + self._prepare(messages), + tools=tools, + add_generation_prompt=add_generation_prompt, + ) + + token_ids: list[int] = [] + message_indices: list[int] = [] + sampled_mask: list[bool] = [] + is_content: list[bool] = [] + for piece in pieces: + encoded = self._encode(piece) + token_ids += encoded + message_indices += [piece.msg_idx] * len(encoded) + sampled_mask += [piece.sampled] * len(encoded) + # is_content is NOT sampled_mask: a user or tool message's body is content + # the caller supplied even though the model never sampled it. Only the + # header scaffolding is excluded on every role. + is_content += [piece.is_body] * len(encoded) + + return RenderedTokens( + token_ids=token_ids, + message_indices=message_indices, + sampled_mask=sampled_mask, + is_content=is_content, + message_roles=[m.get("role") for m in messages], + message_tool_names=[m.get("name") for m in messages], + multi_modal_data=None, + ) + + def bridge_to_next_turn( + self, + previous_prompt_ids: list[int], + previous_completion_ids: list[int], + new_messages, + *, + tools=None, + ) -> RenderedTokens | None: + """Extend prompt + sampled completion with the next turn, without re-rendering. + + Re-rendering the previous turn would round-trip the completion through parse and + back, which can change its tokenization. Keeping the sampled ids verbatim and + appending only the new messages avoids that drift, so the completion the trainer + sees stays bitwise what the generator produced. + + Returns ``None`` -- the caller then re-renders -- when the extension cannot be + appended safely: no prior prompt, nothing to add, an assistant turn in the + extension (its terminator depends on the following message's role, which the + bridge cannot see), or a prior turn with no ``<|eot|>`` to attach to. + + The output is a prompt, so nothing in it is sampled. ``message_indices`` follows + the library convention: -1 over the carried-forward prefix and the trailing + generation prompt, and the index into ``new_messages`` elsewhere. + """ + if not previous_prompt_ids or not new_messages: + return None + if reject_assistant_in_extension(new_messages): + return None + # Under a retention policy that drops history at user-query boundaries, the next + # prompt is not a suffix of this one, so the bridge cannot extend it. + if should_rerender_for_thinking_retention( + self.effective_thinking_retention, new_messages + ): + return None + + # A completion truncated at max_tokens has no terminator; synthesizing <|eot|> + # closes the turn the same way the template would. + previous_ids = trim_to_turn_close( + previous_prompt_ids, + previous_completion_ids, + {EOT_ID}, + synthesize_close=EOT_ID, + ) + if previous_ids is None: + return None + + prepared = self._prepare(list(new_messages)) + # _build always emits the leading bos, and synthesises a default system block + # when no system message is present. Both already exist in previous_ids, so the + # extension starts after them. + pieces = self._build(prepared, tools=tools, add_generation_prompt=True) + skip_synthesized_system = not any(m.get("role") == "system" for m in prepared) + + ext: list[int] = [] + ext_indices: list[int] = [] + ext_content: list[bool] = [] + blocks_seen = 0 + for piece in pieces: + if piece.token_id == START_ID: + blocks_seen += 1 + if piece.token_id == self._tok.bos_token_id and piece.text == self._bos: + continue + if skip_synthesized_system and blocks_seen == 1: + continue + encoded = self._encode(piece) + ext += encoded + ext_indices += [-1 if piece.is_generation_prompt else piece.msg_idx] * len( + encoded + ) + ext_content += [piece.is_body] * len(encoded) + + total = len(previous_ids) + len(ext) + return RenderedTokens( + token_ids=previous_ids + ext, + message_indices=[-1] * len(previous_ids) + ext_indices, + sampled_mask=[False] * total, + is_content=[False] * len(previous_ids) + ext_content, + message_roles=[m.get("role") or "" for m in new_messages], + message_tool_names=extract_message_tool_names(new_messages), + multi_modal_data=None, + ) + + # ------------------------------------------------------------------ parsing + + def parse_response(self, token_ids, *, tools=None) -> ParsedResponse: + """Parse a Muse Glimmer completion into (content, reasoning_content, tool_calls). + + Muse Glimmer emits harmony channels `to=<|message|>`: + - recipient == "self" -> reasoning + - body has an ATEM block -> tool call(s) (env.step reads these) + - otherwise ("user", ...) -> final answer content + Validated against real generations. + """ + text = self._tok.decode(token_ids, skip_special_tokens=False) + + reasoning_parts, content_parts, tool_calls = [], [], [] + for recipient, body in _parse_channels(text): + atem = parse_atem_tool_calls(body) + if atem: + tool_calls += [ + ParsedToolCall( + raw=render_atem_tool_call(c["name"], c["arguments"]), + name=c["name"], + arguments=c["arguments"], + ) + for c in atem + ] + elif recipient == "self": + reasoning_parts.append(body) + else: + # strip any stray ATEM remnants; keep the plain answer text + content_parts.append(_FUNCTION_CALLS_BLOCK.sub("", body).strip()) + + content = "\n".join(p for p in content_parts if p).strip() + reasoning = "\n".join(reasoning_parts).strip() or None + # Opt-in: recover an answer from reasoning when the model produced no + # user-facing channel and no tool call. See the config field for why this is + # off by default. + if ( + self._config.answer_from_reasoning_fallback + and not content + and not tool_calls + and reasoning + ): + lines = [ln.strip() for ln in reasoning.splitlines() if ln.strip()] + if lines: + content = lines[-1] + return ParsedResponse( + content=content, + reasoning_content=reasoning, + tool_calls=tool_calls, + ) + + +def register() -> None: + """Install the muse_glimmer renderer into the ``renderers`` library registry. + + Uses the library's public extension surface -- implement the ``Renderer`` + Protocol, then add the class to ``RENDERER_REGISTRY`` and its config to + ``_CONFIG_BY_NAME`` -- so ``create_renderer(config_from_name("muse_glimmer"))`` + resolves it. Also maps the ``muse_glimmer`` TorchTitan model name to it, which is what + ``RendererConfig(name="muse_glimmer")`` looks up. + + Idempotent. Delete this once the renderer is upstreamed to + PrimeIntellect-ai/renderers (only the _RENDERER_BY_MODEL entry stays). + """ + from renderers import base as renderers_base, configs as renderers_configs + + from torchtitan.experiments.rl.renderer import _RENDERER_BY_MODEL + + # Populate the library's built-ins first: _populate_registry() early-returns if + # RENDERER_REGISTRY is already non-empty, so registering before it runs would + # suppress every built-in renderer. + renderers_base._populate_registry() + + renderers_configs._CONFIG_BY_NAME.setdefault( + RENDERER_NAME, MuseGlimmerRendererConfig + ) + renderers_base.RENDERER_REGISTRY[RENDERER_NAME] = MuseGlimmerRenderer + _RENDERER_BY_MODEL["muse_glimmer"] = RENDERER_NAME diff --git a/torchtitan/experiments/rl/renderer.py b/torchtitan/experiments/rl/renderer.py index 7c5c0ccf72..013f945ae7 100644 --- a/torchtitan/experiments/rl/renderer.py +++ b/torchtitan/experiments/rl/renderer.py @@ -23,6 +23,12 @@ "qwen3_vl": "qwen3-vl", "gpt_oss": "gpt-oss", "deepseek_v3": "deepseek-v3", + # TODO: upstream the Muse Glimmer renderer to PrimeIntellect-ai/renderers, then + # delete its `register()` and point this at the library's name (hyphenated, like + # the entries above). It ships in torchtitan and self-registers only because the + # library has no Muse Glimmer renderer yet; every other model here resolves to one + # the library owns. See rl/models/muse_glimmer/renderer.py. + "muse_glimmer": "muse_glimmer", "default": "default", # llama3 "auto": "auto", # ignores knobs, resolves from tokenizer, } @@ -73,6 +79,21 @@ def build(self, *, tokenizer_path: str) -> Renderer: # `name=None` (or "auto") -> let `create_renderer` resolve from the tokenizer. renderer_name = _RENDERER_BY_MODEL.get(self.name, self.name) + if renderer_name == "muse_glimmer": + # TODO: temporary. Delete this block once the Muse Glimmer renderer is + # upstreamed to PrimeIntellect-ai/renderers -- the library registers its + # own renderers in _populate_registry(), so no torchtitan-side hook is + # needed for any other model here. + # + # Until then it has to live in build(), not in the config registry: the + # renderer is constructed inside a Monarch-spawned RolloutWorker, which + # only receives the serialized config and never imports the recipe module, + # so a register() call there never runs in that process. + from torchtitan.experiments.rl.models.muse_glimmer import ( + renderer as _muse_glimmer_renderer, + ) + + _muse_glimmer_renderer.register() renderer_config = config_from_name(renderer_name) if renderer_name else None if renderer_config is None: return create_renderer(tokenizer, None) diff --git a/torchtitan/experiments/rl/tests/test_muse_glimmer_renderer.py b/torchtitan/experiments/rl/tests/test_muse_glimmer_renderer.py new file mode 100644 index 0000000000..9a0410d845 --- /dev/null +++ b/torchtitan/experiments/rl/tests/test_muse_glimmer_renderer.py @@ -0,0 +1,516 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Parity between the native Muse Glimmer renderer and the published chat template. + +The renderer builds harmony/ATEM prompts in Python instead of calling +``apply_chat_template``. These tests are the spec for that: every case asserts the +native render is byte-identical to what the checkpoint's own ``chat_template.jinja`` +produces, so a template change upstream fails here first. + +Needs a Muse Glimmer tokenizer; point ``MUSE_GLIMMER_TOKENIZER`` at a local checkpoint +directory or leave it unset to pull the public one from the Hub. + + pytest torchtitan/experiments/rl/tests/test_muse_glimmer_renderer.py -v +""" + +from __future__ import annotations + +import os + +import pytest + +from torchtitan.experiments.rl.models.muse_glimmer.renderer import ( + EOM_ID, + EOT_ID, + MESSAGE_ID, + MuseGlimmerRenderer, + MuseGlimmerRendererConfig, + START_ID, +) + +DEFAULT_TOKENIZER = "meta-models/Muse-Glimmer-30B" + +TOOLS = [ + { + "type": "function", + "function": { + "name": "search", + "description": "Run a web search", + "parameters": { + "type": "object", + "properties": { + "query": {"type": "string"}, + "topk": {"type": "integer"}, + }, + "required": ["query"], + }, + }, + } +] +# Namespaced names exercise the "# Valid recipients" and "// Tool metadata" grouping. +NAMESPACED_TOOLS = [ + { + "type": "function", + "function": { + "name": "web.search", + "description": 'Search with a "quoted" phrase', + "parameters": {"type": "object", "properties": {"q": {"type": "string"}}}, + }, + }, + { + "type": "function", + "function": { + "name": "web.fetch", + "description": "Fetch a URL", + "parameters": {"type": "object", "properties": {"url": {"type": "string"}}}, + }, + }, +] + + +def _tool_call(name="search", **arguments): + return {"type": "function", "function": {"name": name, "arguments": arguments}} + + +# (id, messages, tools) -- each rendered both with and without a generation prompt. +CASES = [ + ("user_only", [{"role": "user", "content": "who wrote Blade Runner"}], None), + ("user_with_tools", [{"role": "user", "content": "who wrote Blade Runner"}], TOOLS), + ( + "namespaced_tools", + [{"role": "user", "content": "search please"}], + NAMESPACED_TOOLS, + ), + ( + "explicit_system", + [ + {"role": "system", "content": "You are terse."}, + {"role": "user", "content": "hi"}, + ], + TOOLS, + ), + ( + "system_declaring_reasoning", + [ + {"role": "system", "content": "Be terse.\n\nReasoning strength: low."}, + {"role": "user", "content": "hi"}, + ], + None, + ), + ( + "system_saying_reasoning_effort", + [ + {"role": "system", "content": "Be terse. Reasoning effort: low."}, + {"role": "user", "content": "hi"}, + ], + None, + ), + ( + "assistant_answer", + [ + {"role": "user", "content": "who wrote Blade Runner"}, + {"role": "assistant", "content": "Philip K. Dick"}, + ], + None, + ), + ( + "assistant_reasoning_then_answer", + [ + {"role": "user", "content": "2+2?"}, + { + "role": "assistant", + "reasoning_content": "simple arithmetic", + "content": "4", + }, + ], + None, + ), + ( + "tool_call_then_output", + [ + {"role": "user", "content": "who wrote Blade Runner"}, + { + "role": "assistant", + "reasoning_content": "I should search.", + "tool_calls": [_tool_call(query="Blade Runner author", topk=3)], + }, + {"role": "tool", "name": "search", "content": "Philip K. Dick"}, + {"role": "assistant", "content": "Philip K. Dick"}, + ], + TOOLS, + ), + ( + "parallel_tool_calls", + [ + {"role": "user", "content": "two things"}, + { + "role": "assistant", + "tool_calls": [_tool_call(query="a"), _tool_call(query="b")], + }, + ], + TOOLS, + ), + ( + "tool_args_coercion", + [ + {"role": "user", "content": "coerce"}, + { + "role": "assistant", + "tool_calls": [ + _tool_call( + flag=True, + off=False, + nothing=None, + obj={"b": 1, "a": 2}, + arr=[1, "x"], + text="plain", + ) + ], + }, + ], + TOOLS, + ), + ( + "consecutive_same_role", + [ + {"role": "user", "content": "first"}, + {"role": "user", "content": "second"}, + ], + None, + ), + ( + "unicode_content", + [{"role": "user", "content": "café naïve 你好"}], + None, + ), + ( + "multimodal_parts", + [ + { + "role": "user", + "content": [ + {"type": "text", "text": "describe "}, + {"type": "image"}, + {"type": "text", "text": " please"}, + ], + } + ], + None, + ), +] + +CONFIGS = [ + ("defaults", {}), + ("low_reasoning", {"reasoning_strength": "low"}), + ("pinned_dates", {"knowledge_cutoff": "2025-01-01", "current_date": "2026-02-03"}), +] + + +@pytest.fixture(scope="module") +def tokenizer(): + transformers = pytest.importorskip("transformers") + path = os.environ.get("MUSE_GLIMMER_TOKENIZER", DEFAULT_TOKENIZER) + try: + tok = transformers.AutoTokenizer.from_pretrained(path) + except Exception as exc: # offline, or no access to the Hub + pytest.skip(f"Muse Glimmer tokenizer unavailable at {path!r}: {exc}") + if not tok.chat_template: + pytest.skip(f"tokenizer at {path!r} ships no chat template to compare against") + return tok + + +def _renderer(tokenizer, **overrides): + return MuseGlimmerRenderer(tokenizer, MuseGlimmerRendererConfig(**overrides)) + + +def _template_kwargs(config: MuseGlimmerRendererConfig) -> dict: + """Only forward knobs the caller set, so the template applies its own defaults.""" + return { + name: value + for name, value in ( + ("reasoning_strength", config.reasoning_strength), + ("knowledge_cutoff", config.knowledge_cutoff), + ("current_date", config.current_date), + ) + if value is not None + } + + +@pytest.mark.parametrize("config_id,overrides", CONFIGS, ids=[c[0] for c in CONFIGS]) +@pytest.mark.parametrize("add_generation_prompt", [False, True]) +@pytest.mark.parametrize("case_id,messages,tools", CASES, ids=[c[0] for c in CASES]) +def test_matches_chat_template( + tokenizer, case_id, messages, tools, add_generation_prompt, config_id, overrides +): + """The native render is byte-identical to apply_chat_template.""" + renderer = _renderer(tokenizer, **overrides) + # current_date defaults to today in both implementations; pin it so a run that + # straddles midnight cannot produce a spurious mismatch. + config = renderer._config + kwargs = _template_kwargs(config) + if config.current_date is None: + import datetime + + kwargs["current_date"] = datetime.datetime.now().strftime("%Y-%m-%d") + renderer = _renderer( + tokenizer, **overrides, current_date=kwargs["current_date"] + ) + + expected = tokenizer.apply_chat_template( + messages, + tools=tools, + add_generation_prompt=add_generation_prompt, + tokenize=False, + **kwargs, + ) + actual = renderer._render_text( + messages, tools=tools, add_generation_prompt=add_generation_prompt + ) + assert actual == expected + + +@pytest.mark.parametrize("case_id,messages,tools", CASES, ids=[c[0] for c in CASES]) +def test_render_ids_match_encoding_whole_string(tokenizer, case_id, messages, tools): + """Per-span encoding concatenates to the same ids as encoding the text at once. + + This is what makes the loss mask trustworthy: splitting on special tokens must not + change tokenization. + """ + renderer = _renderer(tokenizer, current_date="2026-02-03") + text = renderer._render_text(messages, tools=tools, add_generation_prompt=True) + assert renderer.render_ids( + messages, tools=tools, add_generation_prompt=True + ) == tokenizer.encode(text, add_special_tokens=False) + + +@pytest.mark.parametrize("case_id,messages,tools", CASES, ids=[c[0] for c in CASES]) +def test_render_mask_aligns_with_ids(tokenizer, case_id, messages, tools): + """render() returns one mask entry per token and the same ids as render_ids().""" + renderer = _renderer(tokenizer, current_date="2026-02-03") + rendered = renderer.render(messages, tools=tools, add_generation_prompt=False) + assert rendered.token_ids == renderer.render_ids( + messages, tools=tools, add_generation_prompt=False + ) + n = len(rendered.token_ids) + assert len(rendered.sampled_mask) == n + assert len(rendered.message_indices) == n + assert len(rendered.is_content) == n + + +def test_mask_excludes_scaffolding_and_covers_assistant_body(tokenizer): + """Scaffolding is never trained on; the assistant body and terminator are.""" + renderer = _renderer(tokenizer, current_date="2026-02-03") + messages = [ + {"role": "user", "content": "who wrote Blade Runner"}, + {"role": "assistant", "content": "Philip K. Dick"}, + ] + rendered = renderer.render(messages) + sampled = [ + tid for tid, keep in zip(rendered.token_ids, rendered.sampled_mask) if keep + ] + + assert sampled, "the assistant turn should contribute trainable tokens" + # Header scaffolding must never be trainable. + assert START_ID not in sampled + assert MESSAGE_ID not in sampled + # The terminator the model itself emits is trainable. + assert sampled[-1] == EOT_ID + # The body round-trips to the assistant text. + assert "Philip K. Dick" in tokenizer.decode(sampled) + # Nothing before the assistant turn is trainable. + first = rendered.sampled_mask.index(True) + assert not any(rendered.sampled_mask[:first]) + + +def test_is_content_covers_non_assistant_bodies(tokenizer): + """is_content marks caller-supplied bodies on every role, not just assistant. + + It is deliberately NOT a copy of sampled_mask: a user or tool message's body is + content the caller supplied even though the model never sampled it. Only header + scaffolding is excluded on every role. + """ + renderer = _renderer(tokenizer, current_date="2026-02-03") + messages = [ + {"role": "user", "content": "who wrote Blade Runner"}, + { + "role": "assistant", + "tool_calls": [_tool_call(query="Blade Runner author")], + }, + {"role": "tool", "name": "search", "content": "Philip K. Dick wrote it"}, + {"role": "assistant", "content": "Philip K. Dick"}, + ] + rendered = renderer.render(messages, tools=TOOLS) + content = [ + tid for tid, keep in zip(rendered.token_ids, rendered.is_content) if keep + ] + text = tokenizer.decode(content) + + # Bodies from every role are present... + assert "who wrote Blade Runner" in text + assert "Philip K. Dick wrote it" in text + # ...and header scaffolding is not content on any role. + assert START_ID not in content + assert MESSAGE_ID not in content + # is_content must differ from sampled_mask now (non-assistant bodies included). + assert rendered.is_content != rendered.sampled_mask + # Every sampled token is still content (assistant bodies + their terminators). + assert all(c for c, s in zip(rendered.is_content, rendered.sampled_mask) if s) + + +def test_bridge_declines_when_retention_requires_rerender(tokenizer): + """thinking_retention='tool_cycle' must stop the bridge at a new user query. + + The default is 'all' (the template renders prior reasoning unconditionally), so + the bridge extends; an explicit tool_cycle policy has to decline instead. + """ + messages = [{"role": "user", "content": "q"}] + keep_all = _renderer(tokenizer, current_date="2026-02-03") + assert keep_all.effective_thinking_retention == "all" + + prompt_ids = keep_all.render_ids(messages, add_generation_prompt=True) + new_user = [{"role": "user", "content": "a second question"}] + tool_only = [{"role": "tool", "name": "search", "content": "x"}] + + # Default policy: a new user query is still bridgeable. + assert keep_all.bridge_to_next_turn(prompt_ids, [EOT_ID], new_user) is not None + + cycle = _renderer( + tokenizer, current_date="2026-02-03", thinking_retention="tool_cycle" + ) + assert cycle.effective_thinking_retention == "tool_cycle" + # A new user query crosses the boundary -> must re-render. + assert cycle.bridge_to_next_turn(prompt_ids, [EOT_ID], new_user) is None + # Staying inside the tool cycle is still fine. + assert cycle.bridge_to_next_turn(prompt_ids, [EOT_ID], tool_only) is not None + + +def test_reasoning_channel_is_trainable_and_ends_with_eom(tokenizer): + """A to=self channel is sampled and closes with <|eom|>, not <|eot|>.""" + renderer = _renderer(tokenizer, current_date="2026-02-03") + rendered = renderer.render( + [ + {"role": "user", "content": "2+2?"}, + {"role": "assistant", "reasoning_content": "arithmetic", "content": "4"}, + ] + ) + sampled = [ + tid for tid, keep in zip(rendered.token_ids, rendered.sampled_mask) if keep + ] + assert EOM_ID in sampled, "the reasoning channel terminator should be trainable" + assert sampled[-1] == EOT_ID + + +def test_drop_reasoning_history(tokenizer): + """retain_reasoning_in_history=False removes the to=self channel from history.""" + messages = [ + {"role": "user", "content": "2+2?"}, + {"role": "assistant", "reasoning_content": "arithmetic", "content": "4"}, + ] + kept = _renderer(tokenizer, current_date="2026-02-03")._render_text(messages) + dropped = _renderer( + tokenizer, current_date="2026-02-03", retain_reasoning_in_history=False + )._render_text(messages) + assert "to=self" in kept + assert "to=self" not in dropped + + +def test_bridge_extends_without_retokenizing_the_completion(tokenizer): + """The bridge preserves sampled ids verbatim and appends the new turn.""" + renderer = _renderer(tokenizer, current_date="2026-02-03") + prompt_ids = renderer.render_ids( + [{"role": "user", "content": "who wrote Blade Runner"}], + tools=TOOLS, + add_generation_prompt=True, + ) + completion_ids = tokenizer.encode( + ' to=search<|message|>\n\n' + 'Blade Runner\n' + "\n<|eot|>", + add_special_tokens=False, + ) + bridged = renderer.bridge_to_next_turn( + prompt_ids, + completion_ids, + [{"role": "tool", "name": "search", "content": "Philip K. Dick"}], + tools=TOOLS, + ) + assert bridged is not None + # The caller reads `.token_ids`, so the bridge must return RenderedTokens, not a list. + carried = len(prompt_ids) + len(completion_ids) + assert bridged.token_ids[:carried] == prompt_ids + completion_ids + suffix = tokenizer.decode(bridged.token_ids[carried:]) + assert suffix.startswith("<|start|>tool search<|message|>") + assert suffix.endswith("<|start|>assistant") + # The bridge must not re-emit the system preamble already in prompt_ids. + assert "Valid recipients" not in suffix + + # A bridge produces a prompt, so nothing in it is trainable, and every parallel + # array must line up with token_ids or the trainer will index past the end. + n = len(bridged.token_ids) + assert len(bridged.sampled_mask) == n + assert len(bridged.message_indices) == n + assert len(bridged.is_content) == n + assert not any(bridged.sampled_mask) + # -1 over the carried prefix and the trailing generation prompt; 0 for new_messages[0]. + assert bridged.message_indices[:carried] == [-1] * carried + assert set(bridged.message_indices[carried:]) <= {0, -1} + assert bridged.message_indices[-1] == -1 + assert bridged.message_roles == ["tool"] + + +def test_bridge_returns_none_when_it_cannot_extend_safely(tokenizer): + """Cases the bridge must decline so the caller re-renders instead.""" + renderer = _renderer(tokenizer, current_date="2026-02-03") + prompt_ids = renderer.render_ids( + [{"role": "user", "content": "q"}], add_generation_prompt=True + ) + tool_msg = [{"role": "tool", "name": "search", "content": "x"}] + + assert renderer.bridge_to_next_turn([], [EOT_ID], tool_msg) is None + assert renderer.bridge_to_next_turn(prompt_ids, [EOT_ID], []) is None + # An assistant turn's terminator depends on the *next* message's role, which the + # bridge cannot see, so it must decline rather than guess. + assert ( + renderer.bridge_to_next_turn( + prompt_ids, [EOT_ID], [{"role": "assistant", "content": "hi"}] + ) + is None + ) + + +def test_bridge_synthesizes_a_close_for_a_truncated_completion(tokenizer): + """A completion cut off at max_tokens has no terminator; the bridge adds one.""" + renderer = _renderer(tokenizer, current_date="2026-02-03") + prompt_ids = renderer.render_ids( + [{"role": "user", "content": "q"}], add_generation_prompt=True + ) + truncated = tokenizer.encode( + " to=user<|message|>cut off mid", add_special_tokens=False + ) + bridged = renderer.bridge_to_next_turn( + prompt_ids, truncated, [{"role": "tool", "name": "search", "content": "x"}] + ) + assert bridged is not None + carried = len(prompt_ids) + len(truncated) + assert bridged.token_ids[:carried] == prompt_ids + truncated + assert bridged.token_ids[carried] == EOT_ID # synthesized turn close + + +def test_rejects_incompatible_tokenizer(tokenizer): + """Special-token ids are load-bearing, so a mismatch fails loudly.""" + + class Shifted: + def __init__(self, inner): + self._inner = inner + self.bos_token = inner.bos_token + self.bos_token_id = inner.bos_token_id + + def convert_tokens_to_ids(self, token): + return 1 + self._inner.convert_tokens_to_ids(token) + + with pytest.raises(ValueError, match="not compatible"): + MuseGlimmerRenderer(Shifted(tokenizer), MuseGlimmerRendererConfig())