123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208 |
- # Copyright (c) Microsoft Corporation.
- # SPDX-License-Identifier: Apache-2.0
- # DeepSpeed Team
- import torch
- from torch.nn import functional as F
- from deepspeed.ops.sparse_attention import BertSparseSelfAttention, SparsityConfig
- '''
- This file contains few utility functions to handle adapting pretrained model with sparse self-attention module.
- '''
- class SparseAttentionUtils:
- """This class provides some utility functions that are use integrating sparse attention into transformer models.
- Such utilities include extending position embeddings, replacing current self-attention layer with sparse attention, padding sequences to multiple of block size, etc.
- """
- @staticmethod
- def extend_position_embedding(model, max_position):
- """This function extends the position embedding weights of a model loaded from a checkpoint.
- It assumes the new max position is bigger than the original max length.
- Arguments:
- model: required: a transformer model
- max_position: required: an integer determining new position embedding size
- Return:
- model: updated model; in which position embedding weights have been extended based on new size
- """
- if hasattr(model, 'bert'):
- original_max_position = model.bert.embeddings.position_embeddings.weight.size(0)
- assert max_position > original_max_position
- extend_multiples = max(1, max_position // original_max_position)
- model.bert.embeddings.position_embeddings.weight.data = model.bert.embeddings.position_embeddings.weight.repeat(
- extend_multiples, 1)
- elif hasattr(model, 'roberta'):
- # RoBERTa has positions 0 & 1 reserved, so embedding size is max position + 2
- original_max_position, embed_size = model.roberta.embeddings.position_embeddings.weight.shape
- original_max_position -= 2
- extend_multiples = max(1, max_position // original_max_position)
- assert max_position > original_max_position
- max_position += 2
- extended_position_embedding = model.roberta.embeddings.position_embeddings.weight.new_empty(
- max_position, embed_size)
- k = 2
- for i in range(extend_multiples):
- extended_position_embedding[k:(
- k + original_max_position)] = model.roberta.embeddings.position_embeddings.weight[2:]
- k += original_max_position
- model.roberta.embeddings.position_embeddings.weight.data = extended_position_embedding
- else:
- raise ValueError(
- 'Please extend \"extend_position_embedding\" function to support your model type. It currently only supports \"bert\" & \"roberta\"!'
- )
- model.config.max_position_embeddings = max_position
- print(f'Extended position embeddings to {original_max_position * extend_multiples}')
- return model
- @staticmethod
- def update_tokenizer_model_max_length(tokenizer, max_position):
- """This function updates the position embedding length of a tokenizer to a new max position.
- Arguments:
- tokenizer: required: a transformer tokenizer
- max_position: required: an integer determining new position embedding size
- Return:
- tokenizer: updated tokenizer; in which model maximum length has been extended based on new size
- """
- tokenizer.model_max_length = max_position
- tokenizer.init_kwargs['model_max_length'] = max_position
- print(f'updated tokenizer model max imum length to {max_position}')
- return tokenizer
- @staticmethod
- def replace_model_self_attention_with_sparse_self_attention(
- model,
- max_position,
- # SparsityConfig parameters needs to be set accordingly
- sparsity_config=SparsityConfig(num_heads=4)):
- """This function replaces the self attention layers in model encoder with sparse self attention.
- It currently supports bert and roberta model and can be easily extended to any other models following similar steps here.
- For sparsityConfig, refer to the config class.
- Arguments:
- model: required: a transformer model
- max_position: required: an integer determining new position embedding size
- sparsity_config: optional: this parameter determines sparsity pattern configuration; it is based on SparsityConfig class
- Return:
- model: updated model; in which self attention layer has been replaced with DeepSpeed Sparse Self Attention layer.
- """
- if hasattr(model, 'bert'):
- model.config.max_position_embeddings = max_position
- model.replace_self_attention_layer_with_sparse_self_attention_layer(model.config, model.bert.encoder.layer,
- sparsity_config)
- elif hasattr(model, 'roberta'):
- model.config.max_position_embeddings = max_position + 2
- model.replace_self_attention_layer_with_sparse_self_attention_layer(model.config,
- model.roberta.encoder.layer,
- sparsity_config)
- else:
- raise ValueError(
- 'Please extend \"update_model_self_attention_to_sparse_self_attention\" function to support \
- your model type. It currently only supports \"bert\" & \"roberta\"!')
- return model
- @staticmethod
- def replace_self_attention_layer_with_sparse_self_attention_layer(
- config,
- layers,
- # SparsityConfig parameters needs to be set accordingly
- sparsity_config=SparsityConfig(num_heads=4)):
- """This function replaces the self attention layers in attention layer with sparse self attention.
- For sparsityConfig, refer to the config class.
- Arguments:
- config: required: transformer model config
- layers: required: transformer model attention layers
- sparsity_config: optional: this parameter determines sparsity pattern configuration; it is based on SparsityConfig class
- Return:
- layers: updated attention layers; in which self attention layers have been replaced with DeepSpeed Sparse Self Attention layer.
- """
- for layer in layers:
- deepspeed_sparse_self_attn = BertSparseSelfAttention(config, sparsity_config)
- deepspeed_sparse_self_attn.query = layer.attention.self.query
- deepspeed_sparse_self_attn.key = layer.attention.self.key
- deepspeed_sparse_self_attn.value = layer.attention.self.value
- layer.attention.self = deepspeed_sparse_self_attn
- return layers
- @staticmethod
- def pad_to_block_size(block_size, input_ids, attention_mask, token_type_ids, position_ids, inputs_embeds,
- pad_token_id, model_embeddings):
- """This function pads input tokens and attention mask on sequence length dimension to be multiple of block size.
- This is a requirement for Sparse Transformer in which the self attention layer works on sequences of length multiple of block size.
- It needs to be called in your model, such as BertModel, right before you calculate the embedding outputs.
- Note)
- 1- instead of passing your embedding layer to this function, you can simply add this function to your model. It can be more simplified if given attention_mask and/or token_type_ids are none.
- 2- you need to call unpad function before returning your model output to unpad the encoder sequence output.
- Arguments:
- block_size: required: an integer determining the block size of sparsity config.
- pad_token_id: required: an integer determining the pad token from the model config; such as bert.config.pad_token_id.
- input_ids: a torch.LongTensor of shape [batch_size, sequence_length] with the word token indices in the vocabulary
- attention_mask: a torch.LongTensor of shape [batch_size, sequence_length] with indices selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max input sequence length in the current batch. It's the mask that we typically use for attention when a batch has varying length sentences.
- token_type_ids: a torch.LongTensor of shape [batch_size, sequence_length] with the token types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to a `sentence B` token (see BERT paper for more details).
- position_ids: a torch.LongTensor of shape [batch_size, sequence_length] with the indices of positions of each input sequence tokens in the position embeddings.
- inputs_embeds: an optional torch.FloatTensor of shape [batch_size, sequence_length, hidden_size] that contains embedded representation and can be passed instead of input_ids directly.
- model_embeddings: an optional object. If inputs_embeds are not none, this will be your model embeddings such as BertEmbeddings from your model such as BertModel. You can move this function inside your model and use self.embeddings instead of passing this parameter.
- Return:
- pad_len: an integer determining how much inputs have been padded to transfer sequence length dimension to multiple of block size.
- input_ids: if input_ids are not none padded input_ids otherwise none.
- attention_mask: if attention_mask is not none padded attention_mask otherwise none.
- token_type_ids: if token_type_ids are not none padded token_type_ids otherwise none.
- position_ids: if position_ids are not none padded position_ids otherwise none.
- inputs_embeds: if inputs_embeds are not none padded inputs_embeds otherwise none.
- """
- batch_size, seq_len = input_ids.shape if input_ids is not None else inputs_embeds.shape[:-1]
- pad_len = (block_size - seq_len % block_size) % block_size
- if pad_len > 0:
- if inputs_embeds is not None:
- pad_input_ids = inputs_embeds.new_full((batch_size, pad_len), pad_token_id, dtype=torch.long)
- pad_inputs_embeds = model_embeddings(pad_input_ids)
- inputs_embeds = torch.cat([inputs_embeds, pad_inputs_embeds], dim=-2)
- # may not be needed as input_ids are not used if inputs_embeds are given
- if input_ids is not None:
- input_ids = F.pad(input_ids, (0, pad_len), value=pad_token_id)
- if position_ids is not None:
- # pad position_id with pad_token_id
- position_ids = F.pad(position_ids, (0, pad_len), value=pad_token_id)
- # pad attention mask without attention on the padding tokens
- attention_mask = F.pad(attention_mask, (0, pad_len), value=False)
- # pad token_type_ids with token_type_id = 0
- token_type_ids = F.pad(token_type_ids, (0, pad_len), value=0)
- return pad_len, input_ids, attention_mask, token_type_ids, position_ids, inputs_embeds
- @staticmethod
- def unpad_sequence_output(pad_len, sequence_output):
- """This function unpads sequence output if inputs of the model were padded.
- This is a requirement for Sparse Transformer in which the self attention layer works on sequences of length multiple of block size.
- It needs to be called in your model, such as BertModel, right before you return the model outputs.
- Arguments:
- pad_len: required: an integer determining how much model inputs have been padded to transfer sequence length dimension to multiple of block size.
- sequence_output: required: sequence output of the encoder layer.
- Return:
- sequence_output: unpaded sequence output of the encoder layer.
- """
- if (pad_len > 0):
- sequence_output = sequence_output[:, :-pad_len]
- return sequence_output
|