feat: SFT 增加 position_ids 边界处理

- SFTProcessor 输出 per-sample position_ids（torch.arange），每个样本从 0 开始
- position_ids 与 sequence/loss_mask 一同打包，边界处自然重置
- PT 路径不生成 position_ids
- 新增 TestPositionIds 测试及 SFTProcessor 相关测试

tests/test_io.py

@@ -7,7 +7,7 @@ import torch
 import h5py
 from pathlib import Path
 
-from pipeline.io import FileScanner, HDF5Handler
+from pipeline.io import FileScanner, HDF5Handler, cache_jsonl
 
 
 class TestFileScanner:
@@ -137,3 +137,48 @@ class TestHDF5Handler:
             h5_path = os.path.join(tmpdir, "meta.h5")
             metadata = HDF5Handler.get_metadata(h5_path)
             assert metadata["data"] == 5
+
+
+class DummyTokenizer:
+    im_end = "<|im_end|>"
+
+    def encode(self, text: str, add_special_tokens: bool = False):
+        return [ord(c) for c in text]
+
+    def apply_chat_template(
+        self, messages, add_generation_prompt=True, tokenize=True
+    ):
+        text = ""
+        for m in messages:
+            text += f"<|im_start|>{m['role']}\n{m['content']}<|im_end|>\n"
+        if add_generation_prompt:
+            text += "<|im_start|>assistant\n"
+        return self.encode(text) if tokenize else text
+
+
+class TestPositionIds:
+    def test_example_specific_position_ids(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            jsonl_path = os.path.join(tmpdir, "data.jsonl")
+            with open(jsonl_path, "w") as f:
+                f.write('{"messages": [{"role": "user", "content": "a"}, {"role": "assistant", "content": "bc"}]}\n')
+                f.write('{"messages": [{"role": "user", "content": "def"}, {"role": "assistant", "content": "g"}]}\n')
+
+            from pipeline.processors import SFTProcessor
+
+            processor = SFTProcessor(DummyTokenizer())
+            out_dir = os.path.join(tmpdir, "cached")
+            cache_jsonl([jsonl_path], out_dir, processor, pack_size=-1)
+
+            h5_path = os.path.join(out_dir, "data.h5")
+            loaded = HDF5Handler.load(h5_path, share_memory=False)
+
+            assert "position_ids" in loaded
+            assert len(loaded["position_ids"]) == 2
+            assert len(loaded["position_ids"]) == len(loaded["sequence"])
+
+            for i, pos in enumerate(loaded["position_ids"]):
+                seq_len = len(loaded["sequence"][i])
+                assert len(pos) == seq_len
+                assert pos[0].item() == 0
+                assert (pos == torch.arange(seq_len, dtype=torch.int32)).all()

tests/test_processors.py

@@ -13,9 +13,21 @@ from pipeline.processors import (
 
 
 class DummyTokenizer:
+    im_end = "<|im_end|>"
+
     def encode(self, text: str, add_special_tokens: bool = False):
         return [ord(c) for c in text]
 
+    def apply_chat_template(
+        self, messages, add_generation_prompt=True, tokenize=True
+    ):
+        text = ""
+        for m in messages:
+            text += f"<|im_start|>{m['role']}\n{m['content']}<|im_end|>\n"
+        if add_generation_prompt:
+            text += "<|im_start|>assistant\n"
+        return self.encode(text) if tokenize else text
+
 
 class TestBaseProcessor:
     def test_abstract_class_cannot_be_instantiated(self):
@@ -41,16 +53,18 @@ class TestPreTrainProcessor:
 
 class TestSFTProcessor:
     def test_output_keys(self):
-        assert SFTProcessor(DummyTokenizer()).output_keys == ["sequence", "loss_mask"]
+        keys = SFTProcessor(DummyTokenizer()).output_keys
+        assert "sequence" in keys
+        assert "loss_mask" in keys
+        assert "position_ids" in keys
 
-    def test_process_returns_both_keys(self):
+    def test_process_returns_all_keys(self):
         result = SFTProcessor(DummyTokenizer()).process(
             {"query": "hello", "response": "world"}
         )
-        assert "sequence" in result
-        assert "loss_mask" in result
-        assert isinstance(result["sequence"], torch.Tensor)
-        assert isinstance(result["loss_mask"], torch.Tensor)
+        for key in ["sequence", "loss_mask", "position_ids"]:
+            assert key in result
+            assert isinstance(result[key], torch.Tensor)
 
     def test_loss_mask_correct_length(self):
         result = SFTProcessor(DummyTokenizer()).process(
@@ -64,6 +78,73 @@ class TestSFTProcessor:
         )
         assert result["loss_mask"].dtype == torch.bool
 
+    def test_position_ids_start_from_zero(self):
+        result = SFTProcessor(DummyTokenizer()).process(
+            {"query": "abc", "response": "de"}
+        )
+        seq_len = len(result["sequence"])
+        expected = torch.arange(seq_len, dtype=torch.int32)
+        assert torch.equal(result["position_ids"], expected)
+
+    def test_messages_single_turn(self):
+        result = SFTProcessor(DummyTokenizer()).process({
+            "messages": [
+                {"role": "user", "content": "hi"},
+                {"role": "assistant", "content": "bye"},
+            ]
+        })
+        for key in ["sequence", "loss_mask", "position_ids"]:
+            assert key in result
+        assert len(result["sequence"]) == len(result["loss_mask"])
+
+    def test_messages_loss_on_last_assistant_only(self):
+        result = SFTProcessor(DummyTokenizer()).process({
+            "messages": [
+                {"role": "user", "content": "q1"},
+                {"role": "assistant", "content": "a1"},
+                {"role": "user", "content": "q2"},
+                {"role": "assistant", "content": "a2"},
+            ]
+        })
+        mask = result["loss_mask"]
+        first_true = mask.tolist().index(True)
+        assert not mask[:first_true].any()
+        assert mask[-1].item() is True
+
+    def test_messages_with_system_prompt(self):
+        result = SFTProcessor(DummyTokenizer()).process({
+            "messages": [
+                {"role": "system", "content": "You are helpful."},
+                {"role": "user", "content": "hi"},
+                {"role": "assistant", "content": "hello"},
+            ]
+        })
+        assert "sequence" in result
+
+    def test_messages_empty_raises(self):
+        with pytest.raises(ValueError, match="Messages list is empty"):
+            SFTProcessor(DummyTokenizer()).process({"messages": []})
+
+    def test_messages_last_not_assistant_raises(self):
+        with pytest.raises(ValueError, match="Last message must"):
+            SFTProcessor(DummyTokenizer()).process({
+                "messages": [{"role": "user", "content": "hi"}]
+            })
+
+    def test_missing_fields_raises(self):
+        with pytest.raises(KeyError):
+            SFTProcessor(DummyTokenizer()).process({"foo": "bar"})
+
+    def test_position_ids_start_from_zero(self):
+        result = SFTProcessor(DummyTokenizer()).process(
+            {"query": "hi", "response": "ok"}
+        )
+        pos_ids = result["position_ids"]
+        assert pos_ids.dtype == torch.int32
+        assert len(pos_ids) == len(result["sequence"])
+        assert pos_ids[0].item() == 0
+        assert (pos_ids == torch.arange(len(pos_ids))).all()
+
 
 class TestDPOProcessor:
     def test_output_keys(self):