feat: SFT 增加 position_ids 边界处理

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

pipeline/io/export.py

@@ -5,7 +5,6 @@ import logging
 import os
 from pathlib import Path
 from typing import Any, Callable, Dict, List, Optional, Union
-
 from datasets import Dataset
 from tqdm import tqdm
 

pipeline/processors/sft.py

@@ -15,15 +15,20 @@ from pipeline.processors.factory import ProcessorFactory
 class SFTProcessor(BaseProcessor):
     """Supervised fine-tuning data processor.
 
-    Processes query-response pairs into tokenized sequences with loss masks.
+    Supports two input formats:
-
+      1. messages (recommended):
-    Input schema:
+         ``{"messages": [{"role": "user", "content": "..."},
-        - query: str - User query/prompt
+                         {"role": "assistant", "content": "..."}]}``
-        - response: str - Assistant response
+         Multi-turn and system prompts are supported.
+         The tokenizer's ``apply_chat_template`` is used for rendering.
+      2. legacy query/response:
+         ``{"query": "...", "response": "..."}``
+         Falls back to the configured PromptStrategy (ChatML by default).
 
     Output schema:
-        - sequence: int32 tensor - Combined token IDs (query + response)
+        - sequence: int32 tensor - Combined token IDs (prompt + response)
         - loss_mask: bool tensor - True for response tokens (compute loss)
+        - position_ids: int32 tensor - Per-sample position IDs starting from 0
     """
 
     def __init__(
@@ -32,28 +37,71 @@ class SFTProcessor(BaseProcessor):
         strategy: Optional[PromptStrategy] = None,
     ):
         self.tokenizer = tokenizer
-        self.strategy = strategy or ChatMLStrategy(tokenizer)
+        self.strategy = strategy
 
     @property
     def schema(self) -> ProcessorSchema:
         return ProcessorSchema(
-            input_fields={"query": str, "response": str},
+            input_fields={
+                "messages": list,
+                "query": str,
+                "response": str,
+            },
             output_fields={
                 "sequence": torch.int32,
                 "loss_mask": torch.bool,
+                "position_ids": torch.int32,
             },
         )
 
     def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
+        if "messages" in input_dict:
+            return self._process_messages(input_dict["messages"])
+        if "query" in input_dict and "response" in input_dict:
+            return self._process_legacy(input_dict)
+        raise KeyError(
+            "Input must contain 'messages' or 'query'/'response' pair"
+        )
+
+    def _process_messages(self, messages: List[Dict[str, str]]) -> Dict[str, Tensor]:
+        if not messages:
+            raise ValueError("Messages list is empty")
+        if messages[-1]["role"] != "assistant":
+            raise ValueError("Last message must have role 'assistant'")
+
+        last_asst_idx = max(
+            i for i, m in enumerate(messages) if m["role"] == "assistant"
+        )
+
+        prompt_tokens = self.tokenizer.apply_chat_template(
+            messages[:last_asst_idx],
+            add_generation_prompt=True,
+            tokenize=True,
+        )
+
+        resp_content = messages[last_asst_idx]["content"]
+        im_end = getattr(self.tokenizer, "im_end", "<|im_end|>")
+        resp_tokens = self.tokenizer.encode(
+            f"{resp_content}{im_end}\n", add_special_tokens=False
+        )
+
+        tokens, loss_mask = encode_with_mask(prompt_tokens, resp_tokens)
+        position_ids = torch.arange(len(tokens), dtype=torch.int32)
+        return {"sequence": tokens, "loss_mask": loss_mask, "position_ids": position_ids}
+
+    def _process_legacy(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
+        strategy = self.strategy or ChatMLStrategy(self.tokenizer)
+
         query_tokens = self.tokenizer.encode(input_dict["query"])
         response_tokens = self.tokenizer.encode(input_dict["response"])
 
-        prompt = self.strategy.assemble_prompt(query_tokens)
+        prompt = strategy.assemble_prompt(query_tokens)
-        response = self.strategy.assemble_response(response_tokens)
+        response = strategy.assemble_response(response_tokens)
 
         tokens, loss_mask = encode_with_mask(prompt, response)
-        return {"sequence": tokens, "loss_mask": loss_mask}
+        position_ids = torch.arange(len(tokens), dtype=torch.int32)
+        return {"sequence": tokens, "loss_mask": loss_mask, "position_ids": position_ids}
 
     @property
     def output_keys(self) -> List[str]:
-        return ["sequence", "loss_mask"]
+        return ["sequence", "loss_mask", "position_ids"]

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
+        for key in ["sequence", "loss_mask", "position_ids"]:
-        assert "loss_mask" in result
+            assert key in result
-        assert isinstance(result["sequence"], torch.Tensor)
+            assert isinstance(result[key], torch.Tensor)
-        assert isinstance(result["loss_mask"], 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):