From ff65c4b278b8f632202061eff55255a06479fdc5 Mon Sep 17 00:00:00 2001 From: Sami Jaghouar Date: Tue, 7 Jul 2026 09:44:41 -0700 Subject: [PATCH 01/12] feat(sft): renderer-only SFT tokenization SFT tokenization now goes exclusively through the renderers library: - build_training_sample with a hand-coded renderer replaces the incremental apply_chat_template masking path, which corrupted loss masks under position-dependent chat templates (e.g. Qwen3 stripping past thinking blocks across user turns). DefaultRenderer is rejected in setup_renderer; fake data skips renderer construction entirely. - [renderer] config defaults to auto-resolution from the tokenizer; template controls (e.g. enable_thinking) are set run-wide via the typed renderer config. - Non-assistant roles opt into the loss via the renderer's body-only path (content_sft_roles). - _drop_null_fields strips Arrow phantom nulls before tool-call deserialization so genuine nulls inside argument strings survive. - Null-check rather than key-check for messages vs prompt/completion resolution (Arrow schema union adds messages: null to prompt/ completion rows). - seq_len default 256 (FakeDataConfig keeps 128); loss_impl defaults to liger_fused; CI reverse-text configs pin [renderer] name=qwen3. Split from #2942 (part 1/7 of #2485), renderer-only without data_files. --- .../integration/reverse_text_sft/resume.toml | 3 + .../integration/reverse_text_sft/start.toml | 3 + .../reverse_text_sft_lora/resume.toml | 3 + .../reverse_text_sft_lora/start.toml | 3 + docs/training.md | 16 ++- examples/reverse_text/sft.toml | 3 + examples/wordle/sft.toml | 3 + .../src/prime_rl/configs/sft.py | 27 ++-- src/prime_rl/trainer/sft/data.py | 136 +++++++++--------- src/prime_rl/trainer/sft/train.py | 28 ++-- src/prime_rl/utils/chat_template.py | 86 +---------- tests/unit/train/sft/test_sft_dataset.py | 46 +++++- 12 files changed, 171 insertions(+), 186 deletions(-) diff --git a/configs/ci/integration/reverse_text_sft/resume.toml b/configs/ci/integration/reverse_text_sft/resume.toml index 0aecee8e49..ba662f4e9e 100644 --- a/configs/ci/integration/reverse_text_sft/resume.toml +++ b/configs/ci/integration/reverse_text_sft/resume.toml @@ -6,6 +6,9 @@ resume_step = 5 [model] name = "PrimeIntellect/Qwen3-0.6B" +[renderer] +name = "qwen3" + [data] name = "PrimeIntellect/Reverse-Text-SFT" batch_size = 4 diff --git a/configs/ci/integration/reverse_text_sft/start.toml b/configs/ci/integration/reverse_text_sft/start.toml index 405d7609e4..4d3250a8c8 100644 --- a/configs/ci/integration/reverse_text_sft/start.toml +++ b/configs/ci/integration/reverse_text_sft/start.toml @@ -5,6 +5,9 @@ max_steps = 5 [model] name = "PrimeIntellect/Qwen3-0.6B" +[renderer] +name = "qwen3" + [data] name = "PrimeIntellect/Reverse-Text-SFT" batch_size = 4 diff --git a/configs/ci/integration/reverse_text_sft_lora/resume.toml b/configs/ci/integration/reverse_text_sft_lora/resume.toml index 67a6ec12af..13269a683f 100644 --- a/configs/ci/integration/reverse_text_sft_lora/resume.toml +++ b/configs/ci/integration/reverse_text_sft_lora/resume.toml @@ -9,6 +9,9 @@ save_adapter_separately = true [model] name = "PrimeIntellect/Qwen3-0.6B" +[renderer] +name = "qwen3" + [model.lora] rank = 8 target_modules = [ diff --git a/configs/ci/integration/reverse_text_sft_lora/start.toml b/configs/ci/integration/reverse_text_sft_lora/start.toml index 28fb516f2e..700a6d6b44 100644 --- a/configs/ci/integration/reverse_text_sft_lora/start.toml +++ b/configs/ci/integration/reverse_text_sft_lora/start.toml @@ -8,6 +8,9 @@ save_adapter_separately = true [model] name = "PrimeIntellect/Qwen3-0.6B" +[renderer] +name = "qwen3" + [model.lora] rank = 8 target_modules = [ diff --git a/docs/training.md b/docs/training.md index 0372fd68c1..d8d4e7e9b8 100644 --- a/docs/training.md +++ b/docs/training.md @@ -144,12 +144,20 @@ Two accepted layouts: If both columns are present, `messages` takes precedence. -**Tool definitions.** For tool-use SFT, add a `tools` column (OpenAI function-calling format) or `tool_defs` ([`verifiers`](https://github.com/PrimeIntellect-ai/verifiers) rollout format). Each row's value can be either a list of dicts or a JSON-encoded string of a list — both are accepted, and `tool_defs` rows are auto-converted to OAI shape before being passed into the chat template's `tools=...` argument. The `chat_template_kwargs` column, if present, is forwarded verbatim into `apply_chat_template`. +**Tool definitions and renderer controls.** For tool-use SFT, add a `tools` column (OpenAI function-calling format) or `tool_defs` ([`verifiers`](https://github.com/PrimeIntellect-ai/verifiers) rollout format). Each row's value can be either a list of dicts or a JSON-encoded string of a list — both are accepted, and `tool_defs` rows are auto-converted to OAI shape before being passed into the renderer. -**Position-dependent chat templates.** Multi-turn SFT under the default tokenization path (`build_incremental_token_mask`) requires that tokenizing the first _k_ turns of a conversation be a strict prefix of tokenizing all _n ≥ k_ turns. Qwen3's upstream template _violates_ this — it strips past `` blocks across user turns, silently corrupting the loss mask. Two fixes: +Renderer-backed SFT reads template controls from the typed `[renderer]` config in the SFT TOML. For example: + +```toml +[renderer] +name = "qwen3" +enable_thinking = false +``` + +If a model needs another template control, add it to that model's renderer config in `renderers` (for example a new field on the relevant `*RendererConfig`) and consume it in the renderer implementation. + +**Renderer-backed tokenization.** SFT tokenization is renderer-only. The [`renderers`](algorithms.md#renderers) package owns message-to-token conversion and loss attribution end-to-end, so position-dependent chat templates (for example templates that strip past `—` blocks across user turns) do not corrupt the loss mask. `[renderer]` defaults to `name = "auto"`; set a typed renderer config only when you need model-specific template controls. Hand-coded renderers ship for Qwen3, Qwen3.5, GLM-5, GLM-4.5, Kimi K2/K2.5, MiniMax M2, DeepSeek V3, Nemotron 3, GPT-OSS. -- **Enable the renderer** (set a typed `[renderer]` config, e.g. `name = "qwen3"`, recommended; defaults to `"auto"` for RL). The [`renderers`](algorithms.md#renderers) package owns tokenization end-to-end and is robust to position-dependent templates. Hand-coded renderers ship for Qwen3, Qwen3.5, GLM-5, GLM-4.5, Kimi K2/K2.5, MiniMax M2, DeepSeek V3, Nemotron 3, GPT-OSS. Not supported for VLMs. -- **Patched chat template** — the prime-rl–patched checkpoints (e.g. `PrimeIntellect/Qwen3-0.6B`, used in `examples/reverse_text/sft.toml`) ship a chat template that preserves thinking. Or supply your own. See [Algorithms § Multi-Turn Trajectories](algorithms.md#multi-turn-trajectories) for the full picture. diff --git a/examples/reverse_text/sft.toml b/examples/reverse_text/sft.toml index 5b21cee566..02abffaab5 100644 --- a/examples/reverse_text/sft.toml +++ b/examples/reverse_text/sft.toml @@ -5,6 +5,9 @@ max_steps = 100 [model] name = "PrimeIntellect/Qwen3-0.6B" +[renderer] +name = "qwen3" + [data] name = "willcb/R1-reverse-wikipedia-paragraphs-v1-1000" seq_len = 4096 diff --git a/examples/wordle/sft.toml b/examples/wordle/sft.toml index 156c5ff5c6..591b2057f7 100644 --- a/examples/wordle/sft.toml +++ b/examples/wordle/sft.toml @@ -5,6 +5,9 @@ max_steps = 20 [model] name = "PrimeIntellect/Qwen3-1.7B" +[renderer] +name = "qwen3" + [data] name = "willcb/V3-wordle" seq_len = 1024 diff --git a/packages/prime-rl-configs/src/prime_rl/configs/sft.py b/packages/prime-rl-configs/src/prime_rl/configs/sft.py index 6352bbd875..f1fb610069 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/sft.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/sft.py @@ -3,7 +3,7 @@ from typing import Annotated, Literal, TypeAlias from pydantic import Field, model_validator -from renderers import RendererConfig +from renderers import AutoRendererConfig, RendererConfig from prime_rl.configs.shared import ( EnvVars, @@ -30,7 +30,7 @@ class BaseDataConfig(BaseConfig): batch_size: int = Field(128, ge=1) """Global batch size.""" - seq_len: int = Field(128, ge=1) + seq_len: int = Field(256, ge=1) """Sequence length.""" pack_function: Literal["cat", "stack"] = "cat" @@ -51,6 +51,12 @@ def validate_batch_size(self): class FakeDataConfig(BaseDataConfig): type: Literal["fake"] = "fake" + seq_len: int = Field(128, ge=1) + """Sequence length.""" + + pack_function: Literal["cat", "stack"] = "cat" + """Sample packing strategy.""" + length: Literal["fixed", "variable"] = "fixed" """Use fixed-length samples or variable-length samples.""" @@ -175,13 +181,8 @@ class SFTConfig(BaseConfig): tokenizer: TokenizerConfig = TokenizerConfig() - renderer: RendererConfig | None = None - """Typed renderer config (``renderers.RendererConfig`` discriminated - union). When set, SFT tokenizes samples through the ``renderers`` - library (single ``render()`` + ``message_indices`` mask) instead of - the default ``build_incremental_token_mask`` path. Required for chat - templates that render position-dependently (e.g. Qwen3, Qwen3.5). - ``None`` (default) uses the legacy tokenization path.""" + renderer: RendererConfig = AutoRendererConfig() + """Renderer config. Defaults to auto-selecting from the tokenizer model name.""" data: DataConfig = SFTDataConfig() @@ -225,8 +226,8 @@ class SFTConfig(BaseConfig): dist_timeout_seconds: int = 3600 """Timeout in seconds for torch distributed ops.""" - loss_impl: Literal["liger", "torch", "liger_fused", "quack_fused"] = "torch" - """Cross-entropy loss implementation. ``liger_fused`` fuses the lm_head projection with the CE loss to avoid materializing full logits. ``quack_fused`` uses quack-kernels for chunked linear + CE with CuTe DSL CUDA kernels.""" + loss_impl: Literal["liger", "torch", "liger_fused", "quack_fused"] = "liger_fused" + """Cross-entropy loss implementation. Defaults to fused Liger loss to avoid materializing full logits.""" heartbeat: HeartbeatConfig | None = None """BetterStack heartbeat configuration for monitoring training progress.""" @@ -319,9 +320,9 @@ def dont_do_massive_traces(self): @model_validator(mode="after") def validate_renderer_vs_vlm(self): - if self.renderer is not None and self.model.vlm is not None: + if self.model.vlm is not None: raise ValueError( - "renderer is not supported for VLMs in SFT. The renderer tokenizes " + "renderer-only SFT does not support VLMs yet. The renderer tokenizes " "text-only message dicts client-side and cannot handle image inputs." ) return self diff --git a/src/prime_rl/trainer/sft/data.py b/src/prime_rl/trainer/sft/data.py index d0edd13787..0dc707a05a 100644 --- a/src/prime_rl/trainer/sft/data.py +++ b/src/prime_rl/trainer/sft/data.py @@ -1,7 +1,7 @@ import json import uuid from collections import defaultdict -from typing import Literal, TypedDict, cast +from typing import Any, Literal, TypedDict, cast import torch from datasets import Dataset, interleave_datasets, load_dataset @@ -15,13 +15,7 @@ from prime_rl.configs.sft import DataConfig, LossMaskConfig, SFTDataConfig from prime_rl.trainer.world import get_world -from prime_rl.utils.chat_template import ( - IncrementalTokenizationError, - build_incremental_token_mask, - deserialize_tool_calls, - normalize_messages, - strip_message_content, -) +from prime_rl.utils.chat_template import deserialize_tool_calls, normalize_messages from prime_rl.utils.logger import get_logger STACKING_DATASET_BUCKET_TIMEOUT = 10 @@ -114,6 +108,23 @@ def __iter__(self): yield fake_sample +def _drop_null_fields(value: Any) -> Any: + """Recursively strip ``None``-valued keys from dict structures. + + PyArrow's JSON loader unifies schemas across rows, so heterogeneous + OAI content blocks (text vs image_url) end up with all union keys + filled with ``None`` where absent. That confuses permissive + content-type predicates inside renderers (e.g. ``"image_url" in item`` + returns ``True`` even when the value is null). Strip the noise before + handing messages off to the renderer. + """ + if isinstance(value, dict): + return {k: _drop_null_fields(v) for k, v in value.items() if v is not None} + if isinstance(value, list): + return [_drop_null_fields(v) for v in value] + return value + + class SFTDataset(StatefulIterableDataset): """A dataset wrapping a HF SFT dataset with prompt/completion or raw messages format.""" @@ -121,6 +132,7 @@ def __init__( self, dataset: Dataset, tokenizer: PreTrainedTokenizer | None, + renderer: Renderer | None = None, shuffle: bool = True, seed: int = 0, seq_len: int = 128, @@ -128,7 +140,6 @@ def __init__( loss_mask_config: LossMaskConfig = LossMaskConfig(), max_examples: int | None = None, max_epochs: int | None = None, - renderer: Renderer | None = None, ): super().__init__() self.logger = get_logger() @@ -144,9 +155,6 @@ def __init__( self.renderer = renderer self._warned_chat_template_kwargs = False - if self.tokenizer is None: - self.logger.warning("No tokenizer provided, will not process examples") - # If specified, select a subset of the dataset if self.max_examples is not None: self.num_examples = min(self.num_examples, self.max_examples) @@ -163,16 +171,19 @@ def __init__( self.data_world_size = get_world().world_size // non_dp_size * num_workers def _process(self, example: dict) -> dict | None: - # Skip processing if no tokenizer was provided if self.tokenizer is None: return example + if self.renderer is None: + raise ValueError("SFT processing requires a renderer.") def resolve_messages(example: dict) -> list[dict]: # `messages` takes precedence over explicit split fields and is interpreted - # as a whole-chat training sample with an empty prompt. - if "messages" in example: + # as a whole-chat training sample with an empty prompt. Null-check rather + # than key-check: Arrow schema union adds `messages: null` to + # prompt/completion rows whenever other rows have a `messages` column. + if example.get("messages") is not None: messages = normalize_messages(example["messages"], default_role="assistant") - elif "prompt" in example and "completion" in example: + elif example.get("prompt") is not None and example.get("completion") is not None: messages = normalize_messages(example["prompt"], default_role="user") + normalize_messages( example["completion"], default_role="assistant" ) @@ -182,22 +193,17 @@ def resolve_messages(example: dict) -> list[dict]: "or both 'prompt' and 'completion' columns for SFT" ) - # Deserialize tool call arguments from message list, if present - assumes OAI format - # Reference: https://platform.openai.com/docs/guides/function-calling#handling-function-calls - messages = deserialize_tool_calls(messages) - - # Strip content from all messages so that incremental tokenization works - # NOTE: This has the side effect that we do never train on leading or trailing whitespace - return strip_message_content(messages) + # Strip nulls before deserializing so genuine nulls inside tool-call + # argument strings survive. + messages = [_drop_null_fields(m) for m in messages] + return deserialize_tool_calls(messages) messages = resolve_messages(example) - # Parse available tools, if present - assumes OAI format - # Reference: https://platform.openai.com/docs/guides/function-calling#function-tool-example - # Accepts either `tools` or `tool_defs` (the verifiers rollout format), - # as either a JSON-encoded string of a list or a list of dicts. Tools - # arriving in the verifiers shape are converted to OAI form so any - # downstream chat template can consume them. + # Parse available tools, if present - assumes OAI format. Accepts either + # `tools` or `tool_defs` (the verifiers rollout format), as either a + # JSON-encoded string of a list or a list of dicts; verifiers-shaped + # tools are converted to OAI form for the chat template. raw_tools = example.get("tools", example.get("tool_defs")) if not raw_tools: tools = [] @@ -223,45 +229,36 @@ def should_mask(message: dict) -> bool: assert "role" in message, "Message must have a role" match message["role"]: case "user": - return True if self.loss_mask_config.user else False + return self.loss_mask_config.user case "assistant": - return True if self.loss_mask_config.assistant else False + return self.loss_mask_config.assistant case "system": - return True if self.loss_mask_config.system else False + return self.loss_mask_config.system case "tool": - return True if self.loss_mask_config.tool else False + return self.loss_mask_config.tool case _: raise ValueError(f"Invalid message role: {message['role']}") - if self.renderer is not None: - if example.get("chat_template_kwargs") and not self._warned_chat_template_kwargs: - self.logger.warning( - "Example carries chat_template_kwargs but a renderer is configured; " - "renderers don't forward chat_template_kwargs (model-specific " - "renderers bake their template behavior in). These kwargs will " - "be ignored. Further warnings suppressed for this dataset." - ) - self._warned_chat_template_kwargs = True - - input_ids, loss_mask = build_training_sample( - self.renderer, - messages, - role_to_mask=should_mask, - tools=tools, + if example.get("chat_template_kwargs") and not self._warned_chat_template_kwargs: + self.logger.warning( + "Ignoring per-example chat_template_kwargs; renderers only take " + "template kwargs run-wide via the [renderer] config." ) - else: - try: - input_ids, loss_mask = build_incremental_token_mask( - self.tokenizer, - messages, - role_to_mask=should_mask, - tools=tools, - chat_template_kwargs=example.get("chat_template_kwargs", {}), - collapse_consecutive_tool_messages=True, - ) - except IncrementalTokenizationError as e: - self.logger.warning(f"Skipping example {example.get('__index', '')}: {e}") - return None + self._warned_chat_template_kwargs = True + + # Non-assistant roles are opted into the loss via the renderer's + # body-only path: the message content is trained, not the role + # scaffolding (e.g. <|im_start|>assistant) the harness emits. + content_sft_roles = {role for role in ("user", "system", "tool") if getattr(self.loss_mask_config, role)} + sample = build_training_sample( + self.renderer, + messages, + role_to_mask=should_mask, + tools=tools, + content_sft_roles=content_sft_roles or None, + ) + input_ids = list(sample.token_ids) + loss_mask = list(sample.loss_mask) # If EOS token is not found, manually append it if not self.tokenizer.eos_token_id in input_ids: @@ -271,7 +268,7 @@ def should_mask(message: dict) -> bool: input_ids.append(cast(int, self.tokenizer.eos_token_id)) loss_mask.append(True) - # Prepare inputs + # Causal shift: model predicts next token from current. target_ids = input_ids.copy()[1:] loss_mask = loss_mask[1:] input_ids = input_ids[:-1] @@ -280,7 +277,7 @@ def should_mask(message: dict) -> bool: self.logger.warning( f"Skipping example {example.get('__index', '')} because no trainable tokens were found within the context window ({self.seq_len}). This is to prevent NaN loss." ) - return + return None assert len(input_ids) == len(loss_mask) == len(target_ids), ( f"input_ids, loss_mask and target_ids must have the same length, but got {len(input_ids)=}, {len(loss_mask)=}, {len(target_ids)=}" @@ -288,7 +285,6 @@ def should_mask(message: dict) -> bool: assert sum(loss_mask) > 0, "There are no tokens in this sample that contribute to the loss" assert self.tokenizer.eos_token_id in target_ids, "EOS token ID must be present in target_ids" - # Create sample (with one fake target for the last token) return { "input_ids": input_ids, "target_ids": target_ids, @@ -297,9 +293,6 @@ def should_mask(message: dict) -> bool: } def __iter__(self): - """ - Apply chat template and tokenize a single example in prompt + completion format (https://github.com/huggingface/trl/blob/de27d612b026526ba39b88eee348994d7636e033/trl/trainer/sft_trainer.py#L661) - """ dataset = self.dataset.shuffle(seed=self.epoch + self.seed) if self.shuffle else self.dataset while True: self.step += 1 @@ -585,21 +578,26 @@ def setup_dataset( ) -> StatefulIterableDataset: if config.type == "fake": return FakeDataset( - vocab_size=tokenizer.vocab_size, seq_len=config.seq_len, length=config.length, input_ids=config.input_ids + vocab_size=tokenizer.vocab_size, + seq_len=config.seq_len, + length=config.length, + input_ids=config.input_ids, ) elif config.type == "sft": + if renderer is None: + raise ValueError("SFT data requires a renderer.") if raw_dataset is None: raw_dataset = load_sft_dataset(config) return SFTDataset( raw_dataset, tokenizer, + renderer=renderer, shuffle=config.shuffle, seed=config.seed, seq_len=config.seq_len, loss_mask_config=config.loss_mask, non_dp_size=non_dp_size, max_epochs=max_epochs, - renderer=renderer, ) else: raise ValueError(f"Invalid dataset type: {config.type}") diff --git a/src/prime_rl/trainer/sft/train.py b/src/prime_rl/trainer/sft/train.py index 2abb73ec0f..0468d36a7f 100644 --- a/src/prime_rl/trainer/sft/train.py +++ b/src/prime_rl/trainer/sft/train.py @@ -58,6 +58,19 @@ from torchtitan.distributed.utils import clip_grad_norm_ +def setup_renderer(tokenizer, config): + """Create the SFT renderer, rejecting the DefaultRenderer fallback.""" + renderer = create_renderer(tokenizer, config) + if isinstance(renderer, DefaultRenderer): + raise ValueError( + f"SFT renderer for {tokenizer.name_or_path!r} resolved to DefaultRenderer. " + "SFT is renderer-only and requires a hand-coded renderer for stable " + "message-to-token attribution. Use a model with a hand-coded renderer " + "(see renderers.base.MODEL_RENDERER_MAP), or set [renderer] name= explicitly." + ) + return renderer + + @clean_exit def train(config: SFTConfig): # Setup world and logger @@ -162,17 +175,12 @@ def train(config: SFTConfig): logger.info(f"Initializing tokenizer ({config.tokenizer})") tokenizer = setup_tokenizer(config.tokenizer) + # Fake data never renders messages, so a model without a hand-coded renderer + # can still be used to benchmark step time / memory. Validation data is + # always real, so it needs the renderer even when training data is fake. renderer = None - if config.renderer is not None: - renderer = create_renderer(tokenizer, config.renderer) - if isinstance(renderer, DefaultRenderer): - raise ValueError( - f"renderer set for {config.tokenizer.name!r} resolved to DefaultRenderer. " - "DefaultRenderer falls back to incremental apply_chat_template and does NOT " - "fix position-dependent chat templates — the bug the renderer client is meant to solve. " - "Either use a model with a hand-coded renderer (see renderers.base.MODEL_RENDERER_MAP), " - "set [renderer] name= explicitly, or remove the [renderer] block." - ) + if config.data.type != "fake" or config.val is not None: + renderer = setup_renderer(tokenizer, config.renderer) logger.info(f"Initialized {type(renderer).__name__} for {config.tokenizer.name}") # Set up the optimizer diff --git a/src/prime_rl/utils/chat_template.py b/src/prime_rl/utils/chat_template.py index 400c97c657..fde9f1e92b 100644 --- a/src/prime_rl/utils/chat_template.py +++ b/src/prime_rl/utils/chat_template.py @@ -1,13 +1,5 @@ import json -from typing import Any, Callable - -from transformers.tokenization_utils import PreTrainedTokenizer - - -class IncrementalTokenizationError(ValueError): - """Raised when incremental tokenization produces inconsistent token prefixes.""" - - pass +from typing import Any def normalize_messages(messages: Any, default_role: str) -> list[dict[str, Any]]: @@ -56,79 +48,3 @@ def _deserialize_tool_call(tool_call: dict[str, Any]) -> dict[str, Any]: ) return deserialized_messages - - -def strip_message_content(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: - def _strip(message: dict[str, Any]) -> dict[str, Any]: - content = message.get("content") - if isinstance(content, str): - return {**message, "content": content.strip()} - return message - - return [_strip(message) for message in messages] - - -def should_add_generation_prompt(messages: list[dict[str, Any]], idx: int) -> bool: - role = messages[idx].get("role") - if role not in ("user", "tool"): - return False - if idx + 1 >= len(messages): - return False - return messages[idx + 1].get("role") == "assistant" - - -def render_messages( - tokenizer: PreTrainedTokenizer, - messages: list[dict[str, Any]], - *, - tools: list[dict[str, Any]] | None = None, - chat_template_kwargs: dict[str, Any] | None = None, - add_generation_prompt: bool = False, -) -> list[int]: - kwargs = dict(chat_template_kwargs or {}) - kwargs["add_generation_prompt"] = add_generation_prompt - if tools is not None: - kwargs["tools"] = tools - kwargs["return_dict"] = False - return list(tokenizer.apply_chat_template(messages, **kwargs)) - - -def build_incremental_token_mask( - tokenizer: PreTrainedTokenizer, - messages: list[dict[str, Any]], - *, - role_to_mask: Callable[[dict[str, Any]], bool], - tools: list[dict[str, Any]] | None = None, - chat_template_kwargs: dict[str, Any] | None = None, - collapse_consecutive_tool_messages: bool = False, -) -> tuple[list[int], list[bool]]: - token_mask: list[bool] = [] - prev_ids: list[int] = [] - prev_len = 0 - - for idx, message in enumerate(messages): - role = message.get("role") - if collapse_consecutive_tool_messages and role == "tool" and idx + 1 < len(messages): - if messages[idx + 1].get("role") == "tool": - continue - - cur_ids = render_messages( - tokenizer, - messages[: idx + 1], - tools=tools, - chat_template_kwargs=chat_template_kwargs, - add_generation_prompt=should_add_generation_prompt(messages, idx), - ) - - if prev_ids != cur_ids[:prev_len]: - raise IncrementalTokenizationError( - f"Mismatch in incremental tokenization with chat template at message {idx} (role={role}). " - "This usually means the chat template is not stable under incremental application. " - "The sample will be skipped." - ) - - token_mask.extend([role_to_mask(message)] * (len(cur_ids) - prev_len)) - prev_ids = cur_ids - prev_len = len(cur_ids) - - return prev_ids, token_mask diff --git a/tests/unit/train/sft/test_sft_dataset.py b/tests/unit/train/sft/test_sft_dataset.py index b8465e59d7..4d6de1fcb6 100644 --- a/tests/unit/train/sft/test_sft_dataset.py +++ b/tests/unit/train/sft/test_sft_dataset.py @@ -2,6 +2,7 @@ import pytest from datasets import Dataset, interleave_datasets +from renderers import create_renderer from transformers import AutoTokenizer from prime_rl.trainer.sft.data import SFTDataset @@ -24,7 +25,7 @@ def test_raise_error_if_no_prompt_and_completion(build_dummy_dataset): """Tests that an error is raised if no supported SFT message fields are provided.""" dataset = build_dummy_dataset("a", 1) tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B") - sft_dataset = SFTDataset(dataset, tokenizer=tokenizer) + sft_dataset = SFTDataset(dataset, tokenizer=tokenizer, renderer=create_renderer(tokenizer)) with pytest.raises(ValueError): next(iter(sft_dataset)) @@ -194,7 +195,7 @@ def test_multiturn_loss_mask(): ] ) tokenizer = AutoTokenizer.from_pretrained("PrimeIntellect/Qwen3-0.6B") # Properly handles multi-turn think - dataset = SFTDataset(dataset, tokenizer=tokenizer, max_examples=1) + dataset = SFTDataset(dataset, tokenizer=tokenizer, renderer=create_renderer(tokenizer), max_examples=1) sample = next(iter(dataset)) print_sample(sample["input_ids"], sample["loss_mask"], tokenizer) @@ -257,7 +258,7 @@ def test_multiturn_loss_mask_with_tools(): dataset = Dataset.from_list([tool_example]) tokenizer = AutoTokenizer.from_pretrained("PrimeIntellect/Qwen3-0.6B") # Properly handles multi-turn think - dataset = SFTDataset(dataset, tokenizer=tokenizer, max_examples=1) + dataset = SFTDataset(dataset, tokenizer=tokenizer, renderer=create_renderer(tokenizer), max_examples=1) sample = next(iter(dataset)) print_sample(sample["input_ids"], sample["loss_mask"], tokenizer) @@ -282,10 +283,16 @@ def test_messages_rows_are_equivalent_to_empty_prompt_completion(): ] tokenizer = AutoTokenizer.from_pretrained("PrimeIntellect/Qwen3-0.6B") - messages_dataset = SFTDataset(Dataset.from_list([{"messages": messages}]), tokenizer=tokenizer, max_examples=1) + messages_dataset = SFTDataset( + Dataset.from_list([{"messages": messages}]), + tokenizer=tokenizer, + renderer=create_renderer(tokenizer), + max_examples=1, + ) split_dataset = SFTDataset( Dataset.from_list([{"prompt": [], "completion": messages}]), tokenizer=tokenizer, + renderer=create_renderer(tokenizer), max_examples=1, ) @@ -304,11 +311,40 @@ def test_messages_take_precedence_over_prompt_and_completion(): "completion": [{"role": "assistant", "content": "Ignored completion"}], } - messages_dataset = SFTDataset(Dataset.from_list([row]), tokenizer=tokenizer, max_examples=1) + messages_dataset = SFTDataset( + Dataset.from_list([row]), + tokenizer=tokenizer, + renderer=create_renderer(tokenizer), + max_examples=1, + ) expected_dataset = SFTDataset( Dataset.from_list([{"prompt": [], "completion": row["messages"]}]), tokenizer=tokenizer, + renderer=create_renderer(tokenizer), max_examples=1, ) assert next(iter(messages_dataset)) == next(iter(expected_dataset)) + + +def test_null_messages_falls_back_to_prompt_and_completion(): + # Arrow schema union adds `messages: None` to prompt/completion rows when + # other rows in the file have a `messages` column + tokenizer = AutoTokenizer.from_pretrained("PrimeIntellect/Qwen3-0.6B") + prompt = [{"role": "user", "content": "What is 2+2?"}] + completion = [{"role": "assistant", "content": "4"}] + + mixed_row_dataset = SFTDataset( + Dataset.from_list([{"messages": None, "prompt": prompt, "completion": completion}]), + tokenizer=tokenizer, + renderer=create_renderer(tokenizer), + max_examples=1, + ) + expected_dataset = SFTDataset( + Dataset.from_list([{"prompt": prompt, "completion": completion}]), + tokenizer=tokenizer, + renderer=create_renderer(tokenizer), + max_examples=1, + ) + + assert next(iter(mixed_row_dataset)) == next(iter(expected_dataset)) From e1f5fd3bbcf98510caa43f148bcabc9475c2bcaf Mon Sep 17 00:00:00 2001 From: hubert-marek Date: Tue, 7 Jul 2026 22:00:39 +0000 Subject: [PATCH 02/12] fix(sft): allow explicit default renderer for checkpoints with their own chat template PrimeIntellect/Qwen3-0.6B and -1.7B ship a chat_template.jinja distinct from Qwen3's, so pinning [renderer] name="qwen3" trained them with the wrong template. setup_renderer now honors an explicit name="default" (apply_chat_template with the checkpoint's own template) and only rejects the silent auto-resolution fallback. Co-Authored-By: Claude Fable 5 --- .../ci/integration/reverse_text_sft/resume.toml | 3 ++- .../ci/integration/reverse_text_sft/start.toml | 3 ++- .../reverse_text_sft_lora/resume.toml | 3 ++- .../reverse_text_sft_lora/start.toml | 3 ++- examples/reverse_text/sft.toml | 3 ++- examples/wordle/sft.toml | 3 ++- src/prime_rl/trainer/sft/train.py | 17 ++++++++++++----- 7 files changed, 24 insertions(+), 11 deletions(-) diff --git a/configs/ci/integration/reverse_text_sft/resume.toml b/configs/ci/integration/reverse_text_sft/resume.toml index ba662f4e9e..77a939039c 100644 --- a/configs/ci/integration/reverse_text_sft/resume.toml +++ b/configs/ci/integration/reverse_text_sft/resume.toml @@ -6,8 +6,9 @@ resume_step = 5 [model] name = "PrimeIntellect/Qwen3-0.6B" +# PrimeIntellect/Qwen3-0.6B ships its own chat template (distinct from Qwen3's) [renderer] -name = "qwen3" +name = "default" [data] name = "PrimeIntellect/Reverse-Text-SFT" diff --git a/configs/ci/integration/reverse_text_sft/start.toml b/configs/ci/integration/reverse_text_sft/start.toml index 4d3250a8c8..5d228a4892 100644 --- a/configs/ci/integration/reverse_text_sft/start.toml +++ b/configs/ci/integration/reverse_text_sft/start.toml @@ -5,8 +5,9 @@ max_steps = 5 [model] name = "PrimeIntellect/Qwen3-0.6B" +# PrimeIntellect/Qwen3-0.6B ships its own chat template (distinct from Qwen3's) [renderer] -name = "qwen3" +name = "default" [data] name = "PrimeIntellect/Reverse-Text-SFT" diff --git a/configs/ci/integration/reverse_text_sft_lora/resume.toml b/configs/ci/integration/reverse_text_sft_lora/resume.toml index 13269a683f..9c4e37ebd1 100644 --- a/configs/ci/integration/reverse_text_sft_lora/resume.toml +++ b/configs/ci/integration/reverse_text_sft_lora/resume.toml @@ -9,8 +9,9 @@ save_adapter_separately = true [model] name = "PrimeIntellect/Qwen3-0.6B" +# PrimeIntellect/Qwen3-0.6B ships its own chat template (distinct from Qwen3's) [renderer] -name = "qwen3" +name = "default" [model.lora] rank = 8 diff --git a/configs/ci/integration/reverse_text_sft_lora/start.toml b/configs/ci/integration/reverse_text_sft_lora/start.toml index 700a6d6b44..eb24c18ae9 100644 --- a/configs/ci/integration/reverse_text_sft_lora/start.toml +++ b/configs/ci/integration/reverse_text_sft_lora/start.toml @@ -8,8 +8,9 @@ save_adapter_separately = true [model] name = "PrimeIntellect/Qwen3-0.6B" +# PrimeIntellect/Qwen3-0.6B ships its own chat template (distinct from Qwen3's) [renderer] -name = "qwen3" +name = "default" [model.lora] rank = 8 diff --git a/examples/reverse_text/sft.toml b/examples/reverse_text/sft.toml index 02abffaab5..345870650b 100644 --- a/examples/reverse_text/sft.toml +++ b/examples/reverse_text/sft.toml @@ -5,8 +5,9 @@ max_steps = 100 [model] name = "PrimeIntellect/Qwen3-0.6B" +# PrimeIntellect/Qwen3-0.6B ships its own chat template (distinct from Qwen3's) [renderer] -name = "qwen3" +name = "default" [data] name = "willcb/R1-reverse-wikipedia-paragraphs-v1-1000" diff --git a/examples/wordle/sft.toml b/examples/wordle/sft.toml index 591b2057f7..8d00ef37b3 100644 --- a/examples/wordle/sft.toml +++ b/examples/wordle/sft.toml @@ -5,8 +5,9 @@ max_steps = 20 [model] name = "PrimeIntellect/Qwen3-1.7B" +# PrimeIntellect/Qwen3-1.7B ships its own chat template (distinct from Qwen3's) [renderer] -name = "qwen3" +name = "default" [data] name = "willcb/V3-wordle" diff --git a/src/prime_rl/trainer/sft/train.py b/src/prime_rl/trainer/sft/train.py index 0468d36a7f..51bb088c9b 100644 --- a/src/prime_rl/trainer/sft/train.py +++ b/src/prime_rl/trainer/sft/train.py @@ -4,6 +4,7 @@ from contextlib import nullcontext from datetime import timedelta +from renderers import DefaultRendererConfig from renderers.base import create_renderer from renderers.default import DefaultRenderer from ring_flash_attn import substitute_hf_flash_attn @@ -59,14 +60,20 @@ def setup_renderer(tokenizer, config): - """Create the SFT renderer, rejecting the DefaultRenderer fallback.""" + """Create the SFT renderer, rejecting a silent DefaultRenderer fallback. + + An explicit ``[renderer] name = "default"`` is honored: checkpoints whose + chat template has no hand-coded renderer (e.g. PrimeIntellect/Qwen3-0.6B, + which ships its own template distinct from Qwen3's) train with the + template they were built with via ``apply_chat_template``. + """ renderer = create_renderer(tokenizer, config) - if isinstance(renderer, DefaultRenderer): + if isinstance(renderer, DefaultRenderer) and not isinstance(config, DefaultRendererConfig): raise ValueError( f"SFT renderer for {tokenizer.name_or_path!r} resolved to DefaultRenderer. " - "SFT is renderer-only and requires a hand-coded renderer for stable " - "message-to-token attribution. Use a model with a hand-coded renderer " - "(see renderers.base.MODEL_RENDERER_MAP), or set [renderer] name= explicitly." + "Use a model with a hand-coded renderer (see renderers.base.MODEL_RENDERER_MAP), " + "set [renderer] name=, or opt into the model's own chat " + 'template explicitly with [renderer] name = "default".' ) return renderer From 4abe2254bbff676e30a2f763e33e1a3e1e418f45 Mon Sep 17 00:00:00 2001 From: hubert-marek Date: Tue, 7 Jul 2026 22:00:39 +0000 Subject: [PATCH 03/12] refactor(sft): validate renderer at SFTDataset construction Renderer pairing is checked once in __init__ (tokenizer=None stays the tests' passthrough mode) instead of raising mid-processing; the chat_template_kwargs backward-compat warning is dropped. Co-Authored-By: Claude Fable 5 --- src/prime_rl/trainer/sft/data.py | 14 ++++---------- 1 file changed, 4 insertions(+), 10 deletions(-) diff --git a/src/prime_rl/trainer/sft/data.py b/src/prime_rl/trainer/sft/data.py index 0dc707a05a..cb28aec7ac 100644 --- a/src/prime_rl/trainer/sft/data.py +++ b/src/prime_rl/trainer/sft/data.py @@ -142,6 +142,10 @@ def __init__( max_epochs: int | None = None, ): super().__init__() + # tokenizer=None puts the dataset in passthrough mode (no processing); + # with a tokenizer, a renderer is required. + if tokenizer is not None and renderer is None: + raise ValueError("SFTDataset requires a renderer when a tokenizer is provided.") self.logger = get_logger() self.dataset = dataset self.num_examples = len(self.dataset) @@ -153,7 +157,6 @@ def __init__( self.max_examples = max_examples self.max_epochs = max_epochs self.renderer = renderer - self._warned_chat_template_kwargs = False # If specified, select a subset of the dataset if self.max_examples is not None: @@ -173,8 +176,6 @@ def __init__( def _process(self, example: dict) -> dict | None: if self.tokenizer is None: return example - if self.renderer is None: - raise ValueError("SFT processing requires a renderer.") def resolve_messages(example: dict) -> list[dict]: # `messages` takes precedence over explicit split fields and is interpreted @@ -239,13 +240,6 @@ def should_mask(message: dict) -> bool: case _: raise ValueError(f"Invalid message role: {message['role']}") - if example.get("chat_template_kwargs") and not self._warned_chat_template_kwargs: - self.logger.warning( - "Ignoring per-example chat_template_kwargs; renderers only take " - "template kwargs run-wide via the [renderer] config." - ) - self._warned_chat_template_kwargs = True - # Non-assistant roles are opted into the loss via the renderer's # body-only path: the message content is trained, not the role # scaffolding (e.g. <|im_start|>assistant) the harness emits. From 57718d97dd15115fe2a92880249f16cb8da2f220 Mon Sep 17 00:00:00 2001 From: hubert-marek Date: Tue, 7 Jul 2026 22:11:37 +0000 Subject: [PATCH 04/12] feat(configs): validate SFT renderer auto-resolution at parse time Mirrors the OrchestratorConfig validator: renderer.name='auto' with a model outside MODEL_RENDERER_MAP now fails config validation (and --dry-run) instead of at trainer startup. The redundant runtime guard in sft/train.py is removed; fake-data runs without validation stay exempt since they build no renderer. Co-Authored-By: Claude Fable 5 --- .../src/prime_rl/configs/sft.py | 26 +++++++++++++++++++ src/prime_rl/trainer/sft/train.py | 23 +--------------- uv.lock | 26 +++++++++---------- 3 files changed, 40 insertions(+), 35 deletions(-) diff --git a/packages/prime-rl-configs/src/prime_rl/configs/sft.py b/packages/prime-rl-configs/src/prime_rl/configs/sft.py index f1fb610069..3274fc5014 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/sft.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/sft.py @@ -4,6 +4,7 @@ from pydantic import Field, model_validator from renderers import AutoRendererConfig, RendererConfig +from renderers.base import MODEL_RENDERER_MAP from prime_rl.configs.shared import ( EnvVars, @@ -272,6 +273,31 @@ def validate_deployment(self): raise ValueError("Must use SLURM for multi-node deployment.") return self + @model_validator(mode="after") + def validate_auto_renderer_resolves(self): + """Reject renderer auto-resolution misses at config time (mirrors the + OrchestratorConfig validator). Resolution is an exact-name lookup, so + it is fully decidable here; fake-data runs without validation need no + renderer and are exempt. + """ + if not isinstance(self.renderer, AutoRendererConfig): + return self + if self.data.type == "fake" and self.val is None: + return self + model_id = self.tokenizer.name or self.model.name + if model_id in MODEL_RENDERER_MAP: + return self + raise ValueError( + f"renderer.name='auto' but {model_id!r} is not in " + f"renderers.base.MODEL_RENDERER_MAP, so it would silently fall back to " + f"DefaultRenderer. Pick one: " + f"(a) [renderer] name='default' — for fine-tunes / vendored mirrors with " + f"custom chat templates (DefaultRenderer calls apply_chat_template). " + f"(b) [renderer] name= — if {model_id!r} is " + f"template-identical to a mapped family (and ideally also add it upstream " + f"to renderers.base.MODEL_RENDERER_MAP)." + ) + @model_validator(mode="after") def validate_pack_function(self): if self.model.cp > 1: diff --git a/src/prime_rl/trainer/sft/train.py b/src/prime_rl/trainer/sft/train.py index 51bb088c9b..692b5488ea 100644 --- a/src/prime_rl/trainer/sft/train.py +++ b/src/prime_rl/trainer/sft/train.py @@ -4,9 +4,7 @@ from contextlib import nullcontext from datetime import timedelta -from renderers import DefaultRendererConfig from renderers.base import create_renderer -from renderers.default import DefaultRenderer from ring_flash_attn import substitute_hf_flash_attn from torch.nn import CrossEntropyLoss @@ -59,25 +57,6 @@ from torchtitan.distributed.utils import clip_grad_norm_ -def setup_renderer(tokenizer, config): - """Create the SFT renderer, rejecting a silent DefaultRenderer fallback. - - An explicit ``[renderer] name = "default"`` is honored: checkpoints whose - chat template has no hand-coded renderer (e.g. PrimeIntellect/Qwen3-0.6B, - which ships its own template distinct from Qwen3's) train with the - template they were built with via ``apply_chat_template``. - """ - renderer = create_renderer(tokenizer, config) - if isinstance(renderer, DefaultRenderer) and not isinstance(config, DefaultRendererConfig): - raise ValueError( - f"SFT renderer for {tokenizer.name_or_path!r} resolved to DefaultRenderer. " - "Use a model with a hand-coded renderer (see renderers.base.MODEL_RENDERER_MAP), " - "set [renderer] name=, or opt into the model's own chat " - 'template explicitly with [renderer] name = "default".' - ) - return renderer - - @clean_exit def train(config: SFTConfig): # Setup world and logger @@ -187,7 +166,7 @@ def train(config: SFTConfig): # always real, so it needs the renderer even when training data is fake. renderer = None if config.data.type != "fake" or config.val is not None: - renderer = setup_renderer(tokenizer, config.renderer) + renderer = create_renderer(tokenizer, config.renderer) logger.info(f"Initialized {type(renderer).__name__} for {config.tokenizer.name}") # Set up the optimizer diff --git a/uv.lock b/uv.lock index f838b922aa..996221eb1f 100644 --- a/uv.lock +++ b/uv.lock @@ -423,7 +423,7 @@ dependencies = [ requires-dist = [ { name = "httpx" }, { name = "openai" }, - { name = "verifiers", specifier = ">=0.1.15.dev422" }, + { name = "verifiers", specifier = ">=0.1.15.dev411" }, ] [[package]] @@ -929,7 +929,7 @@ dependencies = [ requires-dist = [ { name = "datasets" }, { name = "openai" }, - { name = "verifiers", specifier = ">=0.1.15.dev422" }, + { name = "verifiers", specifier = ">=0.1.15.dev411" }, ] [[package]] @@ -1531,7 +1531,7 @@ dependencies = [ requires-dist = [ { name = "datasets" }, { name = "openai" }, - { name = "verifiers", specifier = ">=0.1.15.dev422" }, + { name = "verifiers", specifier = ">=0.1.15.dev411" }, ] [[package]] @@ -1814,7 +1814,7 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "datasets" }, - { name = "verifiers", specifier = ">=0.1.15.dev422" }, + { name = "verifiers", specifier = ">=0.1.15.dev411" }, ] [[package]] @@ -3343,7 +3343,7 @@ dependencies = [ requires-dist = [ { name = "datasets" }, { name = "openai" }, - { name = "verifiers", specifier = ">=0.1.15.dev422" }, + { name = "verifiers", specifier = ">=0.1.15.dev411" }, ] [[package]] @@ -3362,7 +3362,7 @@ requires-dist = [ { name = "datasets" }, { name = "openai" }, { name = "python-dateutil" }, - { name = "verifiers", specifier = ">=0.1.15.dev422" }, + { name = "verifiers", specifier = ">=0.1.15.dev411" }, ] [[package]] @@ -3444,7 +3444,7 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "datasets" }, - { name = "verifiers", specifier = ">=0.1.15.dev422" }, + { name = "verifiers", specifier = ">=0.1.15.dev411" }, ] [[package]] @@ -3662,7 +3662,7 @@ dependencies = [ requires-dist = [ { name = "datasets" }, { name = "openai" }, - { name = "verifiers", specifier = ">=0.1.15.dev422" }, + { name = "verifiers", specifier = ">=0.1.15.dev411" }, ] [[package]] @@ -4764,7 +4764,7 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "datasets" }, - { name = "verifiers", specifier = ">=0.1.15.dev422" }, + { name = "verifiers", specifier = ">=0.1.15.dev411" }, ] [[package]] @@ -4967,7 +4967,7 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "huggingface-hub" }, - { name = "verifiers", specifier = ">=0.1.15.dev422" }, + { name = "verifiers", specifier = ">=0.1.15.dev411" }, ] [[package]] @@ -5132,7 +5132,7 @@ dependencies = [ requires-dist = [ { name = "datasets" }, { name = "openai" }, - { name = "verifiers", specifier = ">=0.1.15.dev422" }, + { name = "verifiers", specifier = ">=0.1.15.dev411" }, ] [[package]] @@ -5149,7 +5149,7 @@ dependencies = [ requires-dist = [ { name = "datasets" }, { name = "openai" }, - { name = "verifiers", specifier = ">=0.1.15.dev422" }, + { name = "verifiers", specifier = ">=0.1.15.dev411" }, ] [[package]] @@ -5488,7 +5488,7 @@ dependencies = [ requires-dist = [ { name = "audioop-lts", marker = "python_full_version >= '3.13'" }, { name = "tau2", git = "https://github.com/sierra-research/tau2-bench.git?rev=337326e" }, - { name = "verifiers", specifier = ">=0.1.15.dev424" }, + { name = "verifiers", specifier = ">=0.1.15.dev411" }, ] [[package]] From 67e2a4f7e1b3a84972c7d67d7aee213ff8b0c573 Mon Sep 17 00:00:00 2001 From: hubert-marek Date: Thu, 9 Jul 2026 21:00:06 +0000 Subject: [PATCH 05/12] Remove tokenizer path in SFT, renderer is only one --- src/prime_rl/trainer/sft/data.py | 32 ++---- tests/unit/train/sft/test_sft_dataset.py | 121 ++++++++++++++--------- uv.lock | 26 ++--- 3 files changed, 95 insertions(+), 84 deletions(-) diff --git a/src/prime_rl/trainer/sft/data.py b/src/prime_rl/trainer/sft/data.py index cb28aec7ac..5c3ffdca2b 100644 --- a/src/prime_rl/trainer/sft/data.py +++ b/src/prime_rl/trainer/sft/data.py @@ -131,8 +131,7 @@ class SFTDataset(StatefulIterableDataset): def __init__( self, dataset: Dataset, - tokenizer: PreTrainedTokenizer | None, - renderer: Renderer | None = None, + renderer: Renderer, shuffle: bool = True, seed: int = 0, seq_len: int = 128, @@ -142,21 +141,16 @@ def __init__( max_epochs: int | None = None, ): super().__init__() - # tokenizer=None puts the dataset in passthrough mode (no processing); - # with a tokenizer, a renderer is required. - if tokenizer is not None and renderer is None: - raise ValueError("SFTDataset requires a renderer when a tokenizer is provided.") self.logger = get_logger() self.dataset = dataset self.num_examples = len(self.dataset) - self.tokenizer = tokenizer + self.renderer = renderer self.shuffle = shuffle self.seed = seed self.seq_len = seq_len self.loss_mask_config = loss_mask_config self.max_examples = max_examples self.max_epochs = max_epochs - self.renderer = renderer # If specified, select a subset of the dataset if self.max_examples is not None: @@ -174,9 +168,6 @@ def __init__( self.data_world_size = get_world().world_size // non_dp_size * num_workers def _process(self, example: dict) -> dict | None: - if self.tokenizer is None: - return example - def resolve_messages(example: dict) -> list[dict]: # `messages` takes precedence over explicit split fields and is interpreted # as a whole-chat training sample with an empty prompt. Null-check rather @@ -244,23 +235,13 @@ def should_mask(message: dict) -> bool: # body-only path: the message content is trained, not the role # scaffolding (e.g. <|im_start|>assistant) the harness emits. content_sft_roles = {role for role in ("user", "system", "tool") if getattr(self.loss_mask_config, role)} - sample = build_training_sample( + input_ids, loss_mask = build_training_sample( self.renderer, messages, role_to_mask=should_mask, tools=tools, content_sft_roles=content_sft_roles or None, ) - input_ids = list(sample.token_ids) - loss_mask = list(sample.loss_mask) - - # If EOS token is not found, manually append it - if not self.tokenizer.eos_token_id in input_ids: - self.logger.warning( - f"Did not find EOS token ID {self.tokenizer.eos_token_id} in input_ids. Is something wrong with the chat template? Manually appending EOS token..." - ) - input_ids.append(cast(int, self.tokenizer.eos_token_id)) - loss_mask.append(True) # Causal shift: model predicts next token from current. target_ids = input_ids.copy()[1:] @@ -277,7 +258,9 @@ def should_mask(message: dict) -> bool: f"input_ids, loss_mask and target_ids must have the same length, but got {len(input_ids)=}, {len(loss_mask)=}, {len(target_ids)=}" ) assert sum(loss_mask) > 0, "There are no tokens in this sample that contribute to the loss" - assert self.tokenizer.eos_token_id in target_ids, "EOS token ID must be present in target_ids" + assert set(self.renderer.get_stop_token_ids()) & set(target_ids), ( + "A renderer stop token must be present in target_ids" + ) return { "input_ids": input_ids, @@ -584,8 +567,7 @@ def setup_dataset( raw_dataset = load_sft_dataset(config) return SFTDataset( raw_dataset, - tokenizer, - renderer=renderer, + renderer, shuffle=config.shuffle, seed=config.seed, seq_len=config.seq_len, diff --git a/tests/unit/train/sft/test_sft_dataset.py b/tests/unit/train/sft/test_sft_dataset.py index 4d6de1fcb6..23e0b76c82 100644 --- a/tests/unit/train/sft/test_sft_dataset.py +++ b/tests/unit/train/sft/test_sft_dataset.py @@ -3,64 +3,102 @@ import pytest from datasets import Dataset, interleave_datasets from renderers import create_renderer +from renderers.base import RenderedTokens from transformers import AutoTokenizer from prime_rl.trainer.sft.data import SFTDataset from prime_rl.trainer.utils import print_sample +_BOS_TOKEN_ID = 0 +_STOP_TOKEN_ID = 1 + + +def _sample_token_ids(value: str) -> list[int]: + return [ord(char) + 2 for char in value] + + +class _DummyRenderer: + def render(self, messages, **kwargs): + content_ids = _sample_token_ids(messages[-1]["content"]) + token_ids = [_BOS_TOKEN_ID, *content_ids, _STOP_TOKEN_ID] + return RenderedTokens( + token_ids=token_ids, + message_indices=[-1, *([len(messages) - 1] * (len(content_ids) + 1))], + sampled_mask=[False, *([True] * (len(content_ids) + 1))], + ) + + def get_stop_token_ids(self): + return [_STOP_TOKEN_ID] + @pytest.fixture(scope="module") def build_dummy_dataset(): - return lambda letter, num_examples: Dataset.from_list([{"text": f"{letter}{i}"} for i in range(num_examples)]) + return lambda letter, num_examples: Dataset.from_list( + [{"messages": [{"role": "assistant", "content": f"{letter}{i}"}]} for i in range(num_examples)] + ) + + +@pytest.fixture +def dummy_renderer(): + return _DummyRenderer() -def test_init_sft_dataset(build_dummy_dataset): +def test_init_sft_dataset(build_dummy_dataset, dummy_renderer): """Tests basic initialization.""" dataset = build_dummy_dataset("a", 1) - sft_dataset = SFTDataset(dataset, tokenizer=None) + sft_dataset = SFTDataset(dataset, dummy_renderer) assert sft_dataset is not None def test_raise_error_if_no_prompt_and_completion(build_dummy_dataset): """Tests that an error is raised if no supported SFT message fields are provided.""" - dataset = build_dummy_dataset("a", 1) + dataset = Dataset.from_list([{"text": "a0"}]) tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B") - sft_dataset = SFTDataset(dataset, tokenizer=tokenizer, renderer=create_renderer(tokenizer)) + sft_dataset = SFTDataset(dataset, create_renderer(tokenizer)) with pytest.raises(ValueError): next(iter(sft_dataset)) @pytest.mark.parametrize("max_epochs", [1, 2, 4]) -def test_sft_first_exhausted(build_dummy_dataset, max_epochs: int): +def test_sft_first_exhausted(build_dummy_dataset, dummy_renderer, max_epochs: int): a = build_dummy_dataset("a", 1) b = build_dummy_dataset("b", 2) ds = [a, b] dataset = interleave_datasets(ds, stopping_strategy="first_exhausted") - dataset = SFTDataset(dataset, tokenizer=None, shuffle=False, max_epochs=max_epochs) + dataset = SFTDataset(dataset, dummy_renderer, shuffle=False, max_epochs=max_epochs) num_samples = 0 sampling_order = [] for x in dataset: - sampling_order.append(x["text"]) + sampling_order.append(x["target_ids"][:-1]) num_samples += 1 assert num_samples == max_epochs * min([len(d) for d in ds]) * len(ds) - assert sampling_order == ["a0", "b0"] * max_epochs + assert sampling_order == [_sample_token_ids("a0"), _sample_token_ids("b0")] * max_epochs @pytest.mark.parametrize("max_epochs", [1, 2, 4]) -def test_sft_all_exhausted(build_dummy_dataset, max_epochs: int): +def test_sft_all_exhausted(build_dummy_dataset, dummy_renderer, max_epochs: int): a = build_dummy_dataset("a", 1) b = build_dummy_dataset("b", 2) ds = [a, b] dataset = interleave_datasets(ds, stopping_strategy="all_exhausted") - dataset = SFTDataset(dataset, tokenizer=None, shuffle=False, max_epochs=max_epochs) + dataset = SFTDataset(dataset, dummy_renderer, shuffle=False, max_epochs=max_epochs) num_samples = 0 sampling_order = [] for x in dataset: - sampling_order.append(x["text"]) + sampling_order.append(x["target_ids"][:-1]) num_samples += 1 assert num_samples == max_epochs * max([len(d) for d in ds]) * len(ds) print(sampling_order) - assert sampling_order == ["a0", "b0", "a0", "b1"] * max_epochs + assert ( + sampling_order + == [ + _sample_token_ids("a0"), + _sample_token_ids("b0"), + _sample_token_ids("a0"), + _sample_token_ids("b1"), + ] + * max_epochs + ) @pytest.mark.parametrize( @@ -71,21 +109,21 @@ def test_sft_all_exhausted(build_dummy_dataset, max_epochs: int): pytest.param((9 / 10, 1 / 10), id="high_low_probs"), ], ) -def test_sft_all_exhausted_with_probs(build_dummy_dataset, probs: list[float]): +def test_sft_all_exhausted_with_probs(build_dummy_dataset, dummy_renderer, probs: list[float]): """Tests that the ratio of samples from different datasets is as specified, in expectation.""" a = build_dummy_dataset("a", int(1e3)) b = build_dummy_dataset("b", int(10e3)) ds = [a, b] dataset = interleave_datasets(ds, stopping_strategy="all_exhausted", probabilities=probs) - dataset = SFTDataset(dataset, tokenizer=None, shuffle=False, max_epochs=1) + dataset = SFTDataset(dataset, dummy_renderer, shuffle=False, max_epochs=1) num_samples = 0 sampling_freq = [] for x in dataset: - sampling_freq.append(x["text"][0]) + sampling_freq.append(x["target_ids"][0]) num_samples += 1 sampling_freq = Counter(sampling_freq) - ratio_a = sampling_freq["a"] / num_samples - ratio_b = sampling_freq["b"] / num_samples + ratio_a = sampling_freq[ord("a") + 2] / num_samples + ratio_b = sampling_freq[ord("b") + 2] / num_samples assert ratio_a > probs[0] * 0.8 and ratio_a < probs[0] * 1.2, ( f"Expected frequency of samples from a to be between {probs[0] * 0.8} and {probs[0] * 1.2}, but got {ratio_a}" ) @@ -94,10 +132,10 @@ def test_sft_all_exhausted_with_probs(build_dummy_dataset, probs: list[float]): ) -def test_sft_dataset_state(build_dummy_dataset): +def test_sft_dataset_state(build_dummy_dataset, dummy_renderer): """Tests the state of the dataset within and across epochs.""" dataset = build_dummy_dataset("", 4) - dataset = SFTDataset(dataset, tokenizer=None, shuffle=False, max_epochs=2) + dataset = SFTDataset(dataset, dummy_renderer, shuffle=False, max_epochs=2) dataiter = iter(dataset) # Initial state @@ -106,24 +144,24 @@ def test_sft_dataset_state(build_dummy_dataset): # Epoch 1 for i in range(4): sample = next(dataiter) - assert sample["text"] == str(i) + assert sample["target_ids"][:-1] == _sample_token_ids(str(i)) assert dataset.state_dict() == {"epoch": 0, "step": i + 1} # Epoch 2 for i in range(4): sample = next(dataiter) - assert sample["text"] == str(i) + assert sample["target_ids"][:-1] == _sample_token_ids(str(i)) assert dataset.state_dict() == {"epoch": 1, "step": 4 + i + 1} with pytest.raises(StopIteration): next(dataiter) -def test_sft_dataset_state_resume(build_dummy_dataset): +def test_sft_dataset_state_resume(build_dummy_dataset, dummy_renderer): """Tests resuming the dataset from checkpoint in between epochs.""" dataset = SFTDataset( build_dummy_dataset("", 4), - tokenizer=None, + dummy_renderer, shuffle=False, max_epochs=2, ) @@ -135,8 +173,7 @@ def test_sft_dataset_state_resume(build_dummy_dataset): # Epoch 1 for i in range(4): sample = next(dataiter) - print(sample["text"]) - assert sample["text"] == str(i) + assert sample["target_ids"][:-1] == _sample_token_ids(str(i)) assert dataset.state_dict() == {"epoch": 0, "step": i + 1} # Resuming from checkpoint cross epoch @@ -144,7 +181,7 @@ def test_sft_dataset_state_resume(build_dummy_dataset): del dataset dataset = SFTDataset( build_dummy_dataset("", 4), - tokenizer=None, + dummy_renderer, shuffle=False, max_epochs=2, ) @@ -154,8 +191,7 @@ def test_sft_dataset_state_resume(build_dummy_dataset): # Epoch 2.1 for i in range(2): sample = next(dataiter) - print(sample["text"]) - assert sample["text"] == str(i) + assert sample["target_ids"][:-1] == _sample_token_ids(str(i)) assert dataset.state_dict() == {"epoch": 1, "step": 4 + i + 1} # Resuming from checkpoint mid epoch @@ -163,7 +199,7 @@ def test_sft_dataset_state_resume(build_dummy_dataset): del dataset dataset = SFTDataset( build_dummy_dataset("", 4), - tokenizer=None, + dummy_renderer, shuffle=False, max_epochs=2, ) @@ -173,8 +209,7 @@ def test_sft_dataset_state_resume(build_dummy_dataset): # Epoch 2.2 for i in range(2, 4): sample = next(dataiter) - print(sample["text"]) - assert sample["text"] == str(i) + assert sample["target_ids"][:-1] == _sample_token_ids(str(i)) assert dataset.state_dict() == {"epoch": 1, "step": 4 + i + 1} with pytest.raises(StopIteration): @@ -195,7 +230,7 @@ def test_multiturn_loss_mask(): ] ) tokenizer = AutoTokenizer.from_pretrained("PrimeIntellect/Qwen3-0.6B") # Properly handles multi-turn think - dataset = SFTDataset(dataset, tokenizer=tokenizer, renderer=create_renderer(tokenizer), max_examples=1) + dataset = SFTDataset(dataset, create_renderer(tokenizer), max_examples=1) sample = next(iter(dataset)) print_sample(sample["input_ids"], sample["loss_mask"], tokenizer) @@ -258,7 +293,7 @@ def test_multiturn_loss_mask_with_tools(): dataset = Dataset.from_list([tool_example]) tokenizer = AutoTokenizer.from_pretrained("PrimeIntellect/Qwen3-0.6B") # Properly handles multi-turn think - dataset = SFTDataset(dataset, tokenizer=tokenizer, renderer=create_renderer(tokenizer), max_examples=1) + dataset = SFTDataset(dataset, create_renderer(tokenizer), max_examples=1) sample = next(iter(dataset)) print_sample(sample["input_ids"], sample["loss_mask"], tokenizer) @@ -285,14 +320,12 @@ def test_messages_rows_are_equivalent_to_empty_prompt_completion(): tokenizer = AutoTokenizer.from_pretrained("PrimeIntellect/Qwen3-0.6B") messages_dataset = SFTDataset( Dataset.from_list([{"messages": messages}]), - tokenizer=tokenizer, - renderer=create_renderer(tokenizer), + create_renderer(tokenizer), max_examples=1, ) split_dataset = SFTDataset( Dataset.from_list([{"prompt": [], "completion": messages}]), - tokenizer=tokenizer, - renderer=create_renderer(tokenizer), + create_renderer(tokenizer), max_examples=1, ) @@ -313,14 +346,12 @@ def test_messages_take_precedence_over_prompt_and_completion(): messages_dataset = SFTDataset( Dataset.from_list([row]), - tokenizer=tokenizer, - renderer=create_renderer(tokenizer), + create_renderer(tokenizer), max_examples=1, ) expected_dataset = SFTDataset( Dataset.from_list([{"prompt": [], "completion": row["messages"]}]), - tokenizer=tokenizer, - renderer=create_renderer(tokenizer), + create_renderer(tokenizer), max_examples=1, ) @@ -336,14 +367,12 @@ def test_null_messages_falls_back_to_prompt_and_completion(): mixed_row_dataset = SFTDataset( Dataset.from_list([{"messages": None, "prompt": prompt, "completion": completion}]), - tokenizer=tokenizer, - renderer=create_renderer(tokenizer), + create_renderer(tokenizer), max_examples=1, ) expected_dataset = SFTDataset( Dataset.from_list([{"prompt": prompt, "completion": completion}]), - tokenizer=tokenizer, - renderer=create_renderer(tokenizer), + create_renderer(tokenizer), max_examples=1, ) diff --git a/uv.lock b/uv.lock index 996221eb1f..f838b922aa 100644 --- a/uv.lock +++ b/uv.lock @@ -423,7 +423,7 @@ dependencies = [ requires-dist = [ { name = "httpx" }, { name = "openai" }, - { name = "verifiers", specifier = ">=0.1.15.dev411" }, + { name = "verifiers", specifier = ">=0.1.15.dev422" }, ] [[package]] @@ -929,7 +929,7 @@ dependencies = [ requires-dist = [ { name = "datasets" }, { name = "openai" }, - { name = "verifiers", specifier = ">=0.1.15.dev411" }, + { name = "verifiers", specifier = ">=0.1.15.dev422" }, ] [[package]] @@ -1531,7 +1531,7 @@ dependencies = [ requires-dist = [ { name = "datasets" }, { name = "openai" }, - { name = "verifiers", specifier = ">=0.1.15.dev411" }, + { name = "verifiers", specifier = ">=0.1.15.dev422" }, ] [[package]] @@ -1814,7 +1814,7 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "datasets" }, - { name = "verifiers", specifier = ">=0.1.15.dev411" }, + { name = "verifiers", specifier = ">=0.1.15.dev422" }, ] [[package]] @@ -3343,7 +3343,7 @@ dependencies = [ requires-dist = [ { name = "datasets" }, { name = "openai" }, - { name = "verifiers", specifier = ">=0.1.15.dev411" }, + { name = "verifiers", specifier = ">=0.1.15.dev422" }, ] [[package]] @@ -3362,7 +3362,7 @@ requires-dist = [ { name = "datasets" }, { name = "openai" }, { name = "python-dateutil" }, - { name = "verifiers", specifier = ">=0.1.15.dev411" }, + { name = "verifiers", specifier = ">=0.1.15.dev422" }, ] [[package]] @@ -3444,7 +3444,7 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "datasets" }, - { name = "verifiers", specifier = ">=0.1.15.dev411" }, + { name = "verifiers", specifier = ">=0.1.15.dev422" }, ] [[package]] @@ -3662,7 +3662,7 @@ dependencies = [ requires-dist = [ { name = "datasets" }, { name = "openai" }, - { name = "verifiers", specifier = ">=0.1.15.dev411" }, + { name = "verifiers", specifier = ">=0.1.15.dev422" }, ] [[package]] @@ -4764,7 +4764,7 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "datasets" }, - { name = "verifiers", specifier = ">=0.1.15.dev411" }, + { name = "verifiers", specifier = ">=0.1.15.dev422" }, ] [[package]] @@ -4967,7 +4967,7 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "huggingface-hub" }, - { name = "verifiers", specifier = ">=0.1.15.dev411" }, + { name = "verifiers", specifier = ">=0.1.15.dev422" }, ] [[package]] @@ -5132,7 +5132,7 @@ dependencies = [ requires-dist = [ { name = "datasets" }, { name = "openai" }, - { name = "verifiers", specifier = ">=0.1.15.dev411" }, + { name = "verifiers", specifier = ">=0.1.15.dev422" }, ] [[package]] @@ -5149,7 +5149,7 @@ dependencies = [ requires-dist = [ { name = "datasets" }, { name = "openai" }, - { name = "verifiers", specifier = ">=0.1.15.dev411" }, + { name = "verifiers", specifier = ">=0.1.15.dev422" }, ] [[package]] @@ -5488,7 +5488,7 @@ dependencies = [ requires-dist = [ { name = "audioop-lts", marker = "python_full_version >= '3.13'" }, { name = "tau2", git = "https://github.com/sierra-research/tau2-bench.git?rev=337326e" }, - { name = "verifiers", specifier = ">=0.1.15.dev411" }, + { name = "verifiers", specifier = ">=0.1.15.dev424" }, ] [[package]] From 36bc39b6f996b168066d7fdba0e77428bb9f8bb9 Mon Sep 17 00:00:00 2001 From: hubert-marek Date: Fri, 10 Jul 2026 02:02:41 +0000 Subject: [PATCH 06/12] Fix build trianing sample --- src/prime_rl/trainer/sft/data.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/prime_rl/trainer/sft/data.py b/src/prime_rl/trainer/sft/data.py index 5c3ffdca2b..85a27363c7 100644 --- a/src/prime_rl/trainer/sft/data.py +++ b/src/prime_rl/trainer/sft/data.py @@ -235,13 +235,15 @@ def should_mask(message: dict) -> bool: # body-only path: the message content is trained, not the role # scaffolding (e.g. <|im_start|>assistant) the harness emits. content_sft_roles = {role for role in ("user", "system", "tool") if getattr(self.loss_mask_config, role)} - input_ids, loss_mask = build_training_sample( + sample = build_training_sample( self.renderer, messages, role_to_mask=should_mask, tools=tools, content_sft_roles=content_sft_roles or None, ) + input_ids = list(sample.token_ids) + loss_mask = list(sample.loss_mask) # Causal shift: model predicts next token from current. target_ids = input_ids.copy()[1:] From d97cf5f0cab9bec89c49e19808c3a8489819fa15 Mon Sep 17 00:00:00 2001 From: hubert-marek Date: Fri, 10 Jul 2026 03:02:38 +0000 Subject: [PATCH 07/12] Handle nested dict tools in SFT datasets --- src/prime_rl/trainer/sft/data.py | 11 +++++++---- tests/unit/train/sft/test_sft_dataset.py | 24 +++++++++++++++++++++++- 2 files changed, 30 insertions(+), 5 deletions(-) diff --git a/src/prime_rl/trainer/sft/data.py b/src/prime_rl/trainer/sft/data.py index 85a27363c7..2e80e3bd55 100644 --- a/src/prime_rl/trainer/sft/data.py +++ b/src/prime_rl/trainer/sft/data.py @@ -108,7 +108,7 @@ def __iter__(self): yield fake_sample -def _drop_null_fields(value: Any) -> Any: +def _drop_null_fields(value: Any, path: tuple[str, ...] = ()) -> Any: """Recursively strip ``None``-valued keys from dict structures. PyArrow's JSON loader unifies schemas across rows, so heterogeneous @@ -116,12 +116,15 @@ def _drop_null_fields(value: Any) -> Any: filled with ``None`` where absent. That confuses permissive content-type predicates inside renderers (e.g. ``"image_url" in item`` returns ``True`` even when the value is null). Strip the noise before - handing messages off to the renderer. + handing messages off to the renderer. Tool-call arguments are opaque + JSON payloads, so preserve their null values. """ + if path[-3:] == ("tool_calls", "function", "arguments"): + return value if isinstance(value, dict): - return {k: _drop_null_fields(v) for k, v in value.items() if v is not None} + return {k: _drop_null_fields(v, (*path, k)) for k, v in value.items() if v is not None} if isinstance(value, list): - return [_drop_null_fields(v) for v in value] + return [_drop_null_fields(v, path) for v in value] return value diff --git a/tests/unit/train/sft/test_sft_dataset.py b/tests/unit/train/sft/test_sft_dataset.py index 23e0b76c82..fc72cb533b 100644 --- a/tests/unit/train/sft/test_sft_dataset.py +++ b/tests/unit/train/sft/test_sft_dataset.py @@ -6,7 +6,7 @@ from renderers.base import RenderedTokens from transformers import AutoTokenizer -from prime_rl.trainer.sft.data import SFTDataset +from prime_rl.trainer.sft.data import SFTDataset, _drop_null_fields from prime_rl.trainer.utils import print_sample _BOS_TOKEN_ID = 0 @@ -43,6 +43,28 @@ def dummy_renderer(): return _DummyRenderer() +@pytest.mark.parametrize( + "arguments", + [ + pytest.param('{"reasoning_effort": null}', id="json-string"), + pytest.param({"reasoning_effort": None}, id="dict"), + ], +) +def test_drop_null_fields_preserves_tool_call_arguments(arguments): + message = { + "role": "assistant", + "content": [{"type": "text", "text": "Calling a tool", "image_url": None}], + "tool_calls": [{"function": {"name": "listReasoningModels", "arguments": arguments}}], + "metadata": {"arguments": {"unrelated_null": None}}, + } + + cleaned = _drop_null_fields(message) + + assert cleaned["tool_calls"][0]["function"]["arguments"] == arguments + assert cleaned["content"] == [{"type": "text", "text": "Calling a tool"}] + assert cleaned["metadata"] == {"arguments": {}} + + def test_init_sft_dataset(build_dummy_dataset, dummy_renderer): """Tests basic initialization.""" dataset = build_dummy_dataset("a", 1) From 6090befbbe91c507abe7b16da73b34156d5bc3eb Mon Sep 17 00:00:00 2001 From: hubert-marek Date: Fri, 10 Jul 2026 20:08:03 +0000 Subject: [PATCH 08/12] feat(sft): require typed renderers via prime-qwen3 Bump renderers to the PrimeIntellect Qwen3 typed renderer, migrate the affected configs off default/qwen3 selections, and reject DefaultRenderer whenever SFT renders real samples. Co-authored-by: Cursor --- .../integration/reverse_text_sft/resume.toml | 2 +- .../integration/reverse_text_sft/start.toml | 2 +- .../reverse_text_sft_lora/resume.toml | 2 +- .../reverse_text_sft_lora/start.toml | 2 +- deps/renderers | 2 +- examples/reverse_text/sft.toml | 2 +- examples/wordle/rl.toml | 7 ++-- examples/wordle/sft.toml | 2 +- .../src/prime_rl/configs/sft.py | 33 +++++++++---------- tests/unit/test_configs.py | 20 +++++++++++ 10 files changed, 44 insertions(+), 30 deletions(-) diff --git a/configs/ci/integration/reverse_text_sft/resume.toml b/configs/ci/integration/reverse_text_sft/resume.toml index 77a939039c..d0431d4f93 100644 --- a/configs/ci/integration/reverse_text_sft/resume.toml +++ b/configs/ci/integration/reverse_text_sft/resume.toml @@ -8,7 +8,7 @@ name = "PrimeIntellect/Qwen3-0.6B" # PrimeIntellect/Qwen3-0.6B ships its own chat template (distinct from Qwen3's) [renderer] -name = "default" +name = "prime-qwen3" [data] name = "PrimeIntellect/Reverse-Text-SFT" diff --git a/configs/ci/integration/reverse_text_sft/start.toml b/configs/ci/integration/reverse_text_sft/start.toml index 5d228a4892..89f3b5d651 100644 --- a/configs/ci/integration/reverse_text_sft/start.toml +++ b/configs/ci/integration/reverse_text_sft/start.toml @@ -7,7 +7,7 @@ name = "PrimeIntellect/Qwen3-0.6B" # PrimeIntellect/Qwen3-0.6B ships its own chat template (distinct from Qwen3's) [renderer] -name = "default" +name = "prime-qwen3" [data] name = "PrimeIntellect/Reverse-Text-SFT" diff --git a/configs/ci/integration/reverse_text_sft_lora/resume.toml b/configs/ci/integration/reverse_text_sft_lora/resume.toml index 9c4e37ebd1..f8ea682d5a 100644 --- a/configs/ci/integration/reverse_text_sft_lora/resume.toml +++ b/configs/ci/integration/reverse_text_sft_lora/resume.toml @@ -11,7 +11,7 @@ name = "PrimeIntellect/Qwen3-0.6B" # PrimeIntellect/Qwen3-0.6B ships its own chat template (distinct from Qwen3's) [renderer] -name = "default" +name = "prime-qwen3" [model.lora] rank = 8 diff --git a/configs/ci/integration/reverse_text_sft_lora/start.toml b/configs/ci/integration/reverse_text_sft_lora/start.toml index eb24c18ae9..f337d3e1c6 100644 --- a/configs/ci/integration/reverse_text_sft_lora/start.toml +++ b/configs/ci/integration/reverse_text_sft_lora/start.toml @@ -10,7 +10,7 @@ name = "PrimeIntellect/Qwen3-0.6B" # PrimeIntellect/Qwen3-0.6B ships its own chat template (distinct from Qwen3's) [renderer] -name = "default" +name = "prime-qwen3" [model.lora] rank = 8 diff --git a/deps/renderers b/deps/renderers index 5904fa24aa..33c8e6f897 160000 --- a/deps/renderers +++ b/deps/renderers @@ -1 +1 @@ -Subproject commit 5904fa24aa73f83c1694fd85a78d0e746d468284 +Subproject commit 33c8e6f897deabaa3b09c9b239afe0acca5fc873 diff --git a/examples/reverse_text/sft.toml b/examples/reverse_text/sft.toml index 345870650b..a2a8f459eb 100644 --- a/examples/reverse_text/sft.toml +++ b/examples/reverse_text/sft.toml @@ -7,7 +7,7 @@ name = "PrimeIntellect/Qwen3-0.6B" # PrimeIntellect/Qwen3-0.6B ships its own chat template (distinct from Qwen3's) [renderer] -name = "default" +name = "prime-qwen3" [data] name = "willcb/R1-reverse-wikipedia-paragraphs-v1-1000" diff --git a/examples/wordle/rl.toml b/examples/wordle/rl.toml index 990fb185fb..db0c7ed584 100644 --- a/examples/wordle/rl.toml +++ b/examples/wordle/rl.toml @@ -30,9 +30,6 @@ max_completion_tokens = 1024 [inference.parallel] dp = 6 -# Qwen3 finetune with the standard PI template patch (byte-identical to -# PrimeIntellect/Qwen3-0.6B base); always re-emits prior blocks. -# Match that with renderer thinking_retention. +# Fine-tune inherits the PrimeIntellect Qwen3 template byte-for-byte. [orchestrator.renderer] -name = "qwen3" -thinking_retention = "all" +name = "prime-qwen3" diff --git a/examples/wordle/sft.toml b/examples/wordle/sft.toml index 8d00ef37b3..ece7d3a9d3 100644 --- a/examples/wordle/sft.toml +++ b/examples/wordle/sft.toml @@ -7,7 +7,7 @@ name = "PrimeIntellect/Qwen3-1.7B" # PrimeIntellect/Qwen3-1.7B ships its own chat template (distinct from Qwen3's) [renderer] -name = "default" +name = "prime-qwen3" [data] name = "willcb/V3-wordle" diff --git a/packages/prime-rl-configs/src/prime_rl/configs/sft.py b/packages/prime-rl-configs/src/prime_rl/configs/sft.py index 3274fc5014..e57c5882f7 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/sft.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/sft.py @@ -3,7 +3,7 @@ from typing import Annotated, Literal, TypeAlias from pydantic import Field, model_validator -from renderers import AutoRendererConfig, RendererConfig +from renderers import AutoRendererConfig, DefaultRendererConfig, RendererConfig from renderers.base import MODEL_RENDERER_MAP from prime_rl.configs.shared import ( @@ -274,28 +274,25 @@ def validate_deployment(self): return self @model_validator(mode="after") - def validate_auto_renderer_resolves(self): - """Reject renderer auto-resolution misses at config time (mirrors the - OrchestratorConfig validator). Resolution is an exact-name lookup, so - it is fully decidable here; fake-data runs without validation need no - renderer and are exempt. - """ - if not isinstance(self.renderer, AutoRendererConfig): - return self + def validate_typed_renderer(self): + """Require a typed renderer whenever SFT renders real samples.""" if self.data.type == "fake" and self.val is None: return self + model_id = self.tokenizer.name or self.model.name - if model_id in MODEL_RENDERER_MAP: + if isinstance(self.renderer, AutoRendererConfig): + if model_id in MODEL_RENDERER_MAP: + return self + reason = f"no typed renderer is registered for {model_id!r}" + elif isinstance(self.renderer, DefaultRendererConfig): + reason = "renderer.name='default' selects DefaultRenderer" + else: return self + raise ValueError( - f"renderer.name='auto' but {model_id!r} is not in " - f"renderers.base.MODEL_RENDERER_MAP, so it would silently fall back to " - f"DefaultRenderer. Pick one: " - f"(a) [renderer] name='default' — for fine-tunes / vendored mirrors with " - f"custom chat templates (DefaultRenderer calls apply_chat_template). " - f"(b) [renderer] name= — if {model_id!r} is " - f"template-identical to a mapped family (and ideally also add it upstream " - f"to renderers.base.MODEL_RENDERER_MAP)." + f"SFT requires a typed renderer with sampled-token and content attribution, but {reason}. " + "Implement and register the renderer in the renderers package, or explicitly select an existing " + "typed renderer only when its template is verified to match." ) @model_validator(mode="after") diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py index 49243a9f98..57687d5666 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -531,6 +531,26 @@ def test_orchestrator_renderer_auto_accepts_mapped_model(): assert config.renderer.name == "auto" +def test_sft_renderer_auto_accepts_prime_qwen_model(): + config = SFTConfig.model_validate({"model": {"name": "PrimeIntellect/Qwen3-0.6B"}}) + assert config.renderer.name == "auto" + + +def test_sft_rejects_default_renderer_for_real_data(): + with pytest.raises(ValidationError, match="requires a typed renderer"): + SFTConfig.model_validate({"renderer": {"name": "default"}}) + + +def test_sft_allows_unused_default_renderer_for_fake_data(): + config = SFTConfig.model_validate( + { + "data": {"type": "fake"}, + "renderer": {"name": "default"}, + } + ) + assert config.renderer.name == "default" + + def test_orchestrator_explicit_renderer_skips_unmapped_check(): """Explicit renderer.name bypasses the auto-resolution check — user opted in.""" config = OrchestratorConfig.model_validate( From 839a543be45300f0f4c41386cfe2ecf2c7da7dfd Mon Sep 17 00:00:00 2001 From: hubert-marek Date: Fri, 10 Jul 2026 21:14:56 +0000 Subject: [PATCH 09/12] chore(configs): drop redundant renderer pins for mapped checkpoints The base PrimeIntellect Qwen3 checkpoints now auto-resolve to the typed renderer; only renamed fine-tunes keep an explicit selection. Co-authored-by: Cursor --- configs/ci/integration/reverse_text_sft/resume.toml | 4 ---- configs/ci/integration/reverse_text_sft/start.toml | 4 ---- configs/ci/integration/reverse_text_sft_lora/resume.toml | 4 ---- configs/ci/integration/reverse_text_sft_lora/start.toml | 4 ---- examples/reverse_text/sft.toml | 4 ---- examples/wordle/sft.toml | 4 ---- 6 files changed, 24 deletions(-) diff --git a/configs/ci/integration/reverse_text_sft/resume.toml b/configs/ci/integration/reverse_text_sft/resume.toml index d0431d4f93..0aecee8e49 100644 --- a/configs/ci/integration/reverse_text_sft/resume.toml +++ b/configs/ci/integration/reverse_text_sft/resume.toml @@ -6,10 +6,6 @@ resume_step = 5 [model] name = "PrimeIntellect/Qwen3-0.6B" -# PrimeIntellect/Qwen3-0.6B ships its own chat template (distinct from Qwen3's) -[renderer] -name = "prime-qwen3" - [data] name = "PrimeIntellect/Reverse-Text-SFT" batch_size = 4 diff --git a/configs/ci/integration/reverse_text_sft/start.toml b/configs/ci/integration/reverse_text_sft/start.toml index 89f3b5d651..405d7609e4 100644 --- a/configs/ci/integration/reverse_text_sft/start.toml +++ b/configs/ci/integration/reverse_text_sft/start.toml @@ -5,10 +5,6 @@ max_steps = 5 [model] name = "PrimeIntellect/Qwen3-0.6B" -# PrimeIntellect/Qwen3-0.6B ships its own chat template (distinct from Qwen3's) -[renderer] -name = "prime-qwen3" - [data] name = "PrimeIntellect/Reverse-Text-SFT" batch_size = 4 diff --git a/configs/ci/integration/reverse_text_sft_lora/resume.toml b/configs/ci/integration/reverse_text_sft_lora/resume.toml index f8ea682d5a..67a6ec12af 100644 --- a/configs/ci/integration/reverse_text_sft_lora/resume.toml +++ b/configs/ci/integration/reverse_text_sft_lora/resume.toml @@ -9,10 +9,6 @@ save_adapter_separately = true [model] name = "PrimeIntellect/Qwen3-0.6B" -# PrimeIntellect/Qwen3-0.6B ships its own chat template (distinct from Qwen3's) -[renderer] -name = "prime-qwen3" - [model.lora] rank = 8 target_modules = [ diff --git a/configs/ci/integration/reverse_text_sft_lora/start.toml b/configs/ci/integration/reverse_text_sft_lora/start.toml index f337d3e1c6..28fb516f2e 100644 --- a/configs/ci/integration/reverse_text_sft_lora/start.toml +++ b/configs/ci/integration/reverse_text_sft_lora/start.toml @@ -8,10 +8,6 @@ save_adapter_separately = true [model] name = "PrimeIntellect/Qwen3-0.6B" -# PrimeIntellect/Qwen3-0.6B ships its own chat template (distinct from Qwen3's) -[renderer] -name = "prime-qwen3" - [model.lora] rank = 8 target_modules = [ diff --git a/examples/reverse_text/sft.toml b/examples/reverse_text/sft.toml index a2a8f459eb..5b21cee566 100644 --- a/examples/reverse_text/sft.toml +++ b/examples/reverse_text/sft.toml @@ -5,10 +5,6 @@ max_steps = 100 [model] name = "PrimeIntellect/Qwen3-0.6B" -# PrimeIntellect/Qwen3-0.6B ships its own chat template (distinct from Qwen3's) -[renderer] -name = "prime-qwen3" - [data] name = "willcb/R1-reverse-wikipedia-paragraphs-v1-1000" seq_len = 4096 diff --git a/examples/wordle/sft.toml b/examples/wordle/sft.toml index ece7d3a9d3..156c5ff5c6 100644 --- a/examples/wordle/sft.toml +++ b/examples/wordle/sft.toml @@ -5,10 +5,6 @@ max_steps = 20 [model] name = "PrimeIntellect/Qwen3-1.7B" -# PrimeIntellect/Qwen3-1.7B ships its own chat template (distinct from Qwen3's) -[renderer] -name = "prime-qwen3" - [data] name = "willcb/V3-wordle" seq_len = 1024 From 592743f857236901bbe76fc0ec3445c2387983f9 Mon Sep 17 00:00:00 2001 From: hubert-marek Date: Fri, 10 Jul 2026 22:09:39 +0000 Subject: [PATCH 10/12] fix(sft): train renderer stop signals for GLM-style templates Defer to the renderer's sampled_mask by default so turn-closing role markers attributed to the next message stay trainable, and append the canonical stop token when a final assistant turn renders without one. Co-authored-by: Cursor --- src/prime_rl/trainer/sft/data.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/src/prime_rl/trainer/sft/data.py b/src/prime_rl/trainer/sft/data.py index 2e80e3bd55..11348ffda0 100644 --- a/src/prime_rl/trainer/sft/data.py +++ b/src/prime_rl/trainer/sft/data.py @@ -234,6 +234,11 @@ def should_mask(message: dict) -> bool: case _: raise ValueError(f"Invalid message role: {message['role']}") + # Defer to the renderer's sampled_mask by default: a role filter would + # drop sampled stop markers attributed to the next message (e.g. GLM's + # turn-closing <|user|> / <|observation|>). + role_to_mask = None if self.loss_mask_config.assistant else should_mask + # Non-assistant roles are opted into the loss via the renderer's # body-only path: the message content is trained, not the role # scaffolding (e.g. <|im_start|>assistant) the harness emits. @@ -241,13 +246,22 @@ def should_mask(message: dict) -> bool: sample = build_training_sample( self.renderer, messages, - role_to_mask=should_mask, + role_to_mask=role_to_mask, tools=tools, content_sft_roles=content_sft_roles or None, ) input_ids = list(sample.token_ids) loss_mask = list(sample.loss_mask) + # Some templates render no terminator after a final assistant turn + # (e.g. GLM); append the canonical stop token as a trainable target. + stop_token_ids = self.renderer.get_stop_token_ids() + last_trainable = next((i for i in range(len(loss_mask) - 1, -1, -1) if loss_mask[i]), None) + ends_with_stop = last_trainable is not None and input_ids[last_trainable] in set(stop_token_ids) + if messages[-1].get("role") == "assistant" and not ends_with_stop: + input_ids.append(stop_token_ids[0]) + loss_mask.append(True) + # Causal shift: model predicts next token from current. target_ids = input_ids.copy()[1:] loss_mask = loss_mask[1:] From b8e69410f8ed2a3d7555f76523e50df56e6f5737 Mon Sep 17 00:00:00 2001 From: hubert-marek Date: Fri, 10 Jul 2026 22:40:41 +0000 Subject: [PATCH 11/12] refactor(sft): delegate final stop handling to build_training_sample Bump renderers to main with ensure_final_stop and drop the trainer-side append. Co-authored-by: Cursor --- deps/renderers | 2 +- src/prime_rl/trainer/sft/data.py | 10 +--------- 2 files changed, 2 insertions(+), 10 deletions(-) diff --git a/deps/renderers b/deps/renderers index 33c8e6f897..fa04acd916 160000 --- a/deps/renderers +++ b/deps/renderers @@ -1 +1 @@ -Subproject commit 33c8e6f897deabaa3b09c9b239afe0acca5fc873 +Subproject commit fa04acd916ce749c302eec905ecf3962bc438429 diff --git a/src/prime_rl/trainer/sft/data.py b/src/prime_rl/trainer/sft/data.py index 11348ffda0..8d0422e246 100644 --- a/src/prime_rl/trainer/sft/data.py +++ b/src/prime_rl/trainer/sft/data.py @@ -249,19 +249,11 @@ def should_mask(message: dict) -> bool: role_to_mask=role_to_mask, tools=tools, content_sft_roles=content_sft_roles or None, + ensure_final_stop=True, ) input_ids = list(sample.token_ids) loss_mask = list(sample.loss_mask) - # Some templates render no terminator after a final assistant turn - # (e.g. GLM); append the canonical stop token as a trainable target. - stop_token_ids = self.renderer.get_stop_token_ids() - last_trainable = next((i for i in range(len(loss_mask) - 1, -1, -1) if loss_mask[i]), None) - ends_with_stop = last_trainable is not None and input_ids[last_trainable] in set(stop_token_ids) - if messages[-1].get("role") == "assistant" and not ends_with_stop: - input_ids.append(stop_token_ids[0]) - loss_mask.append(True) - # Causal shift: model predicts next token from current. target_ids = input_ids.copy()[1:] loss_mask = loss_mask[1:] From 041a1e559851ff82fdbfb7ecd03d4d7513da9c28 Mon Sep 17 00:00:00 2001 From: hubert-marek Date: Fri, 10 Jul 2026 22:40:41 +0000 Subject: [PATCH 12/12] chore(configs): restore seq_len and loss_impl defaults These changes were unrelated to the renderer-only migration. Co-authored-by: Cursor --- .../prime-rl-configs/src/prime_rl/configs/sft.py | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/packages/prime-rl-configs/src/prime_rl/configs/sft.py b/packages/prime-rl-configs/src/prime_rl/configs/sft.py index e57c5882f7..1760ce8aca 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/sft.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/sft.py @@ -31,7 +31,7 @@ class BaseDataConfig(BaseConfig): batch_size: int = Field(128, ge=1) """Global batch size.""" - seq_len: int = Field(256, ge=1) + seq_len: int = Field(128, ge=1) """Sequence length.""" pack_function: Literal["cat", "stack"] = "cat" @@ -52,12 +52,6 @@ def validate_batch_size(self): class FakeDataConfig(BaseDataConfig): type: Literal["fake"] = "fake" - seq_len: int = Field(128, ge=1) - """Sequence length.""" - - pack_function: Literal["cat", "stack"] = "cat" - """Sample packing strategy.""" - length: Literal["fixed", "variable"] = "fixed" """Use fixed-length samples or variable-length samples.""" @@ -227,8 +221,8 @@ class SFTConfig(BaseConfig): dist_timeout_seconds: int = 3600 """Timeout in seconds for torch distributed ops.""" - loss_impl: Literal["liger", "torch", "liger_fused", "quack_fused"] = "liger_fused" - """Cross-entropy loss implementation. Defaults to fused Liger loss to avoid materializing full logits.""" + loss_impl: Literal["liger", "torch", "liger_fused", "quack_fused"] = "torch" + """Cross-entropy loss implementation. ``liger_fused`` fuses the lm_head projection with the CE loss to avoid materializing full logits. ``quack_fused`` uses quack-kernels for chunked linear + CE with CuTe DSL CUDA kernels.""" heartbeat: HeartbeatConfig | None = None """BetterStack heartbeat configuration for monitoring training progress."""