File size: 4,193 Bytes
2ae6691 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | """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
# 524,288 tokens / 4096 = 128 sequences: the most any node count uses in one step.
# Each step owns 128 sample slots; node counts that use fewer just skip the rest.
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()
# each file's windows in one fixed shuffled order (int32 keeps RAM small)
self.perm = [np.random.default_rng(seed + 1000 + f).permutation(n).astype(np.int32)
for f, n in enumerate(self.n_windows)]
# which file each sample of the stage reads, and how many earlier samples used that file
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:]
# document id goes up AFTER each end-of-text, so the EOT stays with its own document
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)
|