philipptrepte's picture
Update BUILD.md with multi-Python build instructions
20daae6 verified
|
Raw
History Blame Contribute Delete
7.43 kB

Building TensorFlow 2.21.0 for DGX Spark (aarch64 / GB10)

Built: 2026-04-07 Target: DGX Spark (GB10 Blackwell, aarch64, SM 10.0 + SM 12.1) HuggingFace: Qanatpharma/tensorflow-2.21.0-cuda13-aarch64 Build time: ~2 hours per Python version

Available Wheels

Python Wheel Size
3.10 tensorflow-2.21.0-cp310-cp310-linux_aarch64.whl 377 MB
3.11 tensorflow-2.21.0-cp311-cp311-linux_aarch64.whl 377 MB
3.12 tensorflow-2.21.0-cp312-cp312-linux_aarch64.whl 377 MB
3.13 tensorflow-2.21.0-cp313-cp313-linux_aarch64.whl 378 MB

Build Configuration

Parameter Value
TensorFlow v2.21.0
Python 3.10 (miniforge3)
CUDA (hermetic) 13.0.2
cuDNN (local) 9.20.0
Compute capability SM 10.0, SM 12.1 (GB10 reports 12.1 under CUDA 13)
Host compiler aarch64-linux-gnu-gcc-13
CPU optimization -march=armv9-a
XLA enabled
Bazel 7.7.0 (via Bazelisk 9.0.1)

The build uses Bazel's hermetic CUDA 13.0.2 for the toolkit (nvcc, headers, libs) but local cuDNN 9.20.0 (hermetic cuDNN lacks CUDA 13 aarch64 packages). A clean cuDNN directory is used to avoid system header contamination.

Prerequisites

Already on the Spark:

  • CUDA 13.0 runtime (/usr/local/cuda-13.0/)
  • cuDNN 9.20.0 (sudo apt-get install cudnn9-cuda-13)
  • GCC 13 (aarch64-linux-gnu-gcc-13)
  • miniforge3

Reproducing the Build

1. Install Bazelisk

mkdir -p ~/bin
curl -L https://github.com/bazelbuild/bazelisk/releases/latest/download/bazelisk-linux-arm64 \
  -o ~/bin/bazel
chmod +x ~/bin/bazel
export PATH="$HOME/bin:$PATH"

2. Create build environment

conda create -n tf-build python=3.10 numpy -y
conda activate tf-build
pip install numpy wheel packaging requests

3. Prepare clean cuDNN directory

The hermetic cuDNN redistributions don't include CUDA 13 aarch64 packages, so we use local cuDNN headers/libs. A clean directory avoids system header contamination that breaks the hermetic compiler sandbox.

mkdir -p ~/bioai/tools/tensorflow-build/cudnn/{include,lib}
cp /usr/include/aarch64-linux-gnu/cudnn*.h ~/bioai/tools/tensorflow-build/cudnn/include/
cp -a /usr/lib/aarch64-linux-gnu/libcudnn*.so* ~/bioai/tools/tensorflow-build/cudnn/lib/

4. Clone TensorFlow (if not already present)

cd ~/bioai/tools/tensorflow-build
git clone --depth 1 --branch v2.21.0 https://github.com/tensorflow/tensorflow.git

5. Configure

cd ~/bioai/tools/tensorflow-build/tensorflow
conda activate tf-build
export PATH="$HOME/bin:$PATH"

export PYTHON_BIN_PATH=$(which python)
export PYTHON_LIB_PATH=$(python -c "import site; print(site.getsitepackages()[0])")
export TF_NEED_CUDA=1
export TF_NEED_ROCM=0
export TF_SET_ANDROID_WORKSPACE=0
export TF_CUDA_CLANG=0
export GCC_HOST_COMPILER_PATH=/usr/bin/gcc
export CC_OPT_FLAGS="-march=armv9-a"
export HERMETIC_CUDA_VERSION=13.0.2
export HERMETIC_CUDNN_VERSION=9.20.0
export HERMETIC_CUDA_COMPUTE_CAPABILITIES="10.0,12.1"
export TF_ENABLE_XLA=1

./configure

Accept defaults for all prompts (env vars drive the config).

Then edit .tf_configure.bazelrc to add the local cuDNN path:

echo 'build:cuda --repo_env LOCAL_CUDNN_PATH="/home/philipptrepte/bioai/tools/tensorflow-build/cudnn"' \
  >> .tf_configure.bazelrc

6. Build

The HERMETIC_PYTHON_VERSION flag controls which Python the wheel targets. Each version requires a full rebuild (~2 hours).

# Python 3.10 (default if HERMETIC_PYTHON_VERSION is omitted)
bazel build \
  --config=cuda_wheel \
  --config=opt \
  --copt=-march=armv9-a \
  --verbose_failures \
  --repo_env=HERMETIC_NVSHMEM_VERSION="" \
  --repo_env=HERMETIC_PYTHON_VERSION=3.10 \
  //tensorflow/tools/pip_package:wheel

For other Python versions, change the env and re-configure first:

# Example: Python 3.12
conda create -n tf-build-312 python=3.12 numpy -y
conda activate tf-build-312
pip install numpy wheel packaging requests

# Re-run configure (sets PYTHON_BIN_PATH for the new env)
export PYTHON_BIN_PATH=$(which python)
# ... same env vars as step 5 ...
./configure
echo 'build:cuda --repo_env LOCAL_CUDNN_PATH=".../cudnn"' >> .tf_configure.bazelrc

# Build with matching hermetic Python
bazel build \
  --config=cuda_wheel --config=opt --copt=-march=armv9-a \
  --verbose_failures --repo_env=HERMETIC_NVSHMEM_VERSION="" \
  --repo_env=HERMETIC_PYTHON_VERSION=3.12 \
  //tensorflow/tools/pip_package:wheel

The wheel lands at:

bazel-bin/tensorflow/tools/pip_package/wheel_house/tensorflow-2.21.0-cp3XX-cp3XX-linux_aarch64.whl

Copy to dist:

mkdir -p ../dist
cp bazel-bin/tensorflow/tools/pip_package/wheel_house/*.whl ../dist/

7. Install

pip install ~/bioai/tools/tensorflow-build/dist/tensorflow-2.21.0-cp3XX-cp3XX-linux_aarch64.whl

8. Verify

# Must run from outside the tensorflow source directory
cd /tmp
python -c "
import tensorflow as tf
print('TF version:', tf.__version__)
print('Built with CUDA:', tf.test.is_built_with_cuda())
gpus = tf.config.list_physical_devices('GPU')
print('GPUs:', gpus)
if gpus:
    for gpu in gpus:
        print('Details:', tf.config.experimental.get_device_details(gpu))
    with tf.device('/GPU:0'):
        a = tf.constant([[1.0, 2.0], [3.0, 4.0]])
        b = tf.constant([[5.0, 6.0], [7.0, 8.0]])
        print('GPU matmul:', tf.matmul(a, b).numpy())
"

Expected:

TF version: 2.21.0
Built with CUDA: True
GPUs: [PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]
Details: {'compute_capability': (12, 1), 'device_name': 'NVIDIA GB10'}
GPU matmul: [[19. 22.]
             [43. 50.]]

Compute Capability Note

The NVIDIA GB10 (Blackwell) reports compute capability 12.1 under CUDA 13 (not 10.0 as documented in early CUDA 12 releases). The wheel includes native kernels for both SM 10.0 and SM 12.1 to handle both numbering schemes. Without SM 12.1, TF falls back to JIT-compiling PTX at runtime, adding significant first-call latency.

Issues Resolved During Build

  1. nvshmem unsupported on CUDA 13 aarch64 — Disabled via --repo_env=HERMETIC_NVSHMEM_VERSION="". Not needed (single-GPU system).

  2. cuDNN hermetic lacks CUDA 13 aarch64 — Hermetic cuDNN redistributions only cover cuda11/cuda12 platforms. Solved by using LOCAL_CUDNN_PATH pointing to a clean directory with just cuDNN headers and libs (not /usr, which contaminates the hermetic sandbox with system headers).

  3. CUPTI headers missing at expected path — CUDA 13 moved CUPTI headers to targets/sbsa-linux/include/. Fixed with symlinks:

    sudo mkdir -p /usr/local/cuda-13.0/extras/CUPTI/include
    sudo ln -sf /usr/local/cuda-13.0/targets/sbsa-linux/include/cupti*.h \
      /usr/local/cuda-13.0/extras/CUPTI/include/
    sudo mkdir -p /usr/local/cuda-13.0/extras/CUPTI/lib64
    sudo ln -sf /usr/local/cuda-13.0/targets/sbsa-linux/lib/libcupti* \
      /usr/local/cuda-13.0/extras/CUPTI/lib64/
    
  4. --config=cuda_wheel required — TF 2.21.0 requires cuda_wheel (not cuda) when building the pip wheel target.

  5. GB10 compute capability 12.1 (not 10.0) — CUDA 13 renumbered Blackwell SM architecture. Build targets both SM 10.0 and SM 12.1 for native kernels.