Creating long versions of Albert

Hi,

I'm trying to follow convert_model_to_long script to convert Albert to a long version by replacing the attention with LongformerSelfAttention.

See https://colab.research.google.com/drive/18ECsAeWFEsNxQ8C-8Q_OMbsfH0eC0sCJ?usp=sharing

I think I've got it working with just dropping in the attention and some architecture changes.. I haven't pretrained yet, but cannot be sure as the BCP on eval is quite high ~ 17. But the BCP is quite high for regular Albert as well (but strangely the BCP of AlbertLong is lower than regular Albert which is ~ 19)

One specific question I have is to confirm that what is returned by the LongformerSelfAttention is 3d tensor: batch x sequence length x hidden_dim, correct?

Albert's q,k,v calculation produces a 4d tensor before doing a strange application against a dense tensor (768x768) to get the 3d tensor, and I am not sure if I'm applying LongformerSelfAttention correctly in this case.

context_layer = context_layer.permute(0, 2, 1, 3).contiguous()

# Should find a better way to do this
w = (
self.dense.weight.t()
.view(self.num_attention_heads, self.attention_head_size, self.hidden_size)
.to(context_layer.dtype)
)
b = self.dense.bias.to(context_layer.dtype)

projected_context_layer = torch.einsum("bfnd,ndh->bfh", context_layer, w) + b
projected_context_layer_dropout = self.dropout(projected_context_layer)
layernormed_context_layer = self.LayerNorm(input_ids + projected_context_layer_dropout)
return (layernormed_context_layer, attention_probs) if output_attentions else (layernormed_context_layer,)

I modified to do this instead:

attention_output = self.attention(hidden_states, attention_mask, output_attentions) # head_mask,
context_layer = attention_output[0]

# this is from the albert implementation. Is this equivlaent to the below??
# Should find a better way to do this
#w = (
# self.dense.weight.t()
# .view(self.num_attention_heads, self.attention_head_size, self.hidden_size)
# .to(context_layer.dtype)
#)
#b = self.dense.bias.to(context_layer.dtype)

new_shape = context_layer.size()[:-1] + (-1,)
context_layer = context_layer.view(*new_shape)
context_layer = self.dense(context_layer) # torch.einsum("bfnd,ndh->bfh", context_layer, w) + b
context_layer = self.dropout(context_layer)
context_layer = self.LayerNorm(hidden_states + context_layer) # projected_context_layer_dropout)

Any comments and suggestions is much appreciated.