diff --git a/src/openjarvis/learning/intelligence/orchestrator/sft_trainer.py b/src/openjarvis/learning/intelligence/orchestrator/sft_trainer.py index 2fe6aa83..2c2e0645 100644 --- a/src/openjarvis/learning/intelligence/orchestrator/sft_trainer.py +++ b/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: diff --git a/src/openjarvis/learning/training/lora.py b/src/openjarvis/learning/training/lora.py index 587707ef..90cba870 100644 --- a/src/openjarvis/learning/training/lora.py +++ b/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 diff --git a/tests/learning/training/test_lora.py b/tests/learning/training/test_lora.py index 85cd2aee..58c3557d 100644 --- a/tests/learning/training/test_lora.py +++ b/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 diff --git a/tests/test_orchestrator_learning/test_sft_trainer.py b/tests/test_orchestrator_learning/test_sft_trainer.py index 3e2dfbcf..941427f3 100644 --- a/tests/test_orchestrator_learning/test_sft_trainer.py +++ b/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 = "" + + 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