Limitations of PyTorch on Cerebras

Floating Point Precision

Only mixed precision is supported for training a model on the Cerebras system. This means that weights are stored as float32 but the computation other than the weight update happens using float16. Casts are automatically inserted; you do not need to insert them manually.

Note

There are no plans to support other precision modes at this time.

Static Graphs

As of the 1.6.1. software release, users may not reprogram the Wafer Scale Engine after initial programming. This means that multiple compiles are not supported and therefore the PyTorch compute graph must not change between iterations.

This means that there are a number of caveats as to how the training loop is allowed to be constructed, all of which are already addressed in our custom PyTorch runner classses. Refer to our implementations of the various hooks mentioned in PyTorch Runners.

Modes

Only mode = "train" and mode = "eval" are supported in this release. mode = "train_and_eval" is not supported because it is not permitted to reprogram the fabric after initial programming. So, the fabric cannot be reprogrammed for eval mode after having been programmed for train in the same run.

Note

Even when Cerebras supports reprogramming the fabric after initial programming in the future, it will still be ideal to avoid recompiling if possible, as recompiling and reprogramming the fabric can take a very long time.

Learning Rate Scheduler

Currently, we do not support the typical PyTorch learning rate scheduler paradigm. A typical PyTorch learning scheduler would compute a learning rate scalar and set the values of the learning rates in the optimizer parameter groups. However, due to current limitations of the system requiring static graphs, we cannot support this behavior.

Scheduling Learning Rate

There are two different ways of scheduling a learning rate, depending on whether you are running in pipeline mode or in weight streaming mode.

Layer Pipelined model

We provide custom LR scheduler classes which pass in the scheduler information directly to the system’s communication manager so that the scheduler can be programmed into the fabric.

See modelzoo.common.pytorch.optim.lr_scheduler for more details.

Weight Streaming mode

For weight streaming mode, we must specify the entire learning rate schedule as a function of the global step. This means that the learning rate becomes less of a scalar value and more of a tensor that depends on the value of the global step. See modelzoo.common.pytorch.optim.lr_scheduler for examples of this.

This does also mean that any optimizers being used need to be written in a way such that the learning rate is not treated as a scalar value, but rather as a tensor. See modelzoo.common.pytorch.optim.AdamBase for an example of this.

Eval metrics

Eval metrics are only allowed to return a single result value. Therefore, only the final metric value should be returned. No intermediate state or values can be retrieved at this time.