TransformerLayer input_mask format

[sshleifer]

I am trying to use the DeepSpeedTransformerLayer and wondering what format the attention mask should be for left to right language model training.
From https://github.com/microsoft/DeepSpeed/blob/44bd538b110ce0e8fc69626854631c3aee0dc094/tests/unit/test_cuda_forward.py#L181 , it seems like (bs, 1, seq_len, seq_len) could be correct,

but input_size: torch.Size([1, 501, 512]) and input_mask.shape=[1, 501, 501] raises

            input_mask = torch.cat((input_mask, torch.ones((inp_size[0], input_mask.shape[1], input_mask.shape[2], \>                                           (16 - (inp_size[1] % 16))), device=input_mask.device, dtype=input_mask.dtype) * -10000), 3)
E           IndexError: Dimension out of range (expected to be in range of [-3, 2], but got 3)

There is no docstring so I figured I'd ask. Thanks!

[RezaYazdaniAminabadi]

Hi @sshleifer

The input_mask that we currently support in our implementation must have this shape (batch_size, 1, 1, sequence_length). This means that the mask can be different for different input.
I think in your case the mask is triangular that has a different shape like (1, 1, sequence_length, sequence_length) or (1, sequence_length, softmax_length). This masking is mostly used for GPT2 type of models.

Can I ask what use-case you want to use the transformer layer for?
Thanks,
Reza

[sshleifer]

Use case:

• I'm trying to update working gpt2 training code in fairseq the mask is shaped (seq_len, seq_len). https://github.com/pytorch/fairseq/blob/master/fairseq/modules/transformer_layer.py#L159

[RezaYazdaniAminabadi]

I see. I can add such feature in the transformer layer.
Just give me some time as I am wrapping up some other issues and I can get back to this reasonably soon!

Thanks,
Reza

[sshleifer]

Thanks!

[sshleifer]

Any timeline here @RezaYazdaniAminabadi ?

[RezaYazdaniAminabadi]

I already have the kernel developed, just running some tests with GPT2 finetuning, I hope to have it ready by the end of this week or the beginning of next week at most.
Thanks in advance for your patience.
Reza

[RezaYazdaniAminabadi]

Hi @sshleifer,

I add a PR to work on this. I am still testing this and will add unit test for this type of masking.
One question on the masking side: when the mask dimension is torch.Size([1, 501, 512]), how does this work? Will it still consider the previous tokens when computing the SoftMax of each one? Or, the mask can be arbitrary no matter what the dimension is?

Thanks.
Reza

[sshleifer]

The mask is shaped (seq_len, seq_len)
It looks like below, and is meant to prevent the attention from considering future tokens.

self_attn_mask: tensor([[0., -inf, -inf,  ..., -inf, -inf, -inf],
        [0., 0., -inf,  ..., -inf, -inf, -inf],
        [0., 0., 0.,  ..., -inf, -inf, -inf],
        ...,
        [0., 0., 0.,  ..., 0., -inf, -inf],
        [0., 0., 0.,  ..., 0., 0., -inf],
        [0., 0., 0.,  ..., 0., 0., 0.]]) 

and then gets passed to MHA here.

I am pretty sure I am not answering your question, so I will try a rephrase on a higher level:

The integration is roughly that if you make this change in fairseq, and likely some more changes here then

pytest tests/test_binaries.py::TestLanguageModeling::test_transformer_lm 

should pass. The fairseq API and the huggingface Bart, Pegasus, T5 decoder layers all want "future mask"/"causal mask" formats.

[RezaYazdaniAminabadi]

Thanks @sshleifer

I think we can start with the (seq_len, seq_len) shape and then add the rest as well. Thanks for sharing the links, I will look more into them.

Reza

[afiaka87]

Thanks @sshleifer

I think we can start with the (seq_len, seq_len) shape and then add the rest as well. Thanks for sharing the links, I will look more into them.

Reza

Edit: I have followed up on the linked issue with a series of tests to try and solve this problem - but I'm honestly not certain it's the issue at the point as I'm also using a branch here which allows us to use CUDA 11. I unfortunately don't have the time to test this behavior as it's quite bizarre and would be costly to try to reproduce on CUDA 10.1. We've also recently implemented horovod support so the motivation to fix this is low for me personally.

Original comment:

Hey there - having a similar issue. We're trying to use stage 3 in order to train the project DALLE-pytorch on a distributed system. For this architecture - the transformer needs a mask in the shape [-1,0]. Rather than to prevent attention on future tokens - in our instance we need to prevent text from paying attention to image tokens and vice versa.

# mask logits to make sure text predicts text (except last token), and image predicts image

logits_mask = self.logits_mask[:, :seq_len]
max_neg_value = -torch.finfo(logits.dtype).max
logits.masked_fill_(logits_mask, max_neg_value)

I think the proposed fix you're working should allow for this. Is that correct?

lucidrains/DALLE-pytorch#185

[sshleifer]

Should #933 be runnable from python if I

DS_BUILD_TRANSFORMER=1 pip install . # from the branch

and then
call layer(x, mask) with x shaped (bs, seq_len, d_model) and mask shaped ( seq_len, seq_len)?
do I need to send a new argument to DeepSpeedTransformerConfig?

I tried this approach and got RuntimeError: CUDA error: an illegal memory access was encountered.

[RezaYazdaniAminabadi]

Hi @sshleifer,

Thanks for trying this out.
No, the implementation should find this out by checking the shape of the mask. Can you please tell me how you are creating the layer (mainly the configuration and sequence length that you are using), so that I can try in on my end too?

Thanks.
Reza

[RezaYazdaniAminabadi]

Can you please try this with having the mask shape as (1, 1, Seq, Seq)?

[RezaYazdaniAminabadi]

I did try with some masks on my side and it seems to work. However, I wait to understand your way of doing this to test this further.
Thanks