-
Notifications
You must be signed in to change notification settings - Fork 342
feat(sft): require renderer-based tokenization #2988
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from 10 commits
ff65c4b
e1f5fd3
4abe225
57718d9
67e2a4f
36bc39b
a37949b
d97cf5f
3b5afdb
6090bef
839a543
592743f
b8e6941
041a1e5
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| +1 −1 | README.md | |
| +6 −1 | docs/renderer-config.md | |
| +10 −0 | renderers/__init__.py | |
| +43 −5 | renderers/base.py | |
| +87 −13 | renderers/configs.py | |
| +706 −0 | renderers/hy3.py | |
| +384 −13 | renderers/laguna_xs2.py | |
| +226 −16 | renderers/parsing.py | |
| +694 −0 | renderers/prime_qwen3.py | |
| +4 −3 | tests/conftest.py | |
| +3 −1 | tests/test_bridge.py | |
| +486 −0 | tests/test_hy3.py | |
| +376 −0 | tests/test_laguna_xs21.py | |
| +196 −0 | tests/test_prime_qwen3_parity.py | |
| +35 −1 | tests/test_renderer_config_parity.py | |
| +6 −1 | tests/test_roundtrip.py | |
| +1 −0 | tests/test_tool_arg_type_preservation.py |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -3,7 +3,8 @@ | |
| from typing import Annotated, Literal, TypeAlias | ||
|
|
||
| from pydantic import Field, model_validator | ||
| from renderers import RendererConfig | ||
| from renderers import AutoRendererConfig, DefaultRendererConfig, RendererConfig | ||
| from renderers.base import MODEL_RENDERER_MAP | ||
|
|
||
| from prime_rl.configs.shared import ( | ||
| EnvVars, | ||
|
|
@@ -30,7 +31,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 +52,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.""" | ||
|
hubert-marek marked this conversation as resolved.
Outdated
|
||
|
|
||
| length: Literal["fixed", "variable"] = "fixed" | ||
| """Use fixed-length samples or variable-length samples.""" | ||
|
|
||
|
|
@@ -175,13 +182,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.""" | ||
|
hubert-marek marked this conversation as resolved.
|
||
|
|
||
| data: DataConfig = SFTDataConfig() | ||
|
|
||
|
|
@@ -225,8 +227,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.""" | ||
|
hubert-marek marked this conversation as resolved.
Outdated
|
||
|
|
||
| heartbeat: HeartbeatConfig | None = None | ||
| """BetterStack heartbeat configuration for monitoring training progress.""" | ||
|
|
@@ -271,6 +273,28 @@ def validate_deployment(self): | |
| raise ValueError("Must use SLURM for multi-node deployment.") | ||
| return self | ||
|
|
||
| @model_validator(mode="after") | ||
| 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 isinstance(self.renderer, AutoRendererConfig): | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. in which case is it not instance of AutoRenderConfig ? shouldn;t it be just if self.renderer is not None ?
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. It needs to resolve later by acutally checking if there is an renderer to be selected by auto |
||
| 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"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") | ||
| def validate_pack_function(self): | ||
| if self.model.cp > 1: | ||
|
|
@@ -319,9 +343,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: | ||
|
hubert-marek marked this conversation as resolved.
|
||
| 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 | ||
|
|
||
Uh oh!
There was an error while loading. Please reload this page.