How to use TransformerEngine?

[tylerweitzman]

On the H100 there is a fully fused layer for an entire transformer block called TransformerLayer provided for PyTorch via TransformerEngine. How can I get this layer into TensorRt? Can it be exported via ONNX?

[zerollzeng]

@nvpohanh can you help here? Thanks!

[nvpohanh]

@tylerweitzman There is an ongoing work to add ONNX-exporting support into TransformerEngine. You can find some active changes in TE repo about this, such as: NVIDIA/TransformerEngine#197

The more complete TE -> ONNX -> TRT support will be done in TRT 8.7 release, but you can start trying this workflow now if you are interested.

cc @nzmora-nvidia

[ttyio]

closing since no activity for more than 3 weeks, thanks!