GPU Support

Set up a GPU environment if you are building a model with image or video data

Setting up GPU environment with Docker

Pre Requisites

• Ensure your OS has a CUDA-capable GPU(s)

• The iai client is using PyTorch with CUDA 11.3. Ensure you install a driver version >= 11.3. (The latest driver should be installed, as it's backward compatible. See: CUDA major component versions)

Linux Setup

1. Install CUDA driver or CUDA toolkit:

   • Install cuda toolkit (which includes the driver, but also contains other unnecessary components):

     • If using rpm/deb package manager, you can install the driver only by using sudo apt-get/yum -y install cuda-drivers

   • Install CUDA driver only:

     • If it's a data centre GPU see: https://docs.nvidia.com/datacenter/tesla/tesla-installation-notes/index.html

     • For the other types of graphic cards, go to Nvidia driver page, select the corresponding graphic card and operating system, download and install the driver.

     • In order for the driver to be loaded, a system reboot may be needed.

2. Install Nvidia container toolkit:

   • Follow this guide

   • Ensure nvidia-docker2 can modify the docker configuration file /etc/docker/daemon.json

Windows Setup

1. Ensure that intel VT-x or AMD SVM is enabled in BIOS, check the motherboard manufacture document for exact steps.

2. Install CUDA driver or CUDA toolkit:

   • Install cuda toolkit (which include driver, but also contains other unnecessary components)

     • In the component selection screen, you can choose to install only the CUDA driver

   • Install CUDA driver only

     • Go to Nvidia driver page, select the corresponding graphic card and operating system, download and install the driver.

3. Setup Docker with WSL2

Running Docker Container with GPU device

Add --gpus all option to the docker run command.

Example:

docker run --gpus all -it -d --name $SILO_NAME -v <data_path>:/root/demo 919740802015.dkr.ecr.ca-central-1.amazonaws.com/edge/fl-client:latest