1. 数据加载为什么在大模型场景下成了独立的硬骨头先抛一个可能不太讨喜的判断在多模态大模型预训练和大模型微调里数据加载环节翻车的概率比模型结构写错还要高。很多人第一次跑多模态预训练脚本模型结构、损失函数、优化器都照着开源代码抄得一字不差结果训练三个小时 loss 纹丝不动回头一查是图文对没对齐或者显存莫名爆掉回头一查是 collate_fn 里 padding 的长度算错了。这类问题的共同特征是——它们不会在编译期报错只会在训练曲线里沉默地折磨你。我接触过的几个做多模态预训练的团队几乎都经历过同一个阶段前期花两周调模型后期花四周调数据管道。这不是能力问题而是数据加载在多模态场景下确实变了性质。传统的图像分类任务里一个样本就是一张图加一个整数标签Dataset 返回(tensor, int)collate_fn 默认的torch.stack就能搞定。但到了多模态大模型这里一个样本可能是一段变长文本、一组分辨率不一的图像、若干个区域标注、一个对话轮次结构甚至带上图像在原文中的位置索引。样本从定长张量退化成了异构字典DataLoader 默认的拼装逻辑直接失效。所以这篇内容想解决的不是怎样写一个 Dataset 类这种入门问题而是在多模态预训练和大模型微调这两个具体场景下数据加载到底该怎么设计、每个设计选择背后的取舍是什么、哪里最容易埋雷。适合已经能跑通单卡小 demo、正在往真实训练规模推进的读者也适合带团队做训练基础设施、需要给数据管道定规范的人。下面我会把 Dataset 结构、collate_fn、预训练与微调的分野、分布式下的采样、以及排查链路拆开讲每个环节都给出我实际踩过的坑和对应的处理方式。1.1 从定长张量到异构字典样本形态的根本变化理解这个变化的起点是看清楚一个 batch 的张量到底长什么样。图像分类时代batch 是一个[B, 3, H, W]的规整张量所有样本共享同一套维度DataLoader只需要沿着第 0 维堆叠即可。到了多模态大模型输入变成了类似{input_ids: [B, L], pixel_values: [B, C, H, W], attention_mask: [B, L], labels: [B, L]}这样的字典而且每个字段的 L 和 H、W 都可能因样本而异。文本长度从几十 token 到几千 token 不等图像从 224 到 1024 分辨率混杂如果还指望默认 collate 帮你处理只会得到一个维度不匹配的报错或者更糟——一个能跑但语义错位的张量。这里有个容易被忽略的点多模态的多不只是模态数量多更是异构维度多。文本侧是变长序列视觉侧是变分辨率张量两者在 batch 内的对齐关系必须由你显式维护。比如一批 8 个样本第 3 个样本带了 2 张图、第 5 个样本带了 1 张图那么pixel_values的第 0 维到底是 8 还是 10注意力掩码怎么标记哪张图属于哪个样本这些都不是框架能替你决定的必须由 collate_fn 里的逻辑明确表达出来。我的经验是在写任何一行数据代码之前先手写三个真实样本的完整字典结构打印出每个字段的形状和 dtype然后模拟拼成一个 batch 看形状是否自洽。这一步花十分钟能省掉后面几小时的 debug。很多翻车都源于以为字段是定长的实际不是。1.2 三个反复出现的翻车场景场景一图文错位。多模态预训练最经典的隐性 bug。Dataset 每次__getitem__返回图像和文本时如果两者来自不同的随机读取源或者图像做了随机裁剪、文本做了随机截断但没有共享同一个随机种子就会出现这批图配的是另一批文本的情况。训练照样能收敛loss 也会降但模型学到的是噪声对齐下游任务表现会明显低于预期。判断方法很简单固定一个 batch把图像和文本还原成可读形式人工核对如果对不上问题就坐实了。场景二变长 padding 把有效 token 淹掉。大模型微调时一个 batch 里既有 50 token 的短样本又有 2000 token 的长样本。如果用批内最长做 padding短样本后面全是 pad注意力掩码一旦写错比如 pad 位置也参与 attention模型就会把大量无效位置当成上下文loss 降得慢且不稳定。更隐蔽的是某些实现里 pad token 和 eos token 用了同一个 id模型分不清句子结束和填充生成阶段就会出现莫名的截断。场景三worker 数量一调结果就变。单进程num_workers0时一切正常改成 8 个 worker 后 loss 曲线开始抖或者某些 batch 的样本重复出现。根源通常是 Dataset 里用了全局随机状态或者自定义 Sampler 在 fork 后各自持有一份副本导致每个 worker 的采样序列不一致。这类问题在分布式训练里会进一步放大因为每个 rank、每个 worker 都在独立采样全局的不重复需要显式设计。2. 多模态Dataset的结构设计样本粒度决定一切Dataset 的设计没有唯一正确答案但有一个核心决策点样本粒度到底切到多细。是把一张图 一段描述作为一个样本还是把一次完整对话 若干张图作为一个样本这个选择会直接决定后续 collate_fn 的复杂度、内存占用模式、以及分布式切分是否均匀。我的倾向是预训练阶段按最小可训练单元切微调阶段按一次交互切。预训练数据量巨大样本应该尽量小、尽量均匀这样批次内 padding 浪费少、shuffle 效果好微调数据量小、质量高样本应该保持完整语义单元避免把一个多轮对话拆散导致上下文断裂。这个差异后面第 4 节还会展开这里先把 Dataset 的通用骨架搭起来。2.1 单样本字段布局与占位策略一个设计良好的多模态样本我通常会包含这几类字段输入侧文本 token 序列、图像张量或图像路径、模态类型标记、监督侧labels、loss 掩码、元信息样本 id、来源、原始长度。元信息不参与前向计算但对排查问题至关重要——当 loss 异常时你需要能凭样本 id 回溯到原始数据。这里有个实操细节图像字段到底存张量还是存路径。存张量简单但 Dataset 对象会常驻内存一个百万级样本的预训练集如果每张图都预解码成张量内存直接爆掉。存路径则需要在__getitem__里读盘解码IO 成为瓶颈。折中方案是存路径 在 worker 里解码 用pin_memory加速传输这也是我在实际项目里用得最多的组合。如果你用的是本地高速存储可以在 Dataset 初始化时缓存一部分高频样本的解码结果做一个简单的 LRU如果是对象存储务必把解码放在 worker 内并行做别在主进程里串行读。import torch from torch.utils.data import Dataset from PIL import Image class MultiModalSampleDataset(Dataset): def __init__(self, records, tokenizer, image_size336, max_text_len512): self.records records # 列表每项含 image_path / text / sample_id self.tokenizer tokenizer self.image_size image_size self.max_text_len max_text_len def __len__(self): return len(self.records) def __getitem__(self, idx): rec self.records[idx] # 文本侧截断 特殊 token enc self.tokenizer( rec[text], truncationTrue, max_lengthself.max_text_len, return_tensorsNone, ) # 视觉侧解码 归一化放在 worker 内完成 img Image.open(rec[image_path]).convert(RGB) img img.resize((self.image_size, self.image_size)) pixel torch.from_numpy( __import__(numpy).array(img, dtypefloat32) / 255.0 ).permute(2, 0, 1) return { input_ids: torch.tensor(enc[input_ids], dtypetorch.long), attention_mask: torch.tensor(enc[attention_mask], dtypetorch.long), pixel_values: pixel, sample_id: rec[sample_id], }这段骨架看起来平淡但有几个点值得强调。tokenizer调用时不要在这里做 paddingpadding 是批次级的决策放到 collate_fn 里统一处理Image.open之后必须convert(RGB)否则灰度图和带 alpha 通道的图会变成 1 通道或 4 通道拼 batch 时直接报错归一化放在 Dataset 里还是 collate 里都行但要全局统一最怕的是训练集归一化了一次、验证集忘了指标对不上还查不出来。2.2 惰性读取与预加载的取舍惰性读取每次__getitem__现读和预加载初始化时全部读进内存之间不是二选一而是按数据规模和访问频率分层。经验数值是数据集小于可用内存的 30%可以全量预加载超过这个比例用惰性读取加多 worker如果某些样本被高频访问比如对比学习里的难负样本池单独做一个小缓存。有个反直觉的结论多 worker 下预加载未必更快。因为 fork 之后每个 worker 都持有一份预加载数据的副本8 个 worker 就是 8 倍内存占用很容易触发 OOM 或者被系统 swap 拖慢。用torch.multiprocessing的共享内存或者把数据放到内存映射文件里能缓解这个问题但复杂度上升。我的建议是先按惰性读取实现把 IO 并行度交给 worker 数量实测 IO 确实是瓶颈之后再考虑缓存别一上来就做过度设计。3. collate_fn批次拼装环节最容易失控的地方如果整个数据管道只能挑一个地方重点加固我会选 collate_fn。原因很直接Dataset 的 bug 通常是局部的、单样本可复现的而 collate_fn 的 bug 是批次级的、和 batch 组成强相关的。同一份代码batch size 是 4 的时候跑得好好的改成 16 就出错这种问题排查成本极高。3.1 变长序列的 padding 与 packingpadding 的目标是让 batch 内所有序列等长但往哪边 pad、pad 什么值、mask 怎么建这三个问题必须一起回答。大模型微调里我固定用右侧 padding、pad id 取 tokenizer 的 pad_token_id、attention_mask 对应位置置 0。右侧 padding 的好处是生成任务里不会打乱前缀位置eval 时取最后一个非 pad 位置的 hidden state 也方便。def collate_pad(batch, pad_id0): max_len max(x[input_ids].size(0) for x in batch) input_ids, attn, labels [], [], [] for x in batch: n x[input_ids].size(0) pad_n max_len - n input_ids.append(torch.cat([x[input_ids], torch.full((pad_n,), pad_id, dtypetorch.long)])) attn.append(torch.cat([x[attention_mask], torch.zeros(pad_n, dtypetorch.long)])) labels.append(torch.cat([x[input_ids], torch.full((pad_n,), -100, dtypetorch.long)])) return { input_ids: torch.stack(input_ids), attention_mask: torch.stack(attn), labels: torch.stack(labels), pixel_values: torch.stack([x[pixel_values] for x in batch]), }注意 labels 的 pad 位置填的是-100而不是pad_id。这是交叉熵损失的约定-100的位置不计入 loss。如果这里填了 pad_id模型会被迫学习预测填充符短样本的有效梯度会被稀释。这个坑我见过不止一次症状是短样本占比高的 batch loss 明显偏低看起来学得很好实际是在学 pad。packing 是另一条路把多个短样本拼成一条长序列用位置 id 和注意力掩码区分样本边界。它能把 padding 浪费降到接近零预训练阶段收益明显但实现复杂度高且一旦掩码写错跨样本的注意力泄漏会让模型学到虚假关联。我的建议是预训练可以上 packing微调阶段老老实实 padding因为微调样本少、质量高padding 浪费几个 token 影响不大没必要冒实现风险。3.2 图像张量与文本张量的拼接顺序多模态 batch 里图像和文本的对齐关系必须由 collate_fn 显式保证。最稳妥的做法是让 Dataset 返回的每个样本严格一对一——一个文本对应一组图像collate 时按样本顺序堆叠图像那一维用[B, N, C, H, W]或者把 N 展平后配合一个image_counts列表。后者更灵活能处理每样本图像数不同的情况。def collate_multimodal(batch): out collate_pad(batch) # 展平所有图像并记录每个样本的图像数量 pixels, counts [], [] for x in batch: imgs x[pixel_values] # [N_i, C, H, W] pixels.append(imgs) counts.append(imgs.size(0)) out[pixel_values] torch.cat(pixels, dim0) # [sum(N_i), C, H, W] out[image_counts] torch.tensor(counts, dtypetorch.long) return out这样模型侧拿到的image_counts就是第 i 个样本文本对应第几张到第几张图的依据。顺序必须和 Dataset 返回的顺序完全一致任何额外的排序、过滤都会破坏对齐。我在一个项目里就遇到过 collate 里对图像按面积排序做效率优化结果图文全错位训练了两天才发现教训很痛。4. 预训练和微调两种场景下数据加载逻辑的分野很多人把预训练和微调的数据管道当成同一套代码改改路径结果两边都不顺。这两者的数据特征差异太大用同一套实现往往意味着两边都做了妥协。分开设计反而更省事。4.1 预训练海量、流式、弱监督多模态预训练的数据规模通常在千万到十亿级特点是单样本质量参差、无需精确的标签对齐、对吞吐量极度敏感。对应的设计取向是样本粒度切小、shuffle 用近似而非精确、异常样本直接跳过不阻塞。流式读取在这个场景下几乎必选。把整个数据集索引一次性加载进内存不现实用IterableDataset配合分片文件顺序读是更常见的做法。代价是失去了随机访问能力shuffle 只能靠分片内 shuffle 分片间 shuffle两级来做。这里要注意多个 worker 各自读不同分片时如果分片划分逻辑有重叠会出现样本重复划分有空洞又会漏样本。我通常会在初始化时把分片列表按 worker 数量取模分配并在日志里打印每个 worker 拿到的分片范围方便核对。异常样本的处理策略也很关键。预训练数据里总会有损坏的图片、空文本、超长序列。在__getitem__里 try/except 并返回 Nonecollate 时过滤掉 None是一种做法更稳的是在预处理阶段就把坏样本剔除训练时零容错。我倾向前者加日志计数因为完全剔除需要全量扫描成本高且新数据进来又要重扫。4.2 微调小规模、高质量、严格对齐微调场景完全反过来数据量可能只有几千到几十万条每条都经过人工或模型筛选样本之间的关系比如同一对话的多轮、同一组图的先后顺序必须保留。这里的核心不是吞吐而是不破坏语义结构。具体做法上我会把一次完整交互作为 Dataset 的一个样本collate 时按交互展开成训练序列。比如多轮对话微调Dataset 返回整个对话的 token 序列和每轮的边界索引collate 负责把 batch 内对话 padding 到等长同时把每轮起始位置整理成一个列表传给 loss 计算。这样做的代价是 batch 内长度方差大、padding 浪费多但换来的是上下文完整模型学到的因果关系不会因为样本切断而断裂。还有一个微调特有的问题验证集的数据加载必须和训练集共用同一套 tokenizer 和归一化参数但 shuffle 必须关掉。我见过验证集用了不同的 max_length导致评估指标和训练日志对不上排查了半天以为是模型问题。建议把预处理配置抽成一个显式对象训练和验证都从它取参数杜绝手写两份。5. DataLoader与分布式worker、shard和随机性的真实代价数据加载的性能问题八成出在 DataLoader 的配置上而不是 Dataset 本身。num_workers、pin_memory、persistent_workers、prefetch_factor这几个参数互相影响调错一个就可能让 GPU 利用率从 95% 掉到 20%。5.1 worker 数量、pin_memory 与持久化num_workers的经验起点是CPU 核数的一半到全部但真正的判断依据是GPU 是否在等数据。用nvidia-smi看 GPU 利用率如果长期低于 70% 且波动大说明数据供给不足加 worker如果利用率稳定在 90% 以上再加 worker 只会增加内存和上下文切换开销。pin_memoryTrue让数据放在页锁定内存里配合non_blockingTrue的.to(device)能实现异步拷贝这个组合在多模态场景下收益明显因为图像张量体积大。但要注意页锁定内存是有限资源worker 多、batch 大时可能耗尽表现为系统变慢而不是报错。persistent_workersTrue能避免每个 epoch 重新 fork worker适合 epoch 多、初始化重的场景代价是内存常驻。5.2 分布式采样器的分片逻辑分布式训练里DistributedSampler负责把数据切到各个 rank。它的默认行为是补齐到能整除也就是给数据量少的 rank 重复一些样本。这在微调时问题不大但在预训练时会导致某些样本被多看几次严格来说破坏了每个样本等权的假设。如果在意这一点可以设drop_lastTrue代价是每个 epoch 丢弃少量尾部数据。更隐蔽的问题是每个 rank 内部还有多个 worker。DistributedSampler只管 rank 间的划分rank 内的 worker 由 DataLoader 自己分配索引默认是连续切块。如果数据本身有顺序比如按来源排序的文件连续切块会导致每个 worker 拿到的样本分布偏移进而让 batch 内的数据多样性下降。解决办法是在 Dataset 初始化时对索引做一次全局 shuffle或者在 Sampler 层做 worker 级的交错分配。5.3 随机种子的可复现问题数据增强里的随机性是复现实验的最大障碍。random、numpy.random、torch三套随机源如果不同步同一个 seed 跑两次结果都可能不同。我的做法是在 worker 初始化函数里显式设置每个 worker 的种子种子由基础种子加 worker id 派生这样既保证可复现又保证不同 worker 的增强序列不重复。def worker_init_fn(worker_id): base_seed torch.initial_seed() % 2**32 seed base_seed worker_id random.seed(seed) np.random.seed(seed) torch.manual_seed(seed)配套地DataLoader 的generator参数也要固定否则 sampler 的 shuffle 顺序每次不同。这两处都设好同一个配置跑两次的 loss 曲线才能对得上否则你根本分不清指标波动是模型问题还是数据顺序问题。6. 排查链路数据加载慢、显存溢出、loss不降的定位顺序最后聊聊排查。数据管道的问题很少单独出现往往是显存溢出和loss不降同时发生让人误以为是模型结构的问题。我的经验是永远先排除数据侧再怀疑模型侧因为数据侧的问题更容易验证且修复成本更低。6.1 一条从外到内的排查路径第一步关掉多 worker用num_workers0跑 10 个 batch。如果单进程正常、多进程异常问题就在 worker 相关逻辑种子、共享状态、文件句柄。第二步把 batch size 设成 1打印每个字段的形状和 dtype。形状不自洽的问题在这一步就会暴露。第三步取一个固定 batch把图像还原成图片、文本 decode 成字符串人工核对对齐关系。图文错位只有肉眼能查出来。第四步检查 labels 的 mask 位置确认 pad 处是-100确认监督位置覆盖了该覆盖的 token。第五步才去看模型侧。显存溢出则要区分是稳定溢出还是随 batch 组成波动溢出。前者是容量问题减 batch 或上梯度累积后者通常是某个 batch 里出现了异常长的序列或异常大的图像在 collate_fn 里加长度和尺寸的上限截断并记录被截断的样本 id很快就能定位到脏数据。6.2 我常用的几个自检习惯第一个习惯在训练脚本启动时先跑一个数据 dry run只加载数据、不跑模型统计 100 个 batch 的文本长度分布、图像尺寸分布、每 batch 耗时。这个统计能提前暴露 80% 的数据问题而且不占 GPU。第二个习惯给每个 batch 打上样本 id 列表的日志采样频率不用高每 1000 步记一次就够。出问题时能凭 id 回溯到具体样本比盲猜快得多。第三个习惯把预处理参数集中到一个配置对象里tokenizer 的 max_length、图像的 size 和归一化均值方差、pad 的方向和 id全部从这里取。训练、验证、推理三处共用杜绝某处忘了改的低级错误。这些习惯看起来琐碎但数据管道的可靠性就是靠这些琐碎堆出来的。模型结构可以抄数据管道抄不了因为它和你的数据形态、存储方式、集群环境强绑定。把 collate_fn 和 Dataset 当成需要单测的代码来对待而不是一次写对的脚本这是我在多个项目里最深的体会。多模态预训练跑起来不难难的是跑得稳、跑得可复现而这两个目标的起点就在数据加载这几十行代码里。