多彩编程 多彩编程MZPH · CODE BLOG
ARTICLE DETAIL

文章详情

深耕前端与后端开发技术的一线实战笔记与踩坑复盘。

XTuner 数据拼接(Pack to Max Length)实战指南:原理、配置与源码解析

XTuner 数据拼接(Pack to Max Length)实战指南:原理、配置与源码解析 XTuner 数据拼接Pack to Max Length实战指南原理、配置与源码解析【免费下载链接】xtunerA Next-Generation Training Engine Built for Ultra-Large MoE Models项目地址: https://gitcode.com/GitHub_Trending/xt/xtuner导读本文聚焦 XTuner 训练引擎中的数据拼接Dataset Packing即 pack_to_max_length技术。对于 LLM 微调而言真实训练数据的长度分布严重偏向短序列而 Transformers 模型要求固定长度输入传统做法是用大量无意义的 Pad Token 填充造成显著的计算浪费。本文将从这一背景出发结合 docs/zh_cn/acceleration/pack_to_max_length.rst 中的基准数据与配置说明深入源码剖析 XTuner 中 Packer 的拼接实现原理并给出关闭拼接、配合长度分组采样器LengthGroupedSampler与 Varlen Attention 等多种实战组合方案。读完本文你将能够根据自己的硬件条件与数据长度分布科学地选择拼接长度并完成训练配置调优。什么是数据集拼接对于大型语言模型LLM的输入而言数据集拼接指将多个 token 序列拼接成一个单独的输入。绝大多数开源数据集如 Alpaca、OASST1、法律、医疗等语料都有一个共同特点长度分布严重偏向较短的序列而 Transformers 模型接收的是固定长度的输入。因此在模型训练过程中通常需要将每条数据 Pad 至当前 batch 内最长序列的长度而 Pad Token 往往是某个特定的无意义 token在 XTuner 中由 xtuner/utils/constants.py 中的DEFAULT_PAD_TOKEN_INDEX定义。当 batch 内长短数据混杂时短数据被填充的冗余 token 会占据大量 GPU 显存与算力而这些计算对模型训练毫无贡献。数据拼接将多条数据首尾相连打包成一个固定长度的输入从而彻底消除 Pad Token 带来的无效计算减少计算资源浪费保持模型作为固定大小输入的静态图表示的优点训练流程无需改变显著降低训练 iteration 数一次前向/反向可以处理更多有效 token。下表来自 pack_to_max_length.rst展示了 InternLM2 7B 模型在 Alpaca 数据集上使用不同数据集拼接策略的训练速度对比。可以看到数据集拼接会大幅度提升训练效率拼接策略每秒处理 token 数加速比不使用362.9-拼接至 2k2677.17.38x拼接至 4k3124.38.61x拼接至 8k3173.98.76x拼接至 16k2864.47.89x拼接至 32k2965.48.17x从表中可以观察到两个重要结论拼接带来的收益是数量级的——即使拼接长度只有 2k吞吐也从 362.9 token/s 跃升至 2677.1 token/s加速约 7.38 倍拼接长度并非越大越好——4k8k 区间达到吞吐峰值约 8.7 倍加速继续增大到 16k、32k 后吞吐反而有所回落这通常与序列长度增大导致的注意力计算与显存占用上升有关。因此在实践中应根据显存容量与数据分布选择一个合适的拼接长度。在 XTuner 中启用数据拼接XTuner 提供的 config 文件默认启用数据集拼接功能。可以通过设置max_length字段来调整数据拼接长度。以 internlm2_chat_7b_qlora_alpaca_e3.py 为例其默认配置为# PART 1 Settings max_length 2048 pack_to_max_length True将拼接长度调整为 32k 的修改方式如下diff 形式来自原文档####################################################################### # PART 1 Settings # ####################################################################### - max_length 2048 max_length 32768 pack_to_max_length True ####################################################################### # PART 3 Dataset Dataloader # ####################################################################### train_dataset dict( max_lengthmax_length, pack_to_max_lengthpack_to_max_length, ...)其中两个字段的含义是max_length拼接后每条训练样本的目标 token 长度即 chunk 大小。启用拼接时数据集会被切分为若干条长度均为max_length的样本最后一个不足整块的 chunk 作为残差保留pack_to_max_length是否启用数据拼接的开关取值为布尔值默认True。关闭数据拼接的场景若不想使用数据拼接在 config 中将pack_to_max_length设为False即可max_length 2048 pack_to_max_length False此时max_length字段的含义变为单条数据最长的 token 数超过该长度的数据会被截断整个 batch 会被 pad 成当前 batch 内最长一条数据的长度。这种模式下每条样本独立计算 attention样本之间互不干扰但会引入 Pad 开销吞吐大幅下降对应上表 362.9 token/s 的基线。不拼接时的替代方案LengthGroupedSampler在不使用数据拼接策略时XTuner 还提供了一种数据集采样策略LengthGroupedSampler可以保证一个 batch 内的数据长度尽可能接近从而减少 Pad 对计算资源的浪费。它通过train_dataloader的sampler参数配置详见 docs/zh_cn/acceleration/length_grouped_sampler.rst。配置要点如下from xtuner.dataset.samplers import LengthGroupedSampler 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))其中length_property需要传入获取数据集长度的属性名。这一数值在通过process_hf_dataset构建数据集时会自动设置为length见下文源码分析因此使用自定义数据集类时需确保该属性被正确设置。其实现位于 xtuner/dataset/samplers/length_grouped.py核心函数get_length_grouped_indices先随机打乱索引再将数据划分为多个 megabatch每个 megabatch 内按长度降序排序从而让同一 batch 内的数据长度尽可能接近最后把最长样本所在的 megabatch 放到首位以尽早释放显存。源码级原理Packer 如何拼接数据理解了配置之后我们来深入源码。XTuner 中数据拼接的核心实现是 xtuner/dataset/utils.py 中的Packer类它通过 Hugging Facedatasets的batchedTruemap 操作逐 batch 执行拼接。调用链路拼接发生在process_hf_datasetxtuner/dataset/huggingface.py的数据预处理流水线中process_hf_dataset ├── build_origin_dataset # 加载原始数据集 ├── map_dataset # 应用 dataset_map_fn转换为统一格式 ├── add_template_to_dataset # 套用 prompt template如 internlm2_chat ├── tokenize_dataset # encode_fn 逐条 tokenize得到 input_ids / labels ├── filter # 过滤掉没有有效 label 的数据 └── pack_dataset # Packer(max_length) 按 max_length 拼接 └── dataset.map(Packer, batchedTrue)相关实现要点xtuner/dataset/huggingface.pydef pack_dataset(dataset, max_length, use_varlen_attn, shuffle_before_pack, map_num_proc): if shuffle_before_pack: dataset dataset.shuffle() dataset dataset.flatten_indices(num_procmap_num_proc) dataset dataset.map( Packer(max_length, use_varlen_attnuse_varlen_attn), batchedTrue, num_procmap_num_proc) return datasetshuffle_before_packTrue会在拼接前对数据集进行洗牌保证被打包进同一条样本的多条数据具有随机性。Packer 的残差residual机制Packer的核心设计是跨 batch 的残差缓存见 xtuner/dataset/utils.pyclass Packer: Pack multiple pieces of data into one. def __init__(self, chunk_size2048, use_varlen_attnFalse, drop_lastFalse): self.chunk_size chunk_size self.residual {input_ids: [], labels: []} self.use_varlen_attn use_varlen_attn self.drop_last drop_last if use_varlen_attn: self.residual_cumulative_len [0]其__call__逻辑分三步拼接把当前 map batch 内的所有样本与上一次残留的residual首尾相连形成一条超长序列concatenated_samples { k: v list(chain(*batch[k])) for k, v in self.residual.items() }切块若总长度达到chunk_size则按chunk_size切成整数个完整 chunk每条 chunk 即一条训练样本切不完的尾部保留在residual中等待下一个 batch 的数据来填满if total_length self.chunk_size: chunk_num total_length // self.chunk_size result { k: [ v[i:i self.chunk_size] for i in range( 0, chunk_num * self.chunk_size, self.chunk_size) ] for k, v in concatenated_samples.items() } self.residual { k: v[(chunk_num * self.chunk_size):] for k, v in concatenated_samples.items() }不足一块的处理若最后剩余数据不足一个 chunk则根据drop_last决定是丢弃还是作为一条不足长度的样本输出else: if self.drop_last: result {k: [] for k, v in concatenated_samples.items()} else: result {k: [v] for k, v in concatenated_samples.items()} self.residual {k: [] for k in concatenated_samples.keys()}正是这种拼接 → 切块 → 残留传递的机制保证了最终产出的几乎每一条训练样本都严格等于max_length长度仅数据集末尾可能余下一条不足长度的样本从而在训练时 batch 内完全不需要 Pad。与 Varlen Attention 的配合当use_varlen_attnTrue时Packer还会同步维护residual_cumulative_len记录残差序列中每条原始样本的累计长度边界xtuner/dataset/utils.py并通过get_cumulative_len与get_position_ids为拼接后的样本生成cumulative_len拼接样本内每条原始子样本的起止位置供 Flash Attention Varlen 接口定位每条子序列position_ids每条子样本独立从 0 开始的位置编码避免跨样本的 token 位置互相串扰。这样拼接后虽然物理上是同一条长序列但注意力计算仍按原始子样本边界进行样本间互不干扰。注意 xtuner/dataset/huggingface.py 中的约束use_varlen_attnTrue时强制要求pack_to_max_lengthTrue反之启用拼接时split必须为train或Nonextuner/dataset/huggingface.py。更完整的 Varlen Attention 使用说明可参考 docs/zh_cn/acceleration/varlen_flash_attn.rst。拼接模式下的 collate_fn训练时 xtuner/dataset/collate_fns/default_collate_fn.py 中的default_collate_fn负责把 Packer 产出的样本组装为 batch。启用拼接后由于所有样本长度一致input_ids可以直接torch.stack不再需要pad_sequencextuner/dataset/collate_fns/default_collate_fn.py。此时传入 config 的collate_fn需带use_varlen_attnuse_varlen_attn参数train_dataloader dict( batch_sizebatch_size, num_workersdataloader_num_workers, datasetalpaca_en, samplerdict(typesampler, shuffleTrue), collate_fndict(typedefault_collate_fn, use_varlen_attnuse_varlen_attn))另外值得留意的是 xtuner/dataset/huggingface.py当pack_to_max_lengthTrue且未显式设置remove_unused_columns时XTuner 会打印警告并强制移除多余列因为拼接要求数据仅保留input_ids/labels等必要字段。拼接后数据集长度属性拼接完成后process_hf_dataset会为数据集计算length属性xtuner/dataset/huggingface.pydataset dataset.map(get_lengths, num_procmap_num_proc) setattr(dataset, length, dataset[length])其中get_lengths返回{length: len(example[input_ids])}xtuner/dataset/huggingface.py。这正是上文LengthGroupedSampler中length_propertylength的数据来源——拼接后每条样本长度都是max_length或末尾不足一块的残差长度因此长度分组在拼接模式下天然均匀。完整配置示例下面给出一个完整可运行的拼接配置骨架基于 internlm2_chat_7b_qlora_alpaca_e3.py 精简展示max_length与pack_to_max_length在 config 中的完整联动####################################################################### # PART 1 Settings # ####################################################################### pretrained_model_name_or_path internlm/internlm2-chat-7b use_varlen_attn False max_length 2048 # 拼接目标长度即 chunk size pack_to_max_length True # 是否启用数据拼接 batch_size 1 # per_device accumulative_counts 16 dataloader_num_workers 0 max_epochs 3 ####################################################################### # PART 3 Dataset Dataloader # ####################################################################### from xtuner.dataset import process_hf_dataset from xtuner.dataset.collate_fns import default_collate_fn from mmengine.dataset import DefaultSampler train_dataset dict( typeprocess_hf_dataset, datasetdict(typeload_dataset, pathtatsu-lab/alpaca), tokenizertokenizer, max_lengthmax_length, dataset_map_fnalpaca_map_fn, template_map_fndict( typetemplate_map_fn_factory, templateprompt_template), remove_unused_columnsTrue, shuffle_before_packTrue, pack_to_max_lengthpack_to_max_length, use_varlen_attnuse_varlen_attn) train_dataloader dict( batch_sizebatch_size, num_workersdataloader_num_workers, datasettrain_dataset, samplerdict(typeDefaultSampler, shuffleTrue), collate_fndict(typedefault_collate_fn, use_varlen_attnuse_varlen_attn))需要调整时显存充足、希望追求更高吞吐可将max_length调大至 40968192对应基准表中吞吐峰值区间显存紧张或序列长度需求更高可搭配 docs/zh_cn/acceleration/train_extreme_long_sequence.rst 与 docs/zh_cn/acceleration/deepspeed.rst 中的 DeepSpeed ZeRO 优化与序列并行策略追求极长上下文参考仓库中的相关配置如 deepseek_v2_lite_chat_full_alpaca_e3_32k_varlen.py32k 拼接 Varlen Attention以及 xtuner/configs/llama_speed_benchmark 目录下覆盖 8k 到 1M 长度的吞吐基准配置。注意事项与最佳实践拼接长度与显存/吞吐的权衡基准表表明拼接长度存在最优区间4k8k过长的拼接反而因注意力计算开销增大而降低吞吐应根据实际显存与数据分布实测调优拼接会改变数据语义随机合并的多条样本共享同一条序列跨样本位置编码已通过position_ids隔离但样本间仍会通过 attention mask 的边界设置产生潜在交互对训练效果的实质影响需自行评估若希望每条样本完全独立可关闭拼接并改用LengthGroupedSampler必须设置remove_unused_columns拼接模式强制只保留必要列process_hf_dataset会自动处理但自定义数据类时应遵循同一约定Varlen Attention 与拼接是绑定关系use_varlen_attnTrue时必须pack_to_max_lengthTrue否则会触发 xtuner/dataset/huggingface.py 中的断言错误评估非 trainsplit 不支持拼接启用拼接时split只能为train或None验证集如需评估请使用非拼接分支或专门的评估流程参见 docs/zh_cn/evaluation 下的评估文档。小结数据拼接是 XTuner 默认启用且收益极为显著的一项训练加速技术仅通过max_length与pack_to_max_length两个字段即可控制。其底层Packer通过跨 batch 残差缓存 定长切块的机制将任意长度分布的语料转化为整齐划一的固定长度样本从根源上消除 Pad Token 的无效计算配合LengthGroupedSampler、Varlen Attention 与序列并行sequence_parallel_size等能力可在不同硬件与场景下灵活组合实现接近极限的训练吞吐。建议读者基于本文的基准数据与源码分析在自己的数据集与硬件上实测不同max_length的吞吐表现找到最适合自身场景的拼接配置。【免费下载链接】xtunerA Next-Generation Training Engine Built for Ultra-Large MoE Models项目地址: https://gitcode.com/GitHub_Trending/xt/xtuner创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表