| """Deterministic sampling of training windows from memory-mapped token shards.""" |
|
|
| from __future__ import annotations |
|
|
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
|
|
|
|
| def effective_mix(target: dict[str, float], group_tokens: dict[str, int], total_tokens: int, |
| max_epochs: dict[str, float]) -> dict[str, float]: |
| """Cap groups at `max_epochs` passes over their data; give the freed share to the other groups.""" |
| capped = {} |
| for g, w in target.items(): |
| if g in max_epochs: |
| cap = max_epochs[g] * group_tokens.get(g, 0) / total_tokens |
| if cap < w: |
| capped[g] = cap |
| free = 1.0 - sum(capped.values()) |
| rest = {g: w for g, w in target.items() if g not in capped} |
| rest_total = sum(rest.values()) |
| out = dict(capped) |
| for g, w in rest.items(): |
| out[g] = free * w / rest_total |
| return {g: out[g] for g in target} |
|
|
|
|
| class MixtureLoader: |
| """Samples (x, y) windows; batch contents depend only on (seed, step, micro).""" |
|
|
| def __init__(self, data_dir: Path, split: str, weights: dict[str, float], context_len: int, micro_batch: int, |
| seed: int): |
| self.data_dir = Path(data_dir) |
| self.manifest = json.loads((self.data_dir / "manifest.json").read_text()) |
| self.context_len, self.micro_batch, self.seed = context_len, micro_batch, seed |
| self.groups: dict[str, tuple[list[np.memmap], np.ndarray]] = {} |
| for group, splits in sorted(self.manifest["groups"].items()): |
| if split not in splits: |
| continue |
| arrays, sizes = [], [] |
| for shard in splits[split]["shards"]: |
| if shard["tokens"] < context_len + 1: |
| continue |
| |
| path = self.data_dir / group / Path(shard["path"]).name |
| arrays.append(np.memmap(path, dtype="<u2", mode="r")) |
| sizes.append(shard["tokens"] - context_len) |
| if arrays: |
| self.groups[group] = (arrays, np.asarray(sizes, dtype=np.float64)) |
| usable = {g: w for g, w in weights.items() if w > 0 and g in self.groups} |
| if not usable: |
| raise ValueError(f"no usable shards for weights {weights} in split {split!r}") |
| total = sum(usable.values()) |
| self.group_names = sorted(usable) |
| self.group_probs = np.asarray([usable[g] / total for g in self.group_names]) |
|
|
| def _window(self, rng: np.random.Generator, group: str) -> np.ndarray: |
| arrays, sizes = self.groups[group] |
| k = rng.choice(len(arrays), p=sizes / sizes.sum()) |
| off = int(rng.integers(0, sizes[k])) |
| return np.asarray(arrays[k][off : off + self.context_len + 1], dtype=np.int64) |
|
|
| def _stack(self, rows: list[np.ndarray]) -> tuple[torch.Tensor, torch.Tensor]: |
| batch = torch.from_numpy(np.stack(rows)) |
| return batch[:, :-1].contiguous(), batch[:, 1:].contiguous() |
|
|
| def get_batch(self, step: int, micro: int) -> tuple[torch.Tensor, torch.Tensor]: |
| rng = np.random.default_rng([self.seed, step, micro]) |
| groups = rng.choice(len(self.group_names), size=self.micro_batch, p=self.group_probs) |
| return self._stack([self._window(rng, self.group_names[g]) for g in groups]) |
|
|
| def get_group_batch(self, group: str, index: int) -> tuple[torch.Tensor, torch.Tensor]: |
| gid = sorted(self.groups).index(group) |
| rng = np.random.default_rng([self.seed, 1_000_000_007, gid, index]) |
| return self._stack([self._window(rng, group) for _ in range(self.micro_batch)]) |
|
|
