数据加载 · get_batch 怎么喂饭 · nanoGPT 源码解读

逐行读 get_batch——用 np.memmap 把几百 MB 的 train.bin 当数组用而不撑爆内存，随机取一批窗口，输入 x 与目标 y 只差一个位置（这就是「预测下一个」的全部监督信号），再用 pin_memory 异步搬上 GPU。

数据加载 · get_batch 怎么喂饭

训练循环每一步都要一批新鲜数据。get_batch 就是「食堂打饭窗口」：从数据那站存的 train.bin 里，随机舀一批喂给模型。别看它只有十来行，里头有三个漂亮的工程点。

全代码

train.py 约 116–131 行：

data_dir = os.path.join('data', dataset)
def get_batch(split):
    # We recreate np.memmap every batch to avoid a memory leak, as per
    # https://stackoverflow.com/questions/45132940/...
    if split == 'train':
        data = np.memmap(os.path.join(data_dir, 'train.bin'), dtype=np.uint16, mode='r')
    else:
        data = np.memmap(os.path.join(data_dir, 'val.bin'), dtype=np.uint16, mode='r')
    ix = torch.randint(len(data) - block_size, (batch_size,))
    x = torch.stack([torch.from_numpy((data[i:i+block_size]).astype(np.int64)) for i in ix])
    y = torch.stack([torch.from_numpy((data[i+1:i+1+block_size]).astype(np.int64)) for i in ix])
    if device_type == 'cuda':
        # pin arrays x,y, which allows us to move them to GPU asynchronously (non_blocking=True)
        x, y = x.pin_memory().to(device, non_blocking=True), y.pin_memory().to(device, non_blocking=True)
    else:
        x, y = x.to(device), y.to(device)
    return x, y

亮点一：memmap——不把大文件读进内存

data = np.memmap(os.path.join(data_dir, 'train.bin'), dtype=np.uint16, mode='r')

train.bin 可能几百 MB 甚至几十 GB(OpenWebText)，全读进内存不现实。np.memmap 的绝活是：把磁盘上的文件「映射」成一个数组，但并不真的把它加载进内存——你访问哪一小段，操作系统才把哪一小段按需换页进来。于是你能像操作普通数组一样 data[i:i+block_size]，内存占用却只有实际碰到的那点。这就是几十 GB 数据集能在有限内存里训练的关键。

那句注释也值得一提：每次 get_batch 都重新创建 memmap，是为了绕开一个已知的内存泄漏（注释里附了 StackOverflow 链接）。这是踩过坑才会写的一行——真实工程的味道。

亮点二：随机取窗口，x 与 y 差一个位置

ix = torch.randint(len(data) - block_size, (batch_size,))
x = torch.stack([torch.from_numpy((data[i:i+block_size]).astype(np.int64)) for i in ix])
y = torch.stack([torch.from_numpy((data[i+1:i+1+block_size]).astype(np.int64)) for i in ix])

ix = randint(len(data) - block_size, (batch_size,))：随机挑 batch_size 个起始位置。注意上界是 len(data) - block_size，保证从任何一个起点都能取够一整窗 block_size。训练数据不是顺序喂的，而是每批从整条数据里随机开窗——这样批与批之间更独立，训练更稳。
x：每个起点 i 取 data[i : i+block_size]，一段长 block_size 的 token，batch_size 段叠成一批。
y：每个起点取 data[i+1 : i+1+block_size]——整体往后挪一位。

这个「挪一位」，是整个语言模型监督信号的全部秘密。 把 x 和 y 对齐看：

位置:      0    1    2    3    4
x (输入):  To   be   or   not   to
y (目标):  be   or   not   to    be     ← 每一位都是 x 的下一个

模型在位置 0 看到 To，要预测 be（y 的第 0 位）；在位置 1 看到 To be，要预测 or……「预测下一个 token」这句口号，落到数据上就是「让 y = x 右移一位」。回忆GPT.forward 里的 cross_entropy，比的正是「模型对每个位置的预测」和「这里的 y」。前后完全咬合。

.astype(np.int64)：train.bin 里是 uint16（省空间），但 PyTorch 索引/嵌入要 int64，所以取出来时转一下。

亮点三：pin_memory + 异步搬运

x, y = x.pin_memory().to(device, non_blocking=True), y.pin_memory().to(device, non_blocking=True)

数据在 CPU 上备好，要送上 GPU。这里用了两个加速小技巧：

pin_memory()：把张量放进锁页内存（不会被操作系统换出），这样才能……
.to(device, non_blocking=True)：异步拷贝到 GPU——不必干等拷贝完成，CPU 可以继续干别的。

配合训练循环里「一边算当前批、一边预取下一批」的写法，数据搬运和 GPU 计算就能重叠起来，GPU 不必停下来等饭。CPU 模式则退回普通的 .to(device)。

回头看，数据这条链特别干净：prepare.py 把文本压成一条 uint16 长龙存盘 → get_batch 用 memmap 随机开窗、右移一位造出 (x， y) → forward 拿 x 出 logits、和 y 算 loss。没有花哨的 DataLoader，十来行搞定，却把「大文件、随机采样、异步上卡」全办了。

饭有了，开火

配置、优化器、数据加载——训练的三大件齐了。下一章，我们终于走进 train.py 的主循环：学习率怎么随步数起落、梯度累积怎么用小显存模拟大 batch、混合精度和梯度裁剪各自在防什么，以及一步「前向—反向—更新」到底怎么写。这是第三站的高潮。👉 训练主循环：一步到底怎么走

参考文献

源码：Andrej Karpathy, nanoGPT， train.py 的 get_batch. https://github.com/karpathy/nanoGPT
NumPy memmap 文档：https://numpy.org/doc/stable/reference/generated/numpy.memmap.html

数据加载 · get_batch 怎么喂饭