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/__init__.py
+2 -2
View File
@@ -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
+4
View File
@@ -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
+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]:
astrai/serialization/__init__.py
+2
View File
@@ -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
+50 -3
View File
@@ -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
+5 -3
View File
@@ -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 -1
View File
@@ -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
+13 -3
View File
@@ -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
+6 -6
View File
@@ -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)
sampler2 = ResumableDistributedSampler(dataset, seed=42)
sampler1 = RDSampler(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)
sampler2 = ResumableDistributedSampler(dataset, seed=123)
sampler1 = RDSampler(dataset, seed=42)
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))