DeepSpeed 稀疏注意力

本教程将介绍如何使用 DeepSpeed 稀疏注意力 (SA) 及其构建块内核。使用 SA 最简单的方法是通过 DeepSpeed 启动器。我们将在如何在 DeepSpeed 启动器中使用稀疏注意力一节中通过一个例子进行描述。但在此之前,我们先在下一节介绍 DeepSpeed SA 提供的模块。

注意:目前,DeepSpeed 稀疏注意力只能在 NVIDIA V100 或 A100 GPU 上使用,且需 Torch >= 1.6 和 CUDA 10.1、10.2、11.0 或 11.1。

稀疏注意力模块

  • MatMul:此模块处理块稀疏矩阵-矩阵乘法。目前它支持DeepSpeed 稀疏注意力部分中描述的 SDD、DSD 和 DDS。
  • Softmax:此模块应用块稀疏 softmax。它处理前向和后向传播。
  • SparseSelfAttention:此模块使用 MatMul 和 Softmax 内核,并根据 Query、Keys 和 Values 生成 Context 层输出。它是任何自注意力层中常见操作的简化版本。它还可以应用
    • 相对位置嵌入
    • 注意力掩码
    • Key padding mask 应用于中间注意力分数。有关自注意力的更多详细信息,请查看MultiHeadAttention
  • BertSparseSelfAttention:此模块包含一个简化的 BertSelfAttention 层,可用于替代原始的密集 Bert 自注意力层。我们的实现基于DeepSpeedExample
  • SparseAttentionUtils:此模块提供了一些实用函数,用于处理预训练模型与稀疏注意力的适配
    • replace_model_self_attention_with_sparse_self_attention:如果您当前已加载模型并想用稀疏自注意力替换自注意力模块,您可以直接使用此函数为您处理。它目前处理基于 BERT 和 RoBERTa 的预训练模型,但如果您的模型类型不同于这两种,您可以根据您的模型类型进行扩展。您还需要扩展位置嵌入以处理新的序列长度;这可以通过extend_position_embedding函数完成。
    • update_tokenizer_model_max_length:此函数简单地使用新值更新您的分词器中的最大位置嵌入。
    • extend_position_embedding:此函数根据当前值扩展位置嵌入。例如,如果您有一个最大序列长度为 128 的模型并将其扩展到 1k 序列长度,它会将当前嵌入复制 8 次以初始化新的嵌入。实验表明,这种初始化方式比从头开始初始化效果更好;可以更快地收敛。
    • pad_to_block_size:此函数在序列长度维度上填充输入 token 和注意力掩码,使其成为块大小的倍数;这是 SA 的要求。
    • unpad_sequence_output:如果模型输入已填充,此函数将取消填充序列输出。
  • SparsityConfig:这是一个稀疏性结构的抽象类。任何稀疏性结构都需要扩展此类别并编写自己的稀疏模式构建;即make_layout函数。DeepSpeed 目前提供以下结构,将在如何配置稀疏结构一节中进行描述
    • FixedSparsityConfig
    • BSLongformerSparsityConfig
    • BigBirdSparsityConfig
    • VariableSparsityConfig
    • DenseSparsityConfig

注意:目前 DeepSpeed Transformer Kernels 不支持稀疏注意力。要使用稀疏注意力,您需要禁用 Transformer Kernels!

如何在 DeepSpeed 启动器中使用稀疏注意力

在本节中,我们将描述如何通过我们的 bing_bert 代码使用 DeepSpeed 稀疏注意力。

  • 更新注意力模块:首先,您需要根据稀疏计算更新您的注意力模块。这里,我们使用BertSparseSelfAttention,它是我们bing_bert代码中BertSelfAttention的稀疏版本。它重写了BertSelfAttention,其中它替换了
attention_scores = torch.matmul(query_layer, key_layer)
attention_scores = attention_scores / math.sqrt(
    self.attention_head_size)

# Apply the attention mask is (precomputed for all layers in BertModel forward() function)
attention_scores = attention_scores + attention_mask

pdtype = attention_scores.dtype
# Normalize the attention scores to probabilities.
attention_probs = self.softmax(attention_scores)

# This is actually dropping out entire tokens to attend to, which might
# seem a bit unusual, but is taken from the original Transformer paper.
attention_probs = self.dropout(attention_probs)

context_layer = torch.matmul(attention_probs, value_layer)

context_layer =
  self.sparse_self_attention(
	query_layer,
	key_layer,
	value_layer,
	key_padding_mask=attention_mask)

其中sparse_self_attentionSparseSelfAttention的一个实例。该模块通过稀疏注意力计算注意力上下文,将底层的矩阵乘法和 softmax 替换为等效的稀疏版本。您可以类似地更新任何其他注意力模块。

  • 在模型中设置稀疏注意力配置:您需要设置稀疏注意力配置。在我们的示例中,这是在BertModel中完成的。
self.pad_token_id = config.pad_token_id if hasattr(
   config, 'pad_token_id') and config.pad_token_id is not None else 0
# set sparse_attention_config if it has been selected
self.sparse_attention_config = get_sparse_attention_config(
   args, config.num_attention_heads)
self.encoder = BertEncoder(
   config, args, sparse_attention_config=self.sparse_attention_config)
  • 更新编码器模型:此外,当启用 SA 时,您需要更新您的编码器模型以在注意力层中使用 SA。请查看我们的bing_bert 示例,其中当 SA 启用时,我们使用BertSparseSelfAttention而不是BertSelfAttention
if sparse_attention_config is not None:
    from deepspeed.ops.sparse_attention import BertSparseSelfAttention

    layer.attention.self = BertSparseSelfAttention(
         config, sparsity_config=sparse_attention_config)
  • 填充和取消填充输入数据:您可能还需要填充input_idsattention_mask的序列维度,使其成为稀疏块大小的倍数。如上文模块部分所述,DeepSpeed 提供了用于填充和取消填充的实用函数。请查看我们的bing_bert 示例,了解如何填充和取消填充模型的输入或输出。
if self.sparse_attention_config is not None:
   pad_len, input_ids, attention_mask, token_type_ids, position_ids, inputs_embeds = SparseAttentionUtils.pad_to_block_size(
      block_size=self.sparse_attention_config.block,
      input_ids=input_ids,
      attention_mask=extended_attention_mask,
      token_type_ids=token_type_ids,
      position_ids=None,
      inputs_embeds=None,
      pad_token_id=self.pad_token_id,
      model_embeddings=self.embeddings)
.
.
.
# If BertEncoder uses sparse attention, and input_ids were padded, sequence output needs to be unpadded to original length
if self.sparse_attention_config is not None and pad_len > 0:
   encoded_layers[-1] = SparseAttentionUtils.unpad_sequence_output(
      pad_len, encoded_layers[-1])
  • *启用稀疏注意力:要使用 DeepSpeed 稀疏注意力,您需要在启动脚本中通过deepspeed_sparse_attention参数启用它
--deepspeed_sparse_attention

请查看我们的 bing_bert 运行脚本,作为如何在 DeepSpeed 启动器中启用 SA 的示例。

"sparse_attention": {
    "mode": "fixed",
    "block": 16,
    "different_layout_per_head": true,
    "num_local_blocks": 4,
    "num_global_blocks": 1,
    "attention": "bidirectional",
    "horizontal_global_attention": false,
    "num_different_global_patterns": 4
}

如何使用单个内核

DeepSpeed 稀疏注意力可以作为 DeepSpeed 的一个功能使用,如上所述,也可以简单地作为自注意力模块单独集成到任何 Transformer 模型中。此外,构建块内核,即矩阵乘法和 softmax 可以单独使用。要单独使用稀疏注意力,您只需安装 DeepSpeed 并导入模块部分中描述的任何模块;例如

from deepspeed.ops.sparse_attention import SparseSelfAttention

请参阅 Docstrings 以了解如何单独使用每个模块的详细信息。

如何配置稀疏结构

接下来我们描述支持的稀疏结构、它们的参数集以及在自注意力层上添加任意稀疏模式的灵活性。您可以使用任何支持的稀疏结构更新 DeepSpeed 配置文件并相应地设置参数。

  • SparsityConfig:此模块是所有稀疏性结构的父类,并包含所有稀疏性结构的共享特性。它接受以下参数
    • num_heads:一个整数,表示层的注意力头数量。
    • block:一个整数,表示块大小。当前稀疏自注意力实现基于块稀疏矩阵。此参数定义了此类方块的大小;Block X Block
    • different_layout_per_head:一个布尔值,确定每个头是否应分配不同的稀疏布局;默认为 false,这将根据可用性满足。
  • Fixed (FixedSparsityConfig):此结构基于 OpenAI 的《Generative Modeling with Sparse Transformers》,其中局部和全局注意力由给定参数固定
    • num_local_blocks:一个整数,表示局部注意力窗口中的块数量。如下图所示(改编自原始论文),局部窗口中的 token 会关注它们局部内的所有 token。在自回归模型的情况下,如图所示,token 会关注它们在局部窗口中之前出现的 token。而在 Masked 模型(如 BERT)的情况下,注意力是双向的。
    • num_global_blocks:一个整数,表示局部窗口中有多少个连续的块被用作全局注意力的窗口代表;如下图所示。
    • attention:一个字符串,表示注意力类型。注意力可以是unidirectional(单向),例如自回归模型,其中 token 只关注上下文中它们之前出现的 token。考虑到这一点,注意力矩阵的上三角部分是空的,如上图所示。或者可以是bidirectional(双向),例如 BERT,其中 token 可以关注它们之前或之后出现的任何其他 token。那么,注意力矩阵的上三角部分是下图下三角部分的镜像。
    • horizontal_global_attention:一个布尔值,确定作为局部窗口的全局代表的块是否也关注所有其他块。这仅在注意力类型为bidirectional时有效。从注意力矩阵来看,这意味着全局注意力不仅包括垂直块,还包括水平块。
    • num_different_global_patterns:一个整数,表示不同全局注意力布局的数量。虽然全局注意力可以通过哪些块代表任何局部窗口来固定,但由于有多个头,每个头可以使用不同的全局代表。例如,对于由 4 个块构成局部窗口且全局注意力大小为单个块的情况,我们可以有 4 种不同版本,其中每个局部窗口的第一个、第二个、第三个或第四个块可以作为该窗口的全局代表。此参数确定我们想要多少种这样的模式。当然,这受到num_local_blocksnum_global_blocks的限制。此外,如果将其设置为大于 1,则需要将different_layout_per_head设置为True

Fixed sparsity structure

  • BSLongformer (BSLongformerSparsityConfig):此结构是Longformer: The Long-Document Transformer的修改版本,其中我们提供的是块状 token 稀疏性,而非单个 token 稀疏性。定义此模式的参数有
    • num_sliding_window_blocks:一个整数,表示滑动局部注意力窗口中的块数量。
    • global_block_indices:一个整数列表,确定哪些块被视为全局注意力。给定索引,确定所有其他 token 块都关注的块,并且它们也关注所有其他 token 块。请注意,如果设置了global_block_end_indices参数,则此参数用作每个全局窗口的起始索引。
    • global_block_end_indices:一个整数列表,表示全局窗口块的结束索引。默认情况下不使用此参数。但如果设置了,它必须与global_block_indices参数大小相同,并且结合这两个参数,对于每个索引i,从global_block_indices[i]global_block_end_indices[i](不包括)的块都被视为全局注意力块。
  • BigBird (BigBirdSparsityConfig):此结构基于《Big Bird: Transformers for Longer Sequences》。它以某种方式结合了fixedlongformer模式以及随机注意力的思想。以下参数定义了此结构
    • num_random_blocks:一个整数,表示每行块中随机关注的块的数量。
    • num_sliding_window_blocks:一个整数,表示滑动局部注意力窗口中的块数量。
    • num_global_blocks:一个整数,表示从索引 0 开始,有多少个连续的块被视为全局注意力。全局块 token 将被所有其他块 token 关注,并且也将关注所有其他块 token。
  • Variable (VariableSparsityConfig):此结构也结合了局部、全局和随机注意力的思想。此外,它还具有定义可变大小局部窗口的灵活性。以下是定义此结构的参数列表
    • num_random_blocks:一个整数,表示每行块中随机关注的块的数量。
    • local_window_blocks:一个整数列表,用于确定每个局部注意力窗口中的块数量。它假定第一个数字确定第一个局部窗口中的块数量,第二个数字确定第二个窗口中的块数量,依此类推,最后一个数字确定剩余局部窗口中的块数量。
    • global_block_indices:一个整数列表,确定哪些块被视为全局注意力。给定索引,确定所有其他 token 块都关注的块,并且它们也关注所有其他 token 块。请注意,如果设置了global_block_end_indices参数,则此参数用作每个全局窗口的起始索引。
    • global_block_end_indices:一个整数列表,表示全局窗口块的结束索引。默认情况下不使用此参数。但如果设置了,它必须与global_block_indices参数大小相同,并且结合这两个参数,对于每个索引i,从global_block_indices[i]global_block_end_indices[i](不包括)的块都被视为全局注意力块。
    • attention:一个字符串,表示注意力类型。注意力可以是unidirectional(单向),例如自回归模型,其中 token 只关注上下文中它们之前出现的 token。考虑到这一点,注意力矩阵的上三角部分是空的,如上图所示。或者可以是bidirectional(双向),例如 BERT,其中 token 可以关注它们之前或之后出现的任何其他 token。那么,注意力矩阵的上三角部分是下图下三角部分的镜像。
    • horizontal_global_attention:一个布尔值,确定作为局部窗口全局代表的块是否也关注所有其他块。这仅在注意力类型为bidirectional时有效。从注意力矩阵来看,这意味着全局注意力不仅包括垂直块,还包括水平块。下图展示了variable稀疏性的一个例子,其中蓝色、橙色和绿色块分别表示局部、全局和随机注意力块。

Variable sparsity structure

此外,我们提供了一个dense模式(DenseSparsityConfig),可用于测试,同时它代表了完整的注意力。

如何支持新的用户定义稀疏结构

我们的构建块内核,即基于块的MatMulSoftmax,可以接受任何基于块的稀疏性。这提供了将任何基于块的稀疏模式应用于注意力分数的灵活性。要定义和应用新的稀疏模式,您只需遵循上述任何稀疏结构即可。您需要添加一个新类来扩展SparsityConfig,并根据稀疏模式的结构定义make_layout函数。您可以添加任何您可能需要的额外参数,或者只使用父类的默认参数。

更新日期: