[PyTorch] TransformerLayer: add support for Falcon architecture

[Marks101]

Falcon-40b and Falcon-180b are two exciting publicly available models. Currently, transformer-engine does not support their architecture because in their implementation self attention and mlp are not computed in sequence. Instead, the blocks (layer norm -> self attention) and (layernorm -> mlp) are fed with the input into the layer. In a computational graph these operations are thus in parallel. In the Falcon configs this is denoted as new_decoder_architecture. This PR introduces this feature and thus makes it possible to finetune Falcon models with transformer-engine. We would be really happy if this feature finds it's way into transformer-engine.

Two notes on the implementation:

• I denoted the new option parallel_attention_mlp. I am not sure if this is a perfect naming, this is up for discussion.
• I created a new method _bias_dropout_add() in order to keep the new code clear
• Falcon models do not have a bias, accordingly I think that a pretty clean solution is to set return_bias=False for attention and mlp and then use bias dropout add. For models using parallel_attention_mpl and bias, this might not be optimal.

[timmoon10]

/te-ci

[timmoon10]

LGTM

[Marks101]

Thanks for proceeding so quickly with this PR 🥳
Sadly unittests for jax and paddle failed. Not sure how this could be influence based on my changes.

[ptrendx]

@Marks101 Could you add a test to test_numerics, similar to https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/test_numerics.py#L622? Otherwise LGTM :-). The other frameworks tests seem to have failed because of machine issues, so not related to this PR.

[timmoon10]

/te-ci pytorch

[Marks101]

Thanks for fixing my spelling mistakes ... sorry for that. The tests failed because there was one last occurance of parallel_attention_ml. I fixed that and ran the tests locally. Now it should be fine.

[timmoon10]

/te-ci pytorch

[ptrendx]

/te-ci pytorch

[ptrendx]

Tim's attempt to run the CI failed due to network issue apparently, I just retried it.

[ptrendx]

Merged. Thank you @Marks101 for the contribution!

[Marks101]

Great, thank you for the support!