Varlen Attention

Varlen Attention#

数据集拼接 一节中，我们讨论了“数据集拼接”策略对模型训练效率的显著提升。 理论上，数据集拼接可能会对注意力（Attention）机制的计算过程产生影响。这是因为，在未采用数据拼接策略的情况下， 每条数据在计算注意力时仅与自身相关联。然而，当采用数据拼接策略后，由多条短数据拼接成的长数据在计算注意力时会相互关联。 以一个由若干短数据拼接成长度为 4096 的数据为例，如果不采用变长注意力机制，在注意力计算阶段，每个 token 将会关注全部 4096 个 tokens ，如图左侧所示。

相反，在使用变长注意力机制的情况下，每个 token 在注意力计算阶段仅会关注其所在短数据中的所有 tokens ，如图右侧所示。因此， 变长注意力机制确保了无论是否采用“数据集拼接”策略，模型训练的行为保持一致性。

XTuner
变长注意力计算原理（拷贝自 https://github.com/InternLM/InternEvo/blob/develop/doc/usage.md）

支持列表#

备注

使用变长注意力需要首先安装 flash attn （ 参考 flash attn 安装 ）

模型

Flash Attention 支持情况

baichuan 1/2

❌

chatglm 2/3

❌

deepseek

✅

gemma

❌

internlm 1/2

✅

llama 2

✅

mistral

✅

qwen 1/1.5

❌

starcoder

❌

yi

✅

zephyr

✅

使用变长注意力机制训练#

步骤 1：安装 flash_attn#

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

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

小技巧

更多安装方式请参考 flash attn 安装

步骤 2：查找模板 config#

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

$ xtuner list-cfg -p internlm

小技巧

-p 为模糊查找，若想训练其他模型，可以修改 internlm 为 XTuner 支持的其他模型名称。

步骤 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 。

步骤 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。

步骤 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 内置了多种策略，包括 ZeRO-1、ZeRO-2、ZeRO-3 等。如果用户期望关闭此功能，请直接移除此参数。

步骤 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 路径下。

目录