[v1] Fix device mesh, fix lora for reward model and fix sp (#10555)

This commit is contained in:
jiaqiw09
2026-06-25 20:05:56 +08:00
committed by GitHub
parent 666ee0ca78
commit b7615dbdc9
7 changed files with 58 additions and 20 deletions

View File

@@ -32,6 +32,7 @@ def _make_model_input(length: int, start: int = 0):
"attention_mask": [1] * length,
"labels": input_ids.copy(),
"loss_weights": [1.0] * length,
"position_ids": list(range(1, length + 1)),
}
@@ -62,7 +63,7 @@ def test_padding_free():
assert len(batch) == 1
assert batch[0]["input_ids"].shape == (1, 5)
assert batch[0]["input_ids"].tolist() == [[0, 1, 10, 11, 12]]
assert batch[0]["attention_mask"].tolist() == [[1, 1, 1, 1, 1]]
assert batch[0]["attention_mask"] is None
assert batch[0]["position_ids"].tolist() == [[0, 1, 0, 1, 2]]
assert batch[0]["labels"].tolist() == [[0, 1, IGNORE_INDEX, 11, 12]]
assert batch[0]["loss_weights"].tolist() == [[1.0, 1.0, 0.0, 1.0, 1.0]]
@@ -115,6 +116,8 @@ def test_dynamic_batching():
assert len(batch) == 1
assert batch[0]["input_ids"].shape == (3, 6)
assert batch[0]["input_ids"].tolist()[0] == [0, 1, 2, 0, 0, 0]
assert batch[0]["position_ids"].shape == (3, 6)
assert batch[0]["position_ids"].tolist()[0] == [1, 2, 3, 0, 0, 0]
assert len(buffer) == 3
@@ -201,6 +204,7 @@ def test_normal_batching():
batch = next(iter(batch_generator))
assert len(batch) == 2
assert batch[0]["input_ids"].shape == (4, 10)
assert batch[0]["position_ids"].shape == (4, 10)
def test_dynamic_padding_free():
@@ -280,8 +284,8 @@ def test_dynamic_padding_free():
] # Sample 3
]
# Verify attention_mask
assert packed_batch["attention_mask"].tolist() == [[1] * 15]
# Verify attention_mask: padding-free relies on reset-style position_ids instead of a dense mask.
assert packed_batch["attention_mask"] is None
# Verify position_ids
assert packed_batch["position_ids"].tolist() == [