GPU 在等数据，具身模型训练中常被忽视的 DataLoader

无论训练大语言模型还是具身模型GPU 侧的计算与通信都是首先需要优化的环节。但对以视频为主要训练数据的 VLA 和 WAM 来说仅仅优化 GPU 还不够。它们的训练样本难以全部提前处理并存储所需视频帧往往要在训练过程中按需读取、在线解码。数据准备一旦跟不上计算GPU 就得停下来等。数据量小的时候每步多等一点并不显眼。一轮训练达到几十万步后这些等待累积起来就是可观的算力浪费。因此在优化具身模型训练性能时数据供给也必须作为一个独立环节进行排查。百度百舸针对 DataLoader 做了全链路优化覆盖采样、索引、解码、后处理各个阶段。世界动作模型WAM方面FastWAM 的单样本视频处理耗时从 133.0 毫秒降到 63.3 毫秒优化后的索引耗时在数据集从 1 个文件增长到 27,500 个文件的过程中保持稳定。VLM 方面Qwen3-VL-30B 多模态 SFT 通过采样优化训练吞吐提升至原来的 1.68 倍。1. 组 batch 是什么具身模型难在哪里训练的每一步由两部分组成。一部分在 GPU 上前向、反向、参数更新。另一部分在 GPU 之外把样本从磁盘取出、解码成张量、做后处理、拼成一个 batch 送进显存。本文把后一部分统称为「组 batch」严格来说样本拼接只是其中的最后一步。这一段由 DataLoader 组织消耗的是 CPU 和 I/O。两部分是流水线关系。只要组 batch 跟得上计算GPU 就不会空闲。跟不上GPU 每一步都要停下来等。买到的算力能否真正转化为有效算力关键就看 GPU 需要等待数据多久。大语言模型的文本通常可以在训练前完成分词并以 token 序列形式存储。训练时主要是读取 token 序列、采样和组 batch不需要像视频数据那样进行随机定位和解码数据准备更容易被计算覆盖。以独立图片为输入的视觉语言模型VLM虽然增加了图片解码但整体数据加载流程仍与大语言模型较为接近。VLM 也可以接视频输入但本文的 VLM 案例是图片输入下表按此对比。具身的视觉语言动作模型VLA与世界动作模型WAM则不同。它们的样本是从一段 episode 里切出的时间窗口起点随机、相邻窗口大量重叠一条 episode 能切出成百上千个样本。这些窗口并不会预先切分并存储。把所有可能的窗口提前裁出来解码存好并非做不到但存储会成倍膨胀采样方式也随之固定所以实际训练通常保留原始 episode训练时随机确定采样窗口再在线定位关键帧、解码、逐帧处理。以下对比针对本文涉及的场景VLM图片输入VLA 与 WAM每个样本的视觉输入一到数张彼此独立的图片多相机、多帧几十帧量级时间维无帧要连续且要与动作时间戳对齐存储形态独立图片文件视频流访问方式整张读取无需定位随机访问先定位到最近的关键帧再连续解到目标位置样本形状随文本长度和图片尺寸变化需填充窗口长度、相机数由配置固定天然对齐具体到训练过程两类模型的数据加载方式存在以下差异VLM 的数据准备较轻容易被计算覆盖VLA 与 WAM 的解码和逐帧处理集中发生在训练过程中更容易形成前面所说的 GPU 等待。2. 优化目标与排查方法DataLoader 优化的目标不是把组 batch 压到极致而是让它被计算掩盖。具体说组 batch 耗时小于单步计算耗时即可。批量大小batch size由训练侧先定在显存、收敛和训练策略允许的范围内调到 GPU 利用率饱和单步计算耗时随之确定。在此基础上再通过数据侧优化将组 batch 的实际耗时压到单步计算耗时以内。要系统地排查 DataLoader先把它拆开。具身场景主要有 VLM、VLA、WAM 三类模型。三者的输入组成并不完全相同但从 DataLoader 的角度看核心任务是一致的把存储中的压缩数据变成一步训练能直接使用的张量组。VLM 处理的是视觉与文本VLA 与WAM 还要处理动作和本体状态。这个过程可以拆成五个阶段。采样和索引决定取什么数据解码和后处理决定怎么把数据变成模型输入拼张量把多个样本组织成一次计算的输入。三类模型可以用同一个阶段框架分析差别在各阶段耗时占比。具体的优化手段和正确性校验仍要按模型和数据形态调整。阶段做什么VLM图片输入VLA 与 WAM采样决定这一步看哪些样本、什么顺序、分给哪个 rank样本彼此独立随机排列即可。难点在各 rank 负载均衡采样单位是轨迹上的时间窗口同一轨迹切出的窗口高度相关、相邻窗口内容重叠需要打散多数据源还要设混采权重索引把采样编号翻译成文件、偏移、帧号区间一个样本对应一个行号列式存储直接给出偏移一个样本由轨迹编号、起始时刻、窗口长度共同确定这套映射要自己维护解码把压缩字节转成张量I/O 与 CPU最密集的阶段单张图片解码随机访问的视频解码后处理整理成模型要求的形状、数值范围序列截断与模板拼接多数可离线抽帧、缩放、归一化多数需在线完成拼张量多条样本拼成一步计算的张量组序列长度随内容变化需填充到 batch 内最大长度形状由配置固定拼接简单本文展开前四个阶段。拼张量主要依赖框架已有实现不单独展开。优化本身分两步按成本从低到高。第一步调 DataLoader 参数。PyTorch 已把组 batch 的完整逻辑封装在 DataLoader 里并开放了并行进程数num_workers、预取深度prefetch_factor、锁页内存pin_memory、进程跨 epoch 存活persistent_workers等参数。调参不改训练代码成本最低作用是先确认瓶颈是否存在以及能否靠增加 CPU 核数、存储带宽这类资源直接解决。调参之后瓶颈仍在就进入第二步按上面的阶段逐段排查找到耗时占比最高的阶段做针对性优化。本文后面的实践全部属于第二步。3. 沿数据加载链路逐阶段优化以下实践均在百度百舸平台 hpas.lgn7ib 实例上完成该实例搭载国内主流 GPU 型号。四个案例里采样案例来自 VLM 训练索引、解码、后处理三个案例来自 VLA 与 WAM 训练。除采样一项给出训练吞吐外其余为对应环节的耗时对比。3.1. 采样让各 rank 的计算负载均衡采样决定每一步训练看哪些样本以及样本如何分配到各个 rank。对 VLM 来说样本序列长度随内容变化采样直接决定各 rank 的负载是否均衡。多卡训练中每一步的耗时由最慢的 rank 决定分配不均时处理短样本的 rank 只能在梯度同步处等待。在 LLaMA-Factory 框架上对 Qwen3-VL-30B 做多模态 SFT单机 8 卡数据并行时默认采样器做全局随机排列各 rank拿到的样本文本长度和视觉 token 数差异较大因此出现了负载不均。优化方法是设计一个负载均衡的 Sampler综合考虑输入序列长度和视觉 token 数带来的计算开销。首先按序列长度对样本降序排列让同一个 global step 中的样本长度尽量接近减少 padding。然后逐个分配样本每次将其放入后计算成本增量最小的 rank并保证各 rank 拿到相同数量的样本使各 rank 的总成本尽量接近。最后打乱物理 rank 编号和 global step 顺序避免某个 rank 长期承担高成本样本也避免此类样本集中出现在训练前期。单机 8 卡实测训练吞吐从 27.14 samples/s 提升到 45.48 samples/s提升 1.68 倍。3.2. 索引避免耗时随文件数增长索引把采样给出的编号翻译成文件、偏移和帧号区间。在 VLM 数据组织方式中一个样本通常对应一行记录因此索引开销较小。VLA 与 WAM 的一个样本是轨迹里的一段时间窗口一个样本编号要先换成「哪条轨迹、从哪一帧起、取多长」再换成「哪个文件、哪几行、视频里哪个时间戳区间」。框架不会自动维护从样本编号到文件位置的映射关系需要训练系统自行建立并查询。数据集一大仅查询这层对应关系就可能成为负担。FastWAM 训练使用 RoboTwin 数据集 LeRobot v2.1 格式每条轨迹一个 parquet 文件每个相机一个视频文件一个样本是 33 帧乘 3 相机的窗口加一段动作序列。取样本之前要先从数据集里查出要解码的视频时间戳和状态、动作列。原实现用 Hugging Face datasets 的 select 接口取数这个接口适合构造数据集子集但在逐样本取数的场景中每次都会额外构造一个数据集视图视图的索引结构和描述信息处理开销会随文件数增加。当数据集从 1 个文件长到 27,500 个文件约 607 万行时单次索引耗时从 0.90 毫秒涨到 81.1 毫秒。优化方法是换用同一个库自带的列表索引取数fancy indexing直接按已有位置关系返回目标记录绕开这层额外开销。优化后单次索引耗时保持在 0.61 毫秒左右不再随文件数增长。这个案例也说明数据集的文件组织方式会直接影响索引性能LeRobot v2.1 格式按轨迹拆分 parquet 文件文件数接近轨迹数select构造数据集视图及处理相关索引结构的开销会随文件数增加LeRobot v3.0 格式将多条轨迹组织到更少的 parquet 文件中减少了文件数量从存储组织层面缓解了文件数过多带来的索引开销。3.3. 解码减少构建开销只解必需帧解码是 DataLoader 里 I/O 和 CPU 最密集的阶段也是两类模型差别最大的阶段VLM 解的是独立图片整张读取开销不大VLA 与 WAM 要在视频流里随机访问解码开销明显更高本文后面的解码优化均针对VLA 与 WAM。对 VLA 与 WAM 的视频解码来说一次解码调用的开销可以拆成两部分。第一部分是构建解码对象并定位到关键帧。无论单次需要解码多少帧这部分固定开销都会产生。在随机访问场景中可以看作固定开销大小受视频编码格式和关键帧间隔影响。第二部分是逐帧解压耗时通常随解码帧数增加。一个样本的总解码开销是各次解码调用开销的总和。在本文场景中调用次数主要由相机数和分段数决定。百舸从三个方面优化解码解码库的选择、只解需要的帧、合并同一样本的多次调用。解码库的选择用 UR5e 数据集的训练视频对比两个主流视频解码库 decord 与 torchcodec。逐帧解压效率两者接近解 9 帧分别为 86.25 毫秒和 73.05 毫秒。差别在固定开销每次新建解码对象decord 需要 141.04 毫秒torchcodec 为 2.54 毫秒。随机访问场景下每次打开视频并构建解码对象都会产生一次固定开销一个样本涉及多个相机或多个分段时这笔开销累积得越明显。从 decord 换到 torchcodec 要注意两点解码结果要做像素一致性校验近似定位模式approximate seek在时间戳不规整的视频上可能取到不同的帧。解码库与 PyTorch 版本存在二进制依赖升级时要同步确认。只解需要的帧FastWAM 的视觉窗口包含 33 帧但进入模型前只需均匀采样其中 9 帧以降低视觉编码器的计算量。原有流程中抽帧位于数据处理链路下游解码层不知道下游需要哪些帧仍会将 33 帧全部解出其中 24 帧随后被丢弃。优化方法是将下游确定的 9 帧索引提前传递给解码层只解码实际需要的帧。优化前流程是「解码 33 帧、后处理 33 帧、最后抽出 9 帧」。优化后则是「先确定 9 帧、只解码这 9 帧」。解码耗时从 111.58 毫秒降到 50.97 毫秒加速比为 2.19 倍。若解码耗时完全随帧数线性变化33 帧减少到 9 帧的理论加速上限为 3.67 倍。实际加速比低于这一上限主要是因为构建解码对象和定位关键帧的固定开销不会随解码帧数减少。合并同一样本的多次调用这个案例的训练对象是一个新模型。视觉输入参考了 Pi0.7 的设计将单个时间片段扩展为三个片段当前帧用于表征当前状态过去帧用于补充时序上下文未来帧则作为目标条件用于消除指令歧义。一个样本包含 3 个相机视角原来的 DataLoader 按「片段 × 相机」逐条调用解码3 个片段乘以 3 个相机每个样本需要 9 次调用。对同一个相机而言3 个片段来自同一个视频文件却会分别打开视频并构建解码对象因此每个视频文件被打开 3 次、解码对象被构造 3 次。由于这 3 个片段只是帧的位置不同仍然来自同一个视频文件因此可以合并同一视频所需的帧列表通过一次解码调用取出。这样每个相机只需打开一次视频、只需构建一次解码对象。无缓存时单样本解码从 77.03毫秒降到 49.06 毫秒加速比为 1.57 倍。有缓存时加速比为 1.10 倍耗时降低约 9%。因为文件已经在缓存里打开文件的开销本来就小合并调用主要减少的是解码对象的重复构建收益相应降低。3.4. 后处理避免处理最终不会进入模型的帧后处理把解出的帧转换为模型要求的形状和数值范围。VLM 的后处理多数可以在训练前完成VLA 与 WAM 的抽帧、缩放随随机窗口变化通常需要在线处理每一帧的算子开销都会落在训练过程中。FastWAM 中上一节只解 9 帧的优化留了一个尾巴。下游后处理的接口是按照 33 帧设计。为了不改下游代码解码层把 9 帧复制回填成 33 帧后再交给下游。于是后处理链上的逐帧算子数值转换、缩放、张量规整仍会在 33 帧上执行最后才由链尾的抽帧操作丢弃其中 24 帧这部分计算没有产生有效收益。抽帧只是选择帧不改变帧内容。对于逐帧算子而言先抽帧再处理与先处理再抽帧结果相同因此可以将抽帧移到这些算子之前。与解码优化合并后解码直接输出 9 帧不再回填后处理也只处理这 9 帧。优化前解码 33 帧、逐帧算子处理 33 帧最后抽出 9 帧优化后解码直接输出 9 帧逐帧算子也只处理 9 帧最终不会进入模型的帧不再经过这些处理。两项优化叠加后FastWAM 完整视频处理链的单样本耗时从 133.0 毫秒降到 63.3 毫秒加速比为 2.10 倍。具体来看只解需要的 9 帧但仍保留 33 帧回填时耗时降至 73.5 毫秒加速比 1.81 倍进一步移除回填操作后耗时降至 63.3 毫秒加速比为 2.10 倍。4. 结语别让 GPU 等数据模型和数据集的文件组织方式不同DataLoader 的瓶颈也会变化。优化的关键是找到当前链路中真正让 GPU 等待的环节。本文的这些优化全部在搭载国内主流 GPU 的实例上完成。百度百舸始终立足国内现有的算力供给格局通过系统性的 AI Infra 优化与高扩展性集群架构设计为各类具身智能场景找到高性价比的算力解决方案。