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/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]: