Pytorch: Getting Started

This quick-start guide describes first-time user setup and user workflow for running PyTorch jobs on a Cerebras Wafer-Scale Cluster. Cerebras Wafer-Scale Cluster is composed of CS-2 system(s), MemoryX, SwarmX, management and input worker nodes. The cluster supports two execution modes to enable ML models of different sizes:

  • Layer pipelined: In this mode, all the layers of the network are loaded together onto the Cerebras WSE. This mode is selected for neural network models that fit entirely on the WSE, approximately up to 1B parameters.

  • Weight streaming: In this mode, one layer of the neural network model is loaded at a time. This layer-by-layer mode is used to run extremely large models (>1B parameters).

Perform the following steps to run your PyTorch jobs on the Wafer-Scale Cluster:

  1. Ensure that the admin setup is complete. See Admin setup checklist.

  2. Follow the first-time user setup procedure for PyTorch below. This includes creating and configuring your virtual environment. This step should be done only once.

  3. Wafer-Scale Clusters currently adopt different workflows to launch jobs for Pipelined execution and Weight Streaming execution.

Note

If you are on the Original Cerebras Installation and have not upgraded to the Wafer-Scale Cluster, you can still use Slurm-based workflow to launch jobs for small to medium models with Pipelined execution. Large models with Weight Streaming execution are not supported on the Original Cerebras Installation. To get started on the Original Cerebras Installation, see `PyTorch: Getting Started <https://docs.cerebras.net/en/latest/getting-started/pytorch/index2.html`_.

If you are ready to start developing / adapting your own PyTorch code for CS System

Skip to Workflow for PyTorch on CS for an in-depth development guide using PyTorch for Cerebras.

First-time user setup for PyTorch

The first time you use Wafer-Scale Cluster for your PyTorch runs, you must set up a virtual environment as shown below.

Note

Make sure that you have the TLS Certificate available from your sysadmin. You need this to communicate between the user node and the Wafer-Scale Cluster. Your admin will have shared the path to this file during the setup.

  1. Set up the Python virtual environment using Python 3.7. Create the environment named venv_cerebras_pt using the following command:

    python3.7 -m venv venv_cerebras_pt
    
  2. Cerebras provides three main packages to set up virtual environments: cerebras_appliance software package, the cerebras_tensorflow package if you are using TensorFlow, and the cerebras_pytorch package if you are using PyTorch. To set up your PyTorch environment, you need two out of these three packages. Enter the following commands on the user node to install the required packages (make sure to execute the commands to install the appliance wheel first):

    source venv_cerebras_pt/bin/activate
    
    pip install <path_to_wheels>/cerebras_appliance-<Cerebras release version>_<date>_<build>_<commit>-py3-none-any.whl --find-links=<path_to_wheels>
    pip install <path_to_wheels>/cerebras_pytorch-<Cerebras release version>_<date>_<build>_<commit>-py3-none-any.whl --find-links=<path_to_wheels>
    

Note

With the find-links command, it finds the correct cerebras-appliance version if you place all the wheels in the same directory. If you are only using the Pipelined execution mode, then you shouldn’t need cerebras_pytorch package and can install cerebras_appliance package only.

Note

In the cerebras_appliance wheel, there exists two scripts: csrun_cpu and csrun_wse. These scripts are required for Pipelined execution and serve the same function as the scripts previously available for Slurm. The csrun_cpu is for non-WSE jobs while the csrun_wse is to launch jobs on the Wafer-Scale Engine. For more information, see csrun_cpu and csrun_wse.

Running PyTorch jobs

After you have completed first-time user setup: