fix: rewrite GRPO data pipeline for offline record-level access
- process_list_field returns List[List[int]] preserving per-response boundaries - GRPODataset rewritten to record-level __getitem__ (no windowing/stride) - grpo_collate_fn pads variable-length responses into [B, G, R] tensors - JsonlStore detects nested List[List[int]] and stores List[Tensor] per record - Store._normalize skips nested-list keys from cumsum bookkeeping - Pipeline._flush handles nested lists without cross-record flattening - Export grpo_collate_fn from astrai.dataset - 6 new GRPO tests + 2 updated builder tests, 114 total pass
This commit is contained in:
@@ -1,6 +1,7 @@
|
||||
from astrai.dataset.dataset import (
|
||||
BaseDataset,
|
||||
DatasetFactory,
|
||||
grpo_collate_fn,
|
||||
)
|
||||
from astrai.dataset.sampler import ResumableDistributedSampler
|
||||
from astrai.dataset.storage import (
|
||||
@@ -21,6 +22,7 @@ from astrai.serialization import (
|
||||
__all__ = [
|
||||
"BaseDataset",
|
||||
"DatasetFactory",
|
||||
"grpo_collate_fn",
|
||||
"Store",
|
||||
"StoreFactory",
|
||||
"H5Store",
|
||||
|
||||
+116
-16
@@ -15,6 +15,49 @@ from astrai.dataset.storage import (
|
||||
from astrai.factory import BaseFactory
|
||||
|
||||
|
||||
def grpo_collate_fn(batch: List[Dict[str, Tensor]]) -> Dict[str, Tensor]:
|
||||
"""Collate variable-length GRPO samples into padded 3-D tensors.
|
||||
|
||||
Input: list of dicts, each with:
|
||||
- prompts: [P_i]
|
||||
- responses: list of G tensors, each [R_ij]
|
||||
- masks: list of G tensors, each [R_ij]
|
||||
- rewards: [G]
|
||||
|
||||
Output:
|
||||
- prompts: [B, P_max]
|
||||
- responses: [B, G, R_max]
|
||||
- masks: [B, G, R_max]
|
||||
- rewards: [B, G]
|
||||
"""
|
||||
B = len(batch)
|
||||
G = len(batch[0]["responses"])
|
||||
P_max = max(b["prompts"].size(0) for b in batch)
|
||||
R_max = max(r.size(0) for b in batch for r in b["responses"])
|
||||
|
||||
prompts = torch.zeros(B, P_max, dtype=torch.long)
|
||||
responses = torch.zeros(B, G, R_max, dtype=torch.long)
|
||||
masks = torch.zeros(B, G, R_max, dtype=torch.bool)
|
||||
rewards = torch.zeros(B, G, dtype=torch.float32)
|
||||
|
||||
for i, b in enumerate(batch):
|
||||
p_len = b["prompts"].size(0)
|
||||
prompts[i, :p_len] = b["prompts"]
|
||||
rewards[i, : b["rewards"].size(0)] = b["rewards"]
|
||||
for g in range(min(G, len(b["responses"]))):
|
||||
r_len = b["responses"][g].size(0)
|
||||
responses[i, g, :r_len] = b["responses"][g]
|
||||
if g < len(b["masks"]):
|
||||
masks[i, g, :r_len] = b["masks"][g]
|
||||
|
||||
return {
|
||||
"prompts": prompts,
|
||||
"responses": responses,
|
||||
"masks": masks,
|
||||
"rewards": rewards,
|
||||
}
|
||||
|
||||
|
||||
class BaseDataset(Dataset, ABC):
|
||||
"""Abstract base class for all dataset types.
|
||||
|
||||
@@ -250,28 +293,85 @@ class DPODataset(BaseDataset):
|
||||
|
||||
@DatasetFactory.register("grpo")
|
||||
class GRPODataset(BaseDataset):
|
||||
"""Dataset for Group Relative Policy Optimization training."""
|
||||
"""Dataset for offline Group Relative Policy Optimization.
|
||||
|
||||
Unlike the window-based datasets (SEQ/SFT/DPO), GRPO data is
|
||||
record-structured: each sample is one prompt with its group of
|
||||
responses and scalar rewards. There is no windowing or stride —
|
||||
every record is an independent training unit.
|
||||
|
||||
Expected storage layout (produced by JsonlStore or pre-tokenized):
|
||||
|
||||
- ``prompts``: List[Tensor] — one 1-D token tensor per record
|
||||
- ``responses``: List[List[Tensor]] — G response tensors per record
|
||||
- ``masks``: List[List[Tensor]] — G mask tensors per record
|
||||
- ``rewards``: List[Tensor] — one 1-D float tensor (len G) per record
|
||||
"""
|
||||
|
||||
def __init__(self, window_size: int = 0, stride: int = 0, **kwargs):
|
||||
super().__init__(window_size=window_size, stride=stride or window_size)
|
||||
self._records: List[dict] = []
|
||||
|
||||
@property
|
||||
def required_keys(self) -> List[str]:
|
||||
return ["prompts", "responses", "masks", "rewards"]
|
||||
|
||||
def _fetch_data(self, begin_idx: int, end_idx: int, key: str) -> Tensor:
|
||||
return self.storage.fetch(begin_idx, end_idx, key)
|
||||
def load(self, load_path: str, storage_type: Optional[str] = None, **kwargs):
|
||||
if storage_type is None:
|
||||
storage_type = detect_format(load_path)
|
||||
self.storage = StoreFactory.create(storage_type, **kwargs)
|
||||
self._load_path = load_path
|
||||
self.storage.load(load_path, **kwargs)
|
||||
self._validate_keys()
|
||||
self._build_records()
|
||||
|
||||
def _validate_keys(self):
|
||||
actual_keys = set(self.storage.keys)
|
||||
missing = [k for k in self.required_keys if k not in actual_keys]
|
||||
if missing:
|
||||
raise KeyError(
|
||||
f"GRPODataset requires keys {self.required_keys}, "
|
||||
f"but storage only has {sorted(actual_keys)}. Missing: {missing}"
|
||||
)
|
||||
|
||||
def _build_records(self):
|
||||
"""Unfold segmented storage into per-record lists.
|
||||
|
||||
``prompts`` is a flat list of 1-D tensors (one per record).
|
||||
``responses`` / ``masks`` are nested lists (G tensors per record).
|
||||
``rewards`` is a flat list of 1-D tensors (len G per record).
|
||||
"""
|
||||
prompt_segs = self.storage._data.get("prompts", [])
|
||||
response_segs = self.storage._data.get("responses", [])
|
||||
mask_segs = self.storage._data.get("masks", [])
|
||||
reward_segs = self.storage._data.get("rewards", [])
|
||||
|
||||
n_records = len(prompt_segs)
|
||||
self._records = []
|
||||
for i in range(n_records):
|
||||
self._records.append(
|
||||
{
|
||||
"prompts": prompt_segs[i],
|
||||
"responses": response_segs[i] if i < len(response_segs) else [],
|
||||
"masks": mask_segs[i] if i < len(mask_segs) else [],
|
||||
"rewards": reward_segs[i]
|
||||
if i < len(reward_segs)
|
||||
else torch.tensor([]),
|
||||
}
|
||||
)
|
||||
|
||||
@property
|
||||
def count(self) -> int:
|
||||
return len(self._records)
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self._records)
|
||||
|
||||
def __getitem__(self, index: int) -> Dict[str, Tensor]:
|
||||
begin_idx, end_idx = self.get_index(index)
|
||||
|
||||
prompts = self._fetch_data(begin_idx, end_idx, "prompts").to(dtype=torch.long)
|
||||
responses = self._fetch_data(begin_idx, end_idx, "responses").to(
|
||||
dtype=torch.long
|
||||
)
|
||||
masks = self._fetch_data(begin_idx, end_idx, "masks").to(dtype=torch.bool)
|
||||
rewards = self._fetch_data(begin_idx, end_idx, "rewards")
|
||||
|
||||
rec = self._records[index]
|
||||
return {
|
||||
"prompts": prompts,
|
||||
"responses": responses,
|
||||
"masks": masks,
|
||||
"rewards": rewards,
|
||||
"prompts": rec["prompts"].to(dtype=torch.long),
|
||||
"responses": [r.to(dtype=torch.long) for r in rec["responses"]],
|
||||
"masks": [m.to(dtype=torch.bool) for m in rec["masks"]],
|
||||
"rewards": rec["rewards"].to(dtype=torch.float32),
|
||||
}
|
||||
|
||||
@@ -148,26 +148,37 @@ class Store(ABC):
|
||||
|
||||
return results[0] if len(results) == 1 else torch.cat(results, dim=0)
|
||||
|
||||
def _normalize(self, raw: Dict[str, List[Tensor]]):
|
||||
def _normalize(self, raw: Dict[str, list]):
|
||||
"""Register segments and pre-compute cumulative lengths.
|
||||
|
||||
Does NOT concatenate — segments are kept as-is to avoid OOM on
|
||||
large datasets. Sets ``self._length`` to the minimum total
|
||||
element count across all keys.
|
||||
element count across all flat-tensor keys.
|
||||
|
||||
For GRPO multi-response keys, values may be ``List[List[Tensor]]``
|
||||
(one list of G tensors per record). These are stored as-is and
|
||||
excluded from the cumulative-length bookkeeping since they are
|
||||
accessed record-by-record via ``_data`` rather than via ``fetch``.
|
||||
"""
|
||||
flat_lengths = []
|
||||
for key, tensors in raw.items():
|
||||
self._data[key] = tensors
|
||||
if not tensors:
|
||||
self._cum[key] = []
|
||||
flat_lengths.append(0)
|
||||
continue
|
||||
# Skip nested lists (GRPO responses/masks) — record-level access
|
||||
if isinstance(tensors[0], list):
|
||||
self._cum[key] = []
|
||||
continue
|
||||
cum = []
|
||||
total = 0
|
||||
for t in tensors:
|
||||
total += t.shape[0]
|
||||
cum.append(total)
|
||||
self._cum[key] = cum
|
||||
self._length = (
|
||||
min((cum[-1] if cum else 0) for cum in self._cum.values())
|
||||
if self._cum
|
||||
else 0
|
||||
)
|
||||
flat_lengths.append(cum[-1] if cum else 0)
|
||||
self._length = min(flat_lengths) if flat_lengths else 0
|
||||
|
||||
|
||||
class StoreFactory(BaseFactory["Store"]):
|
||||
@@ -274,7 +285,13 @@ class JsonlStore(Store):
|
||||
for key, ids in result.items():
|
||||
if key not in raw:
|
||||
raw[key] = []
|
||||
raw[key].append(torch.tensor(ids, dtype=self._infer_dtype(ids)))
|
||||
if ids and isinstance(ids[0], list):
|
||||
# GRPO multi-response: List[List[int]] → List[Tensor]
|
||||
raw[key].append(
|
||||
[torch.tensor(sub, dtype=self._infer_dtype(sub)) for sub in ids]
|
||||
)
|
||||
else:
|
||||
raw[key].append(torch.tensor(ids, dtype=self._infer_dtype(ids)))
|
||||
|
||||
for jsonl_path in sorted(root.glob("*.jsonl")):
|
||||
with open(jsonl_path, "r", encoding="utf-8") as f:
|
||||
@@ -314,7 +331,7 @@ class JsonlStore(Store):
|
||||
|
||||
@staticmethod
|
||||
def _primary_ids(result: dict) -> List[int]:
|
||||
"""Return the first integer list in *result* as the primary id sequence."""
|
||||
"""Return the first flat integer list in *result* as the primary id sequence."""
|
||||
for val in result.values():
|
||||
if isinstance(val, list) and val and isinstance(val[0], int):
|
||||
return val
|
||||
|
||||
@@ -95,8 +95,15 @@ class SectionRenderer:
|
||||
return all_ids, loss_mask
|
||||
|
||||
def process_list_field(self, item: dict, sections: list, config, tokenizer):
|
||||
all_ids: list[int] = []
|
||||
loss_mask: list[int] = []
|
||||
"""Tokenize a list-valued field, preserving per-element boundaries.
|
||||
|
||||
Returns ``(list_of_id_lists, list_of_mask_lists)`` where each
|
||||
inner list corresponds to one element of the source list. This
|
||||
is critical for GRPO where each response must stay a separate
|
||||
sequence so the strategy can form a ``[G, R]`` tensor.
|
||||
"""
|
||||
per_item_ids: list[list[int]] = []
|
||||
per_item_masks: list[list[int]] = []
|
||||
|
||||
for sec in sections:
|
||||
field = sec["field"]
|
||||
@@ -108,17 +115,13 @@ class SectionRenderer:
|
||||
continue
|
||||
|
||||
for val in values:
|
||||
ids: list[int] = []
|
||||
mask: list[int] = []
|
||||
if use_template:
|
||||
if isinstance(val, list):
|
||||
wrapper = {field: val}
|
||||
self._append_template(
|
||||
wrapper,
|
||||
field,
|
||||
action,
|
||||
tokenizer,
|
||||
config,
|
||||
all_ids,
|
||||
loss_mask,
|
||||
wrapper, field, action, tokenizer, config, ids, mask
|
||||
)
|
||||
else:
|
||||
wrapper = {field: str(val)}
|
||||
@@ -130,17 +133,19 @@ class SectionRenderer:
|
||||
False,
|
||||
False,
|
||||
config,
|
||||
all_ids,
|
||||
loss_mask,
|
||||
ids,
|
||||
mask,
|
||||
)
|
||||
if ids:
|
||||
max_len = config.preprocessing.max_seq_len
|
||||
ids = ids[:max_len]
|
||||
mask = mask[: len(ids)]
|
||||
per_item_ids.append(ids)
|
||||
per_item_masks.append(mask)
|
||||
|
||||
max_len = config.preprocessing.max_seq_len
|
||||
all_ids = all_ids[:max_len]
|
||||
loss_mask = loss_mask[: len(all_ids)]
|
||||
|
||||
if not all_ids:
|
||||
if not per_item_ids:
|
||||
return None, None
|
||||
return all_ids, loss_mask
|
||||
return per_item_ids, per_item_masks
|
||||
|
||||
@staticmethod
|
||||
def is_value_section(sections: list) -> bool:
|
||||
@@ -282,10 +287,18 @@ class MultiOutputMaskBuilder(BaseMaskBuilder):
|
||||
ids, mask = self.renderer.process_list_field(
|
||||
item, sections, config, tokenizer
|
||||
)
|
||||
else:
|
||||
ids, mask = self.renderer.process_sections(
|
||||
item, sections, config, tokenizer, is_top_level=True
|
||||
)
|
||||
if ids is None:
|
||||
continue
|
||||
# ids is List[List[int]] — preserve per-response structure
|
||||
result[output_key] = ids
|
||||
if mask is not None:
|
||||
result[mask_key] = mask
|
||||
any_output = True
|
||||
continue
|
||||
|
||||
ids, mask = self.renderer.process_sections(
|
||||
item, sections, config, tokenizer, is_top_level=True
|
||||
)
|
||||
|
||||
if ids is None:
|
||||
continue
|
||||
|
||||
@@ -180,9 +180,24 @@ class Pipeline:
|
||||
dt = _STR_TO_DTYPE.get(
|
||||
self.config.output.dtype.get(key, "int32"), torch.int32
|
||||
)
|
||||
tensors[key] = [
|
||||
torch.tensor(list(chain.from_iterable(ids_list)), dtype=dt)
|
||||
]
|
||||
# GRPO multi-response keys store List[List[int]] per record
|
||||
# (responses/masks). Rewards store List[float] per record.
|
||||
# Both produce List[Tensor] (one tensor per record), but
|
||||
# responses need inner flattening while rewards do not.
|
||||
if ids_list and isinstance(ids_list[0], list):
|
||||
tensors[key] = [
|
||||
torch.tensor(
|
||||
list(chain.from_iterable(ids))
|
||||
if ids and isinstance(ids[0], list)
|
||||
else ids,
|
||||
dtype=dt,
|
||||
)
|
||||
for ids in ids_list
|
||||
]
|
||||
else:
|
||||
tensors[key] = [
|
||||
torch.tensor(list(chain.from_iterable(ids_list)), dtype=dt)
|
||||
]
|
||||
|
||||
if mode == "continuous" and original_sequences:
|
||||
pos_ids = self._position_id.generate(keys.get("sequence", []))
|
||||
|
||||
+290
-23
@@ -327,43 +327,82 @@ def test_normalize_mixed_empty_key():
|
||||
|
||||
|
||||
def test_grpo_dataset_dtype(base_test_env):
|
||||
"""GRPO dataset returns correct dtypes for per-record structured data."""
|
||||
from astrai.dataset.dataset import GRPODataset
|
||||
|
||||
test_dir = base_test_env["test_dir"]
|
||||
dummy_data = {
|
||||
"prompts": [torch.randint(0, 100, (100,), dtype=torch.int32)],
|
||||
"responses": [torch.randint(0, 100, (100,), dtype=torch.int32)],
|
||||
"masks": [torch.ones(100, dtype=torch.int32)],
|
||||
"rewards": [torch.ones(100, dtype=torch.float32)],
|
||||
}
|
||||
dataset = _make_seq_dataset(
|
||||
test_dir, "grpo_dtype", train_type="grpo", data=dummy_data, window_size=32
|
||||
)
|
||||
G = 4
|
||||
dataset = GRPODataset()
|
||||
dataset.storage = type(
|
||||
"FakeStore",
|
||||
(),
|
||||
{
|
||||
"keys": ["prompts", "responses", "masks", "rewards"],
|
||||
"_data": {
|
||||
"prompts": [torch.randint(0, 100, (10,), dtype=torch.int32)],
|
||||
"responses": [
|
||||
[torch.randint(0, 100, (5,), dtype=torch.int32) for _ in range(G)]
|
||||
],
|
||||
"masks": [[torch.ones(5, dtype=torch.int32) for _ in range(G)]],
|
||||
"rewards": [torch.rand(G, dtype=torch.float32)],
|
||||
},
|
||||
},
|
||||
)()
|
||||
dataset._build_records()
|
||||
item = dataset[0]
|
||||
|
||||
assert item["prompts"].dtype == torch.long
|
||||
assert item["responses"].dtype == torch.long
|
||||
assert item["masks"].dtype == torch.bool
|
||||
assert all(r.dtype == torch.long for r in item["responses"])
|
||||
assert all(m.dtype == torch.bool for m in item["masks"])
|
||||
assert item["rewards"].dtype == torch.float32
|
||||
|
||||
|
||||
def test_grpo_dataset_load(base_test_env):
|
||||
"""GRPO dataset loads record-structured data with per-response boundaries."""
|
||||
from astrai.dataset.dataset import GRPODataset
|
||||
|
||||
test_dir = base_test_env["test_dir"]
|
||||
dummy_data = {
|
||||
"prompts": [_rand_seq(200)],
|
||||
"responses": [_rand_seq(200)],
|
||||
"masks": [torch.ones(200, dtype=torch.int64)],
|
||||
"rewards": [torch.rand(200, dtype=torch.float32)],
|
||||
}
|
||||
dataset = _make_seq_dataset(
|
||||
test_dir, "grpo_test", train_type="grpo", data=dummy_data
|
||||
)
|
||||
assert len(dataset) > 0
|
||||
G = 3
|
||||
prompt_len = 8
|
||||
resp_lens = [5, 7, 4]
|
||||
dataset = GRPODataset()
|
||||
dataset.storage = type(
|
||||
"FakeStore",
|
||||
(),
|
||||
{
|
||||
"keys": ["prompts", "responses", "masks", "rewards"],
|
||||
"_data": {
|
||||
"prompts": [torch.randint(0, 100, (prompt_len,))],
|
||||
"responses": [[torch.randint(0, 100, (rl,)) for rl in resp_lens]],
|
||||
"masks": [[torch.ones(rl, dtype=torch.int64) for rl in resp_lens]],
|
||||
"rewards": [torch.tensor([0.9, 0.3, 0.7], dtype=torch.float32)],
|
||||
},
|
||||
},
|
||||
)()
|
||||
dataset._build_records()
|
||||
|
||||
assert len(dataset) == 1
|
||||
item = dataset[0]
|
||||
assert "prompts" in item
|
||||
assert "responses" in item
|
||||
assert "masks" in item
|
||||
assert "rewards" in item
|
||||
assert item["prompts"].shape[0] == 64
|
||||
assert item["responses"].shape[0] == 64
|
||||
|
||||
# Prompts is 1-D
|
||||
assert item["prompts"].shape == (prompt_len,)
|
||||
|
||||
# Responses is a list of G tensors with correct lengths
|
||||
assert len(item["responses"]) == G
|
||||
for i, r in enumerate(item["responses"]):
|
||||
assert r.shape == (resp_lens[i],)
|
||||
|
||||
# Masks align with responses
|
||||
assert len(item["masks"]) == G
|
||||
for i, m in enumerate(item["masks"]):
|
||||
assert m.shape == (resp_lens[i],)
|
||||
|
||||
# Rewards has G elements
|
||||
assert item["rewards"].shape == (G,)
|
||||
|
||||
|
||||
def test_detect_format_bin_dir(base_test_env):
|
||||
@@ -621,3 +660,231 @@ def test_jsonl_store_pipeline_config_roundtrip(base_test_env):
|
||||
config = PipelineConfig.from_dict(raw)
|
||||
assert config.output.position_ids_mode == "doc_reset"
|
||||
assert config.preprocessing.max_seq_len == 64
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# GRPO end-to-end: builder → JsonlStore → GRPODataset → collate_fn
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _write_grpo_jsonl(test_dir, tokenizer_path, records):
|
||||
"""Write a GRPO JSONL dataset directory with config."""
|
||||
data_dir = os.path.join(test_dir, "grpo_jsonl")
|
||||
os.makedirs(data_dir, exist_ok=True)
|
||||
|
||||
with open(os.path.join(data_dir, "data.jsonl"), "w", encoding="utf-8") as f:
|
||||
for rec in records:
|
||||
f.write(json.dumps(rec, ensure_ascii=False) + "\n")
|
||||
|
||||
config = {
|
||||
"tokenizer_path": tokenizer_path,
|
||||
"version": 1,
|
||||
"input": {
|
||||
"sources": {
|
||||
"prompts": {
|
||||
"sections": [
|
||||
{
|
||||
"field": "prompt",
|
||||
"action": "mask",
|
||||
"add_special_tokens": True,
|
||||
}
|
||||
]
|
||||
},
|
||||
"responses": {
|
||||
"sections": [{"field": "responses", "action": "train"}],
|
||||
"list_field": True,
|
||||
"mask_key": "masks",
|
||||
},
|
||||
"rewards": {
|
||||
"sections": [{"field": "rewards", "action": "value"}],
|
||||
},
|
||||
}
|
||||
},
|
||||
"mask": {"user": "mask", "assistant": "train"},
|
||||
"mask_default": "mask",
|
||||
"preprocessing": {"max_seq_len": 128},
|
||||
"output": {"position_ids_mode": "none"},
|
||||
}
|
||||
|
||||
with open(
|
||||
os.path.join(data_dir, "dataset_config.json"), "w", encoding="utf-8"
|
||||
) as f:
|
||||
json.dump(config, f, ensure_ascii=False, indent=2)
|
||||
|
||||
return data_dir
|
||||
|
||||
|
||||
def test_grpo_builder_preserves_response_boundaries(base_test_env):
|
||||
"""MultiOutputMaskBuilder with list_field returns List[List[int]] for responses."""
|
||||
from astrai.preprocessing.builder import SectionedMaskBuilder
|
||||
from tests.data.conftest import make_grpo_no_template_config
|
||||
|
||||
tokenizer = base_test_env["tokenizer"]
|
||||
tokenizer_path = _save_test_tokenizer(base_test_env["test_dir"], tokenizer)
|
||||
|
||||
builder = SectionedMaskBuilder()
|
||||
config = make_grpo_no_template_config()
|
||||
config.preprocessing.max_seq_len = 128
|
||||
|
||||
item = {
|
||||
"prompt": "What is 2+2?",
|
||||
"responses": ["4", "four", "2+2=4"],
|
||||
"rewards": [0.9, 0.1, 0.5],
|
||||
}
|
||||
|
||||
result = builder.build(item, config, tokenizer)
|
||||
assert result is not None
|
||||
|
||||
# prompts should be flat list of ints
|
||||
assert isinstance(result["prompts"], list)
|
||||
assert isinstance(result["prompts"][0], int)
|
||||
|
||||
# responses should be list of lists (one per response)
|
||||
assert isinstance(result["responses"], list)
|
||||
assert isinstance(result["responses"][0], list)
|
||||
assert isinstance(result["responses"][0][0], int)
|
||||
assert len(result["responses"]) == 3
|
||||
|
||||
# masks should match responses structure
|
||||
assert isinstance(result["masks"], list)
|
||||
assert len(result["masks"]) == 3
|
||||
for i in range(3):
|
||||
assert len(result["masks"][i]) == len(result["responses"][i])
|
||||
|
||||
# rewards should be flat list of floats
|
||||
assert isinstance(result["rewards"], list)
|
||||
assert all(isinstance(r, float) for r in result["rewards"])
|
||||
assert len(result["rewards"]) == 3
|
||||
|
||||
|
||||
def test_grpo_end_to_end_jsonl(base_test_env):
|
||||
"""Full GRPO pipeline: JSONL → JsonlStore → GRPODataset → collate_fn."""
|
||||
from astrai.dataset.dataset import grpo_collate_fn
|
||||
|
||||
test_dir = base_test_env["test_dir"]
|
||||
tokenizer = base_test_env["tokenizer"]
|
||||
tokenizer_path = _save_test_tokenizer(test_dir, tokenizer)
|
||||
|
||||
records = [
|
||||
{
|
||||
"prompt": "What is 2+2?",
|
||||
"responses": ["4", "four", "The answer is 4"],
|
||||
"rewards": [0.9, 0.1, 0.5],
|
||||
},
|
||||
{
|
||||
"prompt": "Write a haiku",
|
||||
"responses": ["Leaves fall", "Cherry blossoms bloom in spring"],
|
||||
"rewards": [0.3, 0.8],
|
||||
},
|
||||
]
|
||||
|
||||
data_dir = _write_grpo_jsonl(test_dir, tokenizer_path, records)
|
||||
|
||||
dataset = DatasetFactory.load("grpo", data_dir, window_size=0)
|
||||
assert len(dataset) == 2
|
||||
|
||||
# Item 0: 3 responses
|
||||
item0 = dataset[0]
|
||||
assert item0["prompts"].ndim == 1
|
||||
assert len(item0["responses"]) == 3
|
||||
assert len(item0["masks"]) == 3
|
||||
assert item0["rewards"].shape == (3,)
|
||||
for r, m in zip(item0["responses"], item0["masks"]):
|
||||
assert r.shape == m.shape
|
||||
|
||||
# Item 1: 2 responses (different group size)
|
||||
item1 = dataset[1]
|
||||
assert len(item1["responses"]) == 2
|
||||
assert item1["rewards"].shape == (2,)
|
||||
|
||||
# Collate: batch records with same G (item0 has G=3)
|
||||
batch = grpo_collate_fn([item0, item0])
|
||||
assert batch["prompts"].shape[0] == 2
|
||||
assert batch["responses"].ndim == 3
|
||||
assert batch["responses"].shape[0] == 2
|
||||
assert batch["responses"].shape[1] == 3 # G=3
|
||||
assert batch["masks"].shape == batch["responses"].shape
|
||||
assert batch["rewards"].shape == (2, 3)
|
||||
|
||||
|
||||
def test_grpo_collate_variable_lengths():
|
||||
"""collate_fn pads variable-length responses to [B, G, R_max]."""
|
||||
from astrai.dataset.dataset import grpo_collate_fn
|
||||
|
||||
batch = [
|
||||
{
|
||||
"prompts": torch.tensor([1, 2, 3]),
|
||||
"responses": [torch.tensor([4, 5]), torch.tensor([6, 7, 8, 9])],
|
||||
"masks": [torch.tensor([1, 1]), torch.tensor([1, 1, 1, 1])],
|
||||
"rewards": torch.tensor([0.9, 0.1]),
|
||||
},
|
||||
{
|
||||
"prompts": torch.tensor([10, 11]),
|
||||
"responses": [torch.tensor([12]), torch.tensor([13, 14, 15])],
|
||||
"masks": [torch.tensor([1]), torch.tensor([1, 1, 1])],
|
||||
"rewards": torch.tensor([0.5, 0.5]),
|
||||
},
|
||||
]
|
||||
|
||||
result = grpo_collate_fn(batch)
|
||||
|
||||
assert result["prompts"].shape == (2, 3) # B=2, P_max=3
|
||||
assert result["responses"].shape == (2, 2, 4) # B=2, G=2, R_max=4
|
||||
assert result["masks"].shape == (2, 2, 4)
|
||||
assert result["rewards"].shape == (2, 2)
|
||||
|
||||
# Check padding: item 1 prompt is length 2, padded to 3
|
||||
assert result["prompts"][1, 2] == 0
|
||||
|
||||
# Check response content: item 0, response 0 is [4,5] padded to 4
|
||||
assert result["responses"][0, 0, 0] == 4
|
||||
assert result["responses"][0, 0, 1] == 5
|
||||
assert result["responses"][0, 0, 2] == 0 # padded
|
||||
assert result["masks"][0, 0, 2] == False # padded
|
||||
|
||||
# Check response content: item 0, response 1 is [6,7,8,9] no padding
|
||||
assert result["responses"][0, 1, 3] == 9
|
||||
assert result["masks"][0, 1, 3] == True
|
||||
|
||||
|
||||
def test_grpo_multiple_records(base_test_env):
|
||||
"""GRPODataset loads multiple records with correct structure."""
|
||||
from astrai.dataset.dataset import GRPODataset
|
||||
|
||||
G = 4
|
||||
n_records = 5
|
||||
|
||||
dummy_responses = [
|
||||
[torch.randint(0, 100, (np.random.randint(3, 8),)) for _ in range(G)]
|
||||
for _ in range(n_records)
|
||||
]
|
||||
dataset = GRPODataset()
|
||||
dataset.storage = type(
|
||||
"FakeStore",
|
||||
(),
|
||||
{
|
||||
"keys": ["prompts", "responses", "masks", "rewards"],
|
||||
"_data": {
|
||||
"prompts": [torch.randint(0, 100, (10,)) for _ in range(n_records)],
|
||||
"responses": dummy_responses,
|
||||
"masks": [
|
||||
[torch.ones(r.shape[0], dtype=torch.int64) for r in resps]
|
||||
for resps in dummy_responses
|
||||
],
|
||||
"rewards": [
|
||||
torch.rand(G, dtype=torch.float32) for _ in range(n_records)
|
||||
],
|
||||
},
|
||||
},
|
||||
)()
|
||||
dataset._build_records()
|
||||
|
||||
assert len(dataset) == n_records
|
||||
|
||||
for i in range(n_records):
|
||||
item = dataset[i]
|
||||
assert len(item["responses"]) == G
|
||||
assert len(item["masks"]) == G
|
||||
assert item["rewards"].shape == (G,)
|
||||
for g in range(G):
|
||||
assert item["responses"][g].shape == item["masks"][g].shape
|
||||
|
||||
@@ -349,7 +349,17 @@ def test_grpo_basic(chat_tokenizer, builder):
|
||||
assert "responses" in result
|
||||
assert "masks" in result
|
||||
assert "rewards" in result
|
||||
assert len(result["responses"]) == len(result["masks"])
|
||||
|
||||
# responses is List[List[int]] — one per response
|
||||
assert len(result["responses"]) == 4
|
||||
assert all(isinstance(r, list) for r in result["responses"])
|
||||
assert all(isinstance(r[0], int) for r in result["responses"])
|
||||
|
||||
# masks is List[List[int]] — one per response, matching length
|
||||
assert len(result["masks"]) == 4
|
||||
for i in range(4):
|
||||
assert len(result["masks"][i]) == len(result["responses"][i])
|
||||
|
||||
assert result["rewards"] == [1.0, 0.5, 0.8, 0.2]
|
||||
|
||||
|
||||
@@ -362,8 +372,11 @@ def test_grpo_response_tokens_all_trained(chat_tokenizer, builder):
|
||||
}
|
||||
result = builder.build(item, config, chat_tokenizer)
|
||||
masks = result["masks"]
|
||||
assert all(m == 1 for m in masks)
|
||||
assert len(masks) == len(result["responses"])
|
||||
# masks is List[List[int]] — each response's mask should be all 1s
|
||||
assert len(masks) == 2
|
||||
for m in masks:
|
||||
assert all(v == 1 for v in m)
|
||||
assert len(m) == len(result["responses"][masks.index(m)])
|
||||
|
||||
|
||||
def test_grpo_single_reward(chat_tokenizer, builder):
|
||||
|
||||
Reference in New Issue
Block a user