Skip to content

Make distinctive ROCm TE wheels names#441

Merged
ipanfilo merged 7 commits intodevfrom
ipanfilo/wheel_uniq_names
Mar 20, 2026
Merged

Make distinctive ROCm TE wheels names#441
ipanfilo merged 7 commits intodevfrom
ipanfilo/wheel_uniq_names

Conversation

@ipanfilo
Copy link
Collaborator

@ipanfilo ipanfilo commented Feb 3, 2026

Description

Use wheel package names for ROCm TE that are uniq and do not match CUDA TE ones.

Fixes #15360

Type of change

  • Documentation change (change only to the documentation, either a fix or a new content)
  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Infra/Build change
  • Code refactoring

Changes

Please list the changes introduced in this PR:

  • Rename transformer_engine_(torch|jax) to transformer_engine_rocm_(torch|jax)
  • Make transformer_engine_rocm independent, i.e it can install and work w/o transformer_engine metapackage
  • Add configuration parameters to build_wheels.sh

Checklist:

  • I have read and followed the contributing guidelines
  • The functionality is complete
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes

@wenchenvincent
Copy link
Collaborator

wenchenvincent commented Feb 4, 2026

If we use an AMD pypi repo to host the wheels, do we still need the package name changes?

@ipanfilo
Copy link
Collaborator Author

ipanfilo commented Feb 4, 2026 •
edited by Micky774
Loading

If we use an AMD pypi repo to host the wheels, do we still need the package name changes?

Yes, unless we disable using default PyPi index at all

@ipanfilo ipanfilo mentioned this pull request Feb 26, 2026
Merged
wenchenvincent
wenchenvincent reviewed Mar 20, 2026
View reviewed changes
docs/installation.rst Outdated
1. [For PyTorch support] `https://rocm.docs.amd.com/projects/install-on-linux/en/develop/install/3rd-party/pytorch-install.html`__
2. [For JAX support] `https://rocm.docs.amd.com/projects/install-on-linux/en/develop/install/3rd-party/jax-install.html`__

if HIP compiler complains it cannot detect the platform set `HIP_PLATFORM=amd` in the environment.
Copy link
Collaborator

@wenchenvincent wenchenvincent Mar 20, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: If

wenchenvincent
wenchenvincent reviewed Mar 20, 2026
View reviewed changes
docs/installation.rst Outdated
@@ -1,10 +1,132 @@
..
This file was modified to include portability information to AMDGPU.

Copy link
Collaborator

@wenchenvincent wenchenvincent Mar 20, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: remove empty line

wenchenvincent
wenchenvincent approved these changes Mar 20, 2026
View reviewed changes
Copy link
Collaborator

@wenchenvincent wenchenvincent left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM with nits

@ipanfilo ipanfilo merged commit 779f3fd into dev Mar 20, 2026
2 checks passed
ipanfilo added a commit that referenced this pull request Mar 20, 2026
@ipanfilo
Make distinctive ROCm TE wheels names (#441)
383dc90 
* Make distinctive ROCm TE wheels names
* Add ROCm major version to wheel package. Documentation update

(cherry picked from commit 779f3fd)
@ipanfilo ipanfilo deleted the ipanfilo/wheel_uniq_names branch March 20, 2026 17:44
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@wenchenvincent wenchenvincent wenchenvincent approved these changes

@wangye805 wangye805 Awaiting requested review from wangye805 wangye805 is a code owner

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@ipanfilo @wenchenvincent