DeepSpeed稀疏注意力
在本教程中，我们将介绍如何使用DeepSpeed稀疏注意力（SA）及其构建块内核。使用SA最简单的方法是通过DeepSpeed启动器。我们将在如何使用DeepSpeed启动器进行稀疏注意力部分通过一个示例来描述这一点。但在那之前，我们在下一节中介绍了DeepSpeed SA提供的模块。
注意：目前，DeepSpeed稀疏注意力只能在使用Torch >= 1.6和CUDA 10.1、10.2、11.0或11.1的NVIDIA V100或A100 GPU上使用。
稀疏注意力模块
- MatMul：此模块处理块稀疏矩阵乘法。目前它支持SDD、DSD和DDS，如DeepSpeed稀疏注意力部分所述。
- Softmax：此模块应用块稀疏softmax。它处理前向和后向传递。
- SparseSelfAttention：此模块使用MatMul和Softmax内核，并根据查询、键和值生成上下文层输出。它是任何自注意力层中常见操作的简化版本。它还可以应用
相对位置嵌入注意力掩码键填充掩码到中间注意力分数。有关自注意力的更多详细信息，请查看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：此函数在序列长度维度上填充输入标记和注意力掩码，使其成为块大小的倍数；这是SA的要求。unpad_sequence_output：如果模型的输入已填充，则此函数会取消填充序列输出。
- SparsityConfig：这是一个用于稀疏结构的抽象类。任何稀疏结构都需要扩展此类并编写其自己的稀疏模式构造；
make_layout函数。DeepSpeed目前提供以下结构，将在如何配置稀疏结构部分进行描述FixedSparsityConfigBSLongformerSparsityConfigBigBirdSparsityConfigVariableSparsityConfigDenseSparsityConfig
注意：目前DeepSpeed Transformer内核不支持稀疏注意力。要使用稀疏注意力，您需要禁用Transformer内核！
如何使用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_attention是SparseSelfAttention的实例。此模块通过稀疏注意力计算注意力上下文，用其等效的稀疏版本替换底层矩阵乘法和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_ids和attention_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的示例。
- 添加稀疏配置：稀疏配置可以通过DeepSpeed JSON配置文件设置。在此示例中，我们使用了
fixed稀疏模式，将在如何配置稀疏结构部分进行描述。
"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
请参阅文档字符串，了解如何单独使用每个模块的详细信息。
如何配置稀疏结构
接下来，我们将描述支持的稀疏结构、其参数集以及在自注意力层上添加任意稀疏模式的灵活性。您可以使用任何支持的稀疏结构更新DeepSpeed配置文件并相应地设置参数。
- SparsityConfig：此模块是所有稀疏结构的父类，包含所有稀疏结构的共享特性。它采用以下参数
num_heads：一个整数，确定层的注意力头的数量。block：一个整数，确定块大小。稀疏自注意力的当前实现基于块稀疏矩阵。其中此参数定义此类正方形块的大小；块 X 块。different_layout_per_head：一个布尔值，确定是否应为每个头分配不同的稀疏布局；默认为false，这将根据可用性得到满足。
- Fixed（FixedSparsityConfig）：此结构基于OpenAI的使用稀疏Transformer进行生成建模，其中局部和全局注意力由给定参数固定
num_local_blocks：一个整数，确定局部注意力窗口中块的数量。如下图所示（改编自原始论文），局部窗口中的标记会关注所有局部标记。在自回归模型的情况下，如下图所示，标记会关注在局部窗口中之前出现的标记。在BERT等掩码模型的情况下，注意力是双向的。num_global_blocks：一个整数，确定局部窗口中多少个连续块用作窗口的代表以进行全局注意力；如下图所示。attention：一个字符串，确定注意力类型。注意力可以是unidirectional（例如自回归模型），其中标记仅关注在上下文中之前出现的标记。考虑到这一点，注意力矩阵的上三角形为空，如上图所示。或者它可以是bidirectional（例如BERT），其中标记可以关注其之前或之后的任何其他标记。然后，注意力矩阵的上三角部分是上图中下三角形的镜像。horizontal_global_attention：一个布尔值，确定作为局部窗口全局代表的块是否也关注所有其他块。仅当注意力类型为bidirectional时才有效。查看注意力矩阵，这意味着全局注意力不仅包括垂直块，还包括水平块。num_different_global_patterns：一个整数，决定不同全局注意力布局的数量。虽然全局注意力可以通过哪个/哪些块代表任何局部窗口来固定，但由于存在多头，每个头可以使用不同的全局代表。例如，使用4个块构建局部窗口，全局注意力大小为单个块，我们可以有4个不同的版本，其中每个局部窗口的第一个、第二个、第三个或第四个块可以作为该窗口的全局代表。此参数决定我们想要多少种这样的模式。当然，根据num_local_blocks和num_global_blocks存在限制。此外，如果您将其设置为大于1，则需要将different_layout_per_head设置为True。
- BSLongformer (BSLongformerSparsityConfig)：此结构是Longformer: The Long-Document Transformer的修改版本，其中我们提供块级令牌稀疏性，而不是单个令牌级稀疏性。定义此模式的参数为
num_sliding_window_blocks：一个整数，决定滑动局部注意力窗口中的块数。global_block_indices：一个整数列表，决定哪些块被视为全局注意力。给定索引，确定所有其他令牌块都关注的块，以及它们也关注所有其他令牌块。请注意，如果设置了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。它以某种方式结合了
fixed和longformer模式以及随机注意力的思想。以下参数定义此结构num_random_blocks：一个整数，决定每一行块中多少个块被随机关注。num_sliding_window_blocks：一个整数，决定滑动局部注意力窗口中的块数。num_global_blocks：一个整数，决定从索引0开始有多少个连续块被视为全局注意力。全局块令牌将被所有其他块令牌关注，并将关注所有其他块令牌。
- Variable (VariableSparsityConfig)：此结构还结合了局部、全局和随机注意力的思想。此外，它具有定义可变大小局部窗口的灵活性。以下是定义此结构的参数列表
num_random_blocks：一个整数，决定每一行块中多少个块被随机关注。local_window_blocks：一个整数列表，决定每个局部注意力窗口中的块数。它假设第一个数字决定第一个局部窗口中的块数，第二个数字决定第二个窗口中的块数……，最后一个数字决定剩余局部窗口中的块数。global_block_indices：一个整数列表，决定哪些块被视为全局注意力。给定索引，确定所有其他令牌块都关注的块，以及它们也关注所有其他令牌块。请注意，如果设置了global_block_end_indices参数，则此参数用作每个全局窗口的起始索引。global_block_end_indices：一个整数列表，决定全局窗口块的结束索引。默认情况下不使用它。但如果设置了它，它必须与global_block_indices参数具有相同的大小，并且结合这两个参数，对于每个索引i，从global_block_indices[i]到global_block_end_indices[i]（不包括）的块都被视为全局注意力块。attention：一个字符串，确定注意力类型。注意力可以是unidirectional（例如自回归模型），其中标记仅关注在上下文中之前出现的标记。考虑到这一点，注意力矩阵的上三角形为空，如上图所示。或者它可以是bidirectional（例如BERT），其中标记可以关注其之前或之后的任何其他标记。然后，注意力矩阵的上三角部分是上图中下三角形的镜像。horizontal_global_attention：一个布尔值，决定作为局部窗口全局代表的块是否也关注所有其他块。只有当注意力类型为bidirectional时，此选项才有效。查看注意力矩阵，这意味着全局注意力不仅包括垂直块，还包括水平块。下图说明了variable稀疏性的一个示例，其中蓝色、橙色和绿色块分别表示局部、全局和随机注意力块。
此外，我们提供了一个dense模式（DenseSparsityConfig），可以在测试时使用，因为它表示完全注意力。
如何支持新的用户定义稀疏结构
我们的构建块内核，基于块的MatMul和Softmax，可以接受任何基于块的稀疏性。这提供了将任何基于块的稀疏性模式应用于注意力分数的灵活性。要定义和应用新的稀疏性模式，您可以简单地遵循上述任何稀疏性结构。您需要添加一个扩展SparsityConfig的新类，并根据您的稀疏性结构定义make_layout函数。您可以添加任何您可能需要的额外参数，或者只使用父类的默认参数。