Jetson AI Lab - ML DevOps, Containers, Core Inferencing

Autonomous Machines Jetson & Embedded Systems Jetson Projects
,
2

4/2/24 - Rebuild the Jetson AI/ML stack for latest CUDA, Python, ect

  • NVIDIA has started providing CUDA, cuDNN (and soon TensorRT) installers for Jetson that are downloadable through our website and independent from the versions that ship with JetPack. There were many architectural enhancements in L4T and JetPack 6.0 to enable these out-of-band updates (such as upstreaming Orin SoC support into the mainline Linux kernel), ultimately enabling users to keep up-to-date with the latest CUDA versions and transitively leverage enhancements in downstream packages for these (such as PyTorch)

  • Upon installing a newer CUDA version (or Python version), the entire stack downstream also needs rebuilt - which is tedious, time consuming, and error-prone to do manually. For example, if you install the latest CUDA 12.4 and cuDNN 9.0 on top of JetPack 6.0 (which ships with CUDA 12.2 and cuDNN 8.9), PyTorch and everything after it needs recompiled against that new version of CUDA/cuDNN. And sometimes you may need to build a development branch of these packages where support for that newer CUDA is WIP (for example, PyTorch 2.2 won’t build against cuDNN 9.0, but PyTorch 2.3-rc does)

  • jetson-containers has been enhanced so that it will automatically rebuild all needed packages against the desired CUDA version, and resolving dependencies so that the correct versions of downstream packages are selected (for example, the system knows to use PyTorch 2.3 when CUDA 12.4 is selected). This behavior can be enabled by setting environment variables like CUDA_VERSION, CUDNN_VERSION, and PYTHON_VERSION:

    CUDA_VERSION=12.4 PYTHON_VERSION=3.11 \
jetson-containers/build.sh --name my_cu124_container pytorch torchvision transformers ...

    New CUDA/cuDNN/TensorRT package versions still need to be defined in the container configurations (like here for CUDA, so it knows what installers to download). You can also set the default PYTORCH_VERSION like this, and it will do the reverse and select the correct CUDA version for you.

1 Like