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

astrai/preprocessing/builder.py

@@ -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

astrai/preprocessing/pipeline.py

@@ -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", []))