
XTuner LengthGroupedSampler 数据分组采样器告别 Pad 浪费把训练吞吐提上来【免费下载链接】xtunerA Next-Generation Training Engine Built for Ultra-Large MoE Models项目地址: https://gitcode.com/GitHub_Trending/xt/xtuner导读大模型训练数据的长度天然参差不齐而 Transformer 又是定长输入模型batch 内填充Pad带来的算力浪费往往高达三四成。本文基于 XTuner 的LengthGroupedSampler基于数据长度分组的采样器讲清楚它如何通过同批样本长度相近把填充开销压到最低、与数据拼接Pack方案各自的取舍并给出可直接复制的配置修改方法以及源码级的实现原理与多模态场景InternVL下的真实用法。读完你就能在自己的 SFT / LoRA / 多模态微调配置里一键切换长度分组采样立刻看到训练吞吐的提升。问题背景不定长数据与 Pad 浪费生成式大模型如 LLM的训练数据往往是不定长的这导致同一个 batch 内的数据长短不一。为了实现并行化训练一种常见做法是将同一批次的数据填充Pad到该 batch 内最长样本的长度。然而Pad 操作引入的无意义 token会带来训练的低效。原文档给出了一组直观的示意数据假设数据内各样本的长度分别为2、3、7、9期望分为 2 个批次进行训练如果使用默认的随机采样器左数据处理阶段会引入过多的填充数据实际效率只有65.6%如果使用基于数据长度分组的采样器右同样的数据效率可以提升至87.5%。效率的差距来源很清晰随机采样把2、3与7、9混在一起两个 batch 都要按最长样本3和9填充浪费了大量计算而按长度分组后2、3一组、7、9一组每个 batch 内部的长度差被压缩到最小。填充实际发生在 collate 阶段default_collate_fn通过pad_sequence将 batch 内样本填充到ori_length的最大值并用IGNORE_INDEX屏蔽填充位置的 loss见 default_collate_fn.py填充越多有效计算占比就越低。两条技术路线数据拼接 vs 数据分组现阶段有两种技术方案可以解决 / 缓解这一问题两者选其一即可XTuner 官方优先考虑数据拼接技术方案一数据拼接Pack to Max Length利用数据拼接技术将多条数据拼接至训练支持的最大长度。这一做法可以确保同一批次内的数据长度完全一致进而彻底避免填充数据导致的训练效率降低。优点可以合并多个数据样本显著降低训练 iter 数加速效果好缺点随机合并的多个数据样本间会互相影响进而影响训练效果实际影响程度未知数据进行了合并丢失了一定数据随机性。XTuner 的pack_to_max_length参数即对应这一方案相关配置与性能对比可参考 数据拼接文档其中 InternLM2 7B 在 Alpaca 数据上的测试显示拼接策略可带来数倍的每秒 token 处理量提升。方案二数据长度分组采样器本文主题利用基于数据长度分组的采样器在构建批次数据时基于实际长度进行排序确保同一批次内的数据长度尽可能相近进而尽可能减少填充的长度。优点每条数据依然独立存在独立计算 attention避免了数据拼接技术导致的数据样本间的互相影响数据进行了分组丢失了一定数据随机性缺点在数据样本长度比较一致的情况下加速效果一般。一句话总结取舍拼接追求长度绝对一致分组追求长度尽量相近。前者加速上限更高但样本间会互相串扰后者保持样本独立适合对数据独立性有要求、或暂时不想改变数据语义的场景。配置实战让 XTuner 用上 LengthGroupedSamplerXTuner 中基于数据长度分组的采样器实现在 xtuner/dataset/samplers/length_grouped.py用户可以通过在配置文件中修改train_dataloader的sampler参数进行配置。以 internlm2_chat_7b_qlora_oasst1_512_e3.py 配置文件为例注意该 config 中默认使用DefaultSampler且pack_to_max_length False这正是适合启用分组采样的场景我们通过下列修改使其使用基于数据长度分组的采样器- from mmengine.dataset import DefaultSampler from xtuner.dataset.samplers import LengthGroupedSampler batch_size 16 # per_device accumulative_counts 1 train_dataloader dict( batch_sizebatch_size, num_workersdataloader_num_workers, datasettrain_dataset, - samplerdict(typeDefaultSampler, shuffleTrue), samplerdict( typeLengthGroupedSampler, length_propertylength, per_device_batch_sizebatch_size * accumulative_counts), collate_fndict(typedefault_collate_fn, use_varlen_attnuse_varlen_attn))修改完成后直接运行xtuner train即可。训练日志中会出现LengthGroupedSampler is used.与LengthGroupedSampler construction is complete, and the selected attribute is length两条提示源码见 length_grouped.py 与 length_grouped.py用于确认采样器已生效。关键参数说明参数含义说明length_property获取数据集长度的属性名通过process_hf_dataset构建数据集时会自动设置为length自定义数据集需自行保证该属性正确per_device_batch_size单卡每次迭代的样本数应传入batch_size * accumulative_counts单卡 batch 与梯度累积的乘积即虚拟 batch的大小mega_batch_mult分组粒度倍数可选控制每个 megabatch 的大小默认取min(len(dataset) // (total_batch_size * 4), 50)小数据集自动回退为 1seed随机种子可选默认为sync_random_seed()保证分布式下各 rank 一致round_up是否向上取整补足样本可选默认True保证每个 rank 迭代次数一致便于梯度累积对齐length_property 为何是 length注意length_property需要传入获取数据集长度的属性这一数值在通过process_hf_dataset构建数据集时会自动设置为length因此如果使用自定义的数据类请确保这一属性的正确设置。从源码看process_hf_dataset的底层process函数在完成 tokenization、过滤、packing 之后会执行# add length dataset dataset.map(get_lengths, num_procmap_num_proc) setattr(dataset, length, dataset[length])其中get_lengths定义为return {length: len(example[input_ids])}见 huggingface.py 与 huggingface.py即给每条样本打上length列并挂到数据集对象上。LengthGroupedSampler在构造时通过getattr(self.dataset, length_property)读取该属性见 length_grouped.py并断言其类型为list或tuple。若数据集是torch.utils.data.ConcatDataset则会遍历所有子数据集并把各自的length拼接起来。自定义数据集时最容易踩的坑如果你没有走process_hf_dataset而是自定义了数据集类必须手动给数据集对象设置length属性例如dataset.length [len(x[input_ids]) for x in ...]否则采样器会因getattr取不到属性而报错。源码深度解析分组采样是怎么实现的LengthGroupedSampler的完整实现位于 xtuner/dataset/samplers/length_grouped.py核心逻辑集中在get_length_grouped_indices函数与__iter__方法中。三步核心算法随机分块 块内排序 最大块前置get_length_grouped_indices的实现见 length_grouped.py可以拆解为三步随机排列torch.randperm(len(lengths))对全部样本索引做一次随机打乱保证训练仍有足够的随机性切分 megabatch把打乱后的索引按group_batch_size切成若干 megabatchgroup_batch_size mega_batch_mult * total_batch_size其中total_batch_size per_device_batch_size * world_size见 length_grouped.py块内按长度降序排序对每个 megabatch按lengths[i]从大到小排序。这样相邻样本长度接近切出来的每个小 batch 长度分布自然紧凑。排序完成后还有一个细节处理将最长元素所在的 megabatch交换到最前面见 length_grouped.py。这样可以让训练早期就遇到较长的样本同时保证最终索引序列中先长后短有利于流水线 / 梯度累积场景下的资源均衡。多模态场景用正负长度区分两种模态get_length_grouped_indices对正负长度做了特殊处理见 length_grouped.py若所有样本长度同号全正或全负视为单模态直接走统一流程若正负混杂则把length 0的样本视为多模态样本multimodallength 0的样本视为纯语言样本language先各自分组排序再把各自的最后一个不完整 megabatch 合并成最终批次最后对 megabatch 序列整体再随机打乱一次末尾还有断言assert all(leng ! 0 for leng in lengths), Should not have zero length.即不允许出现长度为 0 的样本。这就是为什么 InternVL 系列的配置里length_property用的是modality_length而不是length。以 internvl_v2_internlm2_5_8b_finetune.py 为例train_dataloader dict( batch_sizebatch_size, num_workersdataloader_num_workers, datasetllava_dataset, samplerdict( typeLengthGroupedSampler, length_propertymodality_length, per_device_batch_sizebatch_size * accumulative_counts), collate_fndict(typedefault_collate_fn))在多模态训练中图像 token 多的样本和纯文本样本混在一个 batch 里会显著加剧填充浪费因此按modality_length模态有效长度分组能把含图的长样本和纯文本样本分别聚拢效果尤其明显。事实上XTuner 的 llava 配置目录 与 internvl 配置目录 下的大量微调 config 都已默认启用了LengthGroupedSampler可以直接作为参考。分布式与 epoch 控制LengthGroupedSampler继承自torch.utils.data.Sampler其分布式语义如下构造时通过get_dist_info()获取rank与world_size见 length_grouped.py当round_upTrue默认时num_samples ceil(len(dataset) / world_size / per_device_batch_size) * per_device_batch_sizetotal_size num_samples * world_size即把数据集长度向上取整补足保证每个 rank 的迭代步数一致__iter__中先用torch.Generator().manual_seed(self.seed self.epoch)固定随机状态见 length_grouped.py随后按indices[self.rank:self.total_size:self.world_size]做分片采样见 length_grouped.py保证各 rank 拿到互不重叠的样本且每个 epoch 顺序不同提供set_epoch方法见 length_grouped.py与 XTuner 的DistSamplerSeedHook配置中default_hooks里的sampler_seed配合多轮 epoch 时每轮重新洗牌。与数据拼接、变长注意力如何配合在 XTuner 的加速工具箱里LengthGroupedSampler与另外两个特性是互补关系与 Pack 的关系两者二选一。pack_to_max_length True时走拼接路线此时 batch 内长度天然一致无需分组采样pack_to_max_length False时如 internlm2_chat_7b_qlora_oasst1_512_e3.py数据按原样参与训练此时启用分组采样能把填充浪费降到最低与变长注意力的关系use_varlen_attn True时process_hf_dataset会强制要求pack_to_max_length True断言见 huggingface.py且default_collate_fn要求 batch size 为 1见 default_collate_fn.py此时填充问题被变长注意力从机制上消解同样不需要分组采样。换句话说分组采样是不使用拼接、不使用变长注意力这条路径上的关键优化手段。常见问题与调优建议训练日志没有出现LengthGroupedSampler is used.检查 config 中train_dataloader.sampler.type是否确实改成了LengthGroupedSampler并确认from xtuner.dataset.samplers import LengthGroupedSampler已加入文件头部导入报错Dataset object has no attribute length自定义数据集未设置length属性请在构建数据集后手动挂载长度列表参考process_hf_dataset的setattr(dataset, length, ...)写法报错Should not have zero length.数据集中存在长度为 0 的样本需要在预处理阶段过滤掉空样本加速不明显正如原文档指出的当数据样本长度本身比较一致时分组采样的加速效果有限——此时可以评估是否转向pack_to_max_length拼接方案per_device_batch_size的取值务必传入batch_size * accumulative_counts单卡实际 batch × 梯度累积次数它决定了一个分组内包含的样本规模取小了分组粒度太碎、取大了又接近随机采样两者都会削弱分组效果。小结LengthGroupedSampler是 XTuner 在不改变数据语义的前提下提升训练效率的轻量方案一条sampler配置即可生效底层通过随机分块 块内按长度降序排序 最长块前置的组织方式让每个 batch 内部的长度差最小化从而显著降低 Pad 浪费其正负长度机制还天然适配 InternVL / LLaVA 等多模态训练场景。结合 数据拼接文档、变长注意力文档 一起阅读可以更完整地掌握 XTuner 在数据组织层面的整套加速方法论。【免费下载链接】xtunerA Next-Generation Training Engine Built for Ultra-Large MoE Models项目地址: https://gitcode.com/GitHub_Trending/xt/xtuner创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考