diff --git a/examples/speculative_decoding/requirements.txt b/examples/speculative_decoding/requirements.txt index 0c679b55024..c13df8ca017 100644 --- a/examples/speculative_decoding/requirements.txt +++ b/examples/speculative_decoding/requirements.txt @@ -1,3 +1,3 @@ accelerate>=1.12.0 peft==0.18.1 -transformers>=5.0,<5.4 +transformers>=5.0,<5.13 diff --git a/modelopt/torch/speculative/plugins/hf_streaming_dataset.py b/modelopt/torch/speculative/plugins/hf_streaming_dataset.py index aebb726e27f..01370914912 100644 --- a/modelopt/torch/speculative/plugins/hf_streaming_dataset.py +++ b/modelopt/torch/speculative/plugins/hf_streaming_dataset.py @@ -41,6 +41,7 @@ import base64 import os +import re import time from typing import TypedDict @@ -105,6 +106,22 @@ def _tokenize_with_loss_mask( recovery = None if answer_only_loss and not getattr(tokenizer, "is_fast", False): recovery = get_loss_mask_recovery(tokenizer) + if answer_only_loss and recovery is None: + # Fail loudly on a template without {% generation %} tags: transformers only + # warns and returns an ALL-ZERO assistant mask, which trains every sample at + # zero loss with no other symptom. (The regex matches the tag itself, not the + # unrelated `add_generation_prompt` identifier most templates contain.) + template = getattr(tokenizer, "chat_template", None) or "" + if not re.search(r"\{%-?\s*generation\b", template): + raise RuntimeError( + "answer_only_loss=True needs assistant masks, but the tokenizer's chat " + "template has no {% generation %} tags, so apply_chat_template would " + "return an all-zero mask and training would silently run at zero loss. " + "Pass a tagged template copy via data.chat_template (see " + "examples/MiniMaxAI/MiniMax-M3/m3_chat_template_generation.jinja for a " + "worked example), register a loss-mask recovery " + "(modelopt.torch.utils.loss_mask), or set answer_only_loss=false." + ) out = tokenizer.apply_chat_template( conversations, tokenize=True, diff --git a/modelopt/torch/speculative/plugins/modeling_fakebase.py b/modelopt/torch/speculative/plugins/modeling_fakebase.py index f896abab696..2b5fe989c03 100644 --- a/modelopt/torch/speculative/plugins/modeling_fakebase.py +++ b/modelopt/torch/speculative/plugins/modeling_fakebase.py @@ -204,7 +204,9 @@ def from_source(cls, source: str, trust_remote_code: bool = False) -> "FakeBaseM intermediate_size=getattr(base_cfg, "intermediate_size", None), rms_norm_eps=getattr(base_cfg, "rms_norm_eps", 1e-6), rope_theta=getattr(base_cfg, "rope_theta", None), - final_norm_type=_select_final_norm_type(getattr(base_cfg, "model_type", None)), + final_norm_type=_select_final_norm_type( + getattr(base_cfg, "model_type", None), base_cfg + ), ) model = cls(config) # Load lm_head, embed_tokens, and (for known models) the final norm into the model. diff --git a/modelopt/torch/speculative/plugins/modeling_final_norm.py b/modelopt/torch/speculative/plugins/modeling_final_norm.py index 85852bc2bef..09e7840707b 100644 --- a/modelopt/torch/speculative/plugins/modeling_final_norm.py +++ b/modelopt/torch/speculative/plugins/modeling_final_norm.py @@ -22,6 +22,8 @@ only when we know which one the model uses. """ +from collections.abc import Callable + import torch from transformers.models.llama.modeling_llama import LlamaRMSNorm @@ -41,11 +43,37 @@ def __init__(self, hidden_size, eps=1e-6, dtype=torch.bfloat16): self.to(dtype) +class _FinalGemmaRMSNorm(torch.nn.Module): + """Gemma-style RMSNorm: fp32 normalize, scale by ``(1 + weight)``, multiply-then-cast. + + Mirrors MiniMax M3's ``MiniMaxM3VLRMSNorm`` (``use_gemma_norm=True`` configs): the stored + ``weight`` is a zero-centered delta, the effective scale is ``1 + weight``, and the multiply + happens in fp32 BEFORE casting back (``(x_hat * (1 + w)).type_as(x)``), unlike LlamaRMSNorm's + cast-then-multiply. Reusing ``_FinalRMSNorm`` here would silently drop the ``+1`` and corrupt + the reconstructed logits. ``weight`` is loaded from the base checkpoint. + """ + + def __init__(self, hidden_size, eps=1e-6, dtype=torch.bfloat16): + super().__init__() + self.eps = eps + self.weight = torch.nn.Parameter(torch.zeros(hidden_size, dtype=dtype)) + + def forward(self, x): + output = x.float() + output = output * torch.rsqrt(output.pow(2).mean(-1, keepdim=True) + self.eps) + output = output * (1.0 + self.weight.float()) + return output.type_as(x) + + def extra_repr(self): + return f"{tuple(self.weight.shape)}, eps={self.eps}" + + # Registry of self-implemented final-norm variants. We deliberately reimplement these # (rather than importing the base model's actual module) to keep FakeBaseModel lightweight. # Only a small, explicit set is supported; add a class here when a new type is needed. -_FINAL_NORM_CLASSES = { +_FINAL_NORM_CLASSES: dict[str, Callable[..., torch.nn.Module]] = { "rmsnorm": _FinalRMSNorm, + "gemma_rmsnorm": _FinalGemmaRMSNorm, } # Base ``model_type`` → final-norm type. ONLY listed models get a norm — applying the wrong or @@ -62,17 +90,26 @@ def __init__(self, hidden_size, eps=1e-6, dtype=torch.bfloat16): "deepseek_v3": "rmsnorm", "kimi_k2": "rmsnorm", # Kimi-K2 / K2-Thinking (DeepSeek-V3 arch) report model_type "kimi_k2" "kimi_k25": "rmsnorm", # Kimi-K2.5 / K2.6 / K2.7 all report model_type "kimi_k25" + "minimax_m3_vl_text": "rmsnorm", # overridden to gemma_rmsnorm via use_gemma_norm below # gpt_oss intentionally DISABLED: GptOssRMSNorm uses an fp32 weight + multiply-then-cast, # unlike _FinalRMSNorm's bf16 weight, so reusing it would silently bias reconstructed logits. # Re-enable once a gpt_oss-style class (fp32 weight, multiply-then-cast) is in _FINAL_NORM_CLASSES. } -def _select_final_norm_type(model_type: str | None) -> str | None: +def _select_final_norm_type(model_type: str | None, base_cfg=None) -> str | None: """Return the final-norm type for a base ``model_type``, or ``None`` if unknown. ``None`` means we don't know the model's final norm, so FakeBaseModel builds no norm. + + ``base_cfg`` (the resolved text config, optional) takes precedence over the model_type + table when it carries an explicit ``use_gemma_norm=True`` flag: only MiniMax M2.x/M3 set + it, and their ``model_type`` is unreliable — MiniMax's VL remote code coerces a + model_type-less ``text_config`` to Mixtral (whose table entry is plain ``rmsnorm``), so + keying off model_type alone would silently apply the wrong norm flavor. """ + if base_cfg is not None and getattr(base_cfg, "use_gemma_norm", False): + return "gemma_rmsnorm" return _FINAL_NORM_TYPE_BY_MODEL_TYPE.get(model_type or "") diff --git a/tests/unit/torch/speculative/plugins/test_hf_streaming_dataset.py b/tests/unit/torch/speculative/plugins/test_hf_streaming_dataset.py index 0cb581f0296..94f89691445 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_streaming_dataset.py +++ b/tests/unit/torch/speculative/plugins/test_hf_streaming_dataset.py @@ -484,3 +484,49 @@ def handler(request: httpx.Request) -> httpx.Response: ) ds[0] assert seen_ports == {advertised_port} + + +# --------------------------------------------------------------------------- +# answer_only_loss template guard +# --------------------------------------------------------------------------- + + +def _fast_tokenizer_with_template(template: str, seq: int = 8) -> MagicMock: + """Fast-tokenizer mock with a given chat template; returns ids + assistant_masks.""" + tok = MagicMock() + tok.is_fast = True + tok.chat_template = template + tok.apply_chat_template.return_value = { + "input_ids": torch.arange(seq, dtype=torch.long).unsqueeze(0), + "assistant_masks": torch.ones(1, seq, dtype=torch.long), + } + return tok + + +_CONV = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}] + + +def test_answer_only_loss_rejects_template_without_generation_tags(): + """A fast tokenizer whose template lacks {% generation %} tags fails loudly. + + Without the guard transformers only warns and returns an ALL-ZERO assistant + mask -- every sample then trains at zero loss with no other symptom. + """ + tok = _fast_tokenizer_with_template("{{ bos }}{% if add_generation_prompt %}x{% endif %}") + with pytest.raises(RuntimeError, match="generation"): + hf_streaming_dataset._tokenize_with_loss_mask(tok, _CONV, answer_only_loss=True) + + +def test_answer_only_loss_accepts_tagged_template(): + """Templates carrying {% generation %} (either whitespace-control form) pass.""" + for tag in ("{% generation %}", "{%- generation -%}"): + tok = _fast_tokenizer_with_template("{{ bos }}" + tag + "{{ c }}") + ids, mask = hf_streaming_dataset._tokenize_with_loss_mask(tok, _CONV, answer_only_loss=True) + assert mask.sum() == ids.shape[-1] + + +def test_full_loss_skips_template_guard(): + """answer_only_loss=False never consults the template (mask is all ones).""" + tok = _fast_tokenizer_with_template("{{ bos }}") + ids, mask = hf_streaming_dataset._tokenize_with_loss_mask(tok, _CONV, answer_only_loss=False) + assert mask.sum() == ids.shape[-1] diff --git a/tools/launcher/common/eagle3/train_eagle_streaming.sh b/tools/launcher/common/eagle3/train_eagle_streaming.sh index 2f1e165062c..72b7388c3b1 100755 --- a/tools/launcher/common/eagle3/train_eagle_streaming.sh +++ b/tools/launcher/common/eagle3/train_eagle_streaming.sh @@ -47,6 +47,9 @@ # SERVE_CPU_OFFLOAD_GB GB/GPU offloaded to host RAM (fits big models on too-few GPUs; slower) # SERVE_MAX_MODEL_LEN cap context length (trims KV/activation) # SERVE_MAX_NUM_SEQS cap concurrent sequences (trims KV/activation) +# SERVE_BLOCK_SIZE KV-cache block size (e.g. 128 for MiniMax-M3 MSA sparse attention). +# Needs its own knob: multi-token SERVE_EXTRA_ARGS values are mangled +# by nemo_run's unquoted env export, so "--block-size 128" cannot ride it. # SERVE_HOST single-node: bind/connect host. default 127.0.0.1 # SERVE_GPU single-node: CUDA_VISIBLE_DEVICES for vllm. default "0" # SERVE_TP tensor-parallel size. default 1 single-node / all serve-node GPUs @@ -153,6 +156,7 @@ launch_vllm() { [ -n "${SERVE_CPU_OFFLOAD_GB:-}" ] && opt_args+=(--cpu-offload-gb "$SERVE_CPU_OFFLOAD_GB") [ -n "${SERVE_MAX_MODEL_LEN:-}" ] && opt_args+=(--max-model-len "$SERVE_MAX_MODEL_LEN") [ -n "${SERVE_MAX_NUM_SEQS:-}" ] && opt_args+=(--max-num-seqs "$SERVE_MAX_NUM_SEQS") + [ -n "${SERVE_BLOCK_SIZE:-}" ] && opt_args+=(--block-size "$SERVE_BLOCK_SIZE") # --no-enable-chunked-prefill / --no-enable-prefix-caching: connector captures hidden states during prefill; both skip recomputing cached/partial prefixes, yielding short/empty hidden_states. Required. # --no-enable-flashinfer-autotune: on NVFP4 MoE the autotuner re-tunes on the first serving step and stalls a worker past vLLM's execute-model timeout, killing EngineCore. # Hidden states move serve -> trainer over NIXL RDMA (no disk round-trip): one diff --git a/tools/launcher/examples/MiniMaxAI/MiniMax-M3/hf_streaming_dspark_multi_node.yaml b/tools/launcher/examples/MiniMaxAI/MiniMax-M3/hf_streaming_dspark_multi_node.yaml new file mode 100644 index 00000000000..d3e6290b13a --- /dev/null +++ b/tools/launcher/examples/MiniMaxAI/MiniMax-M3/hf_streaming_dspark_multi_node.yaml @@ -0,0 +1,134 @@ +# DSpark streaming speculative-decoding training for MiniMax-M3 (multi-node). +# DSpark = the DFlash backbone + a lightweight Markov head + a confidence head, +# generating a causal block semi-autoregressively; see dspark.yaml for the head +# and loss config. Runs the shared streaming pipeline +# (common/eagle3/train_eagle_streaming.sh) with the M3-specific base, draft dims, +# mask token and chat template; trained from scratch. A starting point for +# reproduction — tune node counts, batch, steps and serve limits for your cluster. +# +# MiniMax-M3 specifics this yaml encodes (each was a silent failure mode): +# * SERVE_BLOCK_SIZE=128: M3's MSA sparse attention (sparse_block_size=128) +# requires KV block 128 or vLLM dies at engine init ("No common block size"). +# * data.chat_template: M3 ships a FAST tokenizer whose template has no +# {% generation %} tags -> assistant_masks come back ALL-ZERO and +# answer_only_loss training silently runs at zero loss. The tagged template +# copy next to this yaml wraps the assistant turn (think prefix + content + +# tool calls + eos) in {% generation %} tags. +# * The draft does NOT inherit base GQA/FFN dims (set explicitly below), and +# M3's base rope_theta (5e6) is pinned onto the draft. +# * EAGLE_CAPTURE_IDS = draft default target_layer_ids+1 (6 aux) + final (60). +# * M3 is a VLM wrapper (text_config nested): the base's gemma-style final +# norm is selected via the config's use_gemma_norm flag (see +# modeling_final_norm.py) — its text_config coerces to a mixtral model_type, +# so the model_type table alone would pick the wrong norm. +# +# Run ON the cluster login node (paramiko can't reach it through the login proxy): +# export SLURM_HOST=localhost SLURM_ACCOUNT= \ +# SLURM_PARTITION= \ +# SLURM_HF_LOCAL= \ +# SLURM_JOB_DIR= \ +# NEMORUN_HOME=$PWD +# uv run launch.py --yaml examples/MiniMaxAI/MiniMax-M3/hf_streaming_dspark_multi_node.yaml \ +# identity=$HOME/.ssh/id_ecdsa detach=True --yes +# +# The export lands in /scratchspace/export. + +job_name: MiniMax-M3_DSpark_streaming_multi_node +pipeline: + allow_to_fail: false + skip: false + note: + + global_vars: + hf_model: /hf-local/MiniMaxAI/MiniMax-M3 + + # Build /scratchspace/data/train.jsonl. Point data.data_path at the full + # Spec-Decoding-Dataset-v2 corpus to reproduce; eagle_utils also accepts a + # directory of *.jsonl shards directly. + task_0: + script: common/eagle3/make_dataset.sh + args: + - -f modules/Model-Optimizer/examples/dataset/example_data_config.yaml + - --full-conversations + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 8 + container: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc10 + + task_1: + script: common/eagle3/train_eagle_streaming.sh + args: + - --config modules/Model-Optimizer/modelopt_recipes/general/speculative_decoding/dspark.yaml + - model.model_name_or_path=<> + - model.use_fake_base_for_offline=true + - model.trust_remote_code=true + - data.mode=streaming + - data.data_path=/scratchspace/data/train.jsonl + # M3's own template has no {% generation %} tags; without this tagged copy + # answer_only_loss trains on an all-zero mask (see header). + - data.chat_template=examples/MiniMaxAI/MiniMax-M3/m3_chat_template_generation.jinja + - training.output_dir=/scratchspace/dspark + - training.training_seq_len=4096 + - training.disable_tqdm=true + - training.ar_validate_steps=500000 + - training.num_train_epochs=1 + - training.per_device_train_batch_size=4 + - training.gradient_accumulation_steps=1 + - training.save_steps=1000 + - training.logging_steps=20 + - training.learning_rate=1.0e-4 + - training.warmup_steps=2000 + - training.answer_only_loss=true + # The vLLM serve container has no tensorboard -> trainer init crash. + - training.report_to=none + # The DSpark draft does NOT inherit the base GQA/FFN dims, so set them + # explicitly to match the M3 backbone (else a silently wrong-shape draft). + # intermediate_size matches M3's dense FFN (dense_intermediate_size). + - dflash.dflash_architecture_config.num_hidden_layers=6 + - dflash.dflash_architecture_config.num_key_value_heads=8 + - dflash.dflash_architecture_config.intermediate_size=12288 + # Pin the base's rope_theta onto the draft (M3 uses 5e6, not the Qwen3 + # default 1e6; a mismatch trains rope into the weights and caps AL). + - dflash.dflash_architecture_config.rope_theta=5000000 + # Semi-AR generation block (dspark.yaml ships 16; Kimi/M3 runs use 8). + - dflash.dflash_block_size=8 + # M3 has no mask token; vocab 200064, added tokens end at 200060 -> 200063 free. + - dflash.dflash_mask_token_id=200063 + environment: + - HF_MODEL_CKPT: <> + # 6 aux capture ids = the draft's default target_layer_ids+1, plus the true + # final hidden (60). Requires the aux-capture fix vllm#46788 (in-tree in + # recent nightlies); mismatched ids silently skew train vs inference. + - EAGLE_CAPTURE_IDS: "[2,13,24,36,47,58,60]" + - SERVE_NODES: "4" + - SERVE_TP: "8" + - STREAMING_NUM_WORKERS: "4" + # M3's custom-modeling base needs trust_remote_code at export and serve. + - EXPORT_EXTRA_ARGS: "--trust_remote_code" + - SERVE_EXTRA_ARGS: "--trust-remote-code" + # REQUIRED for M3: MSA sparse attention needs KV block 128 (dedicated knob — + # multi-token SERVE_EXTRA_ARGS values are mangled by nemo_run's unquoted + # env export, so "--block-size 128" cannot ride it). + - SERVE_BLOCK_SIZE: "128" + - SERVE_MAX_MODEL_LEN: "4160" + - SERVE_MAX_NUM_SEQS: "32" + - SERVE_GPU_MEM_UTIL: "0.9" + - SERVE_READY_TIMEOUT: "3600" + - VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: "1200" + - VLLM_ENGINE_ITERATION_TIMEOUT_S: "1200" + # RDMA transport is UCX (InfiniBand) by default. On AWS EFA, uncomment — + # and note UCX SEGFAULTS at agent init on EFA nodes (it detects the EFA + # devices), so LIBFABRIC is required there even for single-node runs: + # - NIXL_BACKENDS: "LIBFABRIC" + # - FI_PROVIDER: "efa" + # - NCCL_IB_DISABLE: "1" + slurm_config: + _factory_: "slurm_factory" + nodes: 6 + ntasks_per_node: 1 + gpus_per_node: 8 + # vLLM x86_64 build with native MiniMax-M3 support (vllm/models/minimax_m3) + # and the aux-capture fix (vllm#46788). + container: diff --git a/tools/launcher/examples/MiniMaxAI/MiniMax-M3/m3_chat_template_generation.jinja b/tools/launcher/examples/MiniMaxAI/MiniMax-M3/m3_chat_template_generation.jinja new file mode 100644 index 00000000000..7a32d7b92f7 --- /dev/null +++ b/tools/launcher/examples/MiniMaxAI/MiniMax-M3/m3_chat_template_generation.jinja @@ -0,0 +1,250 @@ +{# ---------- special token variables ---------- #} +{%- set ns_token = ']<]minimax[>[' -%} +{%- set bod_token = ']~!b[' -%} +{%- set bos_token = ']~b]' -%} +{%- set eos_token = '[e~[' -%} +{%- set toolcall_begin_token = ns_token ~ '' -%} +{%- set toolcall_end_token = ns_token ~ '' -%} +{%- set think_begin_token = '' -%} +{%- set think_end_token = '' -%} +{%- set image_token = ']<]image[>[' -%} +{%- set video_token = ']<]video[>[' -%} +{#- Thinking mode: "enabled" / "disabled" / "adaptive" / not defined -#} +{#- Recursive XML renderer for tool_call arguments ======================== -#} +{#- None values are intentionally skipped in mapping iteration so that + `null` (which would round-trip to the literal string "null") + never appears in the rendered tool_call. The convention is: omit the + field entirely. The top-level `_args` loop applies the same rule. + The `val is none` branch below is a safety net only — upstream cleaning + (drop_none_in_tool_arguments) should ensure no None ever reaches here. -#} +{%- macro to_xml(val, ns) -%} +{%- if val is mapping -%} +{%- for k, v in val.items() if v is not none -%} +{{ ns }}<{{ k }}>{{ to_xml(v, ns) }}{{ ns }} +{%- endfor -%} +{%- elif val is iterable and val is not string -%} +{%- for item in val -%} +{{ ns }}{{ to_xml(item, ns) }}{{ ns }} +{%- endfor -%} +{%- elif val is none -%} +{#- Should be unreachable when upstream cleaning is applied. -#} +{%- elif val is boolean -%} +{{ val | tojson }} +{%- else -%} +{{ val }} +{%- endif -%} +{%- endmacro -%} +{#- Tool Rendering Functions ============================================== -#} +{%- macro render_tool_namespace(namespace_name, tool_list) -%} +{%- for tool in tool_list -%} +{{ tool.function | tojson(ensure_ascii=False) }} +{% endfor -%} +{%- endmacro -%} +{%- macro visible_text(content) -%} + {%- if content is string -%} + {{ content }} + {%- elif content is iterable and content is not mapping -%} + {%- for item in content -%} + {%- if item is mapping and item.type == 'text' -%} + {{- item.text }} + {%- elif item is mapping and item.type == 'image' -%} + {{- image_token }} + {%- elif item is mapping and item.type == 'video' -%} + {{- video_token}} + {%- elif item is string -%} + {{- item }} + {%- endif -%} + {%- endfor -%} + {%- elif content is none -%} + {{- '' }} + {%- else -%} + {{- content }} + {%- endif -%} +{%- endmacro -%} +{#- System Message Construction ============================================ -#} +{%- macro build_system_message(system_message) -%} + {%- if system_message and system_message.content -%} + {{- visible_text(system_message.content) }} + {%- else -%} + {{- 'Your model version is MiniMax-M3, developed by MiniMax. Knowledge cutoff: January 2026. Founded in early 2022, MiniMax is a global AI foundation model company committed to advancing the frontiers of AI towards AGI.' }} + {%- endif -%} + + {#- Thinking mode instructions -#} + {{- '\n\n\n' }} + {{- 'You have a thinking capability that allows you to reason step by step before responding. When thinking is enabled, wrap your reasoning in ' ~ think_begin_token ~ think_end_token ~ ' tags before your response. When thinking is disabled, begin your response directly after the ' ~ think_end_token ~ ' prefix. When thinking is adaptive, decide on your own whether to think for the current turn.\n' }} + {%- if thinking_mode is defined -%} + {%- if thinking_mode == "enabled" -%} + {{- 'Current thinking mode: enabled. You MUST think step by step before every response, including after receiving function/tool results.\n' }} + {%- elif thinking_mode == "disabled" -%} + {{- 'Current thinking mode: disabled. Do not output any thinking process.\n' }} + {%- elif thinking_mode == "adaptive" -%} + {{- 'Current thinking mode: adaptive. You are encouraged to think for complex decision-making, multi-step reasoning, or when analyzing function/tool results.\n' }} + {%- endif -%} + {%- else -%} + {{- 'Current thinking mode: adaptive. You are encouraged to think for complex decision-making, multi-step reasoning, or when analyzing function/tool results.\n' }} + {%- endif -%} + {{- '' }} +{%- endmacro -%} +{%- macro build_developer_message(developer_message) -%} + {%- if developer_message and developer_message.content -%} + {{- visible_text(developer_message.content) }} + {%- else -%} + {%- if model_identity is not defined -%} + {%- set model_identity = "You are a helpful assistant." -%} + {%- endif -%} + {{- model_identity }} + {%- endif -%} +{%- endmacro -%} +{#- Main Template Logic ================================================= -#} +{#- Role mapping: root -> system sp (high priority), system/developer -> developer sp (low priority) -#} +{%- set system_message = none -%} +{%- set developer_message = none -%} +{%- set conversation_messages = messages -%} +{%- if messages and messages[0].role == "root" -%} + {%- set system_message = messages[0] -%} + {%- set conversation_messages = messages[1:] -%} + {%- if conversation_messages and conversation_messages[0].role in ["system", "developer"] -%} + {%- set developer_message = conversation_messages[0] -%} + {%- set conversation_messages = conversation_messages[1:] -%} + {%- endif -%} +{%- elif messages and messages[0].role in ["system", "developer"] -%} + {%- set developer_message = messages[0] -%} + {%- set conversation_messages = messages[1:] -%} +{%- endif -%} +{#- Render system sp (higher priority, root role only) -#} +{{- bod_token ~ bos_token ~ 'system' ~ '\n' }} +{{- build_system_message(system_message) }} +{{- eos_token ~ '\n' }} + +{#- Render developer sp (lower priority: system/developer role + tools) -#} +{{- bos_token ~ 'developer' ~ '\n' }} +{{- build_developer_message(developer_message) }} +{%- if tools -%} + {{- '\n\n' ~ '# Tools' ~ '\n' ~ 'You may call one or more tools to assist with the user query.\nHere are the tools available in JSONSchema format:' ~ '\n' }} + {{- '\n' ~ '' ~ '\n' }} + {{- render_tool_namespace("functions", tools) }} + {{- '' ~ '\n\n' }} + {{- 'To call tools, wrap all invocations in a single ' ~ toolcall_begin_token ~ toolcall_end_token ~ ' block. Parameter values containing nested objects or arrays are recursively expanded into XML elements. Example:\n' }} + {{- '\n' ~ toolcall_begin_token ~ '\n' }} + {{- ns_token + '' }} + {{- ns_token + 'value-1' + ns_token + '' }} + {{- ns_token + '' }} + {{- ns_token + '' }} + {{- ns_token + 'val-a' + ns_token + '' }} + {{- ns_token + 'val-b' + ns_token + '' }} + {{- ns_token + '' }} + {{- ns_token + '' }} + {{- ns_token + '\n' }} + {{- ns_token + '' }} + {{- ns_token + 'value-1' + ns_token + '' }} + {{- ns_token + '\n' }} + {{- toolcall_end_token }} +{%- endif -%} +{{- eos_token ~ '\n' }} + +{#- Render messages -#} +{%- set last_tool_call = namespace(name=none) -%} +{%- for message in conversation_messages -%} + {%- if message.role == 'assistant' -%} + {{- bos_token ~ 'ai' ~ '\n' }} + {%- generation -%} + + {%- set reasoning_content = '' %} + {%- set content = visible_text(message.content) %} + {%- if message.reasoning_content is string %} + {%- set reasoning_content = message.reasoning_content %} + {%- else %} + {%- if think_end_token in content %} + {%- set reasoning_content = content.split(think_end_token)[0].strip('\n').split(think_begin_token)[-1].strip('\n') %} + {%- set content = content.split(think_end_token)[-1].strip('\n') %} + {%- endif %} + {%- endif %} + + {%- if reasoning_content -%} + {#- Render thinking for every assistant turn (all-turn visible) -#} + {{- think_begin_token ~ reasoning_content ~ think_end_token }} + {%- else -%} + {#- No thinking rendered → prefix with think_end_token -#} + {{- think_end_token }} + {%- endif -%} + + {%- if content -%} + {{- content }} + {%- endif -%} + {%- if message.tool_calls -%} + {{- toolcall_begin_token ~ '\n' }} + + {%- for tool_call in message.tool_calls -%} + {%- if tool_call.function -%} + {%- set tool_call = tool_call.function -%} + {%- endif -%} +{{- ns_token + '' }} +{%- set _args = tool_call.arguments -%} +{%- for k, v in _args.items() if v is not none %} +{{- ns_token + '<' + k + '>' -}} +{{- to_xml(v, ns_token) -}} +{{- ns_token + '' }} +{%- endfor -%} +{{- ns_token + '' ~ '\n' }} + {%- endfor -%} + + {{- toolcall_end_token }} + {%- if message.tool_calls[-1].function -%} + {%- set last_tool_call.name = message.tool_calls[-1].function.name -%} + {%- else -%} + {%- set last_tool_call.name = message.tool_calls[-1].name -%} + {%- endif -%} + {%- else -%} + {%- set last_tool_call.name = none -%} + {%- endif -%} + {{- eos_token }} + {%- endgeneration -%} + {{- '\n' }} + + {%- elif message.role == 'tool' -%} + {%- if last_tool_call.name is none -%} + {{- raise_exception("Message has tool role, but there was no previous assistant message with a tool call!") }} + {%- endif -%} + {%- if loop.first or (conversation_messages[loop.index0 - 1].role != 'tool') -%} + {{- bos_token ~ 'tool' }} + {%- endif -%} + {{- '\n' }} + {%- if message.content is string -%} + {{- message.content }} + {%- else -%} + {%- for tr in message.content -%} + {%- if tr is mapping and tr.type is defined and tr.type == 'image' -%} + {{- image_token }} + {%- elif tr is mapping and tr.type is defined and tr.type == 'video' -%} + {{- video_token }} + {%- else -%} + {{- tr.output if tr.output is defined else (tr.text if tr.type == 'text' and tr.text is defined else tr) }} + {%- endif -%} + {%- endfor -%} + {%- endif -%} + {{- '' }} + {%- if loop.last or (conversation_messages[loop.index0 + 1].role != 'tool') -%} + {{- eos_token ~ '\n' -}} + {%- endif -%} + + {%- elif message.role == 'user' -%} + {{- bos_token ~ 'user' ~ '\n' }} + {{- visible_text(message.content) }} + {{- eos_token ~ '\n' }} + {%- endif -%} +{%- endfor -%} + +{#- Generation prompt -#} +{%- if add_generation_prompt -%} +{{- bos_token ~ 'ai' ~ '\n' }} +{%- if thinking_mode is defined and thinking_mode == "disabled" -%} + {{- think_end_token }} +{%- elif thinking_mode is defined and thinking_mode == "adaptive" -%} + {#- adaptive: no prefix, let model decide -#} +{%- elif thinking_mode is defined and thinking_mode == "enabled" -%} + {#- enabled or not defined: default to think -#} + {{- think_begin_token }} +{%- else -%} + {#- adaptive: no prefix, let model decide -#} +{%- endif -%} +{%- endif -%}