feat(gemma4): SFT support — loss-mask generator + tests#1
Merged
Conversation
slime's SFT path builds per-token loss masks via
MultiTurnLossMaskGenerator, keyed on --loss-mask-type. Only qwen variants
were supported; Gemma4's chat template uses a different delimiter scheme
(<|turn>{role}\n ... <turn|>\n, assistant rendered as "model", optional
<|channel>thought ... <channel|> reasoning), so SFT on Gemma4 raised
"Unsupported tokenizer type".
Add gen_multi_turn_loss_mask_gemma4 using the same render-then-offset-map
strategy as the Qwen3.5 generator (robust to tokenizer merges across the
content/marker boundary):
- mask each assistant ("model") turn's content plus its <turn|>
terminator so the model learns to end its turn
- leave system/user/tool spans at 0
- when a thinking channel is rendered, exclude it from the loss
- honor step_loss_mask=0 to skip a turn
Wire "gemma4" into the get_loss_mask dispatch.
Tests: tests/utils/test_loss_mask_type_gemma4.py (9 cases) with a
char-level fake tokenizer modeling the Gemma4 template, so it runs in CI
without the checkpoint. Verified against the real gemma-4-31b-it
tokenizer on-cluster: assistant content + <turn|> masked, everything
else at 0.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Layer above the mask-function unit tests: exercises sft_rollout.generate_rollout end-to-end (minus the GPU trainer) with the real Gemma4 tokenizer, pinning the contract the trainer consumes: - sample.tokens is the FULL rendered token sequence - sample.loss_mask is only the tail (loss_mask[-response_length:]) - the tail of tokens, filtered by loss_mask, decodes to exactly the assistant turns — proving the end-alignment assumption that an off slice would silently violate (training on wrong tokens, no crash) Also guards the all-zero-mask failure mode (would make sft_loss zero/nan and silently train on nothing) and the multi-turn case where response_length spans an intervening user turn that must stay loss-0. Skipped when the gemma-4 checkpoint tokenizer is absent (laptop CI); runs on the CPU dev pod. 4/4 pass against gemma-4-31b-it. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The mask generator added in 375e848 wired gemma4 into the get_loss_mask dispatch, but --loss-mask-type has an argparse choices=[...] whitelist that still rejected "gemma4" before training could start (argument --loss-mask-type: invalid choice: 'gemma4'). Add it to the choices list so the flag is accepted. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
… in sync Regression test for the bug where get_loss_mask accepted 'gemma4' but the --loss-mask-type argparse choices=[...] list didn't, so slime rejected the flag at startup. Mutation-tested: removing 'gemma4' from choices makes this test fail. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds supervised fine-tuning (SFT) support for Gemma4, building on the Gemma4 model support in THUDM#1855 (this PR is stacked on
gemma4-pr— review that first; the diff here is SFT-only).slime's SFT path (
--loss-type sft_loss+slime.rollout.sft_rollout) builds per-token loss masks viaMultiTurnLossMaskGenerator, keyed on--loss-mask-type. Only Qwen variants were supported, so SFT on Gemma4 failed. This adds the Gemma4 chat-template mask generator and wires it through.Changes
mask_utils.py—gen_multi_turn_loss_mask_gemma4(). Gemma4's chat template uses a distinct delimiter scheme (<|turn>{role}\n ... <turn|>\n, assistant rendered asmodel, optional<|channel>thought ... <channel|>reasoning). Uses the same render-then-offset-map strategy as the Qwen3.5 generator (robust to tokenizer merges across content/marker boundaries):<turn|>terminator (so the model learns to end its turn)step_loss_mask=0arguments.py— addgemma4to the--loss-mask-typeargparsechoices. (Without this, the flag is rejected at startup even though the dispatch supports it.)Tests
tests/utils/test_loss_mask_type_gemma4.py— 10 cases with a char-level fake tokenizer modeling the Gemma4 template (CI-runnable, no checkpoint). Includes a regression guard asserting the dispatch and argparsechoicesstay in sync (mutation-tested).tests/gemma4/test_gemma4_sft_rollout.py— integration test exercisingsft_rollout.generate_rolloutwith the real Gemma4 tokenizer, pinning the contract the trainer consumes (sample.tokensfull,sample.loss_mask= tail[-response_length:], end-aligned). Skipped when the checkpoint is absent.Validation
Verified on-cluster end-to-end: both Gemma4-31B dense and 26B-A4B MoE SFT bring-ups (300 steps, OpenHermes-2.5) ran to completion — cold-start at iter 0, decreasing loss, checkpoints written, no NaN. The mask masks exactly the assistant content +
<turn|>on real OpenHermes rows.Risk
Additive: a new
gemma4branch in a dispatch the RL path never calls (RL builds masks inline). No behavior change for existing--loss-mask-typevalues.