数据集滑动窗口:喂给训练的 id 流

数据集滑动窗口:喂给训练的 id 流

本节摘要:预训练是「从 token id 到梯度」的函数,本节建喂 id 进去的传送带。训练一次读一批 id 更新模型,批形状由训练契约定:因果语言模型批持 (B, T) 输入 id 与 (B, T) 目标 id(目标为输入左移一位)。本节建这条管线:分词器把文本变扁平 id 列表,滑动窗口切成训练例,自定义 Dataset 把例暴露为张量,DataLoader 用每 epoch 播种的种子洗牌批化。你会学到步幅与重叠的权衡、移位一的位置、确定性洗牌为何关键,以及「每例 T+1 个原始 id」的形状契约。

对应原课程:Phase 19 · Lesson 31 · tokenized-dataset-sliding-window(原英文 phases/19-capstone-projects/31-tokenized-dataset-sliding-window/docs/en.md)。本节属「从零构建 GPT」赛道第二节。

学习目标

阅读完本节,你应当能够:

1. 调一次分词器把原始语流转成 token id 流。
2. 把 id 流切成定长窗口,步幅可配重叠。
3. 构建 PyTorch Dataset,返回下一 token 预测的输入与目标张量。
4. 用每 epoch 播种的确定性洗牌 DataLoader 包数据集。
5. 论证步幅、冗余、有效数据集大小之间的权衡。

一、问题与直觉

因果 LM 吃 (B, T) 形状的 id,B 是批大小,T 是上下文长度。位置 t 的目标是位置 t+1 的输入——每个训练例覆盖 T+1 个原始 id。窗口步幅控制连续例间重叠多少。

切片器从不与语料边界重叠:最后一个窗口若填不满 T+1 位则丢(用 <|pad|> 填尾部也合法但复杂化损失掩码,本节选丢)。

二、从零实现

为何滑动窗口?预训练语料是一条长 id 流。若模型只见不重叠窗口,每个训练例教它同样的 T 边界。调步幅移动边界,模型见更多样的「预测下一 token」任务。步幅 T 产不重叠窗口;T//2 产 50% 重叠、有效数据集翻倍;1 产最大重叠、增 T 倍。代价是每 epoch 更多算力,收益是更多边界多样性。多数预训练用步幅等于上下文长度,因为语料已远大于模型一 epoch 能跑完的量,边界多样性论证较弱。

Dataset 类(两必需方法):

class SlidingWindowDataset(Dataset):
    def __len__(self):  return max(0, 1 + (N - (T+1)) // S)
    def __getitem__(self, i):
        window = self.ids[i*S : i*S + T + 1]      # 在途算窗口起点
        return window[:-1], window[1:]            # (input, target) 移位一
​

__len__ 返例数,__getitem__ 返一对张量。索引时在途算窗口起点,内存成本是一份 id 流副本,不论步幅产多少例。移位一在 __getitem__ 内:input = window[:-1]、target = window[1:],两者皆 PyTorch long 张量,训练循环当真值。

例数公式(可静态算每 epoch 步数):

examples = max(0, 1 + (N - (T + 1)) // S)
# N=id 流长, T=上下文长, S=步幅
​

三、确定性洗牌

shuffle=True 的 DataLoader 读 PyTorch 随机生成器。传一个显式 torch.Generator 每 epoch 播种,重启就拿到同样洗牌。这性质在你比较两个仅差一超参的运行时要紧——无种子,两运行以不同顺序看数据,损失曲线因与改动无关的原因发散。种子契约简单:epoch_seed = base_seed + epoch_index,base 种子构造时传,epoch 索引由训练器每 epoch 顶自增。同 base 种子重跑,每 epoch 同序。

PyTorch 默认采样器无放回均匀随机挑索引,正是预训练要的。DataLoader 调 __getitem__ B 次堆叠成批;因每例构造上等长,无需填充逻辑。本节 num_workers=0 求简——生产运行用 worker 并行 __getitem__,对这条管线基本是空操作(活只是内存张量切片),但同 Dataset API 干净支持 worker。

四、可复用产物

main.py 定义两类一助手:SlidingWindowDataset(PyTorch Dataset)、make_dataloader(返带种子生成器的配置 DataLoader)、_encode_corpus_to_ids(一次性分词调用)。demo 在进程内建小分词器、编码内置语料、构数据集与 dataloader、打一批、断言形状契约。code/tests/test_dataset.py 钉窗口计数公式、移位一性质、确定性洗牌、步幅权衡。把上下文长从 16 调 32,看每 epoch 例数下降——这就是你的每 epoch 步数预算。

五、框架对比

本节的「一次性编码 + 内存扁平张量 + 滑窗切片」是 nanoGPT 风格的最简形态,适合百万 id 级语料(< 100MB)。生产预训练(HuggingFace datasets、Megatron)用磁盘流式与打包(pack),但 Dataset 契约不变——换存储、保契约。多文档边界由构语料时插 <|endoftext|> id 编码,模型学在边界附近预测。本节不流式、不多文档,是有意的简化,下一节在其上叠嵌入。

六、练习

1. 步幅权衡:对同语料,对比步幅 T/T//2/1 的例数与每 epoch 算力,画出权衡曲线。
2. 确定性洗牌:同 base 种子跑两次,确认每 epoch 同序;换种子确认序变。
3. 形状契约:断言每批输入与目标皆为 (B, T) long 张量,且 target[:, :-1] == input[:, 1:]。
4. 尾部丢弃:构一个 N 使最后窗口填不满 T+1,确认被丢而非填充。
5. 上下文长影响:T 从 16 扫到 128,看每 epoch 步数预算如何降。

本节要点回顾

1. 形状契约:批 (B, T) 输入 + (B, T) 目标(左移一),每例 T+1 原始 id。
2. 滑窗切片:步幅控重叠,T 不重叠、T//2 翻倍、1 增 T 倍。
3. 移位一在 getitem:input=window[:-1],target=window[1:]。
4. 确定性洗牌:每 epoch epoch_seed = base + index,重启同序,可比比超参。
5. 例数公式:max(0, 1 + (N-(T+1))//S),可静态算每 epoch 步数。
6. 内存一份 id 副本:不论步幅产多少例,索引在途算窗口起点。

下一节,我们给 id 配上「token 与位置嵌入」——把离散整数变成模型能算的连续向量。