[AMD] install Transformer Engine from prebuilt wheel on rocm720#1692
[AMD] install Transformer Engine from prebuilt wheel on rocm720#1692XinyuJiangCMU wants to merge 1 commit into
Conversation
Gate the TE install on TE_USE_WHEEL, which docker/build.py sets for the rocm720-mi35x variant: install the prebuilt fp8 blockwise wheel from the miles-wheels-rocm release instead of building from source. Other variants keep the source build via the else branch. Co-authored-by: Xinyu Jiang <xinyuj2@andrew.cmu.edu> Co-authored-by: Zhiyao Jiang <jessicajiang324@gmail.com>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a conditional build argument TE_USE_WHEEL in docker/Dockerfile.rocm to allow installing Transformer Engine from a prebuilt wheel instead of building it from source, and enables this flag in docker/build.py for a specific ROCm build configuration. Feedback was provided regarding the use of a wildcard glob (transformer_engine-*.whl) with pip install, which could cause build failures if multiple wheels are present in the directory; using --find-links was suggested as a more robust alternative.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| ARG TE_USE_WHEEL=0 | ||
| RUN pip uninstall -y transformer-engine transformer_engine transformer_engine_torch || true; \ | ||
| if [ "${TE_USE_WHEEL}" = "1" ]; then \ | ||
| pip install /tmp/wheels/transformer_engine-*.whl; \ |
There was a problem hiding this comment.
Using a wildcard glob like transformer_engine-*.whl with pip install can cause the build to fail if the /tmp/wheels directory contains multiple wheels (for example, wheels built for different Python versions or architectures). In such cases, pip will attempt to install all matching wheels, and will fail with an error when encountering an incompatible wheel.
To make this more robust, use --find-links to let pip automatically select and install only the wheel that is compatible with the current Python environment.
pip install --find-links=/tmp/wheels transformer_engine; \
1 participant
Co-authored-with: @JessicaJiang-123
Install Transformer Engine from a prebuilt wheel on
rocm720-mi35xinstead of building from source, gated on a newTE_USE_WHEELarg (docker/build.pysets it to1for that variant; other variants keep the source build). Cuts the rocm720 image build from ~30 min to ~2 min.The wheel (
transformer_engine-2.14.0.dev0, built bybuild_te_wheel.py, in themiles-wheels-rocmrocm720-gfx950-v0.5.14release) comes from the fp8 forkXinyuJiangCMU/TransformerEngine @ miles-dev(based on ROCm/TransformerEngine9c19ec7, "Add gfx1250 support to CK tile group GEMM", 2026-06-21); its Triton blockwise FP8 path for gfx950 (XinyuJiangCMU/TransformerEngine#4) isn't in the stockv2.8_rocmsource. Verified on MI355X: full build succeeds,Float8BlockScalingLinear forward runs on GPU.