使用 DeepSpeed-Ulysses 训练具有超长序列的 Transformer 模型入门

本教程介绍如何启用 DeepSpeed-Ulysses。DeepSpeed-Ulysses 是一种简单但高度通信和内存高效的序列并行机制，用于训练具有海量序列长度的大型 Transformer 模型。它沿着序列维度对输入张量进行分区，并使用通信高效的全对全集体进行分布式注意力计算。此外，DeepSpeed-Ulysses 集成了高级建模和系统优化，例如 Flash Attention、稀疏注意力和 ZeRO 优化器，以优化计算效率和内存使用。使用 DeepSpeed 序列并行进行训练允许模型大小和序列长度几乎无限扩展，不受单个 GPU 内存限制的限制，并且可以达到接近峰值计算性能的水平。目前，DeepSpeed-Ulysses 可以处理长度高达 100 万的序列（是完整哈利波特书籍大小的 10 倍！），在 64 个 A100 GPU 上运行。请阅读我们的DeepSpeed-Ulysses 博客以了解更多信息！

1. 安装

您需要安装 DeepSpeed v0.10.2 或更高版本才能使用 DeepSpeed 序列功能。安装 DeepSpeed 就像pip install deepspeed一样简单，查看更多详细信息。

2. 如何在您的应用程序中使用 DeepSpeed-Ulysses？

将 DS-Seq 集成到您的训练代码中非常简单，在本节中，我们将介绍如何通过我们的Megatron-DeepSpeed 代码库集成 DeepSpeed-Ulysses。

  • 替换注意力模块：首先，您需要使用 DeepSpeed-Ulysses DistributedAttention 更新您的注意力模块。在这里，我们使用来自Megatron-DeepSpeed 的注意力，这是 GPT-3 等模型训练中使用的因果注意力。重写注意力块
def __init__():
    ...
    self.local_attn = CoreAttention(self.layer_number, config, self.attn_mask_type)
    self.core_attention = local_attn
    ...

def forward():
    ...
    context_layer = self.core_attention(
                    query_layer, key_layer, value_layer, attention_mask)
    ...

使用

from deepspeed.sequence.layer import DistributedAttention

def __init__():
    ...
    self.local_attn = CoreAttention(self.layer_number, config, self.attn_mask_type)
    self.dist_attn = DistributedAttention(self.local_attn, parallel_state.get_sequence_parallel_group())
    ...

def forward():
    ...
    context_layer = self.dist_attn(query_layer, key_layer, value_layer, attention_mask)
    ...
  • 添加序列并行通信组：请注意，DistributedAttention 将local_attn 和sequence_parallel_group 作为参数，其中 local_attn 可以是您的原始注意力块。您还需要构建序列并行通信组并将其传递给 DistributedAttention。一种方法是在模型初始化阶段构建序列并行组。
def initialize_model_parallel(
    ...
    sequence_parallel_size,
    ...
):
    ...
    num_sequence_parallel_groups: int = world_size // sequence_parallel_size
    num_sequence_data_parallel_groups: int = world_size // sequence_parallel_size // data_parallel_size
    ...
    global _SEQUENCE_PARALLEL_GROUP
    for i in range(num_sequence_parallel_groups):
        ranks = range(i * sequence_parallel_size,
                      (i + 1) * sequence_parallel_size)
        group = torch.distributed.new_group(ranks)
        if rank in ranks:
            _SEQUENCE_PARALLEL_GROUP = group

def get_sequence_parallel_group():
    """Get the sequence parallel group the caller rank belongs to."""
    return _SEQUENCE_PARALLEL_GROUP

在 Megatron-DeepSpeed 示例中，要启用序列并行，请使用 –ds-sequence-parallel-size 参数设置并行度。您还需要确保注意力头的数量可以被此值整除。我们已为您准备了一些脚本，以便您快速获得一些使用超长序列训练类似 GPT-3 模型的示例

Megatron-DeepSpeed/examples_deepspeed/sequence_parallel$ bash ds_pretrain_gpt_1.3B_seq_parallel_32k.sh
Megatron-DeepSpeed/examples_deepspeed/sequence_parallel$ bash ds_pretrain_gpt_30B_seq_parallel_32k.sh

请注意，我们的序列并行功能目前与 Megatron-LM 的张量或流水线并行不兼容。

3. 如何使用 FlashAttention 启用 DeepSpeed-Ulysses？

DeepSpeed 的序列并行可以与不同类型的注意力实现相结合，以进一步提高长序列训练的内存和计算效率

经典注意力：通过 PyTorch 实现的注意力机制。

FlashAttention：来自FlashAttention：具有 IO 感知的高效且内存高效的精确注意力 的实现。通过--use-flash-attn 启用。

FlashAttention + Triton：FlashAttention 在 Triton 中的实现（已在 triton==2.0.0.dev20221202 上测试）。通过--use-flash-attn-triton 启用。

为了获得最佳性能，我们建议使用 FlashAttention + Triton。以下是安装步骤。请注意，FlashAttention 仅与 NVIDIA Turing、Ampere、Ada 或 Hopper GPU 兼容。

# install triton
git clone -b legacy-backend https://github.com/openai/triton
cd triton/python/
pip install cmake
pip install .

# install
cd ${WORK_DIR}
git clone -b v1.0.4 https://github.com/HazyResearch/flash-attention
cd flash-attention
python setup.py install

您可能还想确保您的模型配置符合 FlashAttention 的要求。例如，为了获得最佳性能，头部大小应为 8 的倍数。有关更多详细信息，请参阅 FlashAttention 的文档。

更新：