repository·master·Indexed 19 days ago
https://github.com/ibm/aihwkitAn open-source Python toolkit for simulating and training neural networks on analog in-memory computing hardware, such as PCM or RRAM crossbar arrays. It provides PyTorch integration with analog-aware modules and optimizers, a CUDA-capable C++ simulator for modeling material characteristics, and utilities for deploying hardware-aware trained networks to Fusion chips via the IBM Analog Composer.
The aihwkit library integrates with PyTorch by providing analog primitives that fit into standard PyTorch workflows. Most high-level features are exposed as PyTorch modules or optimizers.
Key integration modules:
aihwkit.nn: Provides Analog Modules (layers) and Functions.aihwkit.optim: Provides Analog Optimizers.The aihwkit toolkit is designed for exploring in-memory computing for AI and consists of two primary pillars:
Provides primitives to integrate analog hardware simulation into the standard PyTorch workflow:
AnalogLinear, 1d/2d/3d convolution layers, LSTM layers, and sequential containers.AnalogSGD) and in-situ training algorithms (like Tiki-Taka).A high-performance, CUDA-capable C++ simulator that models material characteristics and crossbar configurations. It supports:
The IBM Analog Hardware Acceleration Kit is an open-source Python toolkit designed for exploring and using in-memory computing devices for artificial intelligence. It provides tools to simulate analog hardware behavior within deep learning workflows, specifically focusing on training and inference on analog crossbar arrays.
The toolkit is composed of two primary components:
AnalogLinear, convolution layers), analog-aware optimizers (like AnalogSGD), and hardware-aware training capabilities that include non-idealities and noise in the forward pass.Additional features include a library of device presets calibrated to real hardware, an experiment execution module for high-level use cases, and integration with the AIHW Composer platform for cloud execution.
Analog AI leverages in-memory computing to overcome the von Neumann bottleneck, where computation and memory are traditionally siloed. Instead of moving data between a CPU and memory, computation happens directly within memory devices (such as PCM, RRAM, or Flash) by using their physical properties (charge or conductance states).
In deep learning, this is achieved using crossbar arrays. Each layer of a neural network is represented by a matrix of synaptic weights stored at the crosspoints of these arrays. Matrix multiplications are performed in-place in an analog manner, significantly reducing energy consumption and latency compared to conventional digital hardware.
The aihwkit library is designed as a multi-layered stack that integrates analog simulation capabilities into the PyTorch ecosystem.
Users should primarily interact with the PyTorch layer, which provides high-level primitives like analog-aware layers and optimizers. This layer inherits from standard PyTorch classes and replaces default functionality with calls to a Tiles object in the simulator abstraction layer.
aihwkit.nn, aihwkit.optim.aihwkit.simulator.tiles, aihwkit.simulator.configs, aihwkit.simulator.presets.aihwkit.simulator.rpu_base.tiles, aihwkit.simulator.rpu_base.devices, aihwkit.simulator.rpu_base.parameters.Tiles occur. This layer is not directly accessible from Python but can be used by other C++ programs via provided headers.An analog in-memory computing chip consists of multiple crossbar arrays that communicate with each other.
The aihwkit includes a statistical model to simulate Conductive Metal Oxide/HfOx-ReRAM (CMO-ReRAM) behavior during inference. This model allows developers to evaluate how realistic hardware non-idealities and noise affect neural network inference accuracy.
Key Characteristics:
[-1, 1] range.-1 maps to the High Resistance State (HRS), or $g_{min}$.1 maps to the Low Resistance State (LRS), or $g_{max}$.Beyond core simulation, aihwkit provides:
Drift simulates the conductance relaxation that occurs after programming. In CMO-ReRAM, this relaxation is modeled as a Gaussian distribution:
$g_{drift} = \mathcal{N}(\mu_{drift}, \sigma_{drift})$
Unlike other resistive memory devices, the drift in CMO-ReRAM is independent of the target conductance level. The time dependency is modeled using the logarithm of time ($t$):
In aihwkit, analog simulation is triggered through the interaction of analog-aware layers and specialized optimizers.
When you define an AnalogLinear layer and an AnalogSGD optimizer, the following lifecycle occurs:
Initialization:
AnalogLinear creates a FloatingPointTile (from the Python abstraction layer), which internally instantiates a FloatingPointTile from the Pybind layer, which finally creates the underlying C++ objects.AnalogSGD inspects the AnalogLinear layer to identify which parameters require analog-specific handling during optimization.Forward and Backward Passes:
aihwkit.nn.functions.AnalogFunction.forward() or backward() methods on the layer's FloatingPointTile.Optimization Step:
opt.step() is called, AnalogSGD retrieves the reference to the FloatingPointTile and calls its update() function.model = AnalogLinear(2, 1)
opt = AnalogSGD(model.parameters(), lr=0.5)
...
for epoch in range(100):
pred = model(x_b)
loss = mse_loss(pred, y_b)
loss.backward()
opt.step()When building complex models, standard PyTorch containers (like nn.Sequential) may not correctly propagate certain operations to the parameters and buffers of analog layers.
To ensure operations like .cuda(), .eval(), or .program_analog_weights() are applied correctly to all children, use aihwkit.nn.AnalogSequential. This class acts as a compatible replacement for nn.Sequential and should be used as the superclass when defining custom analog modules.
from aihwkit.nn import AnalogLinear, AnalogSequential
# As a replacement for nn.Sequential
model = AnalogSequential(
AnalogLinear(10, 20)
)
model.cuda()
model.eval()
model.program_analog_weights()