[fsmt] onnx triu workaround
Author: stas00
Pull request: yes

This PR 

* solves 
   ``` 
   RuntimeError: Exporting the operator triu to ONNX opset version 12 is not supported. Please open a bug to request ONNX    export support for the missing operator.
   ```
   as reported in https://github.com/huggingface/transformers/issues/9737.
   It adds a workaround for `triu` not being supported by pytorch's onnx set, as proposed here: https://github.com/pytorch/pytorch/issues/32968#issuecomment-733240232 with some modifications to make it work with transformers. The original workaround couldn't handle a matrix with `-inf`.
* adds an onnx export test

The workaround fix is localized to fsmt, but transformers has a handful of those:

```
src/transformers/pipelines/question_answering.py:        candidates = np.tril(np.triu(outer), max_answer_len - 1)
src/transformers/models/fsmt/modeling_fsmt.py:    causal_mask = torch.triu(fill_with_neg_inf(torch.zeros(tgt_len, tgt_len)), 1).to(
src/transformers/models/transfo_xl/modeling_transfo_xl.py:            dec_attn_mask = (torch.triu(all_ones, 1 + mlen) + torch.tril(all_ones, -mask_shift_len))[:, :, None]  # -1
src/transformers/models/transfo_xl/modeling_transfo_xl.py:            dec_attn_mask = torch.triu(word_emb.new_ones((qlen, klen), dtype=torch.uint8), diagonal=1 + mlen)[
src/transformers/models/xlnet/modeling_xlnet.py:        mask_up = torch.triu(attn_mask, diagonal=1)
src/transformers/models/ctrl/modeling_ctrl.py:        mask = torch.triu(torch.ones(seq_len + past_length, seq_len + past_length), 1).to(inputs_embeds.device)
src/transformers/models/prophetnet/modeling_prophetnet.py:        left_block[stream_idx].triu_(-stream_idx + 1)
src/transformers/models/prophetnet/modeling_prophetnet.py:        causal_mask = torch.triu(causal_mask, 1)
```

Perhaps the workaround wrapper should be somewhere in the common tools and used in other places?

Or merge this to let the user who needed this in first place move forward and then refactor to other places which haven't been spoken for. But actually if you look at https://github.com/pytorch/pytorch/issues/32968 many of the me-too comments talk about `transformers`.

https://github.com/pytorch/pytorch/issues/32968 proposes other solutions too. I tried them and either they didn't work, or were inefficient as spotted by @patrickvonplaten 

Fixes: https://github.com/huggingface/transformers/issues/9737

@LysandreJik, @mfuntowicz 

Comments:
I'm a bit worried that this fix will blow up the memory for very long sequences, *e.g.* applying this to a max_length of 16384 (which we already have for LED) would create a 16384 ** 2 =  1GB tensor

Wouldn't this work-around be better in terms of memory? https://github.com/pytorch/pytorch/issues/32968#issuecomment-733240232

Glad you have the deep understanding of this one, @patrickvonplaten 

wrt, https://github.com/pytorch/pytorch/issues/32968#issuecomment-733240232 - I tried it originally and while onnx was happy, everything else failed with it.

Perhaps it can be re-written not to use `triu`? some suggest using `np.triu`

> Wouldn't this work-around be better in terms of memory? [pytorch/pytorch#32968 (comment)](https://github.com/pytorch/pytorch/issues/32968#issuecomment-733240232)

I managed to fix this more efficient version to work with `-inf`. It's in the PR now.

> Did you run the FSMT integration tests to ensure that it doesn't diverge?

Yes, the function produces the same output as `triu` for the inputs we use.

This workaround seems to only work for squared-sized tensors, while `torch.triu` works for any shape (PyTorch 1.9.0). For example:

```
a = torch.randn(8, 4)
triu_onnx(a)
```
raises an `RuntimeError: The size of tensor a (8) must match the size of tensor b (4) at non-singleton dimension 1`.

Is there any additional workaround when working with non-squared shape tensors?

Thank you and best regards,
Gustavo.

Indeed, the workaround was written for that sort of inputs.

Are you running into this problem with `transformers`? Could you please file a new issue then including the way to reproduce the problem?

But I also see that perhaps we can switch back to pytorch implementation as reported here:
https://github.com/pytorch/pytorch/issues/32968#issuecomment-827054124
it's called `trilu`. I think it should be in pt-1.9.0.

https://github.com/onnx/onnx/blob/29e7aa7048809784465d06e897f043a4600642b2/docs/Operators.md#Trilu

Would you like to experiment with it and see if it solves the problem? If it works you may consider creating a PR that switches to that version instead and we can work together on polishing out the details (as we need to support older pytorch as well).

To clarify: what I'm trying to propose is to try pytorch-1.9.0 and use its built-in `triu` and see if it now works. One way you could test is by reverting my original PR that introduced the workaround and see if it just works. alternatively, if you have the right know-how you can write some test code that tests torch's  `triu` directly with onnx export.


update: 

> 'triu' support added in PT-ONNX exporter in opset14 https://github.com/pytorch/pytorch/pull/59486 

which I suppose should be available in pytorch-1.10 when it comes out as it was merged on july-14. So once it's released we could re-do this workaround and fallback to it for pt<1.10.