Studio: auto-detect completion masking markers, stop silent full-sequence training - #7054
Conversation
Studio's train_on_completions previously relied only on the hardcoded MODEL_TO_TEMPLATE_MAPPER / TEMPLATE_TO_RESPONSES_MAPPER tables and silently disabled masking when a model was not in the table, so unmapped models (LFM2-8B-A1B, DeepSeek, and others) trained on full sequences without telling the user. Several mapped templates (glm, mistral, llama, starling, zephyr, qwen3-thinking) also carried markers that mask every assistant token, which made every row drop in the post-masking filter. Both training callsites (CUDA trainer.py and MLX worker.py) now share utils.datasets.completion_masking.apply_completion_masking: - Try unsloth_zoo chat template auto-detection first; it raises loudly when the template cannot be parsed and never masks the EOS token. - gpt-oss models keep their manual markers so non-final assistant <|end|> tokens stay trained, matching current behavior. - If auto-detection raises, fall back to the template table exactly as before. - If the table also misses, emit an explicit user-visible warning that completion masking could not be applied and full-sequence training will occur, instead of a quiet log line. The >30 percent dropped-rows safety net in trainer.py now guards the auto path as well. Table consumers for inference and chat templates are unchanged. Validated against one representative tokenizer for every template in TEMPLATE_TO_RESPONSES_MAPPER plus the unmapped models: no template regresses; unit tests cover the four decision paths.
|
@codex review |
for more information, see https://pre-commit.ci
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the completion-only masking configuration by introducing a unified apply_completion_masking helper shared between the CUDA trainer and MLX worker. The new policy attempts chat template auto-detection first, falling back to manual template table markers if auto-detection fails, while keeping gpt-oss pinned to its manual markers. A comprehensive test suite has been added to cover these scenarios. There are no review comments, so we have no additional feedback to provide.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d880487f61
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| except Exception as e: | ||
| notify( | ||
| "warning", | ||
| f"Auto-detection of instruction/response markers failed ({e}); " | ||
| f"falling back to the template table", |
There was a problem hiding this comment.
This catch wraps the entire train_fn(trainer, **kwargs) call, not just the chat-template marker detection. If that call raises later while mapping/filtering the dataset (for example a dataset/map worker or tokenization failure) for an unmapped model, the helper reports an auto-detection miss and returns the original trainer, so training proceeds on full sequences instead of surfacing the failed completion-masking step. Please only fall back for the specific auto-detection error and propagate other exceptions so the CUDA/MLX callers can stop rather than silently changing the training objective.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 6e9d215: markers are now detected separately via get_chat_template_parts inside the narrow try, and train_on_responses_only is then called with the detected parts so any failure while applying the masking propagates to the caller, matching the manual path. Added a regression test that an application-time RuntimeError propagates instead of falling back.
The auto branch wrapped the whole train_on_responses_only call, so a real failure while applying the masking (dataset map, tokenization) was treated as a detection miss and training silently proceeded on full sequences. Detect markers separately via get_chat_template_parts (test seam via detect_fn), then apply them with errors propagating, matching the manual path. Tokenizers with preset unsloth marker attrs skip detection and call bare so zoo reuses the stored parts.
|
@codex review |
for more information, see https://pre-commit.ci
|
Codex Review: Something went wrong. Try again later by commenting “@codex review”. ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 97171efcd4
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| except Exception as e: | ||
| logger.warning(f"Failed to apply train on responses only: {e}") |
There was a problem hiding this comment.
When auto-detection chooses markers for a model but train_on_responses_only raises while applying them (for example because truncation or a template mismatch leaves every row fully masked), this catch leaves masking_applied false and training continues on full sequences. That produces a successful run with prompts included in the loss even though the user requested Train on completions; surface the error instead of downgrading the training objective. The MLX worker uses the same catch pattern.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in ad8f6ae: removed the callsite catches in both the CUDA trainer and the MLX worker. The helper falls back internally on detection failures and returns applied=False on a double miss, so any exception reaching the callsites is a real failure applying the masking and now fails the run instead of silently training on full sequences.
The helper already falls back internally on detection failures and returns applied=False on a double miss, so an exception reaching the callsites is a real failure applying the masking. Remove the callsite catches that downgraded it to full-sequence training; the run now fails visibly instead. Also use the explicit re-export alias form in utils/datasets/__init__.py for the two new names, satisfying the import-hoist source lint.
|
@codex review |
|
Codex Review: Didn't find any major issues. You're on a roll. Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
The import-hoist source lint counts only real name loads, so package-level re-exports of the two new names cannot satisfy it. Import apply_completion_masking from utils.datasets.completion_masking directly at both callsites and leave utils/datasets/__init__.py untouched.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 6c5b8c07d7
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| # gpt-oss keeps its manual markers: with them, non-final assistant <|end|> | ||
| # tokens stay trained, whereas auto-detection would mask them. Preserve | ||
| # the current trained behavior. | ||
| if not is_gpt_oss_model_name(model_name): |
There was a problem hiding this comment.
For renamed or private gpt-oss checkpoints such as org/gpt-oss-20b-sft, is_gpt_oss_model_name() returns true here, but lookup_manual_markers() only does an exact MODEL_TO_TEMPLATE_MAPPER lookup and returns no markers. That skips the auto detector and then falls through to applied=False, so Train on completions runs on full sequences despite the tokenizer template being available. The gpt-oss exception should either provide the gpt-oss table markers for name-based matches or only bypass auto-detection when manual markers were actually found.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 8fd232b: renamed or private gpt-oss checkpoints that miss the exact-name table now default to the gpt-oss template markers instead of falling through, so non-final assistant end tokens keep their trained behavior.
| trainer, _masking_applied = apply_completion_masking( | ||
| trainer, | ||
| model_name, | ||
| train_on_responses_only, | ||
| notify = lambda level, message: _send("status", status_message = message), |
There was a problem hiding this comment.
In the MLX path this helper is still invoked whenever train_on_completions is true, even if format_type == "raw" (the function already computes raw_text_mode above, and the CUDA path explicitly skips raw text). With this change, previously-unmapped raw-text runs can now auto-detect chat markers and try to mask arbitrary text examples, dropping or corrupting trainable tokens instead of training all tokens as raw/CPT data should. Please gate this call the same way as the CUDA path, e.g. do not call it when raw_text_mode is true.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 8fd232b: the MLX callsite now carries the same gating as CUDA (skips raw text mode and alpaca format), so both paths apply masking under identical conditions.
| trainer, _masking_applied = apply_completion_masking( | ||
| trainer, | ||
| model_name, | ||
| train_on_responses_only, | ||
| notify = lambda level, message: _send("status", status_message = message), |
There was a problem hiding this comment.
This MLX call also runs when format_type == "alpaca", while the CUDA path excludes dataset_final_format == "alpaca". For Alpaca-formatted runs with a model that was previously unmapped (for example LFM2-8B), the old table lookup skipped masking, but this new auto path derives the tokenizer's chat markers and applies them to text rendered with DEFAULT_ALPACA_TEMPLATE, where those markers are absent; that can leave every row with no response labels. Please mirror the CUDA Alpaca skip before invoking the helper.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 8fd232b with a regression test (test_renamed_gpt_oss_gets_template_markers); full suite 12 passed.
Renamed or private gpt-oss checkpoints are name-detected as gpt-oss but miss the exact-name table; default them to the gpt-oss template markers instead of falling through to full-sequence training. Gate the MLX masking call on not raw_text_mode and format_type != alpaca, mirroring the CUDA path: raw/CPT text has no chat turns to mask and Alpaca-rendered text lacks the tokenizer's chat markers.
for more information, see https://pre-commit.ci
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c7f7a2c4f6
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| # then fall back to the manual template table if auto-detection raises. | ||
| # Mirror the CUDA path's skips: raw/CPT text has no chat turns to mask and | ||
| # Alpaca-rendered text does not contain the tokenizer's chat markers. | ||
| if config.get("train_on_completions", False) and not raw_text_mode and format_type != "alpaca": |
There was a problem hiding this comment.
When MLXTrainingConfig does not expose the optional append_eos field, raw_text_mode is never assigned because it is currently set only inside that feature-detection block above. With train_on_completions enabled, this new condition then raises UnboundLocalError before response masking or training can start, breaking the older-zoo compatibility that the surrounding optional-field checks are meant to preserve.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good catch, fixed in c1ef9be: raw_text_mode is now assigned above the feature-detect block, so the masking gate sees it even when the zoo's MLXTrainingConfig lacks append_eos.
| auto_instruction = auto_response = None | ||
| try: | ||
| if detect_fn is None: | ||
| from unsloth_zoo.dataset_utils import get_chat_template_parts as detect_fn |
There was a problem hiding this comment.
On torch-free Apple/MLX installs, this default detector import pulls in unsloth_zoo.dataset_utils, so the import failure is swallowed as an auto-detection failure and the MLX worker then falls back to the exact-name table. For unmapped models (the main case this change is meant to fix), train_on_completions therefore still runs on full sequences in the MLX path instead of using the MLX trainer's own response-mask/autodetect path.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
There is no torch-free MLX masking path to preserve. Apple Silicon installs include torch (no-torch mode defaults on only for Intel Macs, which are GGUF-only with no MLX training stack), and the MLX train_on_responses_only itself imports unsloth_zoo.dataset_utils (which imports torch at module top) for both its auto-detection and the masking closure, so with torch absent the bare-call path fails identically. The detector import here adds no new torch dependency.
With an older zoo lacking the append_eos config field, the masking gate referenced raw_text_mode before assignment. Hoist the assignment above the feature detection so both consumers see it.
|
@codex review |
1 similar comment
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c1ef9be4fb
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| # then fall back to the manual template table if auto-detection raises. | ||
| # Mirror the CUDA path's skips: raw/CPT text has no chat turns to mask and | ||
| # Alpaca-rendered text does not contain the tokenizer's chat markers. | ||
| if config.get("train_on_completions", False) and not raw_text_mode and format_type != "alpaca": |
There was a problem hiding this comment.
In the MLX path with format_type="auto", the preceding format_and_template_dataset call can return final_format == "alpaca" for Alpaca-shaped data, but this guard only checks the requested format_type. When train_on_completions is enabled for an auto-detectable model that was not in the old manual table, the new helper now applies tokenizer chat-template markers to the rendered Alpaca prompt text; those markers are not present, so response-only masking can leave every row fully masked/filtered instead of training normally. Track and check the formatter's returned final_format here, as the CUDA path does, rather than only the original request value.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 0b188bd: the worker now records the final_format returned by format_and_template_dataset and the masking gate checks it alongside the requested value (skipping alpaca and raw text), matching the CUDA path's dataset_final_format gating.
format_type auto can resolve to alpaca or raw text; the masking skip checked only the requested value, so auto-detected Alpaca data got chat-template markers applied to rendered prompt text. Track the final_format returned by format_and_template_dataset and gate on it, matching the CUDA path.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0b188bde87
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| processor = getattr(trainer, "processing_class", None) or getattr( | ||
| trainer, "tokenizer", None | ||
| ) | ||
| inner = getattr(processor, "tokenizer", processor) |
There was a problem hiding this comment.
In the MLX training path, trainer.tokenizer can be the mlx-lm TokenizerWrapper, whose real Hugging Face tokenizer lives on _tokenizer rather than .tokenizer (the local get_chat_template wrapper handling documents this). This helper only unwraps .tokenizer and then passes the wrapper into get_chat_template_parts, so auto-detection fails for MLX runs on unmapped models such as LiquidAI/LFM2-8B-A1B; because there is no table fallback for those names, train_on_completions silently falls back to full-sequence training on the MLX path. Please also handle _tokenizer here or delegate detection to the MLX masking helper.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in ffae8d1: the helper unwraps a TokenizerWrapper to its _tokenizer before the preset-marker check and detection, like the zoo MLX resolver. Reproduced both halves against real wrapper semantics: preset unsloth* markers on the inner tokenizer are hidden by the wrapper's underscore handling (real miss), and detection on the wrapper only worked through the loader's call patch (now independent of it). Two regression tests added.
| if detect_fn is None: | ||
| from unsloth_zoo.dataset_utils import get_chat_template_parts as detect_fn | ||
| auto_instruction, auto_response = detect_fn(processor) |
There was a problem hiding this comment.
When this shared helper is called from the MLX worker, importing unsloth_zoo.dataset_utils pulls in the CUDA/PyTorch dataset utilities; this repo already keeps that import optional in unsloth/chat_templates.py because MLX hosts can be torch-free. On such MLX installs the ImportError is swallowed as an auto-detection failure, so unmapped models (the main case this patch is trying to fix) have no manual fallback and still train on full sequences. Use the MLX wrapper's native auto-detection path or a torch-free detector for the MLX callsite.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
The MLX wrapper's native auto-detection path is unsloth_zoo.mlx.trainer.train_on_responses_only called bare, and its first statement imports unsloth_zoo.dataset_utils (torch at module top) for both detection and the masking closure, so no torch-free detector exists to delegate to. Apple Silicon installs include torch (no-torch mode defaults on only for Intel Macs, which get no MLX training stack); with torch absent the masking call fails identically before and after this change.
The wrapper delegates plain reads to the wrapped HF tokenizer but hides underscore attrs, so preset unsloth markers were invisible and detection relied on the loader's call patch. Unwrap to the real tokenizer first, as the zoo MLX resolver does.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ffae8d14b3
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| auto_instruction = auto_response = None | ||
| try: | ||
| if detect_fn is None: | ||
| from unsloth_zoo.dataset_utils import get_chat_template_parts as detect_fn |
There was a problem hiding this comment.
In the MLX worker on a torch-free Apple Silicon install, this import pulls in unsloth_zoo.dataset_utils, which the local unsloth/chat_templates.py already treats as torch-dependent to keep MLX hosts importable. When torch is absent, the broad except turns auto-detection into a fallback/table miss, so unmapped models on MLX still train on full sequences despite train_on_completions=True; use an MLX-safe detector path or avoid importing the torch-backed dataset_utils in that worker.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Same premise as the two earlier reports: there is no torch-free MLX masking configuration. unsloth_zoo.mlx.trainer.train_on_responses_only (the MLX train_fn this helper wraps, and the only MLX-native detection path) imports unsloth_zoo.dataset_utils for both detection and the masking closure, so with torch absent masking fails identically with or without this import. Apple Silicon installs include torch; no-torch mode is the Intel Mac GGUF-only path with no MLX training. An MLX-safe detector would require a zoo change, not a Studio one.
|
@codex review |
|
Codex Review: Didn't find any major issues. 🚀 Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
The quantized and BF16 gpt-oss checkpoints ship a chat template without the channel final header, so the pinned manual markers match nothing there and masking trained zero tokens. Auto-detection derives markers from whichever template the checkpoint ships and keeps the final terminator trained; the manual gpt-oss markers remain the detection failure fallback, including for renamed checkpoints.
for more information, see https://pre-commit.ci
|
@codex review |
|
Codex Review: Didn't find any major issues. Breezy! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
|
@codex review |
|
Codex Review: Didn't find any major issues. You're on a roll. Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
1 participant
Summary
Studio's "Train on completions" is on by default for text SFT, but both training paths applied it only through a hardcoded exact-name template table, and on a lookup miss they silently disabled masking. That means models with no table entry, including every DeepSeek model and most LFM sizes (the table has
unsloth/LFM2-1.2Bbut notLiquidAI/LFM2-8B-A1B), trained on the full sequence, prompts included, while the UI toggle said otherwise.Worse, several manual marker sets never match what the tokenizer's chat template actually renders, so those templates masked every assistant token: mistral's response marker
" [/INST]"(leading space) does not occur in the real rendering, and llama-2, starling, glm, and the qwen3-thinking<think>markers have the same class of bug. Those runs either trained on nothing or died on the "over 30 percent of rows dropped" safety net.unsloth_zoo already ships automatic marker detection (
get_chat_template_parts, used bytrain_on_responses_onlywhen no markers are passed) that derives the instruction and response spans from the tokenizer's own chat template and fails loudly when it cannot. Studio just never called it.Changes
studio/backend/utils/datasets/completion_masking.pywithapply_completion_masking(trainer, model_name, train_fn, ...), one policy shared by both training paths:train_on_responses_only).TEMPLATE_TO_RESPONSES_MAPPERmarkers, with a warning.<|end|>tokens keep their current trained behavior (auto would mask them).trainer.py) and MLX path (worker.py) both use the helper; the MLX path surfaces warnings as status messages in the UI. The over-30-percent-dropped safety net is unchanged and now also guards the auto path.Validation
Catalog-wide, token-for-token label comparison (auto vs manual) on a representative tokenizer for every template in the table, plus the previously unmapped models. Two-turn chat fixture; verified user and system tokens masked, all assistant turns trained, and the terminating EOS label is never -100.
Tests: new
tests/test_completion_masking.py(9 tests: auto-first, gpt-oss pinning, fallback on auto failure, double-miss warning without crash, num_proc forwarding, manual-failure propagation). Full training-related suite: 109 passed, 1 skipped.