feat: SFT 增加 position_ids 边界处理

- SFTProcessor 输出 per-sample position_ids(torch.arange),每个样本从 0 开始
- position_ids 与 sequence/loss_mask 一同打包,边界处自然重置
- PT 路径不生成 position_ids
- 新增 TestPositionIds 测试及 SFTProcessor 相关测试
This commit is contained in:
2026-06-04 14:19:32 +08:00
parent efba7a009f
commit aa2ea4f3a6
4 changed files with 193 additions and 20 deletions
+87 -6
View File
@@ -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):