[Out of Box Experience]: ROCm Transformer Engine Should Be Included in AMD Pytorch Images

[OrenLeung]

Suggestion Description

Hi @hliuca ,

On Nvidia NGC Pytorch Containers nvcr.io/nvidia/pytorch:24.xx-py3, Transformer Engine is included out of the box. This leads to less end user installation misconfiguration issues such as not using the correct build flags.

Currently on rocm/pytorch& on rocm/pytorch-nightly, Transformer Engine is not included out of the box.

It would be great to have parity with Nvidia on this.

https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-24-09.html

Operating System

Ubuntu

GPU

MI300X

ROCm Component

No response

[hliuca]

Thank you @OrenLeung again for the suggestion and support. Let me talk with our team about this.

[hliuca]

@OrenLeung we are working on this request to include TE in torch image. Thanks.