build-te.sh

#!/bin/bash
SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
set -eo pipefail

## Parse command-line arguments
usage() {
    echo "Configure, build, and install TransformerEngine"
    echo ""
    echo "  Usage: $0 [OPTIONS]"
    echo ""
    echo "    OPTIONS                        DESCRIPTION"
    echo "    --clean                        Clear build caches under --src-path-te."
    echo "    -h, --help                     Print usage."
    echo "    --ccache                       Use ccache to build TransformerEngine. It will be installed, but the caller should configure it."
    echo "    --no-install                   Only build a wheel; do not install."
    echo "    --src-path-te                  Path to TransformerEngine source code."
    echo "    --src-path-xla                 Path to XLA source code."
    echo "    --sm SM1,SM2,...               Comma-separated list of CUDA SM versions"
    echo "                                   to compile for, e.g. 7.5,8.0 -- PTX will"
    echo "                                   only be emitted for the last one."
    echo "    --sm local                     Compile for the local GPUs."
    echo "    --sm all                       Compile for a default set of SM versions (default)."
    exit $1
}

# Set defaults
CCACHE=0
CLEAN=0
INSTALL=1
SM="all"
SRC_PATH_TE="/opt/transformer-engine"
SRC_PATH_XLA="/opt/xla"

args=$(getopt -o h --long ccache,clean,help,no-install,src-path-te:,src-path-xla:,sm: -- "$@")
if [[ $? -ne 0 ]]; then
    exit 1
fi

eval set -- "$args"
while [ : ]; do
    case "$1" in
        --clean)
            CLEAN=1
            shift 1
            ;;
        -h | --help)
            usage 1
            ;;
        --ccache)
            CCACHE=1
            shift 1
            ;;
        --no-install)
            INSTALL=0
            shift 1
            ;;
        --src-path-te)
            SRC_PATH_TE=$(realpath $2)
            shift 2
            ;;
        --src-path-xla)
            SRC_PATH_XLA=$(realpath $2)
            shift 2
            ;;
        --sm)
            SM=$2
            shift 2
            ;;
        --)
            shift;
            break
            ;;
        *)
            echo "UNKNOWN OPTION $1"
            usage 1
    esac
done

print_var() {
    echo "$1: ${!1}"
}

clean() {
    pushd "${SRC_PATH_TE}"
    rm -rf build/ .eggs/
    popd
}

# This should standardise on 1.2,3.4,5.6 format
if [[ "$SM" == "all" ]]; then
    if [[ -z "${CUDA_ARCH_LIST}" ]]; then
        echo "CUDA_ARCH_LIST was not set; this is defined by the dl-cuda-base image"
        return 1
    fi
    # Infer the compute capabilities from the CUDA_ARCH_LIST variable if it is set;
    # this is in 1.2 3.4 5.6 format
    SM_LIST=${CUDA_ARCH_LIST// /,}
elif [[ "$SM" == "local" ]]; then
    SM_LIST=$("${SCRIPT_DIR}/local_cuda_arch")
    if [[ -z "${SM_LIST}" ]]; then
        echo "Could not determine the local GPU architecture."
        echo "You should pass --sm when compiling on a machine without GPUs."
        nvidia-smi || true
        exit 1
    fi
else
    SM_LIST=${SM}
fi

## Print info
echo "=================================================="
echo "                  Configuration                   "
echo "--------------------------------------------------"
print_var CLEAN
print_var INSTALL
print_var SM
print_var SM_LIST
print_var SRC_PATH_TE
print_var SRC_PATH_XLA
echo "=================================================="

# Parse SM_LIST into the format accepted by TransformerEngine's build system
# "1.2,3.4,5.6" -> "12;34;56". In principle we would like to compile SASS-only
# plus PTX for the highest known architecture, but TransformerEngine's build
# system does not currently handle that.
NVTE_CUDA_ARCHS="${SM_LIST//,/;}"
set -x
export NVTE_CUDA_ARCHS="${NVTE_CUDA_ARCHS//./}"
# Parallelism within nvcc invocations.
export NVTE_BUILD_THREADS_PER_JOB=8
export NVTE_FRAMEWORK=jax
# TransformerEngine needs FFI headers from XLA
export XLA_HOME=${SRC_PATH_XLA}

pushd ${SRC_PATH_TE}
# Install some build dependencies, but avoid installing everything
# (jax, torch, ...) because we do not want to pull in a released version of
# JAX, or the wheel-based installation of CUDA. Note that when we build TE as
# part of building the JAX containers, JAX and XLA are not yet installed.
python - << EOF
import os, subprocess, sys, tomllib
with open("pyproject.toml", "rb") as ifile:
    data = tomllib.load(ifile)
subprocess.run(
    [sys.executable, "-m", "pip", "install"]
    + [r for r in data["build-system"]["requires"]
       if r.startswith("pybind11") or r.startswith("cmake") or r.startswith("ninja")]
    + [f"nvidia-cudnn-frontend=={os.environ['CUDNN_FRONTEND_VERSION']}"]
)
EOF
if [[ "${CCACHE}" == "1" ]]; then
    # Install ccache if not present (needs >= 4.1 for Redis remote storage support)
    if ! command -v ccache &> /dev/null; then
        apt-get update && apt-get install -y --no-install-recommends ccache
    fi
    #export CXX="ccache g++"
    export NVTE_USE_CCACHE=1
    ccache --zero-stats
fi

# The wheel filename includes the TE commit; if this has changed since the last
# incremental build then we would end up with multiple wheels.
rm -fv dist/*.whl
python setup.py bdist_wheel
ls dist/
popd

if [[ "${CCACHE}" == "1" ]]; then
    ccache --show-stats --verbose
fi

## Install the built packages
if [[ "${INSTALL}" == "1" ]]; then
    pip uninstall -y transformer_engine
    pip install ${SRC_PATH_TE}/dist/*.whl
    pip list | grep ^transformer_engine
fi

## Cleanup
if [[ "$CLEAN" == "1" ]]; then
    clean
fi