Supporting PyTorch on the Cerebras Wafer-Scale Engine - Cerebras

This post details the integration of PyTorch support onto the Cerebras Wafer-Scale Engine (WSE-2). It outlines the strategy of adhering to PyTorch's design principles, minimizing user surprises, and integrating with its ecosystem. The implementation leverages a modified XLA backend and MLIR to lower PyTorch graphs to the Cerebras Intermediate Representation (CIRH). Different execution modes (Multi-Replica, Pipelined, Weight Streaming) are described for various model sizes. The Cerebras ML backend performs graph rewrites and automatic kernel generation. The runtime architecture separates the neural network execution on the CS-2 from the chief worker and input workers. User experience is simplified with minimal code modifications under the _cerebras_torch_ namespace.