diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 958f8f4197..96d1b90b16 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -70,7 +70,7 @@ from core.inference.llama_cpp import _hf_offline_if_dns_dead from utils.models import is_vision_model, detect_audio_type from utils.models.model_config import _env_offline from utils.datasets import format_and_template_dataset -from utils.datasets import MODEL_TO_TEMPLATE_MAPPER, TEMPLATE_TO_RESPONSES_MAPPER +from utils.datasets.completion_masking import apply_completion_masking from utils.datasets.iterable import is_streaming_dataset as detect_streaming_dataset from utils.datasets.raw_text import prepare_raw_text_dataset, resolve_column_names from utils.paths import ( @@ -3455,8 +3455,6 @@ class UnslothTrainer: # ========== TRAIN ON RESPONSES ONLY ========== # Raw-text datasets always train on all tokens. - instruction_part = None - response_part = None is_cpt = training_args.get("is_cpt", False) train_on_responses_enabled = ( False @@ -3473,113 +3471,93 @@ class UnslothTrainer: # DeepSeek OCR handles this internally in its collator, so skip # Audio VLM handles label masking in its collator, so skip + # Markers auto-detected from the chat template first, manual table + # as fallback; gpt-oss stays on its manual markers. See + # apply_completion_masking. if ( train_on_responses_enabled and not self.is_audio_vlm and not self.is_audio and not (is_deepseek_ocr or dataset_final_format == "alpaca") ): - try: - logger.info("Configuring train on responses only...\n") + from unsloth.chat_templates import train_on_responses_only - # Template mapping for this model - model_name_lower = self.model_name.lower() + logger.info("Configuring train on responses only...\n") - if model_name_lower in MODEL_TO_TEMPLATE_MAPPER: - template_name = MODEL_TO_TEMPLATE_MAPPER[model_name_lower] - logger.info(f"Detected template: {template_name}\n") + def _notify(level, message): + if level == "warning": + logger.warning(message) + else: + logger.info(f"{message}\n") - if template_name in TEMPLATE_TO_RESPONSES_MAPPER: - instruction_part = TEMPLATE_TO_RESPONSES_MAPPER[template_name][ - "instruction" - ] - response_part = TEMPLATE_TO_RESPONSES_MAPPER[template_name]["response"] + # No try/except: the helper handles detection failures and + # double misses itself, so an exception here is a real masking + # failure that must fail the run, not silently train on full + # sequences. + self.trainer, masking_applied = apply_completion_masking( + self.trainer, + self.model_name, + train_on_responses_only, + num_proc = config_args["dataset_num_proc"], + notify = _notify, + ) - logger.info(f"Instruction marker: {instruction_part[:50]}...\n") - logger.info(f"Response marker: {response_part[:50]}...\n") + if not masking_applied: + train_on_responses_enabled = False + + if masking_applied: + try: + # ── Safety net: check if all samples were filtered out ── + # train_on_responses_only masks non-response tokens with -100; a + # row becomes all -100 (Unsloth drops it) when the response + # template is not found in the formatted text. Usually a + # dataset/template mismatch (already-formatted data, or 'Train on + # completions' on data that doesn't match the model's chat + # template); only sometimes max_seq_length truncating the response + # away. Skip this len()-based check for streaming. + if detect_streaming_dataset(self.trainer.train_dataset): + logger.info("Skipping post-filter length check for streaming dataset\n") else: - logger.info( - f"No response mapping found for template: {template_name}\n" + filtered_len = len(self.trainer.train_dataset) + original_dataset_obj = ( + dataset["dataset"] if isinstance(dataset, dict) else dataset ) - train_on_responses_enabled = False - else: - logger.info(f"No template mapping found for model: {self.model_name}\n") - train_on_responses_enabled = False - - except Exception as e: - logger.warning(f"Could not configure train on responses: {e}") - train_on_responses_enabled = False - - # Apply train on responses only if we have valid parts - if ( - train_on_responses_enabled - and instruction_part - and response_part - and not self.is_audio_vlm - and not self.is_audio - and not (is_deepseek_ocr or dataset_final_format == "alpaca") - ): - try: - from unsloth.chat_templates import train_on_responses_only - - self.trainer = train_on_responses_only( - self.trainer, - instruction_part = instruction_part, - response_part = response_part, - num_proc = config_args["dataset_num_proc"], - ) - logger.info("Train on responses only configured successfully\n") - - # ── Safety net: check if all samples were filtered out ── - # train_on_responses_only masks non-response tokens with -100; - # a row becomes all -100 (and Unsloth drops it) when the response - # template is not found in the formatted text. That is usually a - # dataset/template mismatch (already-formatted data, or 'Train on - # completions' applied to data that doesn't match the model's chat - # template), and only sometimes max_seq_length truncating the - # response away. Skip this len()-based check for streaming. - if detect_streaming_dataset(self.trainer.train_dataset): - logger.info("Skipping post-filter length check for streaming dataset\n") - else: - filtered_len = len(self.trainer.train_dataset) - original_dataset_obj = ( - dataset["dataset"] if isinstance(dataset, dict) else dataset - ) - original_len = len(original_dataset_obj) - dropped = original_len - filtered_len - drop_pct = round(100 * dropped / original_len, 1) if original_len > 0 else 0 - - if filtered_len == 0 or drop_pct > 30: - max_seq = training_args.get("max_seq_length", 2048) - error_msg = ( - f"{dropped}/{original_len} samples ({drop_pct}%) were " - f"dropped after applying 'Train on completions': after " - f"masking, those rows had no trainable response tokens " - f"left. The usual cause is that this model's response " - f"template was not found in the formatted samples, so " - f"every token was masked out. That typically means the " - f"dataset is already formatted, or its structure does " - f"not match the model's chat template, so 'Train on " - f"completions' should be turned off for this dataset. " - f"Less commonly, a max_seq_length ({max_seq}) shorter " - f"than the prompt can truncate the response away; only " - f"raise it if your samples are actually longer than that." + original_len = len(original_dataset_obj) + dropped = original_len - filtered_len + drop_pct = ( + round(100 * dropped / original_len, 1) if original_len > 0 else 0 ) - logger.error(error_msg) - self._update_progress(error = error_msg, is_training = False) - return - if dropped > 0: - logger.info( - f"⚠️ {dropped}/{original_len} samples " - f"({drop_pct}%) were dropped (all labels " - f"masked). {filtered_len} samples remain.\n" - ) - logger.info(f"Post-filter dataset size: {filtered_len} samples\n") + if filtered_len == 0 or drop_pct > 30: + max_seq = training_args.get("max_seq_length", 2048) + error_msg = ( + f"{dropped}/{original_len} samples ({drop_pct}%) were " + f"dropped after applying 'Train on completions': after " + f"masking, those rows had no trainable response tokens " + f"left. The usual cause is that this model's response " + f"template was not found in the formatted samples, so " + f"every token was masked out. That typically means the " + f"dataset is already formatted, or its structure does " + f"not match the model's chat template, so 'Train on " + f"completions' should be turned off for this dataset. " + f"Less commonly, a max_seq_length ({max_seq}) shorter " + f"than the prompt can truncate the response away; only " + f"raise it if your samples are actually longer than that." + ) + logger.error(error_msg) + self._update_progress(error = error_msg, is_training = False) + return - except Exception as e: - logger.warning(f"Failed to apply train on responses only: {e}") - train_on_responses_enabled = False + if dropped > 0: + logger.info( + f"⚠️ {dropped}/{original_len} samples " + f"({drop_pct}%) were dropped (all labels " + f"masked). {filtered_len} samples remain.\n" + ) + logger.info(f"Post-filter dataset size: {filtered_len} samples\n") + + except Exception as e: + logger.warning(f"Post-masking dataset size check failed: {e}") else: if train_on_responses_enabled and is_deepseek_ocr: logger.info("Train on responses handled by DeepSeek OCR collator\n") diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 0ff4d517ed..5fb8fc2bb6 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -1731,6 +1731,7 @@ def _run_mlx_training(event_queue, stop_queue, config): # sharegpt+images) and text (alpaca/sharegpt/chatml → "text" column). format_type = config.get("format_type", "") custom_format_mapping = config.get("custom_format_mapping") + dataset_final_format = "" try: from utils.datasets import format_and_template_dataset def _fmt_progress(status_message = "", **_kw): @@ -1796,6 +1797,7 @@ def _run_mlx_training(event_queue, stop_queue, config): ) if info.get("success", True): dataset = info.get("dataset", dataset) + dataset_final_format = str(info.get("final_format", "") or "").lower() if eval_dataset is not None: ev = format_and_template_dataset( eval_dataset, @@ -1894,6 +1896,9 @@ def _run_mlx_training(event_queue, stop_queue, config): eval_steps = eval_steps_val, ) + # Also gates the masking skip below, so defined outside the feature-detect block. + raw_text_mode = training_type == "Continued Pretraining" or format_type == "raw" + # Feature-detect optional fields so this PR works without the paired zoo bump. _supported_fields = getattr(MLXTrainingConfig, "__dataclass_fields__", {}) if "cast_norm_output_to_input_dtype" in _supported_fields: @@ -1907,7 +1912,6 @@ def _run_mlx_training(event_queue, stop_queue, config): if "max_grad_leaf_norm" in _supported_fields: mlx_config_kwargs["max_grad_leaf_norm"] = max_grad_leaf_norm if "append_eos" in _supported_fields: - raw_text_mode = training_type == "Continued Pretraining" or format_type == "raw" # Studio SFT formatting owns rendered examples; raw/CPT text still # needs MLX to append EOS like the CUDA raw-text path. mlx_config_kwargs["append_eos"] = bool(raw_text_mode) @@ -1928,29 +1932,27 @@ def _run_mlx_training(event_queue, stop_queue, config): _send("eval_configured") # ── 7. Apply train_on_responses_only if requested ── - if config.get("train_on_completions", False): + # Auto-detect markers from the chat template first, manual table as + # fallback. Mirror the CUDA skips: raw/CPT text has no chat turns and + # Alpaca-rendered text lacks the chat markers. Also check the resolved + # format, since format_type="auto" can land on alpaca or raw text. + if ( + config.get("train_on_completions", False) + and not raw_text_mode + and format_type != "alpaca" + and dataset_final_format not in ("alpaca", "raw_text") + ): _send("status", status_message = "Configuring response-only training...") - try: - from utils.datasets import ( - MODEL_TO_TEMPLATE_MAPPER, - TEMPLATE_TO_RESPONSES_MAPPER, - ) - - template_name = MODEL_TO_TEMPLATE_MAPPER.get(model_name.lower()) - markers = TEMPLATE_TO_RESPONSES_MAPPER.get(template_name) if template_name else None - if markers: - trainer = train_on_responses_only( - trainer, - instruction_part = markers["instruction"], - response_part = markers["response"], - ) - else: - _send( - "status", - status_message = f"train_on_completions skipped (no template for {model_name})", - ) - except Exception as e: - _send("status", status_message = f"train_on_completions failed: {e}") + # No catch: the helper handles detection failures and double misses, so + # an exception here is a real masking failure that must fail the run, + # not silently train on full sequences. + from utils.datasets.completion_masking import apply_completion_masking + trainer, _masking_applied = apply_completion_masking( + trainer, + model_name, + train_on_responses_only, + notify = lambda level, message: _send("status", status_message = message), + ) # ── 8. Setup wandb / tensorboard ── wandb_run = None diff --git a/studio/backend/tests/test_completion_masking.py b/studio/backend/tests/test_completion_masking.py new file mode 100644 index 0000000000..be0d8a69bd --- /dev/null +++ b/studio/backend/tests/test_completion_masking.py @@ -0,0 +1,314 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Completion-only masking policy: auto-detect first, manual table fallback. + +Covers utils.datasets.completion_masking.apply_completion_masking, shared by +the CUDA trainer (core/training/trainer.py) and the MLX worker +(core/training/worker.py): + - unmapped models use chat template auto-detection (previously masking was + silently disabled), + - gpt-oss goes auto-first too (its quantized checkpoints ship a template + the manual markers cannot match), + - an auto-detection failure falls back to the template table markers, + - a table miss after an auto failure warns and leaves the trainer unchanged. +""" + +from __future__ import annotations + +import pytest + +from utils.datasets.completion_masking import apply_completion_masking, lookup_manual_markers +from utils.datasets.model_mappings import TEMPLATE_TO_RESPONSES_MAPPER + + +class _Trainer: + """Sentinel trainer; train_fn wraps it in a new object when applied.""" + + +class _Recorder: + """Fake train_on_responses_only that records calls.""" + + def __init__(self): + self.calls = [] + + def __call__(self, trainer, **kwargs): + self.calls.append(kwargs) + wrapped = _Trainer() + wrapped.wrapped_from = trainer + return wrapped + + +def _detect_ok(processor): + return "", "" + + +def _detect_fail(processor): + raise ValueError( + "Unsloth: Could not reliably auto-detect response_part - " + "pass instruction_part and response_part." + ) + + +_AUTO = {"instruction_part": "", "response_part": ""} + + +class _Notes: + def __init__(self): + self.messages = [] + + def __call__(self, level, message): + self.messages.append((level, message)) + + def warnings(self): + return [m for level, m in self.messages if level == "warning"] + + +def test_unmapped_model_uses_auto_detection(): + # Unmapped model: the auto path applies masking (was silently disabled). + trainer = _Trainer() + train_fn = _Recorder() + notes = _Notes() + + result, applied = apply_completion_masking( + trainer, "LiquidAI/LFM2-8B-A1B", train_fn, notify = notes, detect_fn = _detect_ok + ) + + assert applied is True + assert result.wrapped_from is trainer + assert train_fn.calls == [dict(_AUTO)] # applied with the detected markers + assert notes.warnings() == [] + + +def test_mapped_model_prefers_auto_detection(): + trainer = _Trainer() + train_fn = _Recorder() + + _, applied = apply_completion_masking( + trainer, "unsloth/Qwen3-0.6B", train_fn, detect_fn = _detect_ok + ) + + assert applied is True + assert train_fn.calls == [dict(_AUTO)] + + +def test_gpt_oss_uses_auto_detection_first(): + # The quantized gpt-oss checkpoints ship a template without the + # <|channel|>final header, where the manual markers match nothing; auto + # derives markers from the template the checkpoint actually ships. + trainer = _Trainer() + train_fn = _Recorder() + + _, applied = apply_completion_masking( + trainer, "unsloth/gpt-oss-20b", train_fn, detect_fn = _detect_ok + ) + + assert applied is True + assert train_fn.calls == [dict(_AUTO)] + + +def test_gpt_oss_detection_failure_falls_back_to_manual_markers(): + trainer = _Trainer() + train_fn = _Recorder() + + _, applied = apply_completion_masking( + trainer, "unsloth/gpt-oss-20b", train_fn, detect_fn = _detect_fail + ) + + assert applied is True + expected = TEMPLATE_TO_RESPONSES_MAPPER["gpt-oss"] + assert train_fn.calls == [ + { + "instruction_part": expected["instruction"], + "response_part": expected["response"], + } + ] + + +def test_auto_failure_falls_back_to_template_table(): + trainer = _Trainer() + train_fn = _Recorder() + notes = _Notes() + + result, applied = apply_completion_masking( + trainer, "unsloth/Qwen3-0.6B", train_fn, notify = notes, detect_fn = _detect_fail + ) + + assert applied is True + assert result.wrapped_from is trainer + expected = TEMPLATE_TO_RESPONSES_MAPPER["qwen3"] + assert train_fn.calls == [ + { + "instruction_part": expected["instruction"], + "response_part": expected["response"], + }, + ] + assert any("falling back to the template table" in m for m in notes.warnings()) + + +def test_application_failure_propagates_not_fallback(): + # Detection succeeds; a failure while APPLYING the masking must propagate, + # never silently fall back to full-sequence training. + def train_fn(trainer, **kwargs): + raise RuntimeError("dataset map worker crashed") + + with pytest.raises(RuntimeError, match = "dataset map worker crashed"): + apply_completion_masking(_Trainer(), "LiquidAI/LFM2-8B-A1B", train_fn, detect_fn = _detect_ok) + + +def test_preset_tokenizer_markers_used_directly(): + # Preset unsloth marker attrs skip detection; zoo reuses them on a bare call. + class _Tok: + _unsloth_input_part = "" + _unsloth_output_part = "" + + trainer = _Trainer() + trainer.processing_class = _Tok() + train_fn = _Recorder() + + _, applied = apply_completion_masking( + trainer, "LiquidAI/LFM2-8B-A1B", train_fn, detect_fn = _detect_fail + ) + assert applied is True + assert train_fn.calls == [{}] # bare call, stored parts + + +def test_table_miss_warns_and_disables_without_crashing(): + trainer = _Trainer() + train_fn = _Recorder() + notes = _Notes() + + result, applied = apply_completion_masking( + trainer, "some-org/not-in-any-mapper", train_fn, notify = notes, detect_fn = _detect_fail + ) + + assert applied is False + assert result is trainer # unchanged: full sequence training + assert train_fn.calls == [] # detection failed; nothing applied + assert any("could not be applied" in m for m in notes.warnings()) + assert any("full sequences" in m for m in notes.warnings()) + + +def test_num_proc_forwarded_only_when_given(): + # CUDA path passes num_proc; the MLX path omits it. + train_fn = _Recorder() + apply_completion_masking( + _Trainer(), "unsloth/Qwen3-0.6B", train_fn, num_proc = 4, detect_fn = _detect_ok + ) + assert train_fn.calls == [dict(_AUTO, num_proc = 4)] + + train_fn = _Recorder() + apply_completion_masking( + _Trainer(), "unsloth/Qwen3-0.6B", train_fn, num_proc = 4, detect_fn = _detect_fail + ) + assert train_fn.calls[0]["num_proc"] == 4 + + train_fn = _Recorder() + apply_completion_masking(_Trainer(), "unsloth/Qwen3-0.6B", train_fn, detect_fn = _detect_ok) + assert train_fn.calls == [dict(_AUTO)] + + +def test_manual_fallback_failure_propagates_to_caller(): + # Errors while applying the manual fallback must propagate to the caller. + def train_fn(trainer, **kwargs): + raise RuntimeError("boom") + + with pytest.raises(RuntimeError, match = "boom"): + apply_completion_masking(_Trainer(), "unsloth/gpt-oss-20b", train_fn) + + +def test_notify_is_optional(): + train_fn = _Recorder() + _, applied = apply_completion_masking( + _Trainer(), "some-org/not-in-any-mapper", train_fn, detect_fn = _detect_fail + ) + assert applied is False + + +def test_lookup_manual_markers(): + template, instruction, response = lookup_manual_markers("unsloth/Qwen3-0.6B") + assert template == "qwen3" + assert instruction == TEMPLATE_TO_RESPONSES_MAPPER["qwen3"]["instruction"] + assert response == TEMPLATE_TO_RESPONSES_MAPPER["qwen3"]["response"] + + template, instruction, response = lookup_manual_markers("some-org/unknown") + assert (template, instruction, response) == (None, None, None) + + template, instruction, response = lookup_manual_markers(None) + assert (template, instruction, response) == (None, None, None) + + +def test_renamed_gpt_oss_gets_template_markers(): + # Name-detected as gpt-oss but not in the exact-name table: must use the + # gpt-oss markers, not fall through to full-sequence training. + trainer = _Trainer() + train_fn = _Recorder() + + _, applied = apply_completion_masking( + trainer, "some-org/gpt-oss-20b-sft", train_fn, detect_fn = _detect_fail + ) + assert applied is True + expected = TEMPLATE_TO_RESPONSES_MAPPER["gpt-oss"] + assert train_fn.calls == [ + { + "instruction_part": expected["instruction"], + "response_part": expected["response"], + } + ] + + +class _FakeTokenizerWrapper: + """mlx-lm TokenizerWrapper semantics: plain reads delegate to the wrapped + tokenizer, underscore attrs do not (so preset markers are hidden).""" + + def __init__(self, tokenizer): + object.__setattr__(self, "_tokenizer", tokenizer) + + def __getattr__(self, attr): + if attr.startswith("_"): + return object.__getattribute__(self, attr) + return getattr(object.__getattribute__(self, "_tokenizer"), attr) + + +_FakeTokenizerWrapper.__name__ = "TokenizerWrapper" + + +def test_mlx_tokenizer_wrapper_unwrapped_for_preset_markers(): + # Markers live on the inner HF tokenizer that the wrapper hides; the helper + # must unwrap so the preset bare-call path still fires on MLX. + class _Tok: + _unsloth_input_part = "" + _unsloth_output_part = "" + + trainer = _Trainer() + trainer.tokenizer = _FakeTokenizerWrapper(_Tok()) + train_fn = _Recorder() + + _, applied = apply_completion_masking( + trainer, "LiquidAI/LFM2-8B-A1B", train_fn, detect_fn = _detect_fail + ) + assert applied is True + assert train_fn.calls == [{}] # bare call, stored parts + + +def test_mlx_tokenizer_wrapper_unwrapped_for_detection(): + # Detection must see the real tokenizer, not the wrapper, so it does not + # depend on the loader's __call__ patch. + class _Tok: + pass + + inner = _Tok() + trainer = _Trainer() + trainer.tokenizer = _FakeTokenizerWrapper(inner) + train_fn = _Recorder() + seen = [] + + def detect(processor): + seen.append(processor) + return "", "" + + _, applied = apply_completion_masking( + trainer, "LiquidAI/LFM2-8B-A1B", train_fn, detect_fn = detect + ) + assert applied is True + assert seen == [inner] diff --git a/studio/backend/utils/datasets/completion_masking.py b/studio/backend/utils/datasets/completion_masking.py new file mode 100644 index 0000000000..c7c4a474e3 --- /dev/null +++ b/studio/backend/utils/datasets/completion_masking.py @@ -0,0 +1,144 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Completion-only masking policy shared by the CUDA and MLX training paths. + +Decides how train_on_responses_only is applied for a model: chat template +auto-detection first, manual TEMPLATE_TO_RESPONSES_MAPPER markers as the +fallback. gpt-oss included: its quantized checkpoints ship a different +chat template, so only detection from the actual template is reliable. +""" + +from .model_mappings import ( + MODEL_TO_TEMPLATE_MAPPER, + TEMPLATE_TO_RESPONSES_MAPPER, + is_gpt_oss_model_name, +) + + +def lookup_manual_markers(model_name): + """Return (template_name, instruction_part, response_part) from the + manual template table, with None parts when the model or template is + not mapped.""" + template = MODEL_TO_TEMPLATE_MAPPER.get((model_name or "").lower()) + markers = TEMPLATE_TO_RESPONSES_MAPPER.get(template) if template else None + if markers: + return template, markers["instruction"], markers["response"] + return template, None, None + + +def apply_completion_masking( + trainer, + model_name, + train_fn, + num_proc = None, + notify = None, + detect_fn = None, +): + """Apply completion-only masking with auto-detection first and the manual + template table as fallback. + + Args: + trainer: The platform trainer (SFTTrainer or MLXTrainer). + model_name: Model repo id used for table lookup and the gpt-oss + renamed-checkpoint fallback. + train_fn: The platform train_on_responses_only callable. + num_proc: Forwarded to train_fn when not None (CUDA path only). + notify: Optional callback notify(level, message) with level "info" or + "warning" for user-visible progress and warnings. + detect_fn: Marker detector (tokenizer/processor) -> (instruction_part, + response_part). Defaults to unsloth_zoo's get_chat_template_parts, + which raises loudly when the template cannot be parsed. Test seam. + + Returns: + (trainer, applied): the possibly wrapped trainer and whether masking + was applied. When applied is False the trainer is unchanged and + training runs on full sequences. + + Only marker DETECTION failures trigger the table fallback. Exceptions + raised while applying the masking (dataset map, tokenization) propagate + to the caller in both the auto and manual paths, so a real failure stops + the run instead of silently changing the training objective. + """ + if notify is None: + notify = lambda level, message: None + kwargs = {} + if num_proc is not None: + kwargs["num_proc"] = num_proc + + template, instruction_part, response_part = lookup_manual_markers(model_name) + + # gpt-oss goes auto-first: quantized/BF16 checkpoints ship a channel-less + # template, so the manual markers match nothing (zero tokens trained). Auto + # derives markers from whichever template ships, and per the harmony format + # only the final terminator carries stop supervision. Renamed checkpoints + # miss the exact-name table, so give the fallback the gpt-oss markers. + if is_gpt_oss_model_name(model_name) and not (instruction_part and response_part): + markers = TEMPLATE_TO_RESPONSES_MAPPER.get("gpt-oss") + if markers: + template = "gpt-oss" + instruction_part = markers["instruction"] + response_part = markers["response"] + processor = getattr(trainer, "processing_class", None) or getattr(trainer, "tokenizer", None) + # mlx-lm TokenizerWrapper hides underscore attrs, so preset _unsloth_* + # markers are invisible through it. Unwrap to the real tokenizer (as + # zoo's MLX resolver does) before the preset check and detection. + if type(processor).__name__ == "TokenizerWrapper": + wrapped = getattr(processor, "_tokenizer", None) + if wrapped is not None: + processor = wrapped + inner = getattr(processor, "tokenizer", processor) + if hasattr(inner, "_unsloth_input_part") and hasattr(inner, "_unsloth_output_part"): + # Markers preset on the tokenizer; zoo reuses them on a bare call. + trainer = train_fn(trainer, **kwargs) + notify( + "info", + "Train on responses only configured via tokenizer preset markers", + ) + return trainer, True + auto_instruction = auto_response = None + try: + if detect_fn is None: + # Torch-backed import is fine: the MLX train_fn itself requires + # unsloth_zoo.dataset_utils, so a torch-free host cannot mask either way. + from unsloth_zoo.dataset_utils import get_chat_template_parts as detect_fn + auto_instruction, auto_response = detect_fn(processor) + except Exception as e: + notify( + "warning", + f"Auto-detection of instruction/response markers failed ({e}); " + f"falling back to the template table", + ) + if auto_instruction and auto_response: + trainer = train_fn( + trainer, + instruction_part = auto_instruction, + response_part = auto_response, + **kwargs, + ) + notify( + "info", + "Train on responses only configured via chat template auto-detection", + ) + return trainer, True + + if instruction_part and response_part: + trainer = train_fn( + trainer, + instruction_part = instruction_part, + response_part = response_part, + **kwargs, + ) + notify( + "info", + f"Train on responses only configured with template table markers ({template})", + ) + return trainer, True + + notify( + "warning", + f"'Train on completions' could not be applied for {model_name}: no " + f"auto-detected or mapped instruction/response markers. Training " + f"will run on full sequences (prompts included).", + ) + return trainer, False