GitHub Actions (GPU) Setup

Contents

GitHub Actions (GPU) Setup#

This page documents the QuantEcon GPU enabled GitHub Actions workflow configurations.

Our GPU based workflows are now built using RunsOn

Build Images#

There are two types of workflows:

  1. Using our own custom ami that is preloaded with CUDA, CUDANN, and LaTeX

  2. Using the ami provided by Runs On that includes Docker support. This is used by collab.yml workflows or any containerized workflow.

The benefit of using our own custom AMI is that we can choose CUDA, CUDANN versions and ensure we are able to use the latest versions to support lectures on JAX and GPU computing.

It is best to install only system wide software such as GPU frameworks, tool-kits and drivers.

When setting up a custom AMI we should install the NVIDIA support:

  1. nvidia CUDA framework

  2. nvidia drivers

  3. nvidia CUDANN

using the latest versions and then LaTeX

sudo apt-get -qq update
sudo apt-get install -y     \
texlive-latex-recommended \
texlive-latex-extra       \
texlive-fonts-recommended \
texlive-fonts-extra       \
texlive-xetex             \
latexmk                   \
xindy                     \
dvipng                    \
cm-super

Leave anaconda and local installs to the GitHub actions workflow.

While this incurs a small setup cost in terms of execution time, it provides flexibility in upgrading the python / anaconda software stack.

In addition to anaconda you should install gpu enabled python software in the GitHub Action as a setup step

  - name: Install JAX, Numpyro, PyTorch
    shell: bash -l {0}
    run: |
        pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128
        pip install pyro-ppl
        pip install --upgrade "jax[cuda12-local]"
        pip install numpyro  
        python scripts/test-jax-install.py

Warning

PyTorch typically lags JAX in terms of compatibility with nvidia drivers. It is best to install the nightly versions to ensure the best compatibility.

You can check the latest on the PyTorch site with the installation table

Contents