| """Read-once data loader (Stage 2). |
| |
| Each file is cut into windows of seq_len+1 tokens that overlap by exactly one token |
| (x = first seq_len tokens, y = the same window shifted by one), so every token is a |
| prediction target once. Each file's windows are shuffled once with a fixed seed, and a |
| fixed schedule says which file every sample comes from, following the mix weights. |
| |
| Everything is computed from (seed, step, slot), so a restart at any step continues |
| exactly where training stopped, with no saved loader state. A window is only reused |
| after its whole file has been read once. |
| """ |
| import numpy as np |
|
|
| |
| |
| SAMPLES_PER_STEP = 128 |
|
|
|
|
| class ReadOnceDataset: |
| def __init__(self, bin_paths, seq_len, weights, start_step, end_step, eot_token_id, seed=1234): |
| assert end_step > start_step, "empty stage" |
| self.seq_len = seq_len |
| self.start_step = start_step |
| self.eot = eot_token_id |
| self.data = [np.memmap(p, dtype=np.uint32, mode="r") for p in bin_paths] |
| self.n_windows = [(len(d) - 1) // seq_len for d in self.data] |
| assert all(n > 0 for n in self.n_windows), "a file is shorter than one window" |
| w = np.asarray(weights, dtype=np.float64) |
| assert len(w) == len(self.data) and (w >= 0).all() and w.sum() > 0, "bad weights" |
| w = w / w.sum() |
|
|
| |
| self.perm = [np.random.default_rng(seed + 1000 + f).permutation(n).astype(np.int32) |
| for f, n in enumerate(self.n_windows)] |
|
|
| |
| self.total = (end_step - start_step) * SAMPLES_PER_STEP |
| self.src = np.random.default_rng(seed).choice(len(w), size=self.total, p=w).astype(np.int16) |
| self.occ = np.empty(self.total, dtype=np.int64) |
| for f in range(len(w)): |
| idx = np.flatnonzero(self.src == f) |
| self.occ[idx] = np.arange(len(idx)) |
|
|
| @staticmethod |
| def slot(micro, i, rank, world, batch_size): |
| """Position of one sequence inside its step's 128 slots; unique per (micro, i, rank).""" |
| s = (micro * batch_size + i) * world + rank |
| assert s < SAMPLES_PER_STEP, "more than 128 sequences per step: raise SAMPLES_PER_STEP" |
| return s |
|
|
| def window(self, step, slot): |
| """(file index, window index, tokens) for one sample; tokens has seq_len+1 entries.""" |
| k = (step - self.start_step) * SAMPLES_PER_STEP + slot |
| assert 0 <= k < self.total, f"step {step} is outside this stage" |
| f = int(self.src[k]) |
| w = int(self.perm[f][self.occ[k] % self.n_windows[f]]) |
| start = w * self.seq_len |
| tokens = np.asarray(self.data[f][start : start + self.seq_len + 1], dtype=np.int64) |
| return f, w, tokens |
|
|
| def get_numpy(self, step, micro, batch_size, rank, world): |
| """One micro-batch as numpy arrays: inputs, targets, document ids.""" |
| x = np.empty((batch_size, self.seq_len), dtype=np.int64) |
| y = np.empty((batch_size, self.seq_len), dtype=np.int64) |
| seg = np.empty((batch_size, self.seq_len), dtype=np.int32) |
| for i in range(batch_size): |
| _, _, chunk = self.window(step, self.slot(micro, i, rank, world, batch_size)) |
| x[i], y[i] = chunk[:-1], chunk[1:] |
| |
| is_eot = (x[i] == self.eot).astype(np.int32) |
| seg[i] = np.cumsum(is_eot) - is_eot |
| return x, y, seg |
|
|
| def get_batch(self, step, micro, batch_size, rank, world, device="cuda"): |
| """Same as get_numpy, as torch tensors on `device`.""" |
| import torch |
| out = [] |
| for a in self.get_numpy(step, micro, batch_size, rank, world): |
| t = torch.from_numpy(a) |
| if torch.device(device).type == "cuda": |
| t = t.pin_memory().to(device, non_blocking=True) |
| out.append(t) |
| return tuple(out) |
|
|
