Mra models have nan in output last_hidden_states

[zucchini-nlp]

System Info

When adding tests for batching equivalence (PR), it was found that Mra model has nan in outputs. I found the exact place where this happens, after the SparseDenseMatmul in this line, tensors start containing nan values.

When running tests for Mra, some configurations have to be tweaked as noted in this comment

Who can help?

@amyeroberts , tagging you here for tracking

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

Model tests

Expected behavior

Find out why we are having nan

[amyeroberts]

Thanks for opening @zucchini-nlp and digging into this in depth.

In particular, linking to the relevant comment from the tests. Based on that, I don't think there's really much we can do here - having nans was a compromise for MRA in order to have a faster running test suite. With that in mind, I think we can close the issue, and just skip the batching test for MRA, linking to this issue in the reason.