feat(train): hook in assistant_only_loss for TRL SFT path

.github/workflows/ci.yml

@@ -52,6 +52,21 @@ jobs:
       - run: uv run python -m tests.benchmarks.base_vs_instruct.run --help > /dev/null
       - run: uv run python -m tests.experiments.experiment_e2e_pipeline --help > /dev/null
 
+  train-tests:
+    name: train-extra tests
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v6
+      - uses: astral-sh/setup-uv@v8.1.0
+        with:
+          enable-cache: true
+          cache-dependency-glob: uv.lock
+      - run: uv python install 3.12
+      # transformers/trl/torch are train extras, so the default job skips
+      # every assertion in this file.
+      - run: uv sync --frozen --group dev --extra train --python 3.12
+      - run: uv run pytest tests/unit/test_sft_assistant_only_loss.py -v
+
   json-lint:
     name: json lint
     runs-on: ubuntu-latest

lqh/skills/train/SKILL.md

@@ -368,6 +368,7 @@ Defaults live in `lqh/train/defaults.py`; a sweep overrides `learning_rate` /
   for ~20 updates and look like bad data. Two caveats: the 16 floor means a
   few-hundred-row dataset still lands below the target, and the ~100 target is a
   judgement call rather than a measured optimum (see `defaults.py`).
+- **`assistant_only_loss`** (default: false) — text SFT only; labels assistant turns alone (needs `{% generation %}` in the chat template, as in LFM2.5, and is rejected at launch without it). Rows whose assistant turn falls past `max_seq_length` would train on nothing and are dropped with a count.
 - **`num_iterations`** (default: 5) — DPO only.
 - **`dpo_beta`** (default: 0.1) — DPO KL anchor strength.
 

lqh/tools/definitions.py

@@ -1746,6 +1746,14 @@ def _build_all_tools(*, auto_mode: bool = False) -> list[dict]:
                             "\'training isn\'t working\' branch."
                         ),
                     },
+                    "assistant_only_loss": {
+                        "type": "boolean",
+                        "description": (
+                            "Text SFT only: compute the loss on assistant turns "
+                            "alone. Set when user turns dominate the row."
+                        ),
+                        "default": False,
+                    },
                     "num_iterations": {
                         "type": "integer",
                         "description": (

lqh/tools/handlers.py

@@ -5088,6 +5088,41 @@ def _budget_base_model(project_dir: Path, name: str) -> str:
     return current
 
 
+def _assistant_mask_unsupported(project_dir: Path, base_model: str) -> str | None:
+    """Why *base_model* cannot label assistant turns only, or None if it can.
+
+    trl raises during dataset tokenization — after a cloud job has been
+    provisioned — so a chat template that marks no assistant tokens is
+    cheaper to reject at launch. Local checkpoint paths resolve against the
+    project dir. Fails open when the tokenizer will not load here (e.g. no
+    network): the trainer is the backstop then.
+    """
+    try:
+        from transformers import AutoTokenizer
+
+        local = project_dir / base_model
+        tokenizer = AutoTokenizer.from_pretrained(
+            local if local.is_dir() else base_model
+        )
+    except Exception:  # noqa: BLE001 — fail open; the trainer is the backstop
+        return None
+    try:
+        encoded = tokenizer.apply_chat_template(
+            [
+                {"role": "user", "content": "ping"},
+                {"role": "assistant", "content": "pong"},
+            ],
+            tokenize=True,
+            return_dict=True,
+            return_assistant_tokens_mask=True,
+        )
+    except Exception as exc:  # noqa: BLE001 — any template failure is the reason
+        return str(exc).strip() or exc.__class__.__name__
+    if not any(encoded.get("assistant_masks") or []):
+        return "its chat template has no {% generation %} block"
+    return None
+
+
 async def handle_start_training(
     project_dir: Path,
     *,
@@ -5101,6 +5136,7 @@ async def handle_start_training(
     lora: bool = True,
     num_epochs: int | None = None,
     learning_rate: float | None = None,
+    assistant_only_loss: bool = False,
     num_iterations: int = 5,
     dpo_beta: float = 0.1,
     golden_source: str = "dataset",
@@ -5386,6 +5422,20 @@ async def handle_start_training(
     from lqh.models import is_vlm_model_name
 
     is_vision = is_vlm_model_name(base_model)
+    if assistant_only_loss:
+        if is_vision or type != "sft":
+            return ToolResult.fail(
+                "config",
+                "Error: assistant_only_loss applies to text SFT only.",
+            )
+        unsupported = _assistant_mask_unsupported(project_dir, base_model)
+        if unsupported:
+            return ToolResult.fail(
+                "config",
+                f"Error: {base_model} cannot mask assistant tokens: "
+                f"{unsupported}. Drop assistant_only_loss or pick a base "
+                "model whose chat template marks its assistant turns.",
+            )
     if is_vision and type != "sft":
         return ToolResult.fail(
             "validation",
@@ -5474,6 +5524,7 @@ async def handle_start_training(
         "training": {
             **recommended.training_config(),
             "learning_rate": lr,
+            "assistant_only_loss": assistant_only_loss,
         },
         "lora": recommended.lora,
         "manifest": ["base_model", "dataset"],

lqh/train/data_utils.py

@@ -475,3 +475,33 @@ def load_preferences_parquet(
             }
         )
     return result
+
+
+def drop_rows_without_assistant_labels(
+    rows: list[dict[str, Any]],
+    tokenizer: Any,
+    max_length: int | None,
+) -> tuple[list[dict[str, Any]], int]:
+    """Drop rows whose assistant tokens all sit past *max_length*.
+
+    trl builds the assistant mask before truncation, so a row with a long
+    prompt reaches the loss with every label masked out and contributes a
+    zero-or-NaN batch instead of being rejected. The extra tokenization
+    pass is cheap (~0.7 ms/row on the LFM2.5 fast tokenizer).
+    """
+    if not max_length:
+        return rows, 0
+    kept = [
+        row
+        for row in rows
+        if any(
+            tokenizer.apply_chat_template(
+                row["messages"],
+                tools=row.get("tools"),
+                tokenize=True,
+                return_dict=True,
+                return_assistant_tokens_mask=True,
+            )["assistant_masks"][:max_length]
+        )
+    ]
+    return kept, len(rows) - len(kept)

lqh/train/sft.py

@@ -36,6 +36,7 @@
 
 from lqh.train.data_utils import (
     chatml_to_sft_dataset,
+    drop_rows_without_assistant_labels,
     load_chatml_datasets_with_tools,
     load_eval_sources,
     split_train_eval,
@@ -73,6 +74,7 @@ def _write_checkpoint_lineage(
         "per_device_batch_size": training_cfg.get("per_device_batch_size"),
         "gradient_accumulation_steps": training_cfg.get("gradient_accumulation_steps"),
         "effective_batch_size": training_cfg.get("effective_batch_size"),
+        "assistant_only_loss": bool(training_cfg.get("assistant_only_loss", False)),
     }
     if lora_cfg.get("enabled", True):
         hyperparams.update(
@@ -795,14 +797,32 @@ def sft_loop(run_dir: Path, config: dict[str, Any]) -> None:
         if eval_convos:
             eval_dataset = Dataset.from_list(chatml_to_vlm_dataset(eval_convos))
     else:
-        train_dataset = Dataset.from_list(
-            chatml_to_sft_dataset(train_convos, train_tools)
+        train_rows = chatml_to_sft_dataset(train_convos, train_tools)
+        eval_rows = (
+            chatml_to_sft_dataset(eval_convos, eval_tools) if eval_convos else []
         )
-        eval_dataset = None
-        if eval_convos:
-            eval_dataset = Dataset.from_list(
-                chatml_to_sft_dataset(eval_convos, eval_tools)
+        if training_cfg.get("assistant_only_loss"):
+            max_length = training_cfg.get("max_seq_length", 2048)
+            train_rows, dropped = drop_rows_without_assistant_labels(
+                train_rows, tokenizer, max_length
+            )
+            eval_rows, eval_dropped = drop_rows_without_assistant_labels(
+                eval_rows, tokenizer, max_length
             )
+            if dropped or eval_dropped:
+                print(
+                    f"  assistant_only_loss: dropped {dropped} train and "
+                    f"{eval_dropped} eval rows whose assistant turn falls "
+                    f"past max_seq_length={max_length}"
+                )
+            if not train_rows:
+                raise ValueError(
+                    "assistant_only_loss left no trainable rows: every "
+                    f"assistant turn falls past max_seq_length={max_length}. "
+                    "Raise max_seq_length or shorten the prompts."
+                )
+        train_dataset = Dataset.from_list(train_rows)
+        eval_dataset = Dataset.from_list(eval_rows) if eval_rows else None
 
     # Safe batch-size auto-tuning (GPU_TYPE.md §6). Mutates training_cfg
     # in place (per_device_batch_size + gradient_accumulation_steps) so
@@ -958,6 +978,7 @@ def sft_loop(run_dir: Path, config: dict[str, Any]) -> None:
         gradient_checkpointing=training_cfg.get("gradient_checkpointing", True),
         bf16=training_cfg.get("bf16", True),
         max_length=training_cfg.get("max_seq_length", 2048),
+        assistant_only_loss=bool(training_cfg.get("assistant_only_loss", False)),
         dataloader_num_workers=training_cfg.get("dataloader_num_workers", 4),
         dataloader_pin_memory=True,
         ddp_find_unused_parameters=False,

tests/unit/test_sft_assistant_only_loss.py

@@ -0,0 +1,131 @@
+"""assistant_only_loss labels assistant turns and nothing else."""
+
+from __future__ import annotations
+
+import pytest
+
+MESSAGES = [
+    {
+        "role": "user",
+        "content": (
+            "project 01 | tempo 70\nStereo Out | out | vol 6.0 | clip\n"
+            "Why is the master clipping?"
+        ),
+    },
+    {
+        "role": "assistant",
+        "content": (
+            "findings:\n- [critical] Stereo Out | clip: clipping latched at 6 dB\n"
+            '<|tool_call_start|>[set_volume(track="Stereo Out", db=0.0)]<|tool_call_end|>'
+        ),
+    },
+    {"role": "user", "content": "And the bass?"},
+    {"role": "assistant", "content": "Trap Bass sits at -3 dB and does not clip."},
+]
+
+# The shape of LFM2.5's template: the assistant header is outside the block, so
+# only the reply body and its terminator carry a live label.
+TEMPLATE = (
+    "{%- for m in messages -%}"
+    "{{- '<|im_start|>' + m.role + '\n' -}}"
+    "{%- if m.role == 'assistant' -%}{%- generation -%}"
+    "{{- m.content + '<|im_end|>\n' -}}"
+    "{%- endgeneration -%}"
+    "{%- else -%}{{- m.content + '<|im_end|>\n' -}}{%- endif -%}"
+    "{%- endfor -%}"
+)
+
+
+def make_tokenizer(template: str = TEMPLATE):
+    """A character-level tokenizer built in memory: no hub, no cache, exact offsets."""
+    pytest.importorskip("transformers")
+    from tokenizers import Regex, Tokenizer, decoders, models, pre_tokenizers
+    from transformers import PreTrainedTokenizerFast
+
+    chars = [chr(c) for c in range(32, 127)] + ["\n"]
+    inner = Tokenizer(
+        models.WordLevel(
+            vocab={t: i for i, t in enumerate(["<|pad|>", "<|unk|>", *chars])},
+            unk_token="<|unk|>",
+        )
+    )
+    inner.pre_tokenizer = pre_tokenizers.Split(Regex("."), behavior="isolated")
+    inner.decoder = decoders.Fuse()
+    return PreTrainedTokenizerFast(
+        tokenizer_object=inner,
+        pad_token="<|pad|>",
+        eos_token="\n",
+        chat_template=template,
+    )
+
+
+@pytest.fixture(scope="module")
+def tokenizer():
+    return make_tokenizer()
+
+
+def encode(tokenizer, messages=MESSAGES):
+    return tokenizer.apply_chat_template(
+        messages, tokenize=True, return_dict=True, return_assistant_tokens_mask=True
+    )
+
+
+def test_template_marks_only_the_assistant_replies(tokenizer):
+    encoded = encode(tokenizer)
+    ids, mask = encoded["input_ids"], encoded["assistant_masks"]
+    assert len(ids) == len(mask)
+    assert 0 < sum(mask) < len(ids)
+
+    live = tokenizer.decode([i for i, m in zip(ids, mask) if m])
+    assert live == "".join(
+        f"{m['content']}<|im_end|>\n" for m in MESSAGES if m["role"] == "assistant"
+    )
+    masked = tokenizer.decode([i for i, m in zip(ids, mask) if not m])
+    for m in MESSAGES:
+        if m["role"] == "user":
+            assert m["content"] in masked
+    assert "<|im_start|>assistant" in masked
+
+
+def test_collator_masks_every_non_assistant_token(tokenizer):
+    pytest.importorskip("torch")
+    from trl.trainer.sft_trainer import DataCollatorForLanguageModeling
+
+    encoded = encode(tokenizer)
+    ids, mask = encoded["input_ids"], encoded["assistant_masks"]
+    batch = DataCollatorForLanguageModeling(pad_token_id=tokenizer.pad_token_id)(
+        [{"input_ids": ids, "assistant_masks": mask}]
+    )
+    labels = batch["labels"][0]
+    assert (labels == -100).tolist() == [m == 0 for m in mask]
+
+
+def test_rows_truncated_past_their_assistant_turn_are_dropped(tokenizer):
+    from lqh.train.data_utils import drop_rows_without_assistant_labels
+
+    rows = [{"messages": MESSAGES}]
+    first_label = encode(tokenizer)["assistant_masks"].index(1)
+
+    kept, dropped = drop_rows_without_assistant_labels(rows, tokenizer, first_label)
+    assert (kept, dropped) == ([], 1)
+    kept, dropped = drop_rows_without_assistant_labels(rows, tokenizer, first_label + 1)
+    assert (kept, dropped) == (rows, 0)
+    assert drop_rows_without_assistant_labels(rows, tokenizer, None) == (rows, 0)
+
+
+def test_a_template_without_a_generation_block_is_rejected_before_launch(tmp_path):
+    from lqh.tools.handlers import _assistant_mask_unsupported
+
+    blocked = TEMPLATE.replace("{%- generation -%}", "").replace(
+        "{%- endgeneration -%}", ""
+    )
+    make_tokenizer(blocked).save_pretrained(tmp_path / "tok")
+    # base_model given project-relative, as a local checkpoint would be
+    assert _assistant_mask_unsupported(tmp_path, "tok") == (
+        "its chat template has no {% generation %} block"
+    )
+
+    make_tokenizer().save_pretrained(tmp_path / "tok")
+    assert _assistant_mask_unsupported(tmp_path, "tok") is None
+    # unloadable tokenizer fails open: the trainer is the backstop
+    assert _assistant_mask_unsupported(tmp_path, str(tmp_path / "absent")) is None