ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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

Browse Source 
- 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
This commit is contained in:
ViperEkura
2026-07-18 21:02:29 +08:00
parent 28886e4241
commit a74e5b91a3
12 changed files with 334 additions and 92 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/dataset/__init__.py
+4 -2
View File
@@ -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
+96 -52
View File
@@ -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:
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()
def __getitem__(self, index: int):
begin_idx, end_idx = self.get_index(index)
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"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)
rejected = self._fetch_data(begin_idx, end_idx, "rejected").to(dtype=torch.long)
chosen_mask = self._fetch_data(begin_idx, end_idx, "chosen_mask").to(
dtype=torch.bool
)
rejected_mask = self._fetch_data(begin_idx, end_idx, "rejected_mask").to(
dtype=torch.bool
)
@property
def count(self) -> int:
return self.storage.num_records
def __len__(self) -> int:
return self.storage.num_records
def __getitem__(self, index: int) -> Dict[str, Tensor]:
return {
"chosen": chosen,
"rejected": rejected,
"chosen_mask": chosen_mask,
"rejected_mask": rejected_mask,
"chosen": self.storage.fetch_record(index, "chosen").to(dtype=torch.long),
"rejected": self.storage.fetch_record(index, "rejected").to(
dtype=torch.long
),
"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),
"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),
"prompts": prompts.to(dtype=torch.long),
"responses": [r.to(dtype=torch.long) for r in responses],
"masks": [m.to(dtype=torch.bool) for m in masks],
"rewards": rewards.to(dtype=torch.float32),
}
astrai/dataset/sampler.py
+9 -1
View File
@@ -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
+135 -19
View File
@@ -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).
``len(store)`` returns ``self._length`` (explicit, O(1)), the minimum
total element count across all keys.
Stream mode (SEQ/SFT):
``fetch(begin, end, keys)`` slices across concatenated segments,
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]):
"""Register segments and pre-compute cumulative lengths.
@property
def num_records(self) -> int:
return self._num_records
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 flat-tensor keys.
def fetch_record(
self,
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]]``
(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``.
Returns a tensor (flat key) or ``List[Tensor]`` (nested key such as
GRPO ``responses``).
"""
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]: