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:
2026-07-17 14:34:41 +08:00
parent c17aa0dc54
commit a1ea26d367
7 changed files with 502 additions and 75 deletions
+34 -21
View File
@@ -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