[Bigscience] Non-causal Decoder Generation
Author: haileyschoelkopf
Pull request: yes

@thomasw21 @Muennighoff 

As we discussed, here's a quick hack to try out Prefix-LM on BLOOM via swapping out the mask for one that always attends to the first `prefix_length` tokens (as in this figure from the pretraining objectives paper). 

![image](https://user-images.githubusercontent.com/65563625/176474220-0332f2f1-9f74-4111-92b9-f1d96cf65b51.png)

EDIT: I extended this for cleaner interface + drop-in support for `bigscience/lm-eval-harness`

[Minimal script I was using to test outputs using this code](https://gist.github.com/haileyschoelkopf/33b9e41d07b9222e995c6cce155724de)

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_17944). All of your documentation changes will be reflected on that endpoint.

Demonstration still here: https://gist.github.com/haileyschoelkopf/33b9e41d07b9222e995c6cce155724de

PR now updated to include `BloomForPrefixLM` and passing in a mask tensor of size `[batch_size, 1, input_length, input_length]` 

Running the bigscience lm-eval-harness fork with Bloom Prefix LM should be as simple as just swapping the ("bloom", "BloomForCausalLM") entry in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES to ("bloom", "BloomForPrefixLM").

I proposed some adjustments here: https://github.com/haileyschoelkopf/transformers/pull/1
@haileyschoelkopf @thomasw21 

You should create a test that shows that whatever invariances your model has, it should be tested in the test:
 - changing a value in the input part changes all the input logits as well as target.
 - changing a value in the targer only changes logies before and not after (same as language modeling)
 - `generate` actually works with this function.

Thanks for the input. Will make above into tests. 

Thanks for the edits and comments Lintang!

I may not get to testing these today because of the holiday, but I will do so tomorrow if not.

@lintangsutawika I'm happy to also turn your `tests/models/bloom/test_noncausal_attention_bloom.py` test script into unit tests in huggingface/transformers if you'd like.

@younesbelkada since you're the author of the architecture in `transformers`

Update on the status of this: I started looking at Lintang's edits today. Commit [0f692c07788a536fec103d83c8acdd5565cccb05](https://github.com/huggingface/transformers/pull/17944/commits/0f692c07788a536fec103d83c8acdd5565cccb05) from me is the last one I had checked for generation; I will continue editing Lintang's changes as necessary tomorrow.


If necessary, happy to discuss further the merits of using `prefix_length` vs. passing `prefix_mask` to the model as a design choice--Lintang switched it back to `prefix_length`, and I think I agree with this because it'd be easier to just feed a prefix length to the model when training it on a batch of concatenated input-target sequences using a prefix LM objective.

I think we'd be much happier is the `attention_mask` was explicitely passed to the model. Typically one of the thing we're using in MTF right now is packing. I don't think packing exists as a technique under current bloom implementation.

Is the goal for the Transformers version of Bloom to also support packing? My understanding is that even transformers version of T5 doesn't have packing.

Having to input an attention mask means there would be two attention masks to input? The regular attention masks that the PretrainedModel object requires and the NonCausal Attention mask? I figured it would be easier for users to just declare the length of the prefix and let the model build the mask?

I suppose a compromise would be to have a NonCausalAttentionMask argument but allow the model to accept prefix length as well of the former is not provided.

I don't have a strong opinion but setting an attention mask explictely vs having `n` different mechanism to update an attention_mask is annoying. Today we're handling prefix, tomorrow we're handling another weird thing. One use-case I can think of is like `{input} <pad> <pad> <pad> {target} <pad> <pad> <pad>` I think this case can happen when you generate. 

Also @patrickvonplaten if you have any input about this. For context, we're going to train a prefix lm and so naturally we have bidirectional attention in the input, and casaul in the target. The current solution are:
 - pass a prefix length and update the attention mask
 - pass an explicit attention mask everytime.

If there is another attention formula, we could always add another BloomFor<Attention Type>LM? 

In the case for prefix LM, it seems easier to have a prefix length and modify the attention mask. 

I think there's a way to do both. We can allow `prefix_length` as an extra argument to `BloomForPrefixLM`, which will take care of creating the noncausal attention mask if no mask is given as an argument, and pass this mask to the base `BloomModel` forward and use it for attention computation in the forward pass. (`prefix_length` will never be passed to the `BloomModel` this way)

I'll start implementing this if there aren't any objections. 

Nice, I would suggest to do it differently though:
- Just have one additional kwarg in all forward funcs called `causal_mask` next to attention_mask
- If `causal_mask` is None set it to the `torch.tril` default directly in `BloomPreTrainedModel` (i.e. dont recreate it in every layer)
- Have one test checking that if it's set with prefixes it's different than the default like `test_equivalence_prefix_causal_lm`

I think this is all we would need & then the user just creates its own causal mask which has prefixes. There's no need for separate models to allow users to pass prefix masks or bidirectional masks with skipping like `[1,0,1]` or any other mechanism one may come up with like @thomasw21 said. What do you think?



I like the idea of having a single `BloomForCausalLM` that can handle the causal_masks be it causal or non causal. 

I just think that handling the non-causal mask could be done automatically inside the model especially during generation when the intention is to process the input prompt with non-causal mask. 



What if `causal_mask` has 3 possible input options?
- `None`: a causal mask is generated in the model
- `torch.Tensor` with shape same as input_ids: a manual causal mask (which can also be non-causal) is used.
- `torch.Tensor` with shape `batch size * 1`: is a prefix length matrix, a non-causal attention mask is generated in the model

Hi all,

I just went through the comments of this PR but did not checked the modifications yet. We are currently trying to refactor the modeling code in this PR: https://github.com/huggingface/transformers/pull/17866 - one thing that we are doing is to create the mask only once at the `_prepare_attn_mask` function that we pass to all submodules. It might be easier in your use case. I cannot give a fixed timeline on when the refactoring PR will get merged but if you merge this PR before I will take care of refactoring a bit your code ;) ! 

Thanks @Muennighoff for the comments! 

They make sense and I’ll work on addressing them, but I agree with @lintangsutawika that we should have some way of automatically creating a PrefixLM non-causal mask.

If we’re planning on users using Bloom0++ for generation then I think there should be some way of setting causal_mask to a non-causal mask without manually creating it, whether that’s a PrefixLM class or just a flag passed to the CausalLM forward or similar.

Addressed points 1 and 2 from @Muennighoff ! I left the BloomForPrefixLM class for now.

The tests I added (including slow tests locally) pass—it’s unrelated tests that fail currently.

@thomasw21 @Muennighoff 

FYI, I'm in the middle of refactoring this PR after pulling from main. So for now, if you could use commit `6538564ee1e3f1689ab71b01866aa7771b82edc7` that's the last one that still works.



> Nice, I would suggest to do it differently though:
> 
> * Just have one additional kwarg in all forward funcs called `causal_mask` next to attention_mask
> * If `causal_mask` is None set it to the `torch.tril` default directly in `BloomPreTrainedModel` (i.e. dont recreate it in every layer)
> * Have one test checking that if it's set with prefixes it's different than the default like `test_equivalence_prefix_causal_lm`
> 
> I think this is all we would need & then the user just creates its own causal mask which has prefixes. There's no need for separate models to allow users to pass prefix masks or bidirectional masks with skipping like `[1,0,1]` or any other mechanism one may come up with like @thomasw21 said. What do you think?

Very much agree with your idea here @Muennighoff ! Think this is the way to go which will be 100% backwards compatible (we could btw also add this to GPT2 etc...)

Following @Muennighoff thoughts' here I think we should make the behavior crystal clear in the docstring:

```
If 'causal_mask' is set to 'None' the model will automatically create the conventional causal (unidirectional) attention mask to prevent past tokens to attend to future tokens. If you would like to overwrite this behavior, *e.g.* to create a Prefix-LM architecture, please pass a tensor different to 'None'
```

Also cc @sgugger @LysandreJik here to hear their input on this (since this functionality might be extended to GPT2)

I also agree that @Muennighoff options sounds better than adding a new model which is very very similar to the causal LM architecture. It also will make the functionality more accessible (since users often use the auto-classes to load their models).

Thanks for the feedback! Especially since we probably aren't going with Prefix-LM for Bloom-T0, removing the class and just making this an optional argument sounds like a good idea. 

I'll ping you all after I get around to making these changes!

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.