fix(docker): install TransformerEngine from rolling wheels

.github/workflows/docker-build.yml

@@ -6,6 +6,7 @@ on:
       - main
     paths:
       - 'docker/Dockerfile'
+      - 'docker/verify_transformer_engine.py'
       - 'requirements.txt'
 
   schedule:

.github/workflows/pr-test.yml

@@ -83,16 +83,15 @@ jobs:
         shell: bash
         run: |
           git fetch --depth 1 origin "${{ github.event.pull_request.base.sha }}"
-          # Dockerfile.rocm doesn't feed the cu13-x86 build this pipeline produces.
+          # Dockerfile.rocm doesn't feed the cu13 multi-arch build this pipeline produces.
           if git diff --name-only "${{ github.event.pull_request.base.sha }}" HEAD \
-              | grep -v '^docker/Dockerfile\.rocm$' \
-              | grep -qE '^(docker/(Dockerfile|build\.py|patch/)|requirements\.txt)'; then
+              | grep -qE '^(docker/(Dockerfile|build\.py|verify_transformer_engine\.py)|requirements\.txt)$|^docker/patch/'; then
             echo "changed=true" >> "$GITHUB_OUTPUT"
           else
             echo "changed=false" >> "$GITHUB_OUTPUT"
           fi
 
-  # Docker PRs: build the image first and run every suite inside it.
+  # Docker PRs: build the multi-arch image first and run every suite inside it.
   docker-build:
     needs: [docker-paths]
     if: needs.docker-paths.outputs.changed == 'true' && github.event.pull_request.head.repo.full_name == github.repository
@@ -118,7 +117,7 @@ jobs:
           password: ${{ secrets.DOCKERHUB_TOKEN }}
       - name: Build and push PR tag
         run: |
-          python3 docker/build.py --variant cu13-x86 --image-tag custom \
+          python3 docker/build.py --variant cu13 --image-tag custom \
             --custom-tag pr-${{ github.event.pull_request.number }} --push
 
   resolve-ci-image:

docker/Dockerfile

@@ -105,15 +105,22 @@ RUN if ls /tmp/wheels/causal_conv1d-*.whl 2>/dev/null | grep -q . && \
     fi
 
 # transformer_engine
-RUN if [ "${ENABLE_CUDA_13}" = "1" ]; then \
-      pip install --no-deps transformer_engine==2.17.0 && \
-      pip install transformer_engine_cu13==2.17.0 && \
-      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; \
+RUN --mount=type=bind,source=docker/verify_transformer_engine.py,target=/tmp/verify_transformer_engine.py \
+    if [ "${ENABLE_CUDA_13}" = "1" ]; then \
+      TE_CORE_DIST=transformer_engine_cu13; \
     else \
-      pip -v install --no-build-isolation "transformer_engine[pytorch]==2.17.0"; \
-    fi
+      TE_CORE_DIST=transformer_engine_cu12; \
+    fi && \
+    set -- \
+      /tmp/wheels/transformer_engine-2.17.0-py3-none-any.whl \
+      /tmp/wheels/${TE_CORE_DIST}-2.17.0-py3-none-manylinux_2_28_*.whl \
+      /tmp/wheels/transformer_engine_torch-2.17.0-cp312-cp312-linux_*.whl; \
+    if [ "$#" -ne 3 ] || [ ! -f "$1" ] || [ ! -f "$2" ] || [ ! -f "$3" ]; then \
+      echo "expected exactly one TransformerEngine 2.17 wheel for each component (${TE_CORE_DIST})" >&2; exit 1; \
+    fi && \
+    pip uninstall -y transformer-engine transformer-engine-cu12 transformer-engine-cu13 transformer-engine-torch && \
+    pip install --force-reinstall --no-deps "$@" && \
+    python3 /tmp/verify_transformer_engine.py "${TE_CORE_DIST}"
 
 # TE patches (cu13): B300/GB300 FA2 and backward override fixes
 # te_dequantized_backward_override.patch is a hot fix from

docker/verify_transformer_engine.py

@@ -0,0 +1,30 @@
+"""Verify the installed TransformerEngine wheel triplet."""
+
+import importlib.metadata as metadata
+import importlib.util
+import sys
+
+
+VERSION = "2.17.0"
+
+
+def verify(core_dist: str) -> None:
+    core = core_dist.replace("_", "-")
+    expected = {
+        "transformer-engine": VERSION,
+        core: VERSION,
+        "transformer-engine-torch": VERSION,
+    }
+    actual = {name: metadata.version(name) for name in expected}
+    requires = {
+        requirement.lower().replace("_", "-").replace(" ", "")
+        for requirement in (metadata.requires("transformer-engine-torch") or [])
+    }
+
+    assert actual == expected, f"unexpected TransformerEngine versions: {actual}"
+    assert f"{core}=={VERSION}" in requires, f"unexpected TransformerEngine torch requirements: {requires}"
+    assert importlib.util.find_spec("transformer_engine") is not None, "transformer_engine package not found"
+
+
+if __name__ == "__main__":
+    verify(sys.argv[1])

docs/ci/02-docker-build.md

@@ -61,11 +61,11 @@ A multi-arch build (`cu13`) needs Buildx's `docker-container` driver and is push
 
 Dockerfile changes are build-tested on the PR itself, before merge — `docker-build.yml` only runs after a push to `main`, so without this breakage lands on `main` first.
 
-When a PR touches `docker/Dockerfile`, `docker/build.py`, `docker/patch/**`, or `requirements.txt` (detected by the `docker-paths` job), `pr-test.yml` inserts a build in front of the test matrix:
+When a PR touches `docker/Dockerfile`, `docker/build.py`, `docker/verify_transformer_engine.py`, `docker/patch/**`, or `requirements.txt` (detected by the `docker-paths` job), `pr-test.yml` inserts a build in front of the test matrix:
 
 | Job | What it does |
 | --- | --- |
-| `docker-build` | builds `cu13-x86` and pushes a PR-scoped `radixark/miles:pr-<num>` tag (same-repo PRs; fork PRs skip it and test on `dev`) |
+| `docker-build` | builds `cu13` for `linux/amd64` and `linux/arm64`, then pushes one multi-arch PR-scoped `radixark/miles:pr-<num>` tag (same-repo PRs; fork PRs skip it and test on `dev`) |
 | `resolve-ci-image` | waits for the build and resolves the CI image to `pr-<num>`, so **every GPU suite runs inside the freshly built image**; a failed build stops the matrix instead of testing the stale image. The fresh build outranks a `ci-image-tag:` PR-body directive — the directive applies only when no PR image was built (non-docker or fork PRs) |
 | `delete-pr-tag` (`docker-pr-tag-cleanup.yml`) | removes the `pr-<num>` tag when the PR closes; the tag stays available for re-runs while the PR is open |
 
@@ -82,14 +82,14 @@ The only automated builder of `radixark/miles`. Two jobs:
 
 ### Triggers: automatic vs manual
 
-- **Automatic** (no human) — the **schedule** (cron 00:00 / 12:00 UTC, gated by `check-upstream`) and any **push to `main` that touches `docker/Dockerfile` or `requirements.txt`**. Both leave `--variant` empty and build **two images**: `cu13` → `radixark/miles` (multi-arch) and `cu12-x86` → `radixark/miles:dev-cu12`.
+- **Automatic** (no human) — the **schedule** (cron 00:00 / 12:00 UTC, gated by `check-upstream`) and any **push to `main` that touches `docker/Dockerfile`, `docker/verify_transformer_engine.py`, or `requirements.txt`**. Both leave `--variant` empty and build **two images**: `cu13` → `radixark/miles` (multi-arch) and `cu12-x86` → `radixark/miles:dev-cu12`.
 - **Manual** — `workflow_dispatch` (pick one variant — see Trigger a build yourself below) or running `docker/build.py` locally. Only the `rocm-*` images have **no automatic path** (`cu13-x86` / `cu13-aarch64` just rebuild the same `dev` image single-arch).
 
 
 | Trigger                                     | `check-upstream`                   | builds                | `latest` move     | prune      |
 | ------------------------------------------- | ---------------------------------- | --------------------- | ----------------- | ---------- |
 | schedule (cron 00:00 / 12:00 UTC)           | runs; build only if upstream moved | `cu13` + `cu12-x86`   | yes (both)        | yes (both) |
-| push to `main` touching `docker/Dockerfile` or `requirements.txt` | skipped                            | `cu13` + `cu12-x86`   | no                | no         |
+| push to `main` touching `docker/Dockerfile`, `docker/verify_transformer_engine.py`, or `requirements.txt` | skipped                            | `cu13` + `cu12-x86`   | no                | no         |
 | `workflow_dispatch`                         | skipped                            | the one input variant | no                | no         |
 | `workflow_dispatch` + `simulate_schedule`   | runs                               | the one input variant | no                | no         |