无论训练大语言模型还是具身模型,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 优化与高扩展性集群架构设计,为各类具身智能场景找到高性价比的算力解决方案。