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

[functionstackx]

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.

[functionstackx]

@hliuca should we close this issue now that it is in rocm/pytorch-training or is this issue tracking te into rocm/pytorch?

[hliuca]

Right now, it is included in rocm/pytorch-training.

Internally we have ticket and try to include into torch image too, to fix this issue.

[gugarosa]

Is there any update on this issue? TransformerEngine included in rocm/pytorch would be extremely helpful since it is not trivial to build.

[lucbruni-amd]

Hi @gugarosa,

Sorry for the late response. This (being TE included in rocm/pytorch specifically) has been implemented and is set to be released in a future version. I'll update this issue when I have more details.

Thanks for your patience in the matter.