Source-linked AI summary
A flexible and fast PyTorch toolkit for simulating training and inference on analog crossbar arrays
Malte J. Rasch, Diego Moreda, Tayfun Gokmen, Manuel Le Gallo, Fabio Carta, Cindy Goldberg, Kaoutar El Maghraoui, Abu Sebastian, Vijay Narayanan
TL;DR
Analog AI research needs realistic, scalable simulation integrated with modern deep-learning workflows. The toolkit addresses this with configurable analog tiles in PyTorch, GPU-accelerated models, and training and inference features that represent hardware non-idealities. It supports functional accuracy studies for arbitrary ANNs, including accuracy degradation over time, but does not estimate chip runtime, latency, or power.
Problem
Realistic analog-hardware simulation must account for device constraints on realistically sized DNNs and datasets, while existing tools are often not integrated with PyTorch or flexible enough for modern ANN workloads.
Method
The toolkit uses configurable analog tiles integrated into PyTorch, with GPU-accelerated simulation of analog training, inference non-idealities, and hardware-aware training.
Results
The toolkit enables testing arbitrary networks for realistic inference accuracy and its degradation over time.
Takeaways & Limitations
The toolkit supports algorithm development and functional verification for ANN training and inference on emerging analog chips.
Takeaways & Limitations
The toolkit does not estimate hardware-chip runtime, latency, or power performance, which require complementary tools.
Abstract
from arXiv · showhide
We introduce the IBM Analog Hardware Acceleration Kit, a new and first of a kind open source toolkit to simulate analog crossbar arrays in a convenient fashion from within PyTorch (freely available at https://github.com/IBM/aihwkit). The toolkit is under active development and is centered around the concept of an "analog tile" which captures the computations performed on a crossbar array. Analog tiles are building blocks that can be used to extend existing network modules with analog components and compose arbitrary artificial neural networks (ANNs) using the flexibility of the PyTorch framework. Analog tiles can be conveniently configured to emulate a plethora of different analog hardware characteristics and their non-idealities, such as device-to-device and cycle-to-cycle variations, resistive device response curves, and weight and output noise. Additionally, the toolkit makes it possible to design custom unit cell configurations and to use advanced analog optimization algorithms such as Tiki-Taka. Moreover, the backward and update behavior can be set to "ideal" to enable hardware-aware training features for chips that target inference acceleration only. To evaluate the inference accuracy of such chips over time, we provide statistical programming noise and drift models calibrated on phase-change memory hardware. Our new toolkit is fully GPU accelerated and can be used to conveniently estimate the impact of material properties and non-idealities of future analog technology on the accuracy for arbitrary ANNs.
1 Introduction
Analog crossbar arrays could accelerate ANN computation, but realistic non-idealities and hardware constraints require simulation on modern frameworks and realistically sized workloads.
- Memristive crossbars accelerate matrix-vector multiplication and rank updates by storing matrices as analog quantities in resistive elements.
- Analog device noise, nonlinearities, and stationary weights motivate new ANN topologies and algorithms adapted to analog hardware.
- Realistic simulators must model analog constraints and material properties on deep networks and datasets to inform hardware design and expected accuracy.
- Existing simulators are often disconnected from PyTorch, limited to small MNIST examples, or restricted to inference-only analog chips.
2 Toolkit overview
The toolkit provides a PyTorch-integrated, GPU-accelerated software kit whose analog layers support flexible network construction and simulation.
- The toolkit adds PyTorch-integrated analog linear and convolutional layers for defining ANN models.
- Its C++/CUDA RpuCuda core accelerates analog hardware simulation while exposing hardware parameters and functions through Python without recompilation.
- AnalogLinear uses one analog tile for a fully connected layer, and its training simulation is defined similarly to a PyTorch DNN layer.
3 Analog tile model simulation
The analog tile models crossbar computation, configurable non-idealities, analog training behavior, and device-specific update dynamics with GPU acceleration.
- An analog tile represents a 2D non-volatile resistive crossbar weight matrix with configurable pre- and post-processing.
- The simulated matrix-vector product can include configurable noise, conversion resolution, scaling, clipping, and other hardware non-idealities.
- The RPU C++ library implements the tile model with fused GPU kernels, while PyTorch can disable or adapt internal analog-digital processing.
- The framework assumes digital activation functions, parallel readout, and linear analog-digital conversion, while allowing future extensions through PyTorch.
- Backward computation can have independently configured noise and nonlinearity to represent hardware designs with different forward and reverse behavior.
- Stochastic or deterministic pulse trains support sequential analog updates whose material response curves include device-to-device and pulse-to-pulse variation.
- 60s per epoch versus 15s with floating point is reported for VGG-8/CIFAR10 with parallel pulsed updates on a V100 GPU.
4 Advanced device models and optimizers
The toolkit supports richer device configurations, coupled analog optimization schemes, and temporal noise processes for modeling diverse hardware designs.
- Unit cells can contain arbitrary numbers of devices with distinct material properties, updated individually or together.
- Coupled systems of two analog tiles can test analog optimization algorithms that exchange information between tiles.
- Conductance decay, reset, and diffusion model temporal noise with systematic device-to-device variability applied once per mini-batch.
5 Training for analog inference chips
The toolkit supports hardware-aware training for inference chips by modeling analog noise while keeping backward and update passes ideal. It also provides calibrated programming noise, read noise, and drift models to evaluate inference accuracy and its degradation over time.
- Hardware-aware training can configure a tile’s backward and update passes as “perfect” while injecting analog weight noise during training.
- Calibrated inference models add conductance-dependent programming noise, weight read noise, and conductance drift to trained analog weights.The temporal noise profile is calibrated on a 1M phase-change memory device array.
- The inference setting enables testing realistic accuracy and accuracy degradation over time.
- Figure 4 demonstrates the Tiki-taka modified SGD rule for analog training through a more complex device configuration.After defining the rpu_config, DNN training remains identical to Fig. 2.
6 Extendability and open source contributions
The toolkit is designed for extension across Python-level inference models and C++ device-update models, while its open-source workflow encourages community contributions and maintains software quality practices.
- Python-based inference noise models can be extended by deriving new classes for additional drift or noise behavior.
- The C++ RpuCuda core combines speed with modularity, allowing new device models to inherit existing properties and replace material-specific update curves.
- The project uses GitHub-based open-source development with issue discussions, continuous integration, automated testing, code reviews, and documented packaging practices.
7 Conclusion
The paper presents an open-source toolkit for functional simulation of analog AI in ANN training and inference. It targets algorithm and topology development, while hardware performance estimation requires complementary tools.
- The toolkit focuses on functional verification for ANN training and inference on emerging analog chips, rather than estimating runtime, latency, or power.
- Performance characterization is outside the toolkit’s scope and requires dedicated complementary tools.