File size: 6,548 Bytes
c83c617 | 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 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 | """Bucketed task streams with committed cursors, independent of worker prefetch."""
import math
import numpy as np
import torch
from torch.utils.data import DataLoader
from vimeml.training.data import BlockShuffleSampler, LengthBucketBatchSampler
from vimeml.training.train import worker_init
class ConversionWindows:
"""Keep the frozen conversion reader unchanged, including its source hash."""
context_length = 128
epoch_indexed = False
def __init__(self, data):
self.data, self.split = data, data.split
def __len__(self):
return len(self.data)
def __getitem__(self, index):
return self.data[int(index)]
def window_lengths(self, indices):
indices = np.asarray(indices, dtype=np.int64)
# Normal-reading length is a cheap upper bound for the fixed10% empty
# reading views; actual padding/denominators always use returned rows.
return (self.data.offsets[indices + 1] - self.data.offsets[indices] - 1).astype(np.int64)
class TaskCollator:
def __call__(self, samples):
lengths = torch.tensor([len(s["input_ids"]) for s in samples], dtype=torch.long)
width = math.ceil(int(lengths.max()) / 8) * 8
if not 1 <= width <= 128:
raise ValueError("Task sample outside the fixed context")
ids = torch.zeros((len(samples), width), dtype=torch.long)
labels = torch.full_like(ids, -100)
for row, (sample, n) in enumerate(zip(samples, lengths.tolist())):
if len(sample["labels"]) != n:
raise ValueError("Task labels must already be shifted")
ids[row, :n] = torch.tensor(sample["input_ids"], dtype=torch.long)
labels[row, :n] = torch.tensor(sample["labels"], dtype=torch.long)
return {"input_ids": ids, "labels": labels, "lengths": lengths}
class TailLastBucketBatchSampler(LengthBucketBatchSampler):
"""Keep the single partial batch last, including after cursor restoration.
Base LM batching shuffles the tail inside its last pool. C1's exact update
accounting requires all earlier conversion updates to contain256 examples.
"""
def __iter__(self):
base = LengthBucketBatchSampler(self.dataset, self.sampler, self.batch_size,
self.multiplier, self.seed, start_batch=0)
base.set_epoch(self.epoch, 0)
cursor, tail = 0, None
for batch in base:
if len(batch) < self.batch_size:
if tail is not None:
raise ValueError("More than one partial task batch")
tail = batch
else:
if cursor >= self.start_batch:
yield batch
cursor += 1
if tail is not None and cursor >= self.start_batch:
yield tail
def loader_for(data, batch_size, workers, seed, multiplier=64, indices=None, shuffle=True, pin_memory=False):
sampler = (BlockShuffleSampler(len(data), seed) if shuffle else range(len(data))) if indices is None else indices
batches = TailLastBucketBatchSampler(data, sampler, batch_size, multiplier, seed)
options = {"multiprocessing_context": "spawn", "prefetch_factor": 2} if workers else {}
return DataLoader(data, batch_sampler=batches, num_workers=workers,
collate_fn=TaskCollator(), pin_memory=pin_memory,
persistent_workers=workers > 0, worker_init_fn=worker_init,
generator=torch.Generator().manual_seed(seed), **options)
def slice_batch(batch, start, end):
return {name: value[start:end] for name, value in batch.items()}
class CommittedStream:
"""The state always points to the next unconsumed row in a bucketed batch.
Caller copies state only AFTER a successful optimizer update. A partial LM
tail borrows rows from the next LM epoch without dropping/repeating data.
"""
def __init__(self, loader, state=None):
self.loader = loader
self.epoch, self.batch, self.offset = (0, 0, 0) if state is None else (
state["epoch"], state["batch"], state["offset"])
if any(type(x) is not int or x < 0 for x in (self.epoch, self.batch, self.offset)):
raise ValueError("Invalid task cursor")
self.iterator = self.pending = None
self.loader.batch_sampler.set_epoch(self.epoch, self.batch)
def state_dict(self):
return {"epoch": self.epoch, "batch": self.batch, "offset": self.offset}
def next_batch(self):
if self.offset:
raise ValueError("Whole-batch consumption requires an aligned cursor")
if self.iterator is None:
self.iterator = iter(self.loader)
value = next(self.iterator)
self.batch += 1
return value
def next_group(self, size):
result, remaining = [], size
while remaining:
if self.pending is None:
if self.iterator is None:
self.iterator = iter(self.loader)
try:
self.pending = next(self.iterator)
except StopIteration:
if self.batch == 0:
raise ValueError("Empty replay stream")
self.epoch, self.batch, self.offset = self.epoch + 1, 0, 0
self.loader.batch_sampler.set_epoch(self.epoch, 0)
self.iterator = iter(self.loader)
self.pending = next(self.iterator)
available = len(self.pending["lengths"]) - self.offset
if available <= 0:
raise ValueError("Replay cursor outside its saved batch")
take = min(available, remaining)
result.append(slice_batch(self.pending, self.offset, self.offset + take))
remaining -= take
self.offset += take
if self.offset == len(self.pending["lengths"]):
self.offset, self.pending = 0, None
self.batch += 1
return result
def microbatches(batches, size):
for batch in batches:
for start in range(0, len(batch["lengths"]), size):
part = slice_batch(batch, start, start + size)
width = math.ceil(int(part["lengths"].max()) / 8) * 8
# Slicing the time dimension retains pinned storage. .to() handles
# the strided source; never pad a short microbatch to another's max.
part["input_ids"] = part["input_ids"][:, :width]
part["labels"] = part["labels"][:, :width]
yield part
|