diff --git a/.github/workflows/build_kernel_tpu.yaml b/.github/workflows/build_kernel_tpu.yaml new file mode 100644 index 00000000..921b3b7f --- /dev/null +++ b/.github/workflows/build_kernel_tpu.yaml @@ -0,0 +1,45 @@ +name: "Build and test kernel (TPU)" +on: + push: + branches: [main] + pull_request: + branches: [main] + types: [opened, synchronize, reopened] # trigger on PRs + paths-ignore: + - "docs/**" + - "*.md" + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: true + +permissions: {} + +jobs: + build: + permissions: + contents: read + name: Build kernels (TPU) + runs-on: + group: aws-highmemory-32-plus-nix + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + - uses: DeterminateSystems/nix-installer-action@3138316df39ed29be04236d7ffc686fa525866aa # v23 + with: + extra-conf: | + max-jobs = 4 + cores = 12 + sandbox-fallback = false + - uses: cachix/cachix-action@1eb2ef646ac0255473d23a5907ad7b04ce94065c # v17 + with: + name: huggingface + authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}" + env: + USER: runner + - name: Nix info + run: nix-shell -p nix-info --run "nix-info -m" + # For now we only test that there are no regressions in building TPU + # kernels. Also run tests once we have a TPU runner. + - name: Build all TPU example kernels + run: nix build -L ./examples/kernels#ci-build-tpu diff --git a/README.md b/README.md index 91573c62..39d1b804 100644 --- a/README.md +++ b/README.md @@ -81,7 +81,8 @@ different tiers: | Huawei NPU | ✓ | ✗ | ✗ | 3 | | Neuron | ✓ | x | x | 3 | -**Warning:** Neuron and TPU support are experimental and packages are not public yet. +> [!WARNING] +> Neuron support is experimental and packages are not public yet. TPU support is is experimental. ## 📚 Documentation diff --git a/docs/source/builder/build-variants.md b/docs/source/builder/build-variants.md index 064e9a36..a6282a71 100644 --- a/docs/source/builder/build-variants.md +++ b/docs/source/builder/build-variants.md @@ -98,6 +98,10 @@ available. This list will be updated as new PyTorch versions are released. - `torch215-rocm100-x86_64-linux` - `torch215-rocm714-x86_64-linux` +## TPU x86_64-linux + +- `torch-tpu` + ## XPU x86_64-linux - `torch210-xpu20253-x86_64-linux` @@ -116,4 +120,5 @@ one or more of the following variants: - `torch-cuda` - `torch-metal` - `torch-rocm` +- `torch-tpu` - `torch-xpu` diff --git a/examples/kernels/flake.nix b/examples/kernels/flake.nix index 8a136468..403bcc97 100644 --- a/examples/kernels/flake.nix +++ b/examples/kernels/flake.nix @@ -394,6 +394,15 @@ } ]; + # TPU kernels to build in CI. + ciTpuKernels = [ + { + name = "relu-tpu-kernel"; + path = ./relu-tpu; + drv = sys: out: out.packages.${sys}.redistributable.torch-tpu; + } + ]; + # Metal kernels to build in CI. ciMetalKernels = [ { @@ -477,6 +486,7 @@ ]; ciXpuKernelOutputs = mkKernelOutputs' ciXpuKernels; + ciTpuKernelOutputs = mkKernelOutputs' ciTpuKernels; ciMetalKernelOutputs = mkKernelOutputs' ciMetalKernels; ciCpuKernelOutputs = mkKernelOutputs' ciCpuKernels; in @@ -617,6 +627,7 @@ ci-build-cuda = mkCiBuild "ci-kernels-cuda" ciKernelOutputs; ci-build-rocm = mkCiBuild "ci-kernels-rocm" ciRocmKernelOutputs; ci-build-xpu = mkCiBuild "ci-kernels-xpu" ciXpuKernelOutputs; + ci-build-tpu = mkCiBuild "ci-kernels-tpu" ciTpuKernelOutputs; ci-build-metal = mkCiBuild "ci-kernels-metal" ciMetalKernelOutputs; ci-build-cpu = mkCiBuild "ci-kernels-cpu" ciCpuKernelOutputs; in @@ -626,6 +637,7 @@ ci-build-cuda ci-build-rocm ci-build-xpu + ci-build-tpu ci-build-metal ci-build-cpu ; diff --git a/nix-builder/build-variants.json b/nix-builder/build-variants.json index f30e8a88..32f12af1 100644 --- a/nix-builder/build-variants.json +++ b/nix-builder/build-variants.json @@ -90,6 +90,9 @@ "torch215-rocm100-x86_64-linux", "torch215-rocm714-x86_64-linux" ], + "tpu": [ + "torch-tpu" + ], "xpu": [ "torch210-xpu20253-x86_64-linux", "torch211-xpu20253-x86_64-linux", diff --git a/nix-builder/pkgs/python-modules/libtpu/default.nix b/nix-builder/pkgs/python-modules/libtpu/default.nix index fac7af37..bed5288e 100644 --- a/nix-builder/pkgs/python-modules/libtpu/default.nix +++ b/nix-builder/pkgs/python-modules/libtpu/default.nix @@ -1,33 +1,11 @@ { lib, + fetchPypi, buildPythonPackage, requireFile, python, }: -# libtpu is served from Google's gated Artifact Registry, which requires -# an OAuth2 bearer token. That token can't be handed to a build sandbox -# without making the derivation impure, so instead the wheel is -# prefetched out-of-band and handed to `requireFile`, which only checks -# that a file matching `name` + `hash` already exists in the store: -# -# GCLOUD_ACCESS_TOKEN=$(gcloud auth print-access-token) \ -# nix-builder/scripts/helpers/prefetch-tpu-wheel.sh libtpu 0.0.43 -# -# The script prints the `hash` line to paste in below. To bump the -# version, update `version`, then re-run the script with the new -# version. -# -# libtpu ships the `libtpu/libtpu.so` runtime that both torch_tpu and -# jaxlib dlopen. Pinned to 0.0.x to match the jax 0.10.x / torch_tpu -# 0.1.x stack (torch_tpu requires libtpu>=0.0.40; jax's "tpu" extra -# pins libtpu==0.0.43.*), so bump it together with jaxlib/torch_tpu. -# -# NOTE on license: libtpu's wheel METADATA declares its license as -# "Google Cloud Platform Terms of Service" — i.e. unfree, NOT -# Apache-2.0 (unlike torch_tpu). The tpu buildSet therefore sets -# allowUnfree = true (see lib/mk-build-set.nix). - let # The wheel ships per-CPython-ABI builds (cp311..cp314); pick the tag # matching the python this package set is built for. The hash below is @@ -37,19 +15,16 @@ let in buildPythonPackage rec { pname = "libtpu"; - version = "0.0.48"; + version = "0.0.50"; format = "wheel"; - src = requireFile { - name = "libtpu-${version}-${abi}-${abi}-manylinux_2_31_x86_64.whl"; - hash = "sha256-QD0G/ssztiQYMmd/gwI/IXlYnwyZ4lwRocEf/7V7NBY="; - message = '' - libtpu is served from a gated Google Artifact Registry and cannot - be fetched by a pure Nix build. Fetch and register it with: - - GCLOUD_ACCESS_TOKEN=$(gcloud auth print-access-token) \ - nix-builder/scripts/helpers/prefetch-tpu-wheel.sh libtpu ${version} ${abi} - ''; + src = fetchPypi { + inherit format pname version; + dist = abi; + python = abi; + inherit abi; + platform = "manylinux_2_31_x86_64"; + hash = "sha256-vyMsa9hUVwlGJi3269OB5fcJqRdSxNZDNXj0I1lfz1w="; }; pythonImportsCheck = [ "libtpu" ]; diff --git a/nix-builder/pkgs/python-modules/torch/binary/generic.nix b/nix-builder/pkgs/python-modules/torch/binary/generic.nix index 9002cec7..a38d23d0 100644 --- a/nix-builder/pkgs/python-modules/torch/binary/generic.nix +++ b/nix-builder/pkgs/python-modules/torch/binary/generic.nix @@ -6,6 +6,8 @@ symlinkJoin, buildPythonPackage, fetchFromGitHub, + + fetchPypi, fetchurl, requireFile, @@ -45,6 +47,7 @@ numpy, portpicker, pyyaml, + pyzmq, pyzes, requests, setuptools, @@ -157,23 +160,22 @@ let }; # Once we require multiple versions, factor this out. + # The hash is for the cp314 wheel (the nixpkgs default python). torch-tpu = let abi = "cp${lib.versions.major python.version}${lib.versions.minor python.version}"; - version = "0.1.1.dev20260916100444"; in - requireFile { - name = "torch_tpu-${version}-${abi}-${abi}-manylinux_2_31_x86_64.whl"; - hash = "sha256-A8Rw9/QkMCjj9jfb4s+/d04/m/xH5L9rJDa0+giessM="; - message = '' - torch_tpu is served from a gated Google Artifact Registry and - cannot be fetched by a pure Nix build. Fetch and register it with: - - GCLOUD_ACCESS_TOKEN=$(gcloud auth print-access-token) \ - nix-builder/scripts/helpers/prefetch-tpu-wheel.sh torch_tpu ${version} ${abi} - ''; + fetchPypi { + format = "wheel"; + # Wheel file names use an underscore. + pname = "torch_tpu"; + version = "0.2.0"; + dist = abi; + python = abi; + inherit abi; + platform = "manylinux_2_31_x86_64"; + hash = "sha256-xdqU1lQfQchadDw0En1s2huCAqHafZR6OlLFE00Z3ak="; }; - in buildPythonPackage.override { stdenv = effectiveStdenv; } { pname = "torch"; @@ -293,6 +295,7 @@ buildPythonPackage.override { stdenv = effectiveStdenv; } { frozendict immutabledict portpicker + pyzmq tensorboard ] ++ lib.optionals tritonSupport [ diff --git a/nix-builder/scripts/helpers/prefetch-tpu-wheel.sh b/nix-builder/scripts/helpers/prefetch-tpu-wheel.sh deleted file mode 100755 index 5eaeb540..00000000 --- a/nix-builder/scripts/helpers/prefetch-tpu-wheel.sh +++ /dev/null @@ -1,40 +0,0 @@ -#!/usr/bin/env bash -# Fetch a wheel from Google's gated torch-tpu-virtual-registry and -# register it in the Nix store under the exact (name, sha256) that -# `requireFile` expects in pkgs/python-modules/{libtpu,torch_tpu}/default.nix. -# -# The fetch (and the OAuth2 token it needs) happens here, out-of-band, -# so the actual Nix build stays pure and never touches the network. -# -# Usage: -# GCLOUD_ACCESS_TOKEN=$(gcloud auth print-access-token) \ -# scripts/helpers/prefetch-tpu-wheel.sh libtpu 0.0.43 [cp314] -# GCLOUD_ACCESS_TOKEN=$(gcloud auth print-access-token) \ -# scripts/helpers/prefetch-tpu-wheel.sh torch_tpu 0.1.1.dev20260707090224 [cp314] -# -# Prints a `hash = "sha256-...";` line to paste into the matching -# default.nix when bumping a version. The third argument is the -# CPython ABI tag (default cp314, matching the nixpkgs default python -# used by the builder). -# -# Requires `gcloud auth login` once, plus nix (with the `nix-command` -# experimental feature) and jq. -set -euo pipefail - -pkg="${1:?usage: $0 [abi]}" -version="${2:?usage: $0 [abi]}" -abi="${3:-cp314}" -: "${GCLOUD_ACCESS_TOKEN:?set GCLOUD_ACCESS_TOKEN to \$(gcloud auth print-access-token)}" - -registry_pkg="${pkg//_/-}" -name="${pkg}-${version}-${abi}-${abi}-manylinux_2_31_x86_64.whl" -url="https://us-python.pkg.dev/ml-oss-artifacts-transient/torch-tpu-virtual-registry/${registry_pkg}/${name}" -netrc="$(mktemp)" -trap 'rm -f "${netrc}"' EXIT -printf 'machine us-python.pkg.dev\nlogin oauth2accesstoken\npassword %s\n' \ - "${GCLOUD_ACCESS_TOKEN}" >"${netrc}" - -hash="$(nix store prefetch-file --json --hash-type sha256 --name "${name}" \ - --option netrc-file "${netrc}" "${url}" | jq -r .hash)" - -echo "hash = \"${hash}\"; # ${abi}" diff --git a/nix-builder/versions.nix b/nix-builder/versions.nix index d87b69aa..26e95563 100644 --- a/nix-builder/versions.nix +++ b/nix-builder/versions.nix @@ -127,10 +127,7 @@ torchVersion = "2.11"; tpu = true; systems = [ "x86_64-linux" ]; - # Excluded from the shared build cache: torch_tpu's autoPatchelfHook - # step produces an output that is itself a modified copy of Google's - # gated libtpu/torch_tpu wheels, so it can't be published publicly. - bundleBuild = false; + bundleBuild = true; } {