Integrate TransformerEngine

[Quentin-Anthony]

Needed for fp8 training, and adds some nice fp16/bf16 optimizations for Ampere and newer architectures that we can make use of regardless.

https://github.com/EleutherAI/TransformerEngine

[Quentin-Anthony]

Fairly mature implementation at https://github.com/NVIDIA/Megatron-LM

[mkerin]

As discussed on Discord - if you need some extra dev manpower then I'll happily take this one off your hands

[StellaAthena]

As discussed on Discord - if you need some extra dev manpower then I'll happily take this one off your hands

Thank you!

[tf-nv]

Hi, I am curious about the state of the efforts and don't see a related branch. I read on discord that FP8 was working but there were struggles with convergence.

@Quentin-Anthony IIUC you spent some time on this as well, could you tell me more? :)

[mkerin]

Hi - just commenting to say that I’m afraid that I got distracted by other projects, and didn’t make any significant progress on this. Removing myself as assignee as agreed with Quentin so that I don’t block anyone else from picking it up, as I should have done much sooner.

[tf-nv]

There are a few things to unpack here. I had a look at the difference between GPT-NeoX megatron and the upstream megatron, which has a mature implementation as @Quentin-Anthony said. Here's a draft PR with some thoughts on the diff: #1185

It includes a few thoughts on the matter, let's discuss there :)