wehub-resource-sync
53 lines
1.7 KiB
Python
53 lines
1.7 KiB
Python
"""
|
|
Batch iterator for packed SFT data (mirrors ``data_loader.data_loader.get_batch_iterator``
|
|
but carries the per-token loss mask alongside the tokens).
|
|
|
|
The HDF5 file has two aligned datasets, ``tokens`` and ``loss_mask``, both shape
|
|
(N, context_length), produced by ``scripts/prepare_sft_data.py``.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Iterator
|
|
|
|
import h5py
|
|
import numpy as np
|
|
import torch
|
|
|
|
|
|
def get_sft_batch_iterator(
|
|
data_path: str,
|
|
batch_size: int,
|
|
device: str = "cpu",
|
|
*,
|
|
rank: int = 0,
|
|
world_size: int = 1,
|
|
shuffle: bool = True,
|
|
infinite: bool = True,
|
|
) -> Iterator[tuple[torch.Tensor, torch.Tensor, int]]:
|
|
"""
|
|
Yield ``(tokens, loss_mask, epoch)`` batches from a packed SFT HDF5 file.
|
|
|
|
Rows are sharded across ranks (each rank sees a disjoint stride) so DDP data-parallel
|
|
training covers the dataset once per epoch. Set ``infinite=False`` to stop after one
|
|
pass (used for evaluation).
|
|
"""
|
|
with h5py.File(data_path, "r") as f:
|
|
tokens = f["tokens"]
|
|
masks = f["loss_mask"]
|
|
n = tokens.shape[0]
|
|
idxs = np.arange(rank, n, world_size)
|
|
epoch = 0
|
|
rng = np.random.default_rng(1234 + rank)
|
|
while True:
|
|
if shuffle:
|
|
rng.shuffle(idxs)
|
|
for start in range(0, len(idxs) - batch_size + 1, batch_size):
|
|
batch = np.sort(idxs[start: start + batch_size]) # sorted for h5py fancy-index
|
|
tk = torch.tensor(np.asarray(tokens[batch]), dtype=torch.long, device=device)
|
|
mk = torch.tensor(np.asarray(masks[batch]), dtype=torch.long, device=device)
|
|
yield tk, mk, epoch
|
|
epoch += 1
|
|
if not infinite:
|
|
return
|