A while back I found a small bug in TensorFlow. Now that I had some free time, I decided to submit a PR. The bug was fixed quickly, but when I tried to build TensorFlow locally, I ran into quite a few gotchas. Here are my notes.
1. So Many Versions, So Confusing
Before we start building, let’s go over the relevant Nvidia GPU dependencies.
Nvidia has various GPU architectures. To distinguish between them, Nvidia uses Compute Capability. The version is split into x.y. The major version indicates architectural changes (e.g., Pascal, Volta, Ampere) and is incompatible across versions; the minor version indicates differences within the same architecture, with higher versions being backward-compatible with lower ones.
The GPU Driver provides hardware driver support for the OS. Its version can be found with nvidia-smi --query-gpu=driver_version --format=csv. A given driver version supports a range of GPUs with different compute capabilities and architectures.
The CUDA Driver provides the CUDA API on top of the GPU driver. Unlike the GPU driver, which is a kernel-mode device driver, the CUDA driver is a user-mode dynamic shared object (DSO). The CUDA driver version should generally be updated alongside the GPU driver.
The CUDA Toolkit provides the compiler, runtime, and libraries needed to build CUDA programs. Built CUDA applications depend on the interfaces provided by the CUDA driver. Since CUDA is backward-compatible, older CUDA applications can run on newer CUDA drivers. In other words, to run a CUDA program, you need a CUDA driver version higher than a certain threshold.
The word “backward” in backward compatibility refers to compatibility with things earlier in time. This seems to cause a mental-model difference between Chinese and English.
2. Build Configuration
When building TensorFlow with CUDA support, TensorFlow must satisfy compatibility requirements between the CUDA Toolkit and the CUDA driver. As mentioned, for better compatibility we want a lower CUDA Toolkit version. But higher TensorFlow versions may require newer CUDA features. So you need to balance these two factors.
The TensorFlow website provides a list of tested build configurations. If you have no special requirements, follow the versions listed there. Here we’ll use tensorflow-2.18.0 as an example.
First, clone the repository. If you don’t need to modify the source, you can also download the zip archive.
git clone https://github.com/tensorflow/tensorflow.git
cd tensorflow
git checkout r2.18
Setting up the build environment locally is quite tedious. It’s best to use the Docker image tensorflow/build provided by TensorFlow.
docker run --name tf -w /tf/tensorflow -it -d \
--env TF_PYTHON_VERSION=3.10 \
-v "/path/to/tensorflow:/tf/tensorflow" \
tensorflow/build:latest-python3.10 \
bash
Where /path/to/tensorflow is the host path to the TensorFlow source.
Then enter the container:
docker exec -it tf bash
Note that the build process below differs from what’s described here. The latter has been outdated since v2.16 — see #63298.
Inside the container, run the configuration script configure:
./configure
Keep the default Python configuration, using the image’s python3.10.
$ ./configure
You have bazel 6.5.0 installed.
Please specify the location of python. [Default is /usr/bin/python3]:
Found possible Python library paths:
/usr/lib/python3/dist-packages
/usr/local/lib/python3.10/dist-packages
Please input the desired Python library path to use. Default is [/usr/lib/python3/dist-packages]
Don’t select ROCm support; select CUDA support.
Do you wish to build TensorFlow with ROCm support? [y/N]: n
No ROCm support will be enabled for TensorFlow.
Do you wish to build TensorFlow with CUDA support? [y/N]: y
CUDA support will be enabled for TensorFlow.
Next, set the CUDA and cuDNN versions. According to the official build config, for v2.18, it’s best to use CUDA 12.5 and cuDNN 9.3 for building. However, as mentioned, newer CUDA applications can’t run on older CUDA drivers. So the CUDA version set here should be lower than or equal to the target machine’s CUDA driver version — otherwise you’ll need to update the machine’s CUDA driver to match. The cuDNN version should also be adjusted accordingly based on the official build configuration.
Please specify the hermetic CUDA version you want to use or leave empty to use the default version. 12.5.1
Please specify the hermetic cuDNN version you want to use or leave empty to use the default version. 9.3.0
The CUDA and cuDNN versions should specify exact version numbers, e.g.,
12.5.1, not12.5.
Then set the compute capability. Check your GPU’s compute capability on the Nvidia website. Choose a value where the major version matches and the minor version is ≤ your device’s version. My laptop GPU has compute capability 8.9, so I chose 8.6 here.
Note that setting the compute capability too low may cause issues — newer compiler versions may no longer support older compute capabilities. For instance,
clang 18.1.8in the image no longer supports the default3.5.
Please specify a list of comma-separated CUDA compute capabilities you want to build with.
You can find the compute capability of your device at: https://developer.nvidia.com/cuda-gpus. Each capability can be specified as "x.y" or "compute_xy" to include both virtual and binary GPU code, or as "sm_xy" to only include the binary code.
Please note that each additional compute capability significantly increases your build time and binary size, and that TensorFlow only supports compute capabilities >= 3.5 [Default is: 3.5,7.0]: 8.6
Keep the remaining options at their defaults:
Please specify the local CUDA path you want to use or leave empty to use the default version.
Please specify the local CUDNN path you want to use or leave empty to use the default version.
Please specify the local NCCL path you want to use or leave empty to use the default version.
For the compiler, use the recommended clang that’s built into the image:
Do you want to use clang as CUDA compiler? [Y/n]:
Clang will be used as CUDA compiler.
Please specify clang path that to be used as host compiler. [Default is /usr/lib/llvm-18/bin/clang]:
You have Clang 18.1.8 installed.
Then keep the defaults again:
Please specify optimization flags to use during compilation when bazel option "--config=opt" is specified [Default is -Wno-sign-compare]:
Would you like to interactively configure ./WORKSPACE for Android builds? [y/N]:
Not configuring the WORKSPACE for Android builds.
3. Building TensorFlow
After configuration, build with the following options:
$ bazel build //tensorflow/tools/pip_package:wheel --repo_env=WHEEL_NAME=tensorflow --config=cuda --config=cuda_wheel --config=opt --copt=-Wno-gnu-offsetof-extensions
Note the extra
--copt=-Wno-gnu-offsetof-extensions. Without this flag, the build will fail.
The build process consumes significant system resources. You can limit bazel’s resource usage with options like --local_cpu_resources and --local_ram_resources. See the bazel documentation for details.
After the build completes, the *.whl file will be at bazel-bin/tensorflow/tools/pip_package/wheel_house/.
4. Installation
Install the built wheel with pip. If the path is /path/to/tensorflow.whl:
pip3 install '/path/to/tensorflow.whl[and-cuda]'
[and-cuda] is required; otherwise TensorFlow won’t be able to use the GPU.
Then test the CPU setup:
python3 -c "import tensorflow as tf; print(tf.reduce_sum(tf.random.normal([1000, 1000])))"
And test the GPU setup:
python3 -c "import tensorflow as tf; print(tf.config.list_physical_devices('GPU'))"
If you see an error like:
kernel version xxx.yyy.z does not match DSO version uuu.vv.w -- cannot find working devices in this configurationIt means the CUDA version used during compilation is too high and incompatible with the lower-version CUDA driver at runtime. Either lower the CUDA version in the build configuration and rebuild, or upgrade the GPU driver on the target machine so its CUDA driver version exceeds the build-time version.