Run Cerebras Model Zoo on a GPU#
You can run models in the Cerebras Model Zoo on GPUs as well. However, specific packages must be installed to run the model code on a GPU.
Make sure to install these packages in a virtual environment (virtualenv) or a conda environment.
CUDA Requirements#
A CUDA capable GPU is required to run the Cerebras Model Zoo on a GPU.
PyTorch#
Currently, the Cerebras Model Zoo only supports PyTorch version 2.0
Once all the CUDA requirements get installed, create a virtualenv on your system with Python version 3.8 or newer, activate the virtualenv, and install Pytorch using the following commands:
virtualenv -p python3.8 /path/to/venv_gpu
source /path/to/venv_pt/bin/activate
pip install -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cu117
The requirements.txt file is located in the Cerebras Model Zoo. The CUDA related packages will automatically be installed with PyTorch 2.0.
To test whether PyTorch can adequately access the GPU, start a Python session and run the following commands:
>>> import torch
>>> torch.__version__
'2.0.1+cu117-with-pypi-cudnn' # Confirm that the PT version is `2.0.1`
>>> torch.cuda.is_available()
True # Should return `True`
>>> torch.cuda.device_count()
1 # Number of devices present
>>> torch.cuda.get_device_name(0)
# Should return the proper GPU type