Workflow for PyTorch on CS

When you are targeting the Cerebras system for your neural network jobs, start with the high-level workflow described here.

1

Port PyTorch to CS

2

Load the data

3

Compile on CPU

4

Train on CS

Attention

Several PyTorch models, such as the following, and run scripts, are provided in Cerebras Model Zoo:

• PyTorch version of FC-MNIST, and

• PyTorch versions of BERT Base and BERT Large.

1. Port PyTorch to CS

See Porting PyTorch Model to CS for detailed documentation on how to port your PyTorch model to work on a Cerebras System.

2. Load the data

As mentioned in Porting PyTorch Model to CS, you need to provide get_train_dataloader for training and get_eval_dataloader for evaluation. See cerebras.framework.torch.dataloader() for details.

3. Compile on CPU

Before you run it on the Cerebras system, we recommend that you iterate until your model first compiles successfully on a CPU node. Make sure that this CPU node has the Cerebras Singularity container client software.

As described in Porting PyTorch Model to CS, you can verify that your model will compile by passing in the --compile_only flag to the run function’s CLI.

4. Train on the Cerebras system

As described in Porting PyTorch Model to CS, you can train or evaluate your model by providing the IP address of the Cerebras system and the mode you want to run.