From 56852cdaaf62c25b8f86ae1bd10d13c214d2e520 Mon Sep 17 00:00:00 2001 From: MayDomine <1583143678@qq.com> Date: Fri, 11 Sep 2026 15:09:10 +0800 Subject: [PATCH] [bugfix] render chat-template diffs without the generation prompt in SampleBuilder SampleBuilder.append_text diffed successive apply_chat_template renders that each carried add_generation_prompt=True. Such renders are not prefixes of each other, so slicing new_ids[len(old_ids):] corrupts every multi-turn prompt: for a system+user pair the user header is replaced by a stray assistant header while the token count stays the same, so the corruption is silent. Affects dataset/gsm8k.py (system+user, the default grpo_gsm8k recipe) and dataset/s9_math.py; single-message dataset/math.py is unaffected. Fixes #2. build_sample now diffs renders without the generation prompt (those are true prefixes) and appends the generation prompt once at the end, mask 0. append_text is folded into build_sample; its only caller was build_sample and its incremental generation-prompt semantics cannot be made correct. apply_chat_template output is normalized through _render, which accepts both a BatchEncoding and a plain id list across transformers versions. scripts/bench_train_only.py mirrored the same diff logic and is fixed the same way; it now also inserts the generation prompt before each assistant turn, which is what the engine was actually prompted with. Note: prompt token streams for multi-turn samples change, so metrics from runs before this fix are not directly comparable. --- meshy/utils/sample.py | 51 +++++++++++------ scripts/bench_train_only.py | 31 ++++++---- tests/test_rollout_metrics.py | 2 +- tests/test_sample_builder.py | 103 ++++++++++++++++++++++++++++++++++ 4 files changed, 160 insertions(+), 27 deletions(-) create mode 100644 tests/test_sample_builder.py diff --git a/meshy/utils/sample.py b/meshy/utils/sample.py index 0938a71..5503ccb 100644 --- a/meshy/utils/sample.py +++ b/meshy/utils/sample.py @@ -30,9 +30,39 @@ class SampleBuilder: def __init__(self, model_path: str): self.tokenizer = AutoTokenizer.from_pretrained(model_path) + def _render(self, messages: list[dict[str, str]], *, add_generation_prompt: bool) -> list[int]: + if not messages: + return [] + out = self.tokenizer.apply_chat_template( + messages, tokenize=True, add_generation_prompt=add_generation_prompt + ) + # Some transformers versions return a BatchEncoding (a UserDict), others + # a plain list of ids. + return list(out["input_ids"] if hasattr(out, "keys") else out) + + def _spans(self, messages: list[dict[str, str]]) -> list[tuple[str | None, list[int]]]: + """Split the rendered prompt into per-message ``(role, ids)`` spans. + + Spans are diffed on renders *without* the generation prompt: successive + renders are then true prefixes of each other, so slicing by length is + exact. (Diffing renders that carry the generation prompt is wrong: the + prompt sits at the end of each render, so they are not prefix-related.) + The final span is the generation prompt itself, with role ``None``. + """ + spans: list[tuple[str | None, list[int]]] = [] + prev_ids: list[int] = [] + for i, msg in enumerate(messages): + full_ids = self._render(messages[: i + 1], add_generation_prompt=False) + spans.append((msg["role"], full_ids[len(prev_ids):])) + prev_ids = full_ids + tail = self._render(messages, add_generation_prompt=True)[len(prev_ids):] + spans.append((None, tail)) + return spans + def build_sample(self, messages: Iterable[dict[str, str]], logprob: float = 0.0) -> Sample: + messages = [{"role": m["role"], "content": m["content"]} for m in messages] sample = Sample( - messages=[], + messages=messages, tokens=[], logprobs=[], masks=[], @@ -40,21 +70,10 @@ def build_sample(self, messages: Iterable[dict[str, str]], logprob: float = 0.0) reward=None, advantage=None ) - for msg in messages: - sample = self.append_text(sample, msg["role"], msg["content"], logprob) - return sample - - def append_text(self, sample: Sample, role: str, content: str, logprob: float = 0.0) -> Sample: - sample.messages.append({"role": role, "content": content}) - if len(sample.messages) == 1: - ids = self.tokenizer.apply_chat_template(sample.messages, tokenize=True, add_generation_prompt=True)["input_ids"] - else: - old_ids = self.tokenizer.apply_chat_template(sample.messages[:-1], tokenize=True, add_generation_prompt=True)["input_ids"] - new_ids = self.tokenizer.apply_chat_template(sample.messages, tokenize=True, add_generation_prompt=True)["input_ids"] - ids = new_ids[len(old_ids):] - sample.tokens.extend(ids) - sample.logprobs.extend([logprob] * len(ids)) - sample.masks.extend([1 if role == "assistant" else 0] * len(ids)) + for role, ids in self._spans(messages): + sample.tokens.extend(ids) + sample.logprobs.extend([logprob] * len(ids)) + sample.masks.extend([1 if role == "assistant" else 0] * len(ids)) return sample def append_tokens( diff --git a/scripts/bench_train_only.py b/scripts/bench_train_only.py index b59b7b2..28d2a30 100644 --- a/scripts/bench_train_only.py +++ b/scripts/bench_train_only.py @@ -67,9 +67,11 @@ class _Tokenizer: """Rebuild ``tokens / mask_assistant`` from a message list. Non-assistant turns are tokenized through the chat template exactly like - ``SampleBuilder.append_text`` does (template diff with the generation - prompt); assistant turns are ``encode(content) + [eos]`` — what the - inference engine actually produced, without a trailing generation prompt. + ``SampleBuilder.build_sample`` does (template diffs rendered without the + generation prompt, so successive renders are true prefixes); before each + assistant turn the generation prompt the engine was actually prompted with + is appended (mask 0). Assistant turns are ``encode(content) + [eos]`` — + what the inference engine actually produced. """ def __init__(self, model_path: str) -> None: @@ -78,8 +80,12 @@ def __init__(self, model_path: str) -> None: self.tok = AutoTokenizer.from_pretrained(model_path) self.eos = self.tok.eos_token_id - def _template(self, messages: list[dict[str, str]]) -> list[int]: - out = self.tok.apply_chat_template(messages, tokenize=True, add_generation_prompt=True) + def _template(self, messages: list[dict[str, str]], *, add_generation_prompt: bool) -> list[int]: + if not messages: + return [] + out = self.tok.apply_chat_template( + messages, tokenize=True, add_generation_prompt=add_generation_prompt + ) # Newer transformers return a BatchEncoding (a UserDict, not a dict). return list(out["input_ids"] if hasattr(out, "keys") else out) @@ -87,19 +93,24 @@ def __call__(self, messages: list[dict[str, str]]) -> tuple[list[int], list[int] tokens: list[int] = [] masks: list[int] = [] seen: list[dict[str, str]] = [] + prev: list[int] = [] for msg in messages: - seen.append(msg) if msg["role"] == "assistant": + # The engine saw the generation prompt before producing this turn. + gen = self._template(seen, add_generation_prompt=True)[len(prev):] + tokens.extend(gen) + masks.extend([0] * len(gen)) ids = self.tok.encode(msg["content"], add_special_tokens=False) if self.eos is not None: ids = ids + [self.eos] + seen.append(msg) + prev = self._template(seen, add_generation_prompt=False) masks.extend([1] * len(ids)) else: - # Tokens the template adds beyond the previous turns (the - # first turn is taken whole), mirroring SampleBuilder.append_text. - full = self._template(seen) - prev = self._template(seen[:-1]) if len(seen) > 1 else [] + seen.append(msg) + full = self._template(seen, add_generation_prompt=False) ids = full[len(prev):] + prev = full masks.extend([0] * len(ids)) tokens.extend(ids) return tokens, masks diff --git a/tests/test_rollout_metrics.py b/tests/test_rollout_metrics.py index 0a7353f..6e8bd9f 100644 --- a/tests/test_rollout_metrics.py +++ b/tests/test_rollout_metrics.py @@ -31,7 +31,7 @@ def _sample( mask = torch.ones(length) mask[:prompt] = 0 logprobs = -torch.rand(length, generator=g) - logprobs[:prompt] = 0.0 # prompt placeholders, like SampleBuilder.append_text + logprobs[:prompt] = 0.0 # prompt placeholders, like SampleBuilder.build_sample return TensorDict( { "tokens": torch.randint(1, 1000, (length,), generator=g), diff --git a/tests/test_sample_builder.py b/tests/test_sample_builder.py new file mode 100644 index 0000000..42fbf07 --- /dev/null +++ b/tests/test_sample_builder.py @@ -0,0 +1,103 @@ +"""SampleBuilder must produce exactly the tokens a one-shot chat-template +render with the generation prompt would produce. + +The regression guarded here: building the prompt incrementally by diffing +renders that each carry the generation prompt corrupts multi-turn prompts +(the ``<|im_start|>user`` header of the second message is replaced by a +stray ``<|im_start|>assistant`` header), while keeping the token count +identical — a silent failure. Uses a deterministic ChatML-style fake +tokenizer, so no model download is needed. +""" + +from __future__ import annotations + +import types + +import pytest + + +class _ChatMLTokenizer: + """Character-level tokenizer with a ChatML-style chat template.""" + + def _render_text(self, messages, add_generation_prompt): + text = "".join( + f"<|im_start|>{m['role']}\n{m['content']}<|im_end|>\n" for m in messages + ) + if add_generation_prompt: + text += "<|im_start|>assistant\n" + return text + + def apply_chat_template(self, messages, tokenize=True, add_generation_prompt=False): + text = self._render_text(messages, add_generation_prompt) + return {"input_ids": [ord(c) for c in text]} + + def decode(self, tokens, **kw): + return "".join(chr(t) for t in tokens) + + +@pytest.fixture() +def builder(monkeypatch): + import meshy.utils.sample as sample_mod + + monkeypatch.setattr( + sample_mod, + "AutoTokenizer", + types.SimpleNamespace(from_pretrained=lambda *a, **k: _ChatMLTokenizer()), + ) + return sample_mod.SampleBuilder("stub-model") + + +@pytest.mark.parametrize( + "messages", + [ + [{"role": "user", "content": "What is 2+2?"}], + [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "What is 2+2?"}, + ], + [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "What is 2+2?"}, + {"role": "assistant", "content": "4"}, + {"role": "user", "content": "And 3+3?"}, + ], + ], + ids=["user", "system+user", "multi-turn"], +) +def test_build_sample_matches_full_render(builder, messages): + sample = builder.build_sample(messages) + expected = builder.tokenizer.apply_chat_template( + messages, tokenize=True, add_generation_prompt=True + )["input_ids"] + assert sample.tokens == expected + assert len(sample.logprobs) == len(sample.tokens) + assert len(sample.masks) == len(sample.tokens) + + +def test_masks_cover_assistant_spans_only(builder): + messages = [ + {"role": "system", "content": "S"}, + {"role": "user", "content": "U"}, + {"role": "assistant", "content": "A"}, + {"role": "user", "content": "V"}, + ] + sample = builder.build_sample(messages) + tok = builder.tokenizer + masked = tok.decode([t for t, m in zip(sample.tokens, sample.masks) if m == 1]) + assert masked == "<|im_start|>assistant\nA<|im_end|>\n" + # The trailing generation prompt is part of the prompt, not the response. + assert sample.masks[-1] == 0 + + +def test_render_accepts_plain_list_return(builder): + """apply_chat_template returning a bare list (older transformers) works too.""" + tok = builder.tokenizer + orig = tok.apply_chat_template + tok.apply_chat_template = lambda *a, **k: list(orig(*a, **k)["input_ids"]) + messages = [ + {"role": "system", "content": "S"}, + {"role": "user", "content": "U"}, + ] + sample = builder.build_sample(messages) + expected = orig(messages, tokenize=True, add_generation_prompt=True)["input_ids"] + assert sample.tokens == expected