Skip to content

ALiBi Implementation #101

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
thomasw21 merged 7 commits into from
Sep 24, 2021
Merged

ALiBi Implementation #101

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
9aa94ab
alibi implementation
ofirpress Sep 10, 2021
4792218
works but slow
Sep 15, 2021
7ad5b48
efficient implementation
Sep 15, 2021
85c73f2
Update transformer.py
ofirpress Sep 16, 2021
3350322
Update transformer.py
ofirpress Sep 16, 2021
06dc2d7
Remove unnecessary .view()
ofirpress Sep 16, 2021
5b5afb2
bugfix with beta value in attention computation
Sep 16, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions megatron/arguments.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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(),
Expand Down
1 change: 1 addition & 0 deletions megatron/enums.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,3 +30,4 @@ class AttnMaskType(enum.Enum):
class PositionEmbeddingType(enum.Enum):
rotary = 1
absolute = 2
alibi = 3
65 changes: 50 additions & 15 deletions megatron/model/transformer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Comment thread
ibeltagy marked this conversation as resolved.
# hidden_states: [sq, b, h]

# =====================
Expand Down Expand Up @@ -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]]

@thomasw21 thomasw21 Sep 16, 2021

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

https://github.com/bigscience-workshop/Megatron-DeepSpeed/pull/101/files#diff-46c4c76deb18adf1de8e0be6d4229baed5f1f0308e141479f5a993b3d83dd445R307

When using baddmm shouldn't you set beta to 1? otherwise alibi should be ignored no?
https://pytorch.org/docs/stable/generated/torch.baddbmm.html

If you end up modifying beta then you probably have to replace empty with zeros in if alibi is None

@ibeltagy ibeltagy Sep 16, 2021

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I see, you are right. beta here

Megatron-DeepSpeed/megatron/model/transformer.py

Line 307 in 3350322

beta=0.0, alpha=(1.0/self.norm_factor))
should be 1.0 instead

@thomasw21 thomasw21 Sep 16, 2021

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How does the paper handle the normalizing factor? Is it applied after the sum?

@ofirpress ofirpress Sep 16, 2021

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

When using baddmm shouldn't you set beta to 1?

Yes!!! I meant to do that and then totally forgot! Thanks so much for pointing this out!!!

If you end up modifying beta then you probably have to replace empty with zeros in if alibi is None

Wouldn't it be better to do beta = 1 if alibi is not None else 0?

How does the paper handle the normalizing factor? Is it applied after the sum?

Which normalizing factor are you referring to? The softmax? If so, we apply the softmax after adding the ALiBi bias.

@thomasw21 thomasw21 Sep 16, 2021 •
edited
Loading

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I mean the 1.0 / self.norm_factor part. Is it applied only on QK or on the unormalized entire attention matrix?

@ofirpress ofirpress Sep 16, 2021 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah ok, I understand. We apply it before the sum, see https://github.com/ofirpress/attention_with_linear_biases/blob/master/fairseq/modules/multihead_attention.py#L225

(This means alpha will remain unchanged here)

@ofirpress ofirpress Sep 16, 2021

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed now :)


# Rotary embeddings
if self.position_embedding_type == PositionEmbeddingType.rotary:
Expand All @@ -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]
# ==================================================
Expand Down Expand Up @@ -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.
Expand All @@ -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
Expand Down Expand Up @@ -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):

@thomasw21 thomasw21 Sep 16, 2021

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we build a seperate class for this? something like class AlibiPositionEmbedding. There's a position_embedding.py file where there's a rotary embedding implementation.

# 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][
Comment on lines +613 to +614

@thomasw21 thomasw21 Sep 16, 2021

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why is it important to check that n is always a power of 2?

@ofirpress ofirpress Sep 16, 2021

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Explained in the comment here: https://github.com/ofirpress/attention_with_linear_biases/blob/master/fairseq/models/transformer.py#L749

#In the paper, we only train models that have 2^a heads for some a. This function has
#some good properties that only occur when the input is a power of 2. To maintain that even
#when the number of heads is not a power of 2, we use this workaround.

: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)

@thomasw21 thomasw21 Sep 16, 2021

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
alibi = slopes.unsqueeze(1).unsqueeze(1) * torch.arange(max_seq_len).unsqueeze(0).unsqueeze(0).expand(num_attention_heads, -1, -1)
alibi = slopes[:, None, None] * torch.arange(max_seq_len)[None, None, :].expand(num_attention_heads, -1, -1)

Also can you import the same comment as your original implementation? Typically where the matrix here doesn't match the one on the paper?

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,
Expand Down Expand Up @@ -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())

@ibeltagy ibeltagy Sep 16, 2021

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is .to(torch.cuda.current_device()) safe in the multi-gpu model parallelism setup?
@stas00, @slippylolo do you know?

@stas00 stas00 Sep 16, 2021

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This needs to be threaded with care. But on the other hand things quickly break if it's not done right and pytorch let's you know quickly ;)

It should already allocate the tensor on the current device in __init__(). does it break if don't to() explicity, @ofirpress?

Note that matmul_result is created during forward which indeed has the correct device set.

The other approach to not needing to use torch.cuda.current_device() in forward is to use .device of one of the inputs.

@stas00 stas00 Sep 16, 2021

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

They key is to add a good test and we have a multi-gpu CI, so it should be easy to validate.

@ofirpress ofirpress Sep 16, 2021 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

does it break if don't to() explicity, @ofirpress?

Yes it says that the alibi tensor is on the CPU if I don't to() it.

The other approach to not needing to use torch.cuda.current_device() in forward is to use .device of one of the inputs.

But during init do we have any of the inputs yet?

@stas00 stas00 Sep 16, 2021

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In init you typically have other params, but it's probably OK the way you did it.

Usually the model is made of registered params and buffers and those get automatically switched to the right device with the sub-module, so custom tensors not attached to the model are tricky. and typically you have to update them in forward to be on the same device as inputs or params.

I haven't looked at the full context, so let's revisit this when the tests are added if you're running into problems.

@stas00 stas00 Sep 16, 2021

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If it can be done in forward, then that's where we want it done I believe. Because only in forward you know which device you're on. You can't rely on the where it was init'ed. If it makes sense.

@stas00 stas00 Sep 16, 2021

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

for example:
https://github.com/huggingface/transformers/blob/master/src/transformers/models/fsmt/modeling_fsmt.py#L1276

@ofirpress ofirpress Sep 16, 2021 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I guess I'm not sure what the elegant way to do it is.
I guess we can do something like:

if self.alibi.device != some_input.device
    self.alibi.to(some_input.device)

but I'm not sure if that's OK.

@stas00 stas00 Sep 16, 2021

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

there is no need for if. It's a noop if it's already on the right device. So just the to call (plus assignment) as in the example I linked to.

@ibeltagy ibeltagy Sep 16, 2021

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am happy merging the PR as is then fix this later if it ends up being a problem. As @stas00, if it is wrong, pytorch will complain.

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):
Expand All @@ -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

Expand Down Expand Up @@ -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 = []
Expand All @@ -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)
Expand Down