Studio: auto-detect completion masking markers, stop silent full-sequence training

studio/backend/core/training/trainer.py

@@ -70,7 +70,7 @@
 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 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 @@ def audio_vlm_collate_fn(examples):
 
             # ========== 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,94 @@ def audio_vlm_collate_fn(examples):
 
             # DeepSeek OCR handles this internally in its collator, so skip
             # Audio VLM handles label masking in its collator, so skip
+            # Markers are auto-detected from the chat template first
+            # (unsloth_zoo.get_chat_template_parts); the manual template table
+            # is the fallback when auto-detection raises. 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")
-
-                    # Template mapping for this model
-                    model_name_lower = self.model_name.lower()
-
-                    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")
+                from unsloth.chat_templates import train_on_responses_only
 
-                        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"]
+                logger.info("Configuring train on responses only...\n")
 
-                            logger.info(f"Instruction marker: {instruction_part[:50]}...\n")
-                            logger.info(f"Response marker: {response_part[:50]}...\n")
-                        else:
-                            logger.info(
-                                f"No response mapping found for template: {template_name}\n"
-                            )
-                            train_on_responses_enabled = False
+                def _notify(level, message):
+                    if level == "warning":
+                        logger.warning(message)
                     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
+                        logger.info(f"{message}\n")
 
-            # 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")
-            ):
+                masking_applied = False
                 try:
-                    from unsloth.chat_templates import train_on_responses_only
-
-                    self.trainer = train_on_responses_only(
+                    self.trainer, masking_applied = apply_completion_masking(
                         self.trainer,
-                        instruction_part = instruction_part,
-                        response_part = response_part,
+                        self.model_name,
+                        train_on_responses_only,
                         num_proc = config_args["dataset_num_proc"],
+                        notify = _notify,
                     )
-                    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."
-                            )
-                            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")
-
                 except Exception as e:
                     logger.warning(f"Failed to apply train on responses only: {e}")
+
+                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 (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."
+                                )
+                                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")
+
+                    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")

studio/backend/core/training/worker.py

@@ -1928,27 +1928,18 @@ def _fmt_progress(status_message = "", **_kw):
         _send("eval_configured")
 
     # ── 7. Apply train_on_responses_only if requested ──
+    # Auto-detect instruction/response markers from the chat template first,
+    # then fall back to the manual template table if auto-detection raises.
     if config.get("train_on_completions", False):
         _send("status", status_message = "Configuring response-only training...")
         try:
-            from utils.datasets import (
-                MODEL_TO_TEMPLATE_MAPPER,
-                TEMPLATE_TO_RESPONSES_MAPPER,
+            from utils.datasets 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),
             )
-
-            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}")