XTuner LengthGroupedSampler 实战用长度分组采样消除 Padding 浪费提升 LLM 训练吞吐【免费下载链接】xtunerA Next-Generation Training Engine Built for Ultra-Large MoE Models项目地址: https://gitcode.com/GitHub_Trending/xt/xtuner本指南聚焦 XTuner 训练引擎中基于数据长度分组的采样器LengthGroupedSampler讲解它在不定长样本场景下如何通过按长度相近分组来减少 Pad 填充带来的算力浪费。读完本文你将掌握该采样器的适用场景与局限、在配置文件中的完整接入方法含 diff 示例与参数说明、length属性的来源约定以及其底层实现原理与多模态训练的配合方式。问题背景不定长数据与 Padding 导致的双重浪费生成式大模型如 LLM的训练数据往往是不定长的同一批次batch内各样本的 token 数长短不一。为了实现并行化训练一种常见做法是将同一批次的数据统一填充Pad到该批次的最长长度后再喂给模型。然而这一填充操作会引入大量无效计算导致训练低效。以 length_grouped_sampler.rst 文档中的示例为例假设批次内 4 个样本的长度分别为 2、3、7、9期望分成 2 个批次训练。若使用默认的随机采样器文档图示左侧数据会被随机切分处理阶段可能引入过多的填充数据实际训练效率仅为65.6%而使用基于长度分组的采样器后文档图示右侧将长度相近的样本放入同一批次同样的数据效率可提升至87.5%。从源码层面看填充动作发生在数据整理阶段。default_collate_fn.py 中通过pad_sequence(input_ids, batch_firstTrue, padding_valuepad_index)将批内样本对齐到最大长度并基于ori_length [len(ids) for ids in input_ids]构建attention_mask来屏蔽填充位同时用IGNORE_INDEX屏蔽 labels 中的填充位。也就是说批次内最长样本的长度决定了该 batch 的整体计算量样本长短差异越大Pad 浪费的比例越高。两种主流缓解方案及其取舍现阶段有两种技术方案可以解决或缓解上述问题两者选其一即可优先考虑数据拼接技术方案核心思路优点缺点数据拼接技术Pack to Max Length将多条数据拼接至训练支持的最大长度确保同一批次内数据长度完全一致彻底避免填充可合并多个数据样本显著降低训练 iter 数加速效果好随机合并的多个数据样本间会互相影响进而影响训练效果实际影响程度未知数据被合并丢失了一定数据随机性长度分组采样器LengthGroupedSampler本文构建批次时基于实际长度排序确保同一批次内数据长度尽可能相近尽量减少填充每条数据依然独立存在独立计算 attention避免拼接导致的数据样本间互相影响在数据样本长度比较一致的情况下加速效果一般数据进行了分组丢失了一定数据随机性数据拼接技术的完整介绍见 数据拼接文档。本文聚焦第二种方案——LengthGroupedSampler。使用 LengthGroupedSampler配置修改实战XTuner 中基于数据长度分组的采样器实现位于 xtuner/dataset/samplers/length_grouped.py。用户通过在训练配置文件中修改train_dataloader的sampler参数即可启用。以 internlm2_chat_7b_qlora_oasst1_512_e3.py 配置文件为例它默认使用 mmengine 的DefaultSamplershuffleTrue可通过下列修改切换为基于数据长度分组的采样器- 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))注意上例中batch_size * accumulative_counts即每个设备的有效批量大小batch_size是每设备每步的前向批量accumulative_counts是梯度累积步数二者乘积才是参与同一批长度分组的样本数。参数说明length_property指定从数据集对象上读取长度的属性名。通过process_hf_dataset构建数据集时会自动将该属性设置为length。因此如果使用自定义数据集类必须确保该属性的正确设置详见下文length 属性从何而来一节否则采样器无法工作。per_device_batch_size每个 GPU 设备上的有效批量大小用于配合分布式训练计算分组规模。mega_batch_mult可选megabatch 的放大倍数默认值为min(len(dataset) // (total_batch_size * 4), 50)即默认目标为得到 4 个 megabatch上限 50若计算结果为 0极小数据集则回退为 1。seed可选随机种子不传时通过sync_random_seed()同步生成保证分布式各 rank 一致。round_up可选是否将样本数向上取整复制补足默认True保证每个 rank 的迭代数一致、便于分布式同步。与 varlen attention 的关系train_dataloader中collate_fn的use_varlen_attn参数控制是否启用变长注意力。需要说明的是use_varlen_attnTrue时要求pack_to_max_lengthTrue且 batch size 为 1见 default_collate_fn.py 中assert len(instances) 1此时单 batch 内不存在长短不一的问题长度分组采样的意义不大。LengthGroupedSampler主要服务于未启用 varlen attention、需要靠 padding 对齐的传统训练场景如纯 LLM SFT/QLoRA 训练。源码深度剖析分组算法与分布式采样实现LengthGroupedSampler的完整实现位于 length_grouped.py其核心逻辑分为两部分全局索引重排函数get_length_grouped_indices与采样器类LengthGroupedSampler。核心算法 get_length_grouped_indices函数get_length_grouped_indices(lengths, group_batch_size, generator)按以下步骤工作随机打乱indices torch.randperm(len(lengths), generatorgenerator)对全部样本索引做随机排列保留数据随机性。切分 megabatch将打乱后的索引按group_batch_sizemega_batch_mult * total_batch_size其中total_batch_size per_device_batch_size * world_size切成若干个 megabatch。组内降序排序每个 megabatch 内部按样本长度降序排序sorted(megabatch, keylambda i: lengths[i], reverseTrue)使同一 megabatch 内长度相近、且最长样本排在最前。最大批次优先计算每个 megabatch 的首元素即组内最长样本长度将包含全局最长样本的 megabatch 交换到第一个位置。这样安排是为了让最长的 batch 先被处理减少流水线/显存层面的波动。对单模态纯文本数据直接执行上述流程对多模态数据图像 语言源码用正负长度区分模态正长度代表含图像 token 的样本负长度代表纯语言样本见 internvl_dataset.py 中token_length -token_length的负号标记逻辑。此时两种模态的样本被分开各自分组排序最后合并时取abs(lengths[i])作为排序依据并保证末尾不足一组的样本按长度降序归入最后一个 batch。分布式采样与可复现性LengthGroupedSampler.__init__中通过get_dist_info()获取rank与world_size并计算num_samples本 rank 样本数round_upTrue时num_iters ceil(len(dataset) / world_size / per_device_batch_size)num_samples num_iters * per_device_batch_sizetotal_size num_samples * world_size即所有 rank 的总样本数可能通过复制超出原始数据集长度。__iter__中以generator.manual_seed(self.seed self.epoch)构造随机源保证每个 epoch 的采样顺序不同且分布式下可复现调用get_length_grouped_indices得到重排索引若round_upTrue将索引重复扩展至total_size长度indices * int(self.total_size / len(indices) 1)后截断保证每个 rank 迭代轮次一致按 rank 步长切分indices[self.rank:self.total_size:self.world_size]得到本 rank 的样本子集。此外__len__返回num_samplesset_epoch(epoch)供训练循环在每个 epoch 开始前更新self.epoch从而保证各 epoch 的分组顺序不同这是分布式训练中标准DistSamplerSeedHook配合机制。length 属性从何而来LengthGroupedSampler通过getattr(self.dataset, length_property)获取长度列表对ConcatDataset则拼接各子数据集的该属性因此数据集对象上必须存在名为length或你在length_property中指定的名称的属性。在 XTuner 中这一属性由 huggingface.py 的process_hf_dataset自动注入def get_lengths(example): return {length: len(example[input_ids])} ... # add length dataset dataset.map(get_lengths, num_procmap_num_proc) setattr(dataset, length, dataset[length])即数据经过 map_fn 加模板、tokenize、过滤无效样本、可选pack 之后统一计算每条样本input_ids的长度并写入length列同时通过setattr挂到数据集对象上供采样器读取。对自定义数据集的启示若你绕开process_hf_dataset、使用自定义数据集类接入LengthGroupedSampler必须自己维护一个等长的长度列表list 或 tuple并挂载为数据集属性否则getattr会取不到值、采样器无法构建。这是文档中明确强调的注意事项。仓库中的实际应用多模态训练的标配在 XTuner 仓库中LengthGroupedSampler被大量用于多模态LLaVA / InternVL微调配置例如llava_v15_7b_finetune.py、llava_v15_13b_finetune.pyllava_internlm2_chat_7b_clip_vit_large_p14_336_e1_gpu8_finetune.pyinternvl_v1_5_internlm2_2b_finetune.py、internvl_v2_internlm2_5_8b_finetune.py 等这些配置的sampler段典型写法为train_dataloader dict( batch_sizebatch_size, num_workersdataloader_num_workers, datasettrain_dataset, samplerdict( typeLengthGroupedSampler, length_propertylength, per_device_batch_sizebatch_size * accumulative_counts), collate_fndict(typedefault_collate_fn, use_varlen_attnuse_varlen_attn))多模态样本中含图像与纯文本样本的 token 数差异巨大正是长度分组采样器最能发挥价值的场景它结合源码中正负长度分模态排序的设计将图像样本集中、纯文本样本集中显著缩小批内长度方差从而减少 padding 浪费。小结与选型建议LengthGroupedSampler以长度相近的样本同批为原则用很小的改动仅替换sampler配置即可降低 padding 造成的训练效率损耗且不破坏单条样本的独立性。适用与注意要点如下何时选用数据长度分布差异明显的场景尤其多模态数据样本长度比较一致时加速效果一般可维持默认DefaultSampler与数据拼接的关系两者二选一即可官方优先推荐数据拼接技术见 数据拼接文档但若你担心拼接造成样本间相互影响LengthGroupedSampler是更稳妥的折中接入三步① 引入LengthGroupedSampler② 配置length_propertylength与per_device_batch_size③ 确保数据集对象带有length属性走process_hf_dataset则自动满足随机性与可复现分组后数据随机性会有所损失分布式训练下种子由sync_random_seed同步、每个 epoch 由set_epoch更新训练可复现。相关源码与文档索引实现 xtuner/dataset/samplers/length_grouped.py、数据集长度注入 xtuner/dataset/huggingface.py、默认配置示例 internlm2_chat_7b_qlora_oasst1_512_e3.py、完整中文文档 length_grouped_sampler.rst。【免费下载链接】xtunerA Next-Generation Training Engine Built for Ultra-Large MoE Models项目地址: https://gitcode.com/GitHub_Trending/xt/xtuner创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考