Cerebras PyTorch Layer API

Cerebras PyTorch Layer API implements a subset of PyTorch APIs with our custom implementation that takes advantage of our high-performance kernels and provides extra functionalities as compared to the native PyTorch version. The extra functionalities are optional and opt-in; if you don’t use the extra functionalities, then the Layer API is equivalent to the native PyTorch version.

• modelzoo.common.pytorch.layers.MultiheadAttention is the replacement for torch.nn.MultiheadAttention

• modelzoo.common.pytorch.layers.TransformerDecoderLayer is the replacement for torch.nn.TransformerDecoderLayer

• modelzoo.common.pytorch.layers.TransformerDecoder is the replacement for torch.nn.TransformerDecoder

• modelzoo.common.pytorch.layers.TransformerEncoderLayer is the replacement for torch.nn.TransformerEncoderLayer

• modelzoo.common.pytorch.layers.TransformerEncoder is the replacement for torch.nn.TransformerEncoder

Supported PyTorch Ops

If your model implementation requires additional PyTorch Ops beyond the layer APIs above, Cerebras also supports the following PyTorch operations.

Attention

The following list of supported PyTorch ops is very preliminary. We cannot guarantee that mixing and matching them in your models will work. Support is only provided for the way they are used in the Cerebras Model Zoo.

nn

• torch.nn.BCEWithLogitsLoss

• torch.nn.CrossEntropyLoss

  Note: Known limitation: ignore_index can only be -100

• torch.nn.Dropout

• torch.nn.Embedding

  Note: Known limitation: num_embeddings < 65536

• torch.nn.functional.dropout

• torch.nn.functional.gelu

  Note: Known limitation: May have precision issue when approximation !=tanh

• torch.nn.functional.pad

• torch.nn.functional.softmax

• torch.nn.LayerNorm

• torch.nn.Linear

• torch.nn.MSELoss

• torch.nn.NLLLoss

  Note: Known limitation: ignore_index can only be -100

• torch.nn.ReLU

• torch.nn.Softmax

• torch.nn.TanH

Functional

• torch.nn.functional.log_softmax

• torch.nn.functional.relu

• torch.nn.functional.silu

Other ops

• torch.abs

• torch.all

• torch.arange

• torch.broadcast_to

• torch.cat

• torch.einsum

• torch.flatten

• torch.full

• torch.full_like

• torch.gather

• torch.log

• torch.matmul

• torch.min

• torch.ones

• torch.rsqrt

• torch.sigmoid

• torch.sum

• torch.tanh

• torch.where

• torch.zeros

• torch.zeros_like

Layers

• modelzoo.common.pytorch.layers.MultiheadAttention
• modelzoo.common.pytorch.layers.TransformerDecoderLayer
• modelzoo.common.pytorch.layers.TransformerDecoder
• modelzoo.common.pytorch.layers.TransformerEncoderLayer
• modelzoo.common.pytorch.layers.TransformerEncoder