A flexible and fast PyTorch toolkit for simulating training and inference on analog crossbar arrays (2104.02184v1)

Abstract: 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.

• The paper introduces an analog tile model that simulates the processing operations of analog crossbar arrays within PyTorch.
• It integrates device variations and noise modeling to realistically simulate both training and inference processes with GPU acceleration.
• The toolkit offers a user-friendly interface, enabling researchers to explore analog AI innovations and optimize deep learning architectures.

Overview of the IBM Analog Hardware Acceleration Kit

The paper presents a detailed exposition of the IBM Analog Hardware Acceleration Kit, a novel open-source toolkit built for simulating analog crossbar arrays within the PyTorch framework. This toolkit uniquely accommodates the simulation demands of future memristive crossbar architectures used for accelerating artificial neural networks (ANNs). The authors have succeeded in seamlessly integrating the capacities of analog hardware simulation with the established PyTorch environment, thereby broadening the accessibility and applicability of analog computing research in the ANN domain.

Technical Contributions

The toolkit introduces an "analog tile" as its fundamental unit, a versatile construct that simulates the processing operations typical of analog crossbar arrays. Through analog tiles, researchers can create network modules incorporating analog computing elements with varying hardware characteristics encompassing device-to-device variations, noise factors, and material response inconsistencies. This adaptability allows the simulation of both training and inference processes, inclusive of advanced optimization strategies like Tiki-Taka. Importantly, the toolkit offers a GPU-accelerated environment facilitating the rapid simulation of material properties and non-idealities and their impacts on neural network performance.

Comparative Analysis

Within the broader landscape of simulation toolkits focusing on AI workloads on non-volatile memory devices, the IBM toolkit distinguishes itself by providing a user-friendly interface comparable to PyTorch. Unlike existing tools like MLP-NeuroSim and RxNN, which either lack integration with modern AI frameworks or are based on outdated platforms, the IBM toolkit offers compatibility and extension capabilities with prevailing AI methodologies. This ensures that research in analog AI can leverage state-of-the-art AI constructs, augmenting the research potential in this promising field.

Simulation Features

The core functionality of the toolkit is encapsulated in its analog tile model, conforming to the resistive processing unit (RPU) definitions and offering a continuum of noise and non-ideality configurations. Noteworthy is the toolkit's approach to analog training simulations, which, with its CUDA-optimized core, achieves comparably efficient training times relative to conventional floating-point training. The toolkit supports device variation modeling essential for realistic and predictive simulation outcomes, thereby facilitating the exploration of new neural architectures and algorithmic paradigms adapted to analog technologies.

Future Implications

The toolkit is positioned as a cornerstone for both hardware and algorithm-level innovations in analog computing. By catering to the functional verification of analog AI systems rather than performance metrics such as power consumption or latency, the toolkit clears a dedicated path for the algorithmic development side of analog computing. As AI continues its trajectory towards resource-efficient designs, the toolkit promises to catalyze the research needed to optimize deep learning architectures for analog computation paradigms. Moreover, the open-source nature of this toolkit invites community-wide collaboration and enhancement, a critical factor for the maturation and acceptance of analog AI technologies.

Conclusion

The IBM Analog Hardware Acceleration Kit exemplifies a significant advancement in the alignment of analog hardware capabilities with ANN frameworks, championing an important shift toward analog-centric AI insights. By furnishing a scalable and integrable platform within PyTorch, it invites deep learning researchers and hardware developers alike to delve into the analog computing domain, potentially steering future innovations in AI technologies. As AI hardware approaches evolve, this toolkit stands poised to underpin the development of more harmonious and effective AI systems exploiting the potentialities of analog computing.