Efficient Weight Initialization#
Initializing the model parameters on the CPU torch.device, while completely
valid, can be slow for extremely large models, and can cause memory issues as the
parameters may not fit within RAM and may spill over into swap memory or outright
fail to allocate more memory.
To address this, instantiate a cerebras_pytorch.backend that contains
a Cerebras device that can be used as a context manager, much like a
torch.device.
For example,
backend = cstorch.backend("CSX")
with backend.device:
model = Model()
# compile the model the same way as before
compiled_model = cstorch.compile(model, backend)
The parameters get moved to the Cerebras device automatically, which saves the parameter data that will be sent to the Cerebras Wafer-Scale cluster. This frees up memory for subsequent parameters and keeps the overall memory usage low.