fix(learning): exclude padding tokens from SFT loss (#521)

* fix(learning): exclude padding tokens from SFT loss

* test(learning): add regression tests for SFT padding-loss masking

Cover the fix in both trainers (#521):
- OrchestratorSFTDataset.__getitem__: labels are -100 at padded positions,
  equal to input_ids elsewhere, and input_ids is not mutated.
- LoRATrainer._train_step: the labels passed to the model are masked at
  padded positions (captured via an injected model), with input_ids intact.

Both are torch-gated (pytest.importorskip / skipif HAS_TORCH), matching the
project's existing torch test gating, so they skip cleanly in the default CI
env. Verified locally with CPU torch: both PASS against the fix and FAIL
against the pre-fix code.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Jon Saad-Falcon <jonsaadfalcon@gmail.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

src/openjarvis/learning/intelligence/orchestrator/sft_trainer.py

@@ -141,10 +141,18 @@ class OrchestratorSFTDataset:
             return_tensors="pt",
         )
 
+        input_ids = encoding["input_ids"].squeeze(0)
+        attention_mask = encoding["attention_mask"].squeeze(0)
+
+        # Exclude padding positions from the loss (-100 is the
+        # cross-entropy ignore_index).
+        labels = input_ids.clone()
+        labels[attention_mask == 0] = -100
+
         return {
-            "input_ids": encoding["input_ids"].squeeze(0),
-            "attention_mask": encoding["attention_mask"].squeeze(0),
-            "labels": encoding["input_ids"].squeeze(0).clone(),
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "labels": labels,
         }
 
     def _format_conversation(self, conversations: List[Dict[str, str]]) -> str:

src/openjarvis/learning/training/lora.py

@@ -389,10 +389,15 @@ class LoRATrainer:
             [item["attention_mask"] for item in batch_items]
         ).to(self.device)
 
+        # Exclude padding positions from the loss (-100 is the
+        # cross-entropy ignore_index; pad_token == eos_token here).
+        labels = input_ids.clone()
+        labels[attention_mask == 0] = -100
+
         outputs = self.model(
             input_ids=input_ids,
             attention_mask=attention_mask,
-            labels=input_ids,
+            labels=labels,
         )
         loss = outputs.loss
 

tests/learning/training/test_lora.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+from unittest.mock import MagicMock
+
 import pytest
 
 from openjarvis.learning.training.lora import HAS_TORCH, LoRATrainer, LoRATrainingConfig
@@ -145,3 +147,55 @@ class TestLoRATrainerWithTorch:
 
         assert result["status"] == "skipped"
         assert "reason" in result
+
+
+@pytest.mark.skipif(not HAS_TORCH, reason="torch not installed")
+class TestLoRATrainStepMasking:
+    """Regression for #521: _train_step must exclude padding from the loss labels.
+
+    The pre-fix code passed ``labels=input_ids`` (the *same* tensor), so the loss
+    counted every padded EOS position and an in-place mask would also corrupt
+    ``input_ids``. ``_train_step`` must build a masked clone: ``-100`` wherever
+    ``attention_mask == 0``, equal to ``input_ids`` elsewhere, leaving
+    ``input_ids`` untouched.
+    """
+
+    def test_train_step_masks_padding_in_labels(self) -> None:
+        import torch
+
+        # Bypass __init__ (which loads a real model) and inject the minimal
+        # attributes _train_step touches before the forward pass.
+        trainer = LoRATrainer.__new__(LoRATrainer)
+        trainer.device = "cpu"
+
+        captured: dict = {}
+
+        class _StopForward(Exception):
+            pass
+
+        def _capture_model(*, input_ids, attention_mask, labels):
+            captured["labels"] = labels
+            raise _StopForward  # stop before backward()/optimizer.step()
+
+        trainer.model = _capture_model
+
+        batch_items = [
+            {
+                "input_ids": torch.tensor([11, 12, 0, 0]),
+                "attention_mask": torch.tensor([1, 1, 0, 0]),
+            },
+            {
+                "input_ids": torch.tensor([13, 14, 15, 0]),
+                "attention_mask": torch.tensor([1, 1, 1, 0]),
+            },
+        ]
+        with pytest.raises(_StopForward):
+            trainer._train_step(batch_items, optimizer=MagicMock())
+
+        ids = torch.stack([b["input_ids"] for b in batch_items])
+        mask = torch.stack([b["attention_mask"] for b in batch_items])
+        labels = captured["labels"]
+
+        assert (labels[mask == 0] == -100).all()  # padded -> ignored by loss
+        assert (labels[mask == 1] == ids[mask == 1]).all()  # real -> unchanged
+        assert (ids[mask == 0] != -100).all()  # input_ids not mutated in place

tests/test_orchestrator_learning/test_sft_trainer.py

@@ -4,6 +4,8 @@ from __future__ import annotations
 
 from unittest.mock import MagicMock
 
+import pytest
+
 from openjarvis.learning.intelligence.orchestrator.sft_trainer import (
     OrchestratorSFTConfig,
     OrchestratorSFTDataset,
@@ -123,6 +125,56 @@ class TestOrchestratorSFTDataset:
         assert len(batches) == 3  # 2+2+1
 
 
+class TestSFTLabelMasking:
+    """Regression for #521: padding positions must be excluded from the SFT loss.
+
+    With ``padding="max_length"`` and ``pad_token == eos_token``, an unmasked
+    ``labels`` makes the model optimise "predict EOS at a padded position" for
+    the bulk of every example, diluting the gradient on real content and
+    understating the reported loss. ``labels`` must be ``-100`` wherever
+    ``attention_mask == 0`` and equal to ``input_ids`` elsewhere — without
+    mutating ``input_ids`` (the pre-fix code aliased the two).
+    """
+
+    def test_getitem_masks_padding_positions(self, tmp_path):
+        torch = pytest.importorskip("torch")
+        import json
+
+        trace_file = tmp_path / "traces.jsonl"
+        trace_file.write_text(
+            json.dumps(
+                {
+                    "conversations": [
+                        {"role": "user", "content": "hi"},
+                        {"role": "assistant", "content": "yo"},
+                    ]
+                }
+            )
+            + "\n"
+        )
+
+        class _FakeTokenizer:
+            """Returns a fixed padded encoding: 2 real tokens, 6 pad (id 0)."""
+
+            eos_token = "</s>"
+
+            def __call__(self, text, **kwargs):
+                input_ids = torch.tensor([[11, 12, 0, 0, 0, 0, 0, 0]])
+                attention_mask = torch.tensor([[1, 1, 0, 0, 0, 0, 0, 0]])
+                return {"input_ids": input_ids, "attention_mask": attention_mask}
+
+        ds = OrchestratorSFTDataset(
+            trace_path=str(trace_file), tokenizer=_FakeTokenizer()
+        )
+        item = ds[0]
+        ids, mask, labels = item["input_ids"], item["attention_mask"], item["labels"]
+
+        assert (labels[mask == 0] == -100).all()  # padded -> ignored by loss
+        assert (labels[mask == 1] == ids[mask == 1]).all()  # real -> unchanged
+        assert (ids[mask == 0] != -100).all()  # input_ids not mutated in place
+        assert not torch.equal(ids, labels)  # masked clone, not an alias
+
+
 class TestSFTRegistration:
     def test_registered_in_learning_registry(self):
         # Import to trigger registration