PyTorch Compile on CPU

We recommend that you first compile your PyTorch model on a CPU node before running it on the Cerebras system. Note that this CPU node must have the Cerebras Singularity container installed on it.

Setting the flag for compile

To perform compile-only, make sure that compile_only=True is set in the cbtorch.initialize. See below:

import cerebras.framework.torch as cbtorch

cbtorch.initialize(
    ...
    compile_only=True,
    ...
)

Load the model and dataloader

Load the module and dataloader as follows:

model: torch.nn.Module = ...
dataloader: torch.utils.data.DataLoader = ...

model = cbtorch.module(model)
dataloader = cbtorch.dataloader(dataloader)

Compile

To compile, use the cbtorch.compile function as follows:

sample_data = next(iter(dataloader))
outputs = model(sample_data)

cbtorch.compile(
    model,
    outputs,
    mode="train",
)

which will compile assuming a training workflow. Eval workflows are also supported.

Important

Compiling is only supported for training and eval workflows at this time.