XTuner 变长注意力（Varlen Attention）训练指南：原理、模型支持与完整实战流程-平芜编程栈

XTuner 变长注意力（Varlen Attention）训练指南：原理、模型支持与完整实战流程

【免费下载链接】xtunerA Next-Generation Training Engine Built for Ultra-Large MoE Models项目地址: https://gitcode.com/GitHub_Trending/xt/xtuner

变长注意力（Varlen Attention）是 XTuner 针对"数据集拼接（Packing）"训练策略配套提供的一项核心加速机制：它允许将多条短数据直接拼接成长序列喂给模型，同时保证每个 token 只关注其所在原始数据片段内部的 token，从而在享受拼接带来的训练吞吐提升的同时，保持与未拼接训练完全一致的注意力行为。本文将以 docs/zh_cn/acceleration/varlen_flash_attn.rst 为骨架，结合 XTuner 源码与仓库内真实配置，完整讲解变长注意力的工作原理、模型支持范围，以及从安装 Flash Attention 到模型转换的六步实战流程。

变长注意力：为什么"数据拼接"需要它

在 docs/zh_cn/acceleration/pack_to_max_length.rst 中我们讨论过"数据拼接"策略对模型训练效率的显著提升：真实数据集的长度分布严重偏向短序列，而 Transformers 模型接收固定长度输入，传统做法需要用无意义的 Pad Token 把每条数据补齐到 batch 内最长序列，造成大量计算浪费。将多条短数据打包成一条固定长度的长序列，可以彻底去掉 Pad Token，同时保留静态图输入的优势。

但拼接引入了一个理论问题：它可能改变注意力（Attention）机制的计算行为。在未拼接的情况下，每条数据计算注意力时只与自身相关联；而拼接后，由多条短数据组成的长度为 4096（或任意max_length）的长序列中，如果采用普通注意力，每个 token 会关注拼接序列中的全部 4096 个 tokens，即"跨数据片段"发生了信息泄漏。

变长注意力正是为了解决这一问题而生：使用变长注意力后，每个 token 在注意力计算阶段仅关注其所在短数据片段内的所有 tokens。因此，变长注意力机制确保了无论是否采用"数据集拼接"策略，模型训练的行为都保持一致性——既拿到了拼接带来的吞吐收益，又不会引入跨样本的注意力污染。

下图对比了传统带填充拼接（上方）与变长注意力无填充拼接（下方）的差异：传统方式中填充 token 造成大量 GPU 内存浪费；变长方式将偏好对/短数据直接无填充拼成长序列，每个块（由红竖线分隔）内部的 token 仅与同块内 token 计算注意力，实现零填充、零内存浪费：

变长注意力的工作原理：从数据打包到前向计算

要理解use_varlen_attn = True在 XTuner 内部做了什么，可以从数据流水线的三个环节看起。

1. 打包阶段：生成 cumulative_len 与 position_ids

在 xtuner/dataset/utils.py 的Packer类中，当use_varlen_attn=True时，打包器会额外维护residual_cumulative_len残差列表（初始为[0]），记录每个样本在拼接流中的累积长度边界。get_cumulative_len()依据每个chunk_size（即max_length）切分点，通过二分查找（np.searchsorted）还原出每个 chunk 内各条短数据的起止边界cumulative_len；随后get_position_ids()根据cumulative_len为每个 chunk 生成逐片段归零重计的position_ids（即每段短数据内部从 0 开始的位置编码），保证 RoPE 位置编码不会跨片段错乱。

2. 组装阶段：collate_fn 输出变长参数

在 xtuner/dataset/collate_fns/default_collate_fn.py 中，default_collate_fn(..., use_varlen_attn=False)开启变长模式后行为明显不同：

强制assert len(instances) == 1，即每个 batch 只允许一条拼接后的长样本；
不再生成传统attention_mask，而是改为输出cumulative_len（各短数据边界）、position_ids（片段内位置）、以及max_seqlen（本 batch 内最长短片段的 token 数）；
同时拒绝多模态场景（has_image为真时报错），说明当前版本变长注意力仅支持纯文本训练。

3. 前向阶段：Hook 传递参数 + 注意力 Kernel 替换

组装好的变长参数并不会直接进入模型 forward。在 xtuner/engine/hooks/varlen_attn_args_to_messagehub_hook.py 中，VarlenAttnArgsToMessageHubHook在before_train_iter/before_val_iter时从data_batch['data']中弹出cumulative_len与max_seqlen，按 rank 写入 MMEngine 的MessageHub，并在迭代结束后清空。

模型侧，在 xtuner/model/sft.py 的SupervisedFinetune中，use_varlen_attn参数会被传入build_llm_from_cfg，最终调用dispatch_modules(llm, use_varlen_attn=use_varlen_attn)。在 xtuner/model/modules/dispatch/init.py 中，VARLEN_ATTN_DISPATCH_MAPPING将各模型架构的 FlashAttention2 注意力类映射到对应的*_varlen_attn_forward实现。以 xtuner/model/modules/dispatch/llama.py 的llama_varlen_attn_forward为例，它会从MessageHub读取cumulative_len与max_seqlen，然后调用varlen_flash_attn(...)（Flash Attention 2 的_flash_attn_varlen_forward封装），实现"块内注意力"；当无变长参数时则回退到普通flash_attn_wo_mask。

模型支持列表

使用变长注意力需要首先安装 Flash Attention（详见下文"步骤 1"）。当前 XTuner 各模型对变长注意力的支持情况如下表所示：

模型	Flash Attention 支持情况
baichuan 1/2	❌
chatglm 2/3	❌
deepseek	✅
gemma	❌
internlm 1/2	✅
llama 2	✅
mistral	✅
qwen 1/1.5	❌
starcoder	❌
yi	✅
zephyr	✅

从源码看，该表的判断依据与 xtuner/model/modules/dispatch/init.py 中VARLEN_ATTN_DISPATCH_MAPPING的注册情况一致：InternLM2、Llama、Phi3、Mistral、Mixtral、Qwen2/Qwen2Moe、DeepseekV2 等 FlashAttention2 架构都有对应的 varlen forward 实现，而 baichuan、chatglm、gemma、qwen1/1.5、starcoder 等未在映射表中注册。此外，SUPPORT_FLASH2 = is_flash_attn_2_available()，即没有正确安装 flash-attn 时，dispatch_varlen_attn_forward会直接返回、变长注意力无法生效。

使用变长注意力机制训练的完整流程

以下六个步骤以 XTuner 官方流程为准，每一步都可在仓库中找到对应落点。

步骤 1：安装 flash_attn

XTuner 中实现的变长注意力依赖 Flash Attention 2，可通过以下命令安装（需要 CUDA）：

$ MAX_JOBS=4 pip install flash-attn --no-build-isolation

MAX_JOBS=4用于限制编译时的并行任务数，避免 OOM 或编译超时；--no-build-isolation复用当前环境已有的构建依赖以加快安装。安装完成后，可通过 xtuner/model/modules/dispatch/init.py 中的is_flash_attn_2_available()逻辑确认环境就绪（未安装时训练脚本会在check_cfg阶段直接报错）。

步骤 2：查找模板 config

XTuner 提供多个开箱即用的配置文件，可以通过以下命令查看：

$ xtuner list-cfg -p internlm

-p为模糊查找。若想训练其他模型，将internlm替换为 XTuner 支持的其他模型名称即可，例如llama、mistral、yi、deepseek等。

步骤 3：复制 config 文件

导出需要使用的 config：

xtuner copy-cfg ${CONFIG_NAME} ${SAVE_DIR}

例如，将名为internlm_7b_full_oasst1_e3的 config 导出至当前目录下：

$ xtuner copy-cfg internlm_7b_full_oasst1_e3 .

执行后当前目录下会生成一个新的 config：internlm_7b_full_oasst1_e3_copy.py。对应的模板文件位于 xtuner/configs/internlm/internlm_7b/internlm_7b_full_oasst1_e3.py。

步骤 4：修改 config 文件激活变长注意力

将步骤 3 复制得到的 config 文件中的use_varlen_attn属性由False改为True：

... ####################################################################### # PART 1 Settings # ####################################################################### # Model pretrained_model_name_or_path = 'internlm/internlm-7b' - use_varlen_attn = False + use_varlen_attn = True ...

⚠️重要约束：当设置use_varlen_attn = True后，batch_size = 2, max_length = 2k的配置与batch_size = 1, max_length = 4k的配置训练行为是近似的，因此 XTuner 目前只支持batch_size = 1的情况。另外，use_varlen_attn = True时pack_to_max_length也必须设置为True。

上述约束在 xtuner/tools/train.py 的check_cfg中有硬性校验：use_varlen_attn开启时若train_dataloader.batch_size > 1会直接抛出NotImplementedError；同时要求SUPPORT_FLASH2为真、model.llm.attn_implementation为flash_attention_2（或未设置，由 xtuner/model/sft.py 的_prepare_for_flash_attn自动写入），且max_length必须能被 sequence parallel 的 world size 整除。

步骤 5：开始训练

xtuner train ${CONFIG_NAME_OR_PATH}

例如，可以基于步骤 4 修改得到的internlm_7b_full_oasst1_e3_copy.py进行训练：

$ # On a single GPU $ xtuner train internlm_7b_full_oasst1_e3_copy.py --deepspeed deepspeed_zero1 $ # On multiple GPUs (torchrun) $ NPROC_PER_NODE=${GPU_NUM} xtuner train internlm_7b_full_oasst1_e3_copy.py --deepspeed deepspeed_zero1 $ # On multiple GPUs (slurm) $ srun ${SRUN_ARGS} xtuner train internlm_7b_full_oasst1_e3_copy.py --launcher slurm --deepspeed deepspeed_zero1

--deepspeed表示使用 DeepSpeed 来优化训练过程。若未安装 DeepSpeed，可通过pip install deepspeed>=0.12.3安装（xtuner/tools/train.py 中会校验版本不低于 0.12.3）。XTuner 内置了多种策略，包括 ZeRO-1、ZeRO-2、ZeRO-3 等，对应文件位于 xtuner/configs/deepspeed 目录（如deepspeed_zero1.json、deepspeed_zero2.json、deepspeed_zero3.json及其 offload 版本）。如果期望关闭此功能，直接移除该参数即可。

步骤 6：模型转换

训练结束后，将保存的 PTH 模型（如果使用了 DeepSpeed，则是一个文件夹）转换为 HuggingFace 模型：

xtuner convert pth_to_hf ${CONFIG_NAME_OR_PATH} ${PTH} ${SAVE_PATH}

对应上面的例子，模型转换命令为：

xtuner convert pth_to_hf internlm_7b_full_oasst1_e3_copy.py ${PTH} ${SAVE_PATH}

其中${PTH}为训练权重保存的路径，若训练时未指定，默认保存在./work_dirs/internlm_7b_full_oasst1_e3_copy路径下。转换实现的入口可参考 xtuner/tools/model_converters/pth_to_hf.py。

一份开箱即用的变长注意力参考配置

仓库中提供了官方开箱即用的变长注意力配置示例：xtuner/configs/deepseek/deepseek_v2_lite_chat/deepseek_v2_lite_chat_full_alpaca_e3_32k_varlen.py。该配置完整展示了变长注意力的关键参数组合：

pretrained_model_name_or_path = 'deepseek-ai/DeepSeek-V2-Lite-Chat' use_varlen_attn = True data_path = 'tatsu-lab/alpaca' prompt_template = PROMPT_TEMPLATE.deepseek_v2 max_length = 32768 # 拼接目标长度：32k pack_to_max_length = True # 必须为 True sequence_parallel_size = 1 batch_size = 1 # 变长注意力要求 per device batch_size = 1

其中train_dataset的构建需要同时传入pack_to_max_length=pack_to_max_length与use_varlen_attn=use_varlen_attn；train_dataloader的collate_fn配置为dict(type=default_collate_fn, use_varlen_attn=use_varlen_attn)；并且在custom_hooks中追加dict(type=VarlenAttnArgsToMessageHubHook)——正是上文中负责把cumulative_len/max_seqlen送入注意力 kernel 的关键 Hook。

变长注意力的更多应用场景

除了 SFT 全参/微调训练，变长注意力在仓库中还被扩展到偏好对齐类训练：例如 xtuner/configs/dpo/internlm/internlm2_chat_1_8b_dpo_full_varlenattn.py、xtuner/configs/dpo/llama/llama3_8b_instruct_dpo_qlora_varlenattn.py、xtuner/configs/orpo/internlm/internlm2_chat_7b_orpo_qlora_varlenattn_ultrafeedback_e5.py 以及 xtuner/configs/reward_model/internlm/internlm2_chat_1_8b_reward_full_varlenattn_ultrafeedback.py。在这些场景中，chosen/rejected 偏好对或分段数据被无填充地打包进同一序列，每个片段独立计算注意力（即本文首图展示的机制），从而大幅节省 DPO、ORPO 与 Reward Model 训练中因 padding 造成的显存浪费。对应实现可在 xtuner/model/dpo.py、xtuner/model/orpo.py、xtuner/model/reward.py 中查看。

小结

原理：变长注意力让"数据拼接"与"注意力行为"解耦——拼接提速的同时，每个 token 只在所属短数据片段内部做注意力计算，保证训练行为与未拼接时一致。
实现链路：Packer（生成cumulative_len/position_ids）→default_collate_fn（输出变长参数，强制 batch_size=1）→VarlenAttnArgsToMessageHubHook（按 rank 传递参数）→dispatch_modules替换为*_varlen_attn_forward（调用 Flash Attention 2 的 varlen kernel）。
使用前提：先安装 flash-attn（需要 CUDA）；模型需在支持列表内；config 中use_varlen_attn=True、pack_to_max_length=True、batch_size=1，且max_length能被 sequence parallel world size 整除。
完整链路：xtuner list-cfg→xtuner copy-cfg→ 修改 config →xtuner train ... --deepspeed deepspeed_zero1→xtuner convert pth_to_hf，即可完成一次变长注意力训练的全流程。