feat: add record-mode to Store for DPO/GRPO

- Store gains fetch_record/num_records alongside stream fetch/__len__
- save_bin/load_bin support per-record offsets via record_keys param
- H5Store/MmapStore/JsonlStore all support dual stream+record access
- DPODataset/GRPODataset use fetch_record, no cross-record concat
- dpo_collate_fn + collate_fn wired through TrainConfig
- fixes attention context leakage in DPO from windowed concatenation

astrai/serialization/__init__.py

@@ -19,6 +19,7 @@ from astrai.serialization.checkpoint import (
 )
 from astrai.serialization.dataset import (
     load_bin,
+    load_bin_offsets,
     load_h5,
     save_bin,
     save_h5,
@@ -37,6 +38,7 @@ __all__ = [
     "save_safetensors",
     "save_torch",
     "load_bin",
+    "load_bin_offsets",
     "load_h5",
     "save_bin",
     "save_h5",

astrai/serialization/dataset.py

@@ -3,7 +3,7 @@
 import json
 import os
 from pathlib import Path
-from typing import Dict, List
+from typing import Any, Dict, List, Optional
 
 import h5py
 import numpy as np
@@ -50,12 +50,43 @@ def load_h5(file_path: str, share_memory=True) -> Dict[str, List[Tensor]]:
     return tensor_group
 
 
-def save_bin(file_path: str, tensor_group: Dict[str, List[Tensor]]):
+def save_bin(
+    file_path: str,
+    tensor_group: Dict[str, List[Tensor]],
+    record_keys: Optional[List[str]] = None,
+):
+    """Save tensors as memory-mapped binary files.
+
+    When *record_keys* is provided, those keys are written with per-record
+    cumulative offsets in ``meta.json`` so that ``MmapStore.fetch_record``
+    can slice individual records from the concatenated binary without
+    cross-record concatenation.  Keys not in *record_keys* (e.g. SEQ
+    ``sequence``) are written as a single contiguous stream without
+    offsets, preserving backward compatibility.
+
+    Nested keys (``List[List[Tensor]]`` such as GRPO ``responses``) are
+    not supported in bin format — use H5 for those.
+    """
     os.makedirs(file_path, exist_ok=True)
+    record_keys = set(record_keys or [])
     meta = {}
     for key, tensors in tensor_group.items():
+        if tensors and isinstance(tensors[0], list):
+            raise ValueError(
+                f"Nested key '{key}' (List[List[Tensor]]) is not supported "
+                f"in bin format. Use H5 or JSONL storage instead."
+            )
         cat = torch.cat(tensors, dim=0)
-        meta[key] = {"shape": list(cat.shape), "dtype": str(cat.dtype).split(".")[-1]}
+        entry: Dict[str, Any] = {
+            "shape": list(cat.shape),
+            "dtype": str(cat.dtype).split(".")[-1],
+        }
+        if key in record_keys:
+            offsets = [0]
+            for t in tensors:
+                offsets.append(offsets[-1] + t.shape[0])
+            entry["offsets"] = offsets
+        meta[key] = entry
         np.asarray(cat.cpu().numpy()).tofile(os.path.join(file_path, f"{key}.bin"))
     with open(os.path.join(file_path, "meta.json"), "w") as f:
         json.dump(meta, f)
@@ -74,3 +105,19 @@ def load_bin(file_path: str) -> Dict[str, List[Tensor]]:
         )
         segments[key] = [torch.from_numpy(arr)]
     return segments
+
+
+def load_bin_offsets(file_path: str) -> Dict[str, List[int]]:
+    """Read per-record cumulative offsets from ``meta.json``.
+
+    Returns an empty dict when no key has offsets (legacy bin files),
+    in which case record-mode access falls back to per-record segment
+    indexing (H5/JSONL layout).
+    """
+    with open(os.path.join(file_path, "meta.json"), "r") as f:
+        meta = json.load(f)
+    offsets: Dict[str, List[int]] = {}
+    for key, info in meta.items():
+        if "offsets" in info:
+            offsets[key] = info["offsets"]
+    return offsets