diff --git a/megatron/arguments.py b/megatron/arguments.py index de16fd5b3..457c7c18e 100644 --- a/megatron/arguments.py +++ b/megatron/arguments.py @@ -222,7 +222,7 @@ def parse_args(extra_args_provider=None, defaults={}, assert args.encoder_seq_length is not None args.seq_length = args.encoder_seq_length - if args.position_embedding_type == PositionEmbeddingType.absolute: + if args.position_embedding_type == PositionEmbeddingType.absolute or args.position_embedding_type == PositionEmbeddingType.alibi: assert args.max_position_embeddings is not None if args.seq_length is not None: assert args.max_position_embeddings >= args.seq_length @@ -312,7 +312,7 @@ def _add_network_size_args(parser): group.add_argument('--position-embedding-type', type=lambda x: PositionEmbeddingType[x], choices=list(PositionEmbeddingType), default=PositionEmbeddingType.absolute, - help='Define position embedding type ("absolute" | "rotary"). "absolute" by default.' + help='Define position embedding type ("absolute" | "rotary" | "alibi"). "absolute" by default.' ) group.add_argument('--glu-activation', type=str, choices=megatron.model.glu_activations.GLU_ACTIVATIONS.keys(), diff --git a/megatron/enums.py b/megatron/enums.py index 84d9a4ff0..46a6816a1 100644 --- a/megatron/enums.py +++ b/megatron/enums.py @@ -30,3 +30,4 @@ class AttnMaskType(enum.Enum): class PositionEmbeddingType(enum.Enum): rotary = 1 absolute = 2 + alibi = 3 diff --git a/megatron/model/transformer.py b/megatron/model/transformer.py index 473b8e06b..8d141e91f 100644 --- a/megatron/model/transformer.py +++ b/megatron/model/transformer.py @@ -203,7 +203,7 @@ def __init__(self, init_method, self.rotary_emb = RotaryEmbedding(self.hidden_size_per_attention_head, precision=args.params_dtype) def forward(self, hidden_states, attention_mask, layer_past=None, - get_key_value=False, encoder_output=None): + get_key_value=False, encoder_output=None, alibi=None): # hidden_states: [sq, b, h] # ===================== @@ -277,12 +277,15 @@ def forward(self, hidden_states, attention_mask, layer_past=None, output_size[0] * output_size[1], -1) # preallocting result tensor: [b * np, sq, sk] - matmul_result = torch.empty( - output_size[0]*output_size[1], - output_size[2], - output_size[3], - dtype=query_layer.dtype, - device=torch.cuda.current_device()) + if alibi is None: + matmul_result = torch.empty( + output_size[0]*output_size[1], + output_size[2], + output_size[3], + dtype=query_layer.dtype, + device=torch.cuda.current_device()) + else: + matmul_result = alibi[:output_size[0]*output_size[1], :, :output_size[3]] # Rotary embeddings if self.position_embedding_type == PositionEmbeddingType.rotary: @@ -301,11 +304,10 @@ def forward(self, hidden_states, attention_mask, layer_past=None, matmul_result, query_layer.transpose(0, 1), # [b * np, sq, hn] key_layer.transpose(0, 1).transpose(1, 2), # [b * np, hn, sk] - beta=0.0, alpha=(1.0/self.norm_factor)) + beta=0.0 if alibi is None else 1.0, alpha=(1.0/self.norm_factor)) # change view to [b, np, sq, sk] attention_scores = matmul_result.view(*output_size) - # ================================================== # Update attention mask for inference. [b, np, sq, sk] # ================================================== @@ -467,7 +469,7 @@ def __init__(self, init_method, output_layer_init_method, def forward(self, hidden_states, attention_mask, encoder_output=None, enc_dec_attn_mask=None, - layer_past=None, get_key_value=False): + layer_past=None, get_key_value=False, alibi=None): # hidden_states: [b, s, h] # Layer norm at the beginning of the transformer layer. @@ -477,7 +479,8 @@ def forward(self, hidden_states, attention_mask, self.self_attention(layernorm_output, attention_mask, layer_past=layer_past, - get_key_value=get_key_value) + get_key_value=get_key_value, + alibi=alibi) if get_key_value: attention_output, presents = attention_output @@ -594,6 +597,27 @@ def forward(self, inputs, **kwargs): class ParallelTransformer(MegatronModule): """Transformer class.""" + @staticmethod + def _build_alibi_tensor(max_seq_len, num_attention_heads, batch_size): + # Based on https://github.com/ofirpress/attention_with_linear_biases/blob/a35aaca144e0eb6b789dfcb46784c4b8e31b7983/fairseq/models/transformer.py#L742 + """Returns tensor shaped (batch_size * num_attention_heads, 1, max_seq_len)""" + def get_slopes(n): + def get_slopes_power_of_2(n): + start = (2 ** (-2 ** -(math.log2(n) - 3))) + ratio = start + return [start * ratio ** i for i in range(n)] + + if math.log2(n).is_integer(): + return get_slopes_power_of_2(n) + else: + closest_power_of_2 = 2 ** math.floor(math.log2(n)) + return get_slopes_power_of_2(closest_power_of_2) + get_slopes(2 * closest_power_of_2)[0::2][ + :n - closest_power_of_2] + slopes = torch.Tensor(get_slopes(num_attention_heads)) + alibi = slopes.unsqueeze(1).unsqueeze(1) * torch.arange(max_seq_len).unsqueeze(0).unsqueeze(0).expand(num_attention_heads, -1, -1) + alibi = alibi.repeat(batch_size, 1, 1) + return alibi + def __init__(self, init_method, output_layer_init_method, layer_type=LayerType.encoder, self_attn_mask_type=AttnMaskType.padding, @@ -660,11 +684,20 @@ def build_layer(layer_number): get_cuda_rng_tracker = deepspeed.checkpointing.get_cuda_rng_tracker checkpoint = deepspeed.checkpointing.checkpoint + if args.position_embedding_type == PositionEmbeddingType.alibi: + self.alibi = self._build_alibi_tensor(args.seq_length, args.num_attention_heads, args.micro_batch_size).to(torch.cuda.current_device()) + if args.params_dtype == torch.float16: + self.alibi = self.alibi.to(torch.float16) + elif args.params_dtype == torch.bfloat16: + self.alibi = self.alibi.to(torch.bfloat16) + else: + self.alibi = None + def _get_layer(self, layer_number): return self.layers[layer_number] def _checkpointed_forward(self, hidden_states, attention_mask, - encoder_output, enc_dec_attn_mask): + encoder_output, enc_dec_attn_mask, alibi=None): """Forward method with activation checkpointing.""" def custom(start, end): def custom_forward(*inputs): @@ -674,7 +707,7 @@ def custom_forward(*inputs): enc_dec_attn_mask = inputs[3] for index in range(start, end): layer = self._get_layer(index) - x_ = layer(x_, attention_mask, encoder_output, enc_dec_attn_mask) + x_ = layer(x_, attention_mask, encoder_output, enc_dec_attn_mask, alibi=alibi) return x_ return custom_forward @@ -731,7 +764,8 @@ def forward(self, hidden_states, attention_mask, layer_past=None, hidden_states = self._checkpointed_forward(hidden_states, attention_mask, encoder_output, - enc_dec_attn_mask) + enc_dec_attn_mask, + alibi=self.alibi) else: if get_key_value: presents = [] @@ -745,7 +779,8 @@ def forward(self, hidden_states, attention_mask, layer_past=None, encoder_output=encoder_output, enc_dec_attn_mask=enc_dec_attn_mask, layer_past=past, - get_key_value=get_key_value) + get_key_value=get_key_value, + alibi=self.alibi) if get_key_value: hidden_states, present = hidden_states presents.append(present)