xwjim/modelzoo

Cerebras Model Zoo

Introduction

This repository contains examples of common deep learning models that can be trained on Cerebras hardware. These models demonstrate the best practices for coding a model targeted at the Cerebras hardware so that you can take full advantage of this new powerful compute engine.

In order to get started with running your models on a Cerebras system, please refer to the Developer Documentation along with this readme.

NOTE: If you are interested in trying out Cerebras Model Zoo on Cerebras Hardware (CS-2 Systems), we offer the following options:

• Academics - Please fill out our Partner Hardware Access Request form here and we will contact you about gaining access to a system from one of our partners.
• Commercial - Please fill out our Get Demo form here so that our team can provide you with a demo and discuss access to our system.
• For all others - Please contact us at [email protected].

For a list of all supported models, please check models in this repository.

Supported frameworks

We support the models developed in PyTorch and TensorFlow.

Basic workflow

When you are targeting the Cerebras Wafer-Scale Cluster for your neural network jobs, please follow the quick start guide from the developer docs to compile, validate and train the models in this Model Zoo for the framework of your choice.

For advanced use cases and porting your existing code please refer to the developer docs.

Execution modes

On the Cerebras Wafer Scale Cluster you can run neural networks of different model sizes. Cerebras Software supports different execution modes to efficiently run such variety of models.

The execution mode refers to how the Cerebras runtime loads your neural network model onto the Cerebras Wafer Scale Engine (WSE). Two execution modes are supported:

• 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 (with billions to trillions of parameters).
• Layer pipelined: In this mode all the layers of the network are loaded altogether onto the Cerebras WSE. This mode is selected for neural network models of small to medium sized models (with less than a billion parameters).

You can get more information about this on the developer page section on Cerebras Execution Modes

Models in this repository

Model	Layer Pipeline mode	Weight Streaming mode
BERT	TensorFlow code PyTorch code	PyTorch code
BERT (fine-tuning) Classifier	TensorFlow code PyTorch code	-
BERT (fine-tuning) Named Entity Recognition	TensorFlow code PyTorch code	-
BERT (fine-tuning) Summarization	TensorFlow code PyTorch code	-
BERT (fine-tuning) Question Answering	TensorFlow code PyTorch code	-
GPT-2	TensorFlow code PyTorch code	TensorFlow code PyTorch code
GPT-3	-	TensorFlow code PyTorch code
GPT-J	-	TensorFlow code PyTorch code
GPT-NeoX	-	TensorFlow code PyTorch code
GPT-J (fine-tuning) Summarization	-	TensorFlow code
Linformer	TensorFlow code	-
RoBERTa	TensorFlow code PyTorch code	-
T5	TensorFlow code PyTorch code	PyTorch code
Transformer	TensorFlow code PyTorch code	-
MNIST (fully connected)	TensorFlow code PyTorch code	-
UNet	-	PyTorch code

License

Apache License 2.0