第 5 课：数据分片、Shuffle 与精确 Resume

本课目标与完成标准

本课产出：实现无重复的全局样本位置分配，并说明 epoch、seed、world-size 变化下的恢复语义。

通过要求：唯一代码填空题通过给定检查；Q1～Q3 都给出因果链；能指出一个正确性边界和一个性能/运营取舍；总分至少 8/10。

与既有六条主线的边界

train/ 把 batch 当数字；本课关注样本如何在 rank/worker 间唯一分配、何时重复/丢失，以及 checkpoint 如何恢复数据位置。

前置：Python、PyTorch、train 第 1～5 课、CUDA 基础。本课只补 AI Infra 缺口，不重复已经完成的 kernel 或并行算法推导。

核心概念与运行机制

1. 它是什么，解决什么问题

在万亿 Token 级别的大模型预训练与持续多轮微调中，数据加载流水线（Data Pipeline）并非简单的单机 Python 迭代器，而是直接关系到模型收敛确定性与集群万小时级容灾能力的核心分布式子系统。

分布式数据流水线必须解决三大核心痛点：

1. 全局无重复正交分片（Zero-Redundancy Sharding）：如图 5-1 所示，

全量训练数据集 (N 条记录)

Samples [0, 1, 2, ..., N-1]

DistributedSampler 确定性分发

随机种子 + Epoch 序号

确定性伪随机置换序列

Rank 0 数据分片: [idx0, idx4, ...]

Rank 1 数据分片: [idx1, idx5, ...]

图 5-1：分布式数据分片采样与可复现 Shuffle 逻辑映射模型

如何保证数千张卡在同一时刻获取的数据样本互不重叠，且全体样本的并集恰好覆盖全量数据集；

2. 确定性伪随机可复现（Deterministic Reproducibility）：在引入分布式 Shuffle 洗牌打乱数据分布的同时，保证多次运行或断点恢复时数据生成的序列严格一致；

3. 断点精确续训（Exact Resume & Fast-Forward）：如图 5-2 所示，

训练异常中断 / 节点抢占故障

加载恢复 Checkpoint 元数据

恢复 Epoch 计数器、随机种子与已消费 Batch 偏移量

Sampler 快速跳过当前分片已完成样本 (O(1) 索引计算)

无损无缝对齐数据加载管道，避免重复消费或样本遗漏

图 5-2：断点续训 (Resume) 精确状态恢复与跳步定位工作流

当集群遭遇硬件损坏或抢占中断时，如何不从头遍历 TB/PB 级数据集，以 $\mathcal{O}(1)$ 的时间复杂度精准跳过已训练样本，实现无缝对齐恢复，杜绝样本重复吞吐与梯度漂移。

2. 它如何工作

分布式数据分片与精确恢复依托于严格的分层索引映射与状态机持久化设计（如图 5-1 与 图 5-2 所示）：

1. 确定性两级分片映射（图 5-1）：
   • 全局置换生成：DistributedSampler 依据固定公式 $\text{Seed}_{\text{epoch}} = \text{seed} + \text{epoch}$ 初始化确定性伪随机数生成器，对大小为 $N$ 的数据集索引生成全局置换数组 $I = [i_0, i_1, \dots, i_{N-1}]$；
   • 对齐填充约束（Padding Invariant）：若样本总数 $N$ 不能被集群总卡数 $P$ 整除，采样器必须对尾部补全至 $P$ 的整数倍，以确保所有 Rank 在当前 Epoch 拥有完全相等的 Batch 步数，从根本上防止因某张卡少一个 Batch 导致 DDP 全局通信挂起死锁；
   • 跨卡交织切分（Interleaved Slicing）：对于 Rank $r \in [0, P-1]$，其分配到的样本子集严格遵循交错步长采样： $$\text{Indices}(r) = \{I[g] \mid g \pmod P == r\}$$ 全局样本位置 $g$ 与本地样本索引 $l$ 的双向映射为： $$\text{owner}(g) = g \pmod P, \quad \text{local\_idx}(g) = \lfloor g / P \rfloor$$
2. 多进程 Worker 二级分发：
   • 在每个 Rank 内部，PyTorch DataLoader 通常启动 $W$ 个工作进程（num_workers）。
   • 本地样本集进一步在多进程间切分，由各 Worker 进程异步执行数据解码、分词（Tokenization）与张量组装，通过共享内存环形队列传递给主进程。
3. 断点精准快速跳步恢复（图 5-2）：
   • 在保存 Checkpoint 时，引擎不仅持久化模型权重与优化器状态，还原子记录数据进度游标（Data State Cursor）：current_epoch、已完成迭代步数 global_step、以及数据加载器的伪随机种子；
   • 发生故障重启时，采样器重新构建当前 Epoch 的确定性全局置换 $I$；
   • 随后计算已消费的样本偏移量：$\Delta = \text{global\_step} \times \text{batch\_size} \times P$；
   • 采样器直接从索引偏移处切片切分（Fast-Forward 切片切分），零 I/O 开销瞬间跳过数百 GB 的已学数据，实现无缝精准恢复。

3. 正确性条件与常见误区

• DistributedSampler.set_epoch(epoch) 必调契约：
  • 致命陷阱：在每轮 Epoch 训练循环开头忘记显式调用 sampler.set_epoch(epoch)；
  • 底层物理实况：采样器内部的随机数种子将永久固定在初始值 seed + 0。尽管单 Epoch 内各卡数据无重复，但在接下来的数十个 Epoch 中，模型将在完全相同的样本先后顺序和完全相同的局部 Batch 组合下被重复喂养。模型泛化能力严重劣化，过拟合风险剧增，且该 Bug 在单 Epoch 冒烟测试中绝对无法被察觉。
• IterableDataset 与多 Worker 的“八倍数据重复灾难”：
  • 常见误区：直接将流式数据加载器 IterableDataset 配合 num_workers=8 启动训练，误以为框架会自动分片；
  • 实际后果：在 PyTorch 架构中，IterableDataset 默认会在每个 Worker 进程中完整克隆一份实例。若用户未在数据集的 __iter__ 中通过 torch.utils.data.get_worker_info() 主动按 worker_id 切分数据流，8 个 Worker 进程将各自从头开始读取完整数据流，导致该卡上的所有样本被完全重复计算 8 次。
• 弹性伸缩（Elastic Scaling）下的局部 Offset 失效：
  • 若将集群规模从 64 卡动态缩容至 32 卡，原先各卡保存的私有局部步数（local_step）完全失去拓扑意义。必须依赖全局不变样本坐标（Global Sample Index / Token Offset）重新解构分片拓扑。

4. 性能与工程取舍

• Map 索引随机加载 vs. Iterable 流式预取：
  • Map-style Dataset（索引式）：天然支持 $\mathcal{O}(1)$ 确定性断点跳转与灵活多卡切分；但要求底层存储支持高效随机点查（Random Seek），在千亿级文本的小文件随机读取中会打崩网络文件系统（NFS / Lustre）的元数据服务器；
  • IterableDataset（流式打包）：将数据预先分块打包为大尺寸压缩分卷（如 WebDataset / Tar / Parquet），吞吐逼近物理网卡线速；但断点续训状态维护极其复杂，需要显式持久化文件分卷编号与内部分卷 Byte Offset。
• 严格零重复保证 vs. 异步预取缓冲丢弃：
  • 为了维持 GPU 算力饱和，DataLoader 会在后台预取数个 Batch。当触发 Checkpoint 保存时，预取队列中已出队但未实际前向计算的张量，在恢复时必须被严谨对齐；
  • 生产中通常采用“按最后提交的已反向 Step”作为物理截断点，在 Resume 时清空后台预取缓冲并从该 Step 精确重启。

具体演示

设总卡数 $P = 4$，全局生成了一条确定性置换序列： $$I = [\text{idx}_0, \text{idx}_1, \dots, \text{idx}_9, \text{idx}_{10}, \dots]$$

1. 全局位置 $g = 10$ 的归属判定：
   • 宿主卡计算：$\text{owner} = 10 \pmod 4 = 2$（由 Rank 2 消费）；
   • 该卡内的局部序列索引：$\text{local\_idx} = \lfloor 10 / 4 \rfloor = 2$（该样本是 Rank 2 在该 Epoch 分配到的第 3 个样本，从 0 计）；
2. 故障断点恢复推演：
   • 训练进行到第 2 个 Step 中途崩溃，每个 Step 包含全局 Batch 为 4 个样本（每卡 1 个）；
   • Checkpoint 读取到最后提交的已完成 Step 为 1（即前 4 个全局样本 $g \in [0, 3]$ 已完全训练）；
   • 重启时，各卡采样器直接从本地索引 1（即全局位置 4 开始）提取数据，零冗余无缝衔接。

实践任务：唯一代码填空题

补齐全局位置到 rank 的确定性分配，并返回本 rank 的局部序号。

只能修改 TODO/______ 位置，不得删除断言或放宽通过条件。

def owner_and_local_index(global_position, world_size):
    if global_position < 0 or world_size <= 0:
        raise ValueError("invalid shard position")
    # TODO：round-robin 分片，owner=g%P，local=g//P。
    return ______

assert owner_and_local_index(10, 4) == (2, 2)
assert owner_and_local_index(0, 8) == (0, 0)

检查方法

运行本单元格，所有断言必须通过；再补一个边界输入并解释预期。

Q1

忘记 DistributedSampler.set_epoch(epoch) 会怎样？

你的答案：

Q2

IterableDataset + num_workers=8 为什么可能把数据重复 8 次？

你的答案：

Q3

弹性恢复 world size 改变时，怎样减少重复/丢失？

你的答案：

评分规则

• 代码 4 分：正常输入 2 分、边界输入 1 分、解释 1 分；
• Q1～Q3 各 2 分；
• 一票否决：混淆测量与推测、忽略租户/请求隔离、用平均值掩盖尾延迟、删除失败路径。

Tip💡 点击展开参考答案与详细代码实现

参考答案（仅 answer 分支）

先独立完成。核对后请改变一个规模或故障条件重新推演。

def owner_and_local_index(global_position, world_size):
    if global_position < 0 or world_size <= 0:
        raise ValueError("invalid shard position")
    return global_position % world_size, global_position // world_size

assert owner_and_local_index(10, 4) == (2, 2)
assert owner_and_local_index(0, 8) == (0, 0)

Q1 参考答案

忘记显式调用 DistributedSampler.set_epoch(epoch) 会破坏分布式随机洗牌机制，给模型训练埋下极其隐蔽的收敛隐患：

1. 伪随机数种子永久冻结（RNG Seed Freezing）：
   • DistributedSampler 内部的随机置换生成算法严格依赖当前 Epoch 计数：g = torch.Generator(); g.manual_seed(self.seed + self.epoch)。若未在每个 Epoch 循环开头调用 set_epoch(epoch)，self.epoch 将永久保持在默认初始值 0；
2. 训练样本访问轨迹的完全固化（Static Data Trajectory）：
   • 尽管在同一个 Epoch 内部，各个 Rank 之间的数据依然保持不重复切分，但在随后的第 2、第 10 甚至第 100 个 Epoch 中，全量数据集生成的全局置换顺序完全相同；
   • 每张卡在每个 Epoch 获得的数据样本子集及其批次排列顺序被锁死，网络在反复消化完全固定顺序与固定组合的 Mini-batches；
3. 收敛性劣化与测试隐蔽性：
   • 随机梯度下降（SGD）收敛的理论基石要求数据序列具备充分的独立随机性。静态固化的样本轨迹会导致模型对特定样本出现记忆效应，极易引发过拟合与泛化能力断崖式下跌。而该 Bug 在单 Epoch 调试或短周期的冒烟测试中绝对无法被断言拦截，极难排查。

Q2 参考答案

IterableDataset 配合 num_workers=8 导致数据被无意重复 8 次，根源在于 PyTorch 多进程数据加载架构对流式迭代器的处理机制：

1. 独立子进程深拷贝隔离（Process Forking & Memory Independence）：
   • 当 DataLoader 配置 num_workers > 0 时，PyTorch 通过 multiprocessing 衍生出 8 个完全独立的操作系统工作进程。每个 Worker 进程在其独立的虚拟地址空间中各自深拷贝并维护一份专属的 dataset 实例；
2. 流式迭代器的“分片盲区”（Stream Iterator Blindness）：
   • 对于基于索引的传统 Dataset（Map-style），主进程中的 BatchSampler 集中掌控全局索引，向各个 Worker 派发互斥的样本索引切片，各 Worker 仅按 Index 读取，天然无冲突；
   • 而对于流式 IterableDataset，数据是由每个 Worker 内部的 __iter__() 函数自主生成并推送的。若开发者未在 __iter__() 中通过 torch.utils.data.get_worker_info() 探测当前 Worker 的 worker_id 并主动切分数据流（例如按 Worker ID 划分文件列表或按步长交错切片），8 个子进程将全部从头开始完整遍历相同的数据流；
3. 最终后果：8 个子进程各自产出一份完全相同的全量数据并压入共享内存队列，导致单卡内部的数据被机械重复消费了 8 次，不仅浪费了 87.5% 的算力，更破坏了数据集的真实分布。

Q3 参考答案

在弹性训练（Elastic Training）中，当集群发生动态缩容或扩容（如总卡数从 64 卡变更为 32 卡）时，必须通过三项核心机制确保断点续训的零重复与零遗漏：

1. 持久化全局绝对样本游标（Global Invariant Coordinate）：
   • 严禁在 Checkpoint 中仅记录各卡私有的局部步数（local_step）。因为 64 卡下的一步对应消费了 64 个全局样本，拓扑缩容到 32 卡后，原先私有的步数坐标完全丧失物理意义；
   • 必须在 Checkpoint 中统一记录全集群原子提交的全局样本绝对累积位置（Global Consumed Sample Cursor / Token Count）： $$G = \text{global\_step} \times \text{batch\_size} \times P_{\text{old}}$$
2. 基于新拓扑动态重构分片几何（Dynamic Topology Re-sharding）：
   • 重启时，各卡读取全局游标 $G$，依据新的总卡数 $P_{\text{new}}$ 与新 Rank ID 重新推导分片规则；
   • 采样器根据确定性置换序列，直接切除前 $G$ 个已完全提交的全局样本；
   • 新拓扑下的各个 Rank 严格从全局第 $G$ 个样本之后重新交错切分（$\text{new\_local\_start} = \lfloor G / P_{\text{new}} \rfloor$），确保全体活跃卡的数据并集精确从 $G+1$ 处无缝向后延伸；
3. 清空异步预取缓冲并严格以已提交 Step 对齐（Drain Prefetch Buffer）：
   • 训练崩溃往往发生在 Step 内部，各卡内部的 DataLoader 预取管道中可能残留数个未被前向计算读取的临时 Batch；
   • Resume 时必须清空所有 Worker 预取流，强制以“最后一次完整执行了 optimizer.step() 并成功落盘的全局 Step”作为截断点，实现确定性的严格 Exactly-Once 数据一致性。

参考资料

• PyTorch DistributedSampler implementation
• PyTorch IterableDataset Multi-Worker Guidelines
• PyTorch Distributed Checkpoint (DCP) Design
• Megatron-LM Data Loading & Exact Resume Architecture

API 与平台能力会演进；部署前应按目标版本重新核对。

← 上一课
第 4 课：Collective 成本模型与… 下一课 →
第 6 课：存储→CPU→Pinned Me…