feat(glm): make next-message role marker trainable after assistant turns

[hallerite]

Summary

GLM-4.5 / GLM-5 chat templates have no per-turn close token inside the assistant span — the next message's role marker (<|user|> / <|observation|> / <|system|>) doubles as the inference stop signal (see get_stop_token_ids).

Before this change, that role marker was attributed to the next message with is_sampled=False, so SFT never trained the model to emit it. With LossMaskConfig.assistant=True and tool/user/system=False (the defaults), the only training signal saying "stop after </tool_call>" was the data.py EOS-fallback on the very last assistant of the conversation — intermediate turns got nothing. At inference the natural continuation of </tool_call>\n is another <tool_call>, which is exactly what we observe in practice.

Repro

GLM-Air SFT step_250 evaluated on SWE-bench Verified (full, 500 examples × 2 rollouts) emitted 50-283 parallel tool_calls per turn in a single assistant turn. 26% of assistant turns had len(tool_calls) > 1, with median max-tool-calls = 35 for solved rollouts vs 55 for failed. Pass@1 was 25.0%; clear correlation between high-burst rollouts and failure.

Fix

When the previous message is an assistant, attribute that single role-opening token to the assistant's msg_idx with is_sampled=True. The byte stream is unchanged (token_ids still match apply_chat_template); only message_indices and sampled_mask for that one token change. build_training_sample then yields loss_mask=True for it (because the assistant role is trainable by default), so SFT explicitly trains the model to emit the stop token after </tool_call> / content.

Applied symmetrically in both glm45.py and glm5.py for:

• <|observation|> (when next message is tool)
• <|user|> (when next message is user)
• <|system|> (when next message is system)

Test plan

• test_build_training_sample_ids_match still passes (token_ids unchanged)
• New unit test: in a conversation [user, assistant w/ tool_call, tool, user, ...], the <|observation|> and <|user|> tokens that follow the assistant have loss_mask=True
• Re-SFT GLM-Air on INTELLECT-5-SFT-Processed and re-evaluate on SWE-bench Verified; expect multi-tool-call ratio to drop sharply and pass@1 to increase from 25.0%

🤖 Generated with Claude Code

---

Note

Medium Risk
Changes which tokens receive SFT loss for GLM-4.5/5 (mask metadata only, not token_ids) and alters build_training_sample defaults; incorrect caller assumptions could raise ValueError or shift training signal.

Overview
GLM chat templates use the next message’s <|user|> / <|observation|> as inference stop tokens after assistant tool turns, but those tokens were previously marked is_sampled=False, so typical assistant-only SFT never trained the model to emit them after </tool_call>.

In glm45.py and glm5.py, when the prior turn is assistant, the following user/tool role opener is emitted with is_sampled=True (still is_content=False and same token_ids). bridge_to_next_turn keeps those openers is_sampled=False and documents the intentional SFT vs RL split.

In base.py, build_training_sample now accepts optional role_to_mask=None: if the renderer provides sampled_mask, loss follows sampled_mask alone so GLM stop markers on the next message can be trained without a role filter. Callers that omit role_to_mask on renderers without sampled_mask get ValueError instead of a wrong mask.

^{Reviewed by Cursor Bugbot for commit b9fae66. Bugbot is set up for automated code reviews on this repo. Configure here.}

Note

Mark next-message role tokens as trainable after assistant turns in GLM renderers

• In glm45.py and glm5.py, the <|user|> and <|observation|> role opener tokens that immediately follow an assistant message are now emitted with is_sampled=True instead of always False.
• In renderers/base.py, build_training_sample now accepts an optional role_to_mask (defaulting to None) and uses sampled_mask from the renderer when no role filter is supplied, raising ValueError if neither is available.
• Behavioral Change: loss mask semantics for GLM models now include the next-message role marker token in the trainable region when it follows an assistant turn.

^{Macroscope summarized b9fae66.}

[macroscopeapp]

Approvability

Verdict: Needs human review

This PR changes training loss mask behavior to make role-marker tokens trainable after assistant turns, which is a significant runtime change affecting model training. Additionally, there is an unresolved high-severity review comment identifying a potential bug where the intended tokens may still be excluded from the loss mask.

^{You can customize Macroscope's approvability policy. Learn more.}

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit cb07cc4. Configure here.}

[cursor]

Stop tokens still excluded from loss

High Severity

The build_training_sample function incorrectly excludes renderer-generated scaffolding tokens (like GLM's <|observation|>) from the loss_mask. Despite being sampled_mask=True, these tokens are masked out when role_to_mask filters them or when content_sft_roles is active, as the latter prioritizes is_content (which is False for scaffolding). This prevents training on these critical tokens.

 

^{Reviewed by Cursor Bugbot for commit cb07cc4. Configure here.}

[hallerite]

This is the natural consequence of GLM's chat-template design where the per-turn stop signal lives inside the next message's span — there is no dedicated <|im_end|>-style terminator inside the assistant turn itself. We have three knobs and they trade off:

1. msg_idx attribution of <|observation|> / <|user|> after an assistant turn. We keep it attached to the next message (the tool/user message it opens) rather than retro-attributing to the assistant. The byte stream still ends the assistant's turn at that token, but structurally it belongs to the next message — which is what every downstream consumer (message_token_spans, role-based slicing, bridge/extension code) expects.
2. is_sampled=True on that opener — added by this PR. Captures the fact that the model would sample that token at inference to close its turn. This is what makes the default training path (role_to_mask=None) train the model to emit the stop.
3. is_content=False on that opener — preserved from before. The token is a role marker, not message body, so content_mask_for_roles({"tool"}) and content_token_spans_by_role() correctly exclude it.

The two specialized modes you flag fall out of this:

• role_to_mask=lambda m: m["role"] == "assistant" — by construction, "only train on assistant-attributed tokens." Stop opener is attributed to the tool message → excluded. Working as specified; if the caller wants the stop trained too, they pass role_to_mask=None (default) and rely on the sampled-only path.
• content_sft_roles={"tool"} — body-only SFT on tool responses, gated by is_content. Scaffold tokens (including the stop opener) are intentionally is_content=False to keep content_token_spans_by_role honest. In this mixed RL-on-assistant + SFT-on-tool-body mode, the assistant stop-signal supervision belongs to the RL trajectory, not to body-only SFT.

In short, the GLM chat template can't express "this token belongs to message N but counts as message N-1's output for training" cleanly, and the choice we made here (keep structural attribution, use is_sampled as the per-token "model would sample this" signal) is the cleanest fit with the existing build_training_sample contract. Default-mode training works correctly; the filter-mode misses are explicit user opt-ins, not silent bugs.

[mikasenghaas]

comment to code ratio is crazy on this one