[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.

[gabeweisz]

Jax too please!

[zez626]

Is there any update? TE is hard to install.

[hliuca]

@zez626 we have training images with TE, https://hub.docker.com/r/rocm/pytorch-training/tags

Also, the TE will be released with PyTorch images soon. Thank you.

[zez626]

Thank you for your replay, this helps a lot.

[functionstackx]

Hi @hliuca any update on the TE in pytorch images?

came around this issue where they wanted to use torch image and not the torch-training image and took them 8.5 hours to install TE

https://zdtech.substack.com/p/8-bit-training-with-rocm-transformer

[hliuca]

@functionstackx I understand this. Sorry about this.

@wenchenvincent could you please provide some updates? Thank you.

[wenchenvincent]

@functionstackx We will have pytorch TE docker images with ROCm 7.0 release.
In a parallel line, there is a series of unified training docker images published, with Megatron and TE pre-installed: https://hub.docker.com/r/rocm/megatron-lm/tags

AMD blog post for the unified training docker images: https://rocm.blogs.amd.com/software-tools-optimization/amd-optimized-rocm-docker-for-distributed-training/README.html

[hliuca]

Jax too please!

@gabeweisz for now, you may try https://hub.docker.com/r/rocm/jax-training/tags