From 98d15d6785752f8c574911bdef217d7812d07411 Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Fri, 24 Jul 2026 12:14:54 -0700 Subject: [PATCH] docker: install TE from PyPI, drop the prebuilt TE-torch wheel path The wheels release still ships transformer_engine_torch 2.12.0; installing it after the 2.17.0 bump (#1781) downgrades transformer-engine-cu13 via its dependency pin and the TE import sanity check fails the build. TE now comes from PyPI, so remove the wheel branch (and delete the stale wheel from /tmp/wheels so nothing else can pick it up). --- docker/Dockerfile | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/docker/Dockerfile b/docker/Dockerfile index 6e59a8df6b9..7d5d4f11bbd 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -108,12 +108,9 @@ RUN if ls /tmp/wheels/causal_conv1d-*.whl 2>/dev/null | grep -q . && \ RUN if [ "${ENABLE_CUDA_13}" = "1" ]; then \ pip install --no-deps transformer_engine==2.17.0 && \ pip install transformer_engine_cu13==2.17.0 && \ - if ls /tmp/wheels/transformer_engine_torch-*.whl 2>/dev/null | grep -q .; then \ - pip install /tmp/wheels/transformer_engine_torch-*.whl; \ - else \ - pip install nvidia-mathdx==25.6.0 && \ - pip -v install --no-build-isolation transformer_engine_torch==2.17.0; \ - fi; \ + rm -f /tmp/wheels/transformer_engine*.whl && \ + pip install nvidia-mathdx==25.6.0 && \ + pip -v install --no-build-isolation transformer_engine_torch==2.17.0; \ else \ pip -v install --no-build-isolation "transformer_engine[pytorch]==2.17.0"; \ fi