Swap K2V3 TITO tokenizer to IFM template; rename legacy to k2v3_oldbackup (#43)

* Swap TITO tokenizer's K2V3 to IFM template; rename legacy to k2v3_oldbackup

The K2V3 family is migrating to the IFM-style chat template introduced
in bbq-0601 (used by bbq-8b-mid3_v3 and later checkpoints). The new
template namespaces ChatML tokens as <|ifm|im_start|> / <|ifm|im_end|>,
emits no whitespace between messages, and requires assistant messages
to carry a thinking field. The legacy <|im_end|>\n template stays
supported for older K2V3 checkpoints (bbq-8b-mid3-final and earlier)
that haven't migrated yet.

Changes:
  - K2V3TITOTokenizer now targets the IFM template. merge_tokens is
    pure concat — the buffer already matches the canonical render
    (model stops at <|ifm|im_end|> and no trailing whitespace
    follows in the template).
  - Renamed the legacy K2V3TITOTokenizer to K2V3OldBackupTITOTokenizer.
    Its <|im_end|> + \n boundary-fix logic is preserved bit-for-bit.
  - Added TITOTokenizerType.K2V3_OLDBACKUP enum value and registry
    entry. TITOTokenizerType.K2V3 now points at the new IFM class.
  - Both classes hard-assert at __init__ that the loaded tokenizer's
    vocab matches their target template (refuses to load on a
    misconfigured checkpoint, with an error pointing at the right
    --tito-model value).
  - test_tito_k2v3.py rewritten for IFM invariants (no boundary fix,
    BOS prepend, thinking required, hard-assert sanity).
  - Renamed previous test file to test_tito_k2v3_oldbackup.py with
    K2V3OldBackup references.

Breaking change for downstream sbatch:
  --tito-model k2v3 now refers to the IFM template. Legacy checkpoint
  users must update to --tito-model k2v3_oldbackup. Misconfiguration
  raises at init rather than silently producing wrong TITO buffers.

Out of scope (required separately for IFM training):
  - IFM-compatible SGLang reasoning_parser + tool_parser (see
    LLM360/sglang#33).

Verification:
  - tests/fast/.../test_tito_k2v3.py: 43 passed, 12 skipped (skipped =
    SGLang IFM parsers not yet in this container build).
  - tests/fast/.../test_tito_k2v3_oldbackup.py: 54 passed (legacy
    behavior unchanged).

* Use raw-string docstrings to display \n literally in tito K2V3 classes/tests

Docstrings on K2V3TITOTokenizer / K2V3OldBackupTITOTokenizer and the
two K2V3 test files contain visual references to the literal `\n`
escape sequence (the chat-template trailing newline). The previous
\\n escaping renders correctly but reads awkwardly in source. Convert
the affected docstrings to raw strings (r"""...""") so the source
literally contains \n, which is easier to read and write.

No code or test behavior changes.

Tested: 109 passed (55 IFM + 54 oldbackup) inside the agentic-rl
container with sglang PR #33 shadowed for the parser tests.

---------

Co-authored-by: Zhenting Wang <zhenting.wang@mbzuai.ac.ae>

Author: ZhentingWang

miles/utils/chat_template_utils/tito_tokenizer.py

@@ -340,25 +340,90 @@ def merge_tokens(
 
 
 # ---------------------------------------------------------------------------
-# K2V3 family implementation
+# K2V3 family — current (IFM) chat template
 # ---------------------------------------------------------------------------
 
 
 class K2V3TITOTokenizer(TITOTokenizer):
-    """K2V3 family.
+    r"""K2V3 family with the IFM-style chat template (introduced 2026-06-01).
 
-    The chat template emits ``<|im_end|>\\n`` after every message (jinja
-    block whitespace between ``{{- '<|im_end|>' }}`` and the next block
-    is preserved by default ``trim_blocks``), but the model
-    autoregressively stops at ``<|im_end|>`` without generating the
-    trailing ``\\n``. ``merge_tokens`` inserts the missing newline so the
-    pretokenized buffer matches the canonical template output.
+    The current K2V3 chat template (``bbq-0601`` / ``bbq-8b-mid3_v3`` and
+    later) namespaces ChatML tokens as ``<|ifm|im_start|>`` /
+    ``<|ifm|im_end|>`` and emits NO whitespace between
+    ``<|ifm|im_end|>`` and the next ``<|ifm|im_start|>``. The model
+    autoregressively stops at ``<|ifm|im_end|>`` with no trailing byte;
+    the rollout buffer already matches the canonical template render
+    exactly. ``merge_tokens`` therefore needs no boundary fix — it
+    inherits the base ``TITOTokenizer`` concat behavior.
+
+    Empirical sanity check::
+
+        apply_chat_template([user, asst, user], tokenize=False)
+        → '...A1<|ifm|im_end|><|ifm|im_start|>user\n...'
+                              ^^ no \n between messages
+
+    For legacy K2V3 checkpoints (``bbq-8b-mid3-final`` and earlier) whose
+    chat template uses ``<|im_end|>\n`` between messages, use
+    :class:`K2V3OldBackupTITOTokenizer` (``--tito-model k2v3_oldbackup``)
+    instead.
+    """
+
+    _default_assistant_start_str: str = "<|ifm|im_start|>assistant"
+
+    def __init__(
+        self,
+        tokenizer: Any,
+        chat_template_kwargs: dict[str, Any] | None = None,
+        assistant_start_str: str | None = None,
+        allowed_append_roles: list[str] | None = None,
+    ):
+        super().__init__(
+            tokenizer,
+            chat_template_kwargs,
+            assistant_start_str or self._default_assistant_start_str,
+            allowed_append_roles=allowed_append_roles,
+        )
+        # Hard assert against misconfiguration: refuse to load on a legacy
+        # K2V3 checkpoint whose vocab does not have <|ifm|im_end|>.
+        ifm_end_id = tokenizer.convert_tokens_to_ids("<|ifm|im_end|>")
+        unk_id = getattr(tokenizer, "unk_token_id", None)
+        if ifm_end_id is None or ifm_end_id == unk_id:
+            raise ValueError(
+                "K2V3TITOTokenizer (current/IFM chat template) requires "
+                "<|ifm|im_end|> in the tokenizer vocab. The loaded "
+                "tokenizer does not have this token, suggesting you are "
+                "on a legacy K2V3 checkpoint. Use --tito-model "
+                "k2v3_oldbackup for those."
+            )
+        self._im_end_id: int = ifm_end_id
+        self.trailing_token_ids = frozenset({ifm_end_id})
+
+
+# ---------------------------------------------------------------------------
+# K2V3 family — legacy (<|im_end|>\n) chat template
+# ---------------------------------------------------------------------------
+
+
+class K2V3OldBackupTITOTokenizer(TITOTokenizer):
+    r"""K2V3 family with the LEGACY chat template (``<|im_end|>\n``).
+
+    Use this with legacy K2V3 checkpoints (``bbq-8b-mid3-final`` and
+    earlier) whose chat template emits ``<|im_end|>\n`` after every
+    message (jinja block whitespace between ``{{- '<|im_end|>' }}`` and
+    the next block is preserved by default ``trim_blocks``), but where
+    the model autoregressively stops at ``<|im_end|>`` without producing
+    the trailing ``\n``. ``merge_tokens`` inserts the missing newline so
+    the pretokenized buffer matches the canonical template output.
 
     Empirical sanity check::
 
         apply_chat_template([user, assistant, user], tokenize=False)
-        → '...hello<|im_end|>\\n<|im_start|>user\\n...'
+        → '...hello<|im_end|>\n<|im_start|>user\n...'
                           ^^
+
+    For current K2V3 checkpoints (``bbq-8b-mid3_v3`` and later) whose
+    template uses ``<|ifm|im_end|>`` with no trailing ``\n``, use
+    :class:`K2V3TITOTokenizer` (``--tito-model k2v3``) instead.
     """
 
     _default_assistant_start_str: str = "<|im_start|>assistant"
@@ -376,10 +441,22 @@ def __init__(
             assistant_start_str or self._default_assistant_start_str,
             allowed_append_roles=allowed_append_roles,
         )
+        # Hard assert against misconfiguration: refuse to load on a current
+        # K2V3 checkpoint whose vocab does not have <|im_end|>.
+        im_end_id = tokenizer.convert_tokens_to_ids("<|im_end|>")
+        unk_id = getattr(tokenizer, "unk_token_id", None)
+        if im_end_id is None or im_end_id == unk_id:
+            raise ValueError(
+                "K2V3OldBackupTITOTokenizer (legacy chat template) "
+                "requires <|im_end|> in the tokenizer vocab. The loaded "
+                "tokenizer does not have this token, suggesting you are "
+                "on a current K2V3 checkpoint that uses the IFM template. "
+                "Use --tito-model k2v3 for those."
+            )
         nl_ids = tokenizer.encode("\n", add_special_tokens=False)
         assert len(nl_ids) == 1, f"Expected single newline token, got {nl_ids}"
         self._newline_id: int = nl_ids[0]
-        self._im_end_id: int = tokenizer.convert_tokens_to_ids("<|im_end|>")
+        self._im_end_id: int = im_end_id
         self.trailing_token_ids = frozenset({self._newline_id})
 
     def merge_tokens(
@@ -406,13 +483,15 @@ class TITOTokenizerType(str, Enum):
     QWEN3 = "qwen3"
     GLM47 = "glm47"
     K2V3 = "k2v3"
+    K2V3_OLDBACKUP = "k2v3_oldbackup"
 
 
 _TOKENIZER_REGISTRY: dict[TITOTokenizerType, type[TITOTokenizer]] = {
     TITOTokenizerType.DEFAULT: TITOTokenizer,
     TITOTokenizerType.QWEN3: Qwen3TITOTokenizer,
     TITOTokenizerType.GLM47: GLM47TITOTokenizer,
     TITOTokenizerType.K2V3: K2V3TITOTokenizer,
+    TITOTokenizerType.K2V3_OLDBACKUP: K2V3OldBackupTITOTokenizer,
 }