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/__init__.py

@@ -12,7 +12,7 @@ from astrai.config import (
 from astrai.dataset import (
     BaseDataset,
     DatasetFactory,
-    ResumableDistributedSampler,
+    RDSampler,
     Store,
     StoreFactory,
 )
@@ -77,7 +77,7 @@ __all__ = [
     "Pipeline",
     "PipelineConfig",
     "ProtocolHandler",
-    "ResumableDistributedSampler",
+    "RDSampler",
     "SamplingPipeline",
     "SchedulerFactory",
     "Store",

astrai/config/train_config.py

@@ -87,6 +87,10 @@ class TrainConfig(BaseConfig):
     pin_memory: bool = field(
         default=False, metadata={"help": "Pin memory for dataloader."}
     )
+    collate_fn: Optional[Callable[[List[Any]], Any]] = field(
+        default=None,
+        metadata={"help": "Collate function for dataloader (e.g. dpo_collate_fn)."},
+    )
 
     # distributed training
     nprocs: int = field(

astrai/dataset/__init__.py

@@ -1,9 +1,10 @@
 from astrai.dataset.dataset import (
     BaseDataset,
     DatasetFactory,
+    dpo_collate_fn,
     grpo_collate_fn,
 )
-from astrai.dataset.sampler import ResumableDistributedSampler
+from astrai.dataset.sampler import RDSampler
 from astrai.dataset.storage import (
     H5Store,
     JsonlStore,
@@ -22,6 +23,7 @@ from astrai.serialization import (
 __all__ = [
     "BaseDataset",
     "DatasetFactory",
+    "dpo_collate_fn",
     "grpo_collate_fn",
     "Store",
     "StoreFactory",
@@ -33,5 +35,5 @@ __all__ = [
     "load_h5",
     "save_bin",
     "load_bin",
-    "ResumableDistributedSampler",
+    "RDSampler",
 ]

astrai/dataset/dataset.py

@@ -15,6 +15,46 @@ from astrai.dataset.storage import (
 from astrai.factory import BaseFactory
 
 
+def dpo_collate_fn(batch: List[Dict[str, Tensor]]) -> Dict[str, Tensor]:
+    """Collate variable-length DPO samples into padded 2-D tensors.
+
+    Input: list of dicts, each with:
+      - chosen:        [C_i]
+      - rejected:      [R_i]
+      - chosen_mask:   [C_i]
+      - rejected_mask: [R_i]
+
+    Output (padded to the max length across chosen/rejected within the batch):
+      - chosen:        [B, S_max]
+      - rejected:      [B, S_max]
+      - chosen_mask:   [B, S_max]
+      - rejected_mask: [B, S_max]
+    """
+    B = len(batch)
+    S_max = max(b["chosen"].size(0) for b in batch)
+    S_max = max(S_max, max(b["rejected"].size(0) for b in batch))
+
+    chosen = torch.zeros(B, S_max, dtype=torch.long)
+    rejected = torch.zeros(B, S_max, dtype=torch.long)
+    chosen_mask = torch.zeros(B, S_max, dtype=torch.bool)
+    rejected_mask = torch.zeros(B, S_max, dtype=torch.bool)
+
+    for i, b in enumerate(batch):
+        c_len = b["chosen"].size(0)
+        r_len = b["rejected"].size(0)
+        chosen[i, :c_len] = b["chosen"]
+        rejected[i, :r_len] = b["rejected"]
+        chosen_mask[i, :c_len] = b["chosen_mask"]
+        rejected_mask[i, :r_len] = b["rejected_mask"]
+
+    return {
+        "chosen": chosen,
+        "rejected": rejected,
+        "chosen_mask": chosen_mask,
+        "rejected_mask": rejected_mask,
+    }
+
+
 def grpo_collate_fn(batch: List[Dict[str, Tensor]]) -> Dict[str, Tensor]:
     """Collate variable-length GRPO samples into padded 3-D tensors.
 
@@ -262,32 +302,61 @@ class SFTDataset(BaseDataset):
 
 @DatasetFactory.register("dpo")
 class DPODataset(BaseDataset):
-    """Dataset for Direct Preference Optimization training."""
+    """Record-structured dataset for Direct Preference Optimization.
+
+    Each sample is one preference pair (chosen + rejected) and is an
+    independent training unit — no windowing, stride, or cross-record
+    concatenation.  This keeps each sequence self-contained so attention
+    never leaks across preference pairs.
+
+    Delegates record access to ``Store.fetch_record``, which works with
+    any storage backend (H5 per-record datasets, bin+offsets memmap, or
+    JSONL on-the-fly tokenization).
+    """
+
+    def __init__(self, window_size: int = 0, stride: int = 0, **kwargs):
+        super().__init__(window_size=window_size, stride=stride or window_size)
 
     @property
     def required_keys(self) -> List[str]:
         return ["chosen", "rejected", "chosen_mask", "rejected_mask"]
 
-    def _fetch_data(self, begin_idx: int, end_idx: int, key: str) -> Tensor:
+    def load(self, load_path: str, storage_type: Optional[str] = None, **kwargs):
-        return self.storage.fetch(begin_idx, end_idx, key)
+        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()
 
-    def __getitem__(self, index: int):
+    def _validate_keys(self):
-        begin_idx, end_idx = self.get_index(index)
+        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"DPODataset requires keys {self.required_keys}, "
+                f"but storage only has {sorted(actual_keys)}. Missing: {missing}"
+            )
 
-        chosen = self._fetch_data(begin_idx, end_idx, "chosen").to(dtype=torch.long)
+    @property
-        rejected = self._fetch_data(begin_idx, end_idx, "rejected").to(dtype=torch.long)
+    def count(self) -> int:
-        chosen_mask = self._fetch_data(begin_idx, end_idx, "chosen_mask").to(
+        return self.storage.num_records
-            dtype=torch.bool
-        )
-        rejected_mask = self._fetch_data(begin_idx, end_idx, "rejected_mask").to(
-            dtype=torch.bool
-        )
 
+    def __len__(self) -> int:
+        return self.storage.num_records
+
+    def __getitem__(self, index: int) -> Dict[str, Tensor]:
         return {
-            "chosen": chosen,
+            "chosen": self.storage.fetch_record(index, "chosen").to(dtype=torch.long),
-            "rejected": rejected,
+            "rejected": self.storage.fetch_record(index, "rejected").to(
-            "chosen_mask": chosen_mask,
+                dtype=torch.long
-            "rejected_mask": rejected_mask,
+            ),
+            "chosen_mask": self.storage.fetch_record(index, "chosen_mask").to(
+                dtype=torch.bool
+            ),
+            "rejected_mask": self.storage.fetch_record(index, "rejected_mask").to(
+                dtype=torch.bool
+            ),
         }
 
 
@@ -310,7 +379,6 @@ class GRPODataset(BaseDataset):
 
     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]:
@@ -323,7 +391,6 @@ class GRPODataset(BaseDataset):
         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)
@@ -334,44 +401,21 @@ class GRPODataset(BaseDataset):
                 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)
+        return self.storage.num_records
 
     def __len__(self) -> int:
-        return len(self._records)
+        return self.storage.num_records
 
     def __getitem__(self, index: int) -> Dict[str, Tensor]:
-        rec = self._records[index]
+        prompts = self.storage.fetch_record(index, "prompts")
+        responses = self.storage.fetch_record(index, "responses")
+        masks = self.storage.fetch_record(index, "masks")
+        rewards = self.storage.fetch_record(index, "rewards")
         return {
-            "prompts": rec["prompts"].to(dtype=torch.long),
+            "prompts": prompts.to(dtype=torch.long),
-            "responses": [r.to(dtype=torch.long) for r in rec["responses"]],
+            "responses": [r.to(dtype=torch.long) for r in responses],
-            "masks": [m.to(dtype=torch.bool) for m in rec["masks"]],
+            "masks": [m.to(dtype=torch.bool) for m in masks],
-            "rewards": rec["rewards"].to(dtype=torch.float32),
+            "rewards": rewards.to(dtype=torch.float32),
         }

astrai/dataset/sampler.py

@@ -5,7 +5,15 @@ import torch.distributed as dist
 from torch.utils.data import Dataset, Sampler
 
 
-class ResumableDistributedSampler(Sampler[int]):
+class RDSampler(Sampler[int]):
+    """Resumable Distributed Sampler.
+
+    A distributed sampler that supports checkpoint-based resume: iteration
+    state (epoch, position) is tracked so training can continue from the
+    exact sample after a restart.  Shards the dataset across
+    ``dist.world_size`` replicas with optional shuffling.
+    """
+
     def __init__(
         self,
         data_source: Dataset,

astrai/dataset/storage.py

@@ -23,7 +23,7 @@ import json
 import logging
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import Dict, List, Union
+from typing import Dict, List, Optional, Union
 
 import torch
 from torch import Tensor
@@ -34,6 +34,7 @@ from astrai.preprocessing.builder import MaskBuilderFactory
 from astrai.preprocessing.position_id import PositionIdStrategyFactory
 from astrai.serialization import (
     load_bin,
+    load_bin_offsets,
     load_h5,
 )
 from astrai.tokenize import AutoTokenizer
@@ -90,11 +91,18 @@ def detect_format(load_path: str) -> str:
 
 
 class Store(ABC):
-    """String keys -> segmented tensors with ``fetch(begin, end, keys)``.
+    """String keys -> segmented tensors with two access modes.
 
-    Each key maps to one or more tensor segments (no forced concatenation).
+    Stream mode (SEQ/SFT):
-    ``len(store)`` returns ``self._length`` (explicit, O(1)), the minimum
+        ``fetch(begin, end, keys)`` slices across concatenated segments,
-    total element count across all keys.
+        transparently ``torch.cat``-ing across segment boundaries.
+        ``len(store)`` returns total token count.
+
+    Record mode (DPO/GRPO):
+        ``fetch_record(index, keys)`` returns the i-th record without
+        cross-record concatenation.  ``num_records`` returns the record
+        count.  Backed by either per-record segment lists (H5/JSONL) or
+        a single concatenated segment plus per-record offsets (bin).
 
     Subclasses fill ``self._data`` and ``self._cum`` during ``load()``
     via ``_normalize()``.
@@ -103,7 +111,9 @@ class Store(ABC):
     def __init__(self):
         self._data: Dict[str, List[Tensor]] = {}
         self._cum: Dict[str, List[int]] = {}
+        self._offsets: Dict[str, List[int]] = {}
         self._length: int = 0
+        self._num_records: int = 0
 
     @abstractmethod
     def load(self, path: str) -> None:
@@ -148,17 +158,71 @@ class Store(ABC):
 
         return results[0] if len(results) == 1 else torch.cat(results, dim=0)
 
-    def _normalize(self, raw: Dict[str, list]):
+    @property
-        """Register segments and pre-compute cumulative lengths.
+    def num_records(self) -> int:
+        return self._num_records
 
-        Does NOT concatenate — segments are kept as-is to avoid OOM on
+    def fetch_record(
-        large datasets.  Sets ``self._length`` to the minimum total
+        self,
-        element count across all flat-tensor keys.
+        index: int,
+        keys: Union[str, List[str]],
+    ):
+        """Fetch the *index*-th record without cross-record concatenation.
 
-        For GRPO multi-response keys, values may be ``List[List[Tensor]]``
+        Returns a tensor (flat key) or ``List[Tensor]`` (nested key such as
-        (one list of G tensors per record).  These are stored as-is and
+        GRPO ``responses``).
-        excluded from the cumulative-length bookkeeping since they are
+        """
-        accessed record-by-record via ``_data`` rather than via ``fetch``.
+        if not self._data:
+            raise RuntimeError("Store not loaded")
+        if not 0 <= index < self._num_records:
+            raise ValueError(
+                f"Record index out of bounds: {index}, num_records={self._num_records}"
+            )
+        if isinstance(keys, str):
+            return self._fetch_record_key(keys, index)
+        return {k: self._fetch_record_key(k, index) for k in keys}
+
+    def _fetch_record_key(self, key: str, index: int):
+        """Return the *index*-th record for *key*.
+
+        Two storage layouts are supported:
+
+        - **bin + offsets**: ``_data[key]`` is ``[single_long_segment]``;
+          ``_offsets[key]`` holds cumulative per-record offsets.  The record
+          is sliced as ``segment[offsets[i]:offsets[i+1]]``.
+        - **h5 / jsonl**: ``_data[key]`` is ``[t0, t1, ...]`` with one tensor
+          (or nested list of tensors for GRPO) per record.  Direct indexing.
+        """
+        offsets = self._offsets.get(key)
+        if offsets:
+            start = offsets[index]
+            end = (
+                offsets[index + 1]
+                if index + 1 < len(offsets)
+                else self._data[key][0].shape[0]
+            )
+            return self._data[key][0][start:end]
+        return self._data[key][index]
+
+    def _normalize(
+        self,
+        raw: Dict[str, list],
+        offsets: Optional[Dict[str, List[int]]] = None,
+        per_record: bool = False,
+    ):
+        """Register segments and pre-compute indices for both access modes.
+
+        Stream mode: ``_cum[key]`` accumulates per-segment lengths so
+        ``_fetch_key`` can bisect across segments without concatenation.
+
+        Record mode: if *offsets* is provided (bin layout), ``_offsets[key]``
+        stores cumulative per-record offsets into the single concatenated
+        segment.  Otherwise (h5/jsonl layout), ``_data[key]`` is already a
+        per-record list and ``fetch_record`` indexes it directly.
+
+        Nested keys (GRPO ``responses``/``masks`` as ``List[List[Tensor]]``)
+        are stored as-is and excluded from both cumulative bookkeepings —
+        they are only accessed record-by-record.
         """
         flat_lengths = []
         for key, tensors in raw.items():
@@ -180,6 +244,42 @@ class Store(ABC):
             flat_lengths.append(cum[-1] if cum else 0)
         self._length = min(flat_lengths) if flat_lengths else 0
 
+        # Record-mode offsets (bin layout).  Only valid when each key is a
+        # single concatenated segment — multi-shard bin + offsets is not
+        # supported (merge shards or use H5/JSONL instead).
+        valid_offsets: Dict[str, List[int]] = {}
+        if offsets:
+            for key, off in offsets.items():
+                segs = self._data.get(key, [])
+                if len(segs) == 1 and len(off) > 1:
+                    valid_offsets[key] = off
+                elif len(segs) > 1:
+                    logger.warning(
+                        "Key '%s' has %d segments with offsets — record mode "
+                        "disabled for this key (multi-shard bin+offsets not "
+                        "supported). Merge shards or use H5/JSONL.",
+                        key,
+                        len(segs),
+                    )
+        self._offsets = valid_offsets
+        if valid_offsets:
+            record_counts = [len(v) - 1 for v in valid_offsets.values()]
+            self._num_records = min(record_counts) if record_counts else 0
+        elif per_record:
+            # H5/JSONL layout: _data[key] is a per-record list where each
+            # segment is one record.  Even a single segment counts as one
+            # record.
+            per_record_counts = []
+            for key, tensors in self._data.items():
+                if not tensors or isinstance(tensors[0], list):
+                    continue
+                per_record_counts.append(len(tensors))
+            self._num_records = min(per_record_counts) if per_record_counts else 0
+        else:
+            # bin layout without offsets: _data[key] is [concatenated_stream].
+            # Cannot determine record boundaries — stream mode only.
+            self._num_records = 0
+
 
 class StoreFactory(BaseFactory["Store"]):
     """Factory for creating Store instances by type name.
@@ -194,10 +294,15 @@ class StoreFactory(BaseFactory["Store"]):
 
 @StoreFactory.register("h5")
 class H5Store(Store):
-    """HDF5-based storage backend (pre-tokenized data)."""
+    """HDF5-based storage backend (pre-tokenized data).
+
+    Each key is stored as a group of per-record datasets (``data_0``,
+    ``data_1``, …), so record mode indexes ``_data[key]`` directly.
+    Stream mode concatenates across records via ``_cum``.
+    """
 
     def load(self, path: str):
-        self._normalize(load_h5(path))
+        self._normalize(load_h5(path), per_record=True)
 
 
 @StoreFactory.register("bin")
@@ -208,10 +313,15 @@ class MmapStore(Store):
     No per-process memory duplication — all DataLoader workers share the
     same OS page-cache pages.
 
+    When ``meta.json`` contains per-record ``offsets`` for a key (written
+    via ``save_bin(..., record_keys=...)``), record-mode access slices
+    individual records from the concatenated memmap.  Legacy bin files
+    without offsets only support stream mode.
+
     Format on disk::
 
         data_root/
-          meta.json          # {key: {shape, dtype}, ...}
+          meta.json          # {key: {shape, dtype, offsets?}, ...}
           <key>.bin          # raw numpy array, one per key
     """
 
@@ -219,18 +329,24 @@ class MmapStore(Store):
         self._mmap_refs = []
         root = Path(path)
         all_raw: Dict[str, List[Tensor]] = {}
+        all_offsets: Dict[str, List[int]] = {}
         meta_paths = [
             Path(p) for p in glob.glob(str(root / "**" / "meta.json"), recursive=True)
         ]
         for meta_path in meta_paths:
             raw = load_bin(str(meta_path.parent))
+            off = load_bin_offsets(str(meta_path.parent))
             for key, tensors in raw.items():
                 if key not in all_raw:
                     all_raw[key] = []
                 all_raw[key].extend(tensors)
+            for key, o in off.items():
+                if key not in all_offsets:
+                    all_offsets[key] = []
+                all_offsets[key].extend(o)
         if not meta_paths:
             raise FileNotFoundError(f"No meta.json found under {path}")
-        self._normalize(all_raw)
+        self._normalize(all_raw, offsets=all_offsets or None)
         for tensors in self._data.values():
             self._mmap_refs.extend(tensors)
 
@@ -327,7 +443,7 @@ class JsonlStore(Store):
         if pos_ids:
             raw["position_ids"] = [torch.tensor(pos_ids, dtype=torch.int32)]
 
-        self._normalize(raw)
+        self._normalize(raw, per_record=True)
 
     @staticmethod
     def _primary_ids(result: dict) -> List[int]:

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

astrai/trainer/train_context.py

@@ -7,7 +7,7 @@ import torch.nn as nn
 from torch.utils.data import DataLoader, random_split
 
 from astrai.config.train_config import TrainConfig
-from astrai.dataset import ResumableDistributedSampler
+from astrai.dataset import RDSampler
 from astrai.model.components.lora import inject_lora
 from astrai.parallel.executor import BaseExecutor, ExecutorFactory
 from astrai.parallel.setup import get_current_device, get_rank, get_world_size
@@ -141,7 +141,7 @@ class TrainContextBuilder:
             )
 
         sampler_offset = context.consumed_samples // context.world_size
-        sampler = ResumableDistributedSampler(
+        sampler = RDSampler(
             data_source=train_dataset,
             start_epoch=context.epoch,
             start_iter=sampler_offset,
@@ -154,10 +154,11 @@ class TrainContextBuilder:
             num_workers=cfg.num_workers,
             pin_memory=cfg.pin_memory,
             prefetch_factor=cfg.prefetch_factor,
+            collate_fn=cfg.collate_fn,
         )
 
         if val_dataset is not None:
-            val_sampler = ResumableDistributedSampler(
+            val_sampler = RDSampler(
                 data_source=val_dataset,
                 start_epoch=0,
                 start_iter=0,
@@ -171,6 +172,7 @@ class TrainContextBuilder:
                 num_workers=cfg.num_workers,
                 pin_memory=cfg.pin_memory,
                 prefetch_factor=cfg.prefetch_factor,
+                collate_fn=cfg.collate_fn,
             )
 
         context.model, context.optimizer, context.dataloader, context.scheduler = (

scripts/tools/train.py

@@ -8,7 +8,7 @@ import torch.optim as optim
 from torch import Tensor, nn
 
 from astrai.config import AutoRegressiveLMConfig, TrainConfig
-from astrai.dataset import DatasetFactory
+from astrai.dataset import DatasetFactory, dpo_collate_fn, grpo_collate_fn
 from astrai.model import AutoRegressiveLM
 from astrai.model.components.decoder_block import DecoderBlock
 from astrai.trainer import SchedulerFactory, Trainer
@@ -504,6 +504,12 @@ def train(
 
     grad_ckpt_modules = [DecoderBlock] if gradient_checkpointing else []
 
+    collate_fn = None
+    if train_type == "dpo":
+        collate_fn = dpo_collate_fn
+    elif train_type == "grpo":
+        collate_fn = grpo_collate_fn
+
     train_config = TrainConfig(
         model_fn=model_fn,
         strategy=train_type,
@@ -536,6 +542,7 @@ def train(
         executor_kwargs=executor_kwargs,
         extra_kwargs=strategy_kwargs,
         neftune_alpha=neftune_alpha,
+        collate_fn=collate_fn,
     )
 
     trainer = Trainer(train_config)

tests/data/test_dataset.py

@@ -56,6 +56,13 @@ def _write_jsonl_dataset(test_dir, tokenizer_path, records, config_overrides=Non
     return data_dir
 
 
+def _fake_fetch_record(self, idx, keys):
+    """FakeStore.fetch_record matching real Store semantics."""
+    if isinstance(keys, str):
+        return self._data[keys][idx]
+    return {k: self._data[k][idx] for k in keys}
+
+
 def _make_seq_dataset(
     test_dir, name="data", seq_length=200, train_type="seq", data=None, **load_kwargs
 ):
@@ -338,6 +345,7 @@ def test_grpo_dataset_dtype(base_test_env):
         (),
         {
             "keys": ["prompts", "responses", "masks", "rewards"],
+            "num_records": 1,
             "_data": {
                 "prompts": [torch.randint(0, 100, (10,), dtype=torch.int32)],
                 "responses": [
@@ -346,9 +354,9 @@ def test_grpo_dataset_dtype(base_test_env):
                 "masks": [[torch.ones(5, dtype=torch.int32) for _ in range(G)]],
                 "rewards": [torch.rand(G, dtype=torch.float32)],
             },
+            "fetch_record": _fake_fetch_record,
         },
     )()
-    dataset._build_records()
     item = dataset[0]
 
     assert item["prompts"].dtype == torch.long
@@ -371,15 +379,16 @@ def test_grpo_dataset_load(base_test_env):
         (),
         {
             "keys": ["prompts", "responses", "masks", "rewards"],
+            "num_records": 1,
             "_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)],
             },
+            "fetch_record": _fake_fetch_record,
         },
     )()
-    dataset._build_records()
 
     assert len(dataset) == 1
     item = dataset[0]
@@ -864,6 +873,7 @@ def test_grpo_multiple_records(base_test_env):
         (),
         {
             "keys": ["prompts", "responses", "masks", "rewards"],
+            "num_records": n_records,
             "_data": {
                 "prompts": [torch.randint(0, 100, (10,)) for _ in range(n_records)],
                 "responses": dummy_responses,
@@ -875,9 +885,9 @@ def test_grpo_multiple_records(base_test_env):
                     torch.rand(G, dtype=torch.float32) for _ in range(n_records)
                 ],
             },
+            "fetch_record": _fake_fetch_record,
         },
     )()
-    dataset._build_records()
 
     assert len(dataset) == n_records
 

tests/data/test_sampler.py

@@ -1,4 +1,4 @@
-from astrai.dataset import ResumableDistributedSampler
+from astrai.dataset import RDSampler
 
 
 def test_random_sampler_consistency(random_dataset):
@@ -6,8 +6,8 @@ def test_random_sampler_consistency(random_dataset):
     dataset = random_dataset
 
     # Create two samplers with same seed
-    sampler1 = ResumableDistributedSampler(dataset, seed=42)
+    sampler1 = RDSampler(dataset, seed=42)
-    sampler2 = ResumableDistributedSampler(dataset, seed=42)
+    sampler2 = RDSampler(dataset, seed=42)
 
     indices1 = list(iter(sampler1))
     indices2 = list(iter(sampler2))
@@ -20,8 +20,8 @@ def test_random_sampler_different_seeds(random_dataset):
     dataset = random_dataset
 
     # Create two samplers with different seeds
-    sampler1 = ResumableDistributedSampler(dataset, seed=42)
+    sampler1 = RDSampler(dataset, seed=42)
-    sampler2 = ResumableDistributedSampler(dataset, seed=123)
+    sampler2 = RDSampler(dataset, seed=123)
 
     indices1 = list(iter(sampler1))
     indices2 = list(iter(sampler2))
@@ -35,7 +35,7 @@ def test_sampler_across_epochs(random_dataset):
     dataset = random_dataset
     n = len(dataset)
 
-    sampler = ResumableDistributedSampler(dataset, seed=42)
+    sampler = RDSampler(dataset, seed=42)
 
     # Get indices for first epoch
     epoch1_indices = list(iter(sampler))