File size: 5,849 Bytes
22fb4ec |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 |
===============================================
Varlen Attention
===============================================
\ :ref:`数据集拼接 <pack_to_max_length>` \ 一节中，我们讨论了“数据集拼接”策略对模型训练效率的显著提升。
理论上，数据集拼接可能会对注意力（Attention）机制的计算过程产生影响。这是因为，在未采用数据拼接策略的情况下，
每条数据在计算注意力时仅与自身相关联。然而，当采用数据拼接策略后，由多条短数据拼接成的长数据在计算注意力时会相互关联。
以一个由若干短数据拼接成长度为 4096 的数据为例，如果不采用变长注意力机制，在注意力计算阶段，每个 token 将会关注全部 4096 个 tokens ，如图左侧所示。
相反，在使用变长注意力机制的情况下，每个 token 在注意力计算阶段仅会关注其所在短数据中的所有 tokens ，如图右侧所示。因此， **变长注意力机制确保了无论是否采用“数据集拼接”策略，模型训练的行为保持一致性。**
.. raw:: html
<p align="center">
<img src="https://github.com/InternLM/InternLM/assets/41630003/7e0c6a02-a970-4bd3-a10b-8341720bf654" alt="XTuner" width="600"/>
<br />变长注意力计算原理（拷贝自 https://github.com/InternLM/InternEvo/blob/develop/doc/usage.md）<br />
</p>
支持列表
=====================
.. note::
使用变长注意力需要首先安装 `flash attn <https://github.com/Dao-AILab/flash-attention>`_ （
参考 `flash attn 安装 <https://github.com/Dao-AILab/flash-attention?tab=readme-ov-file#installation-and-features>`_ ）
.. list-table::
:widths: 25 50
:header-rows: 1
* - 模型
- 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）：
.. code:: console
$ MAX_JOBS=4 pip install flash-attn --no-build-isolation
.. tip::
更多安装方式请参考 `flash attn 安装 <https://github.com/Dao-AILab/flash-attention?tab=readme-ov-file#installation-and-features>`_
步骤 2：查找模板 config
---------------------------
XTuner 提供多个开箱即用的配置文件，用户可以通过下列命令查看：
.. code-block:: console
$ xtuner list-cfg -p internlm
.. tip::
``-p`` 为模糊查找，若想训练其他模型，可以修改 ``internlm`` 为 XTuner 支持的其他模型名称。
步骤 3：复制 config 文件
-----------------------------
导出需要使用的 config ：
.. code-block:: bash
xtuner copy-cfg ${CONFIG_NAME} ${SAVE_DIR}
例如通过下列命令将名为 ``internlm_7b_full_oasst1_e3`` 的 config 导出至当前目录下：
.. code-block:: console
$ xtuner copy-cfg internlm_7b_full_oasst1_e3 .
.. note::
当前目录下会存在一个新 config
``internlm_7b_full_oasst1_e3_copy.py`` 。
步骤 4：修改 config 文件
-------------------------------
将步骤 3 复制得到的 config 文件中的 ``use_varlen_attn`` 属性由 False 改为 True 即可激活变长注意力训练机制：
.. code-block:: diff
...
#######################################################################
# PART 1 Settings #
#######################################################################
# Model
pretrained_model_name_or_path = 'internlm/internlm-7b'
- use_varlen_attn = False
+ use_varlen_attn = True
...
.. warning::
当设置 ``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：开始训练
-----------------------
.. code-block:: bash
xtuner train ${CONFIG_NAME_OR_PATH}
例如，我们可以基于步骤 4 中修改得到的 `internlm_7b_full_oasst1_e3_copy.py` 进行训练：
.. code-block:: console
$ # 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
.. tip::
``--deepspeed`` 表示使用 `DeepSpeed <https://github.com/microsoft/DeepSpeed>`_ 🚀 来优化训练过程。若未安装 DeepSpeed ，可通过 ``pip install deepspeed>=0.12.3`` 进行安装。XTuner 内置了多种策略，包括 ZeRO-1、ZeRO-2、ZeRO-3 等。如果用户期望关闭此功能，请直接移除此参数。
步骤 6：模型转换
^^^^^^^^^^^^^^^^^^^^^^^^^^^
将保存的 PTH 模型（如果使用的DeepSpeed，则将会是一个文件夹）转换为 HuggingFace 模型：
.. code-block:: bash
xtuner convert pth_to_hf ${CONFIG_NAME_OR_PATH} ${PTH} ${SAVE_PATH}
对应上面的例子，模型转换脚本为：
.. code-block:: bash
xtuner convert pth_to_hf internlm_7b_full_oasst1_e3_copy.py ${PTH} ${SAVE_PATH}
.. note::
其中 ``${PTH}`` 为训练权重保存的路径，若训练时未指定，默认保存在 ``./work_dirs/internlm_7b_full_oasst1_e3_copy`` 路径下。
|