[论文解读] Using the IBM Analog In-Memory Hardware Acceleration Kit for Neural Network Training and Inference

Manuel Le Gallo, Corey Lammie+9 more|arXiv (Cornell University)|Jul 18, 2023

Neural Networks and ApplicationsComputer Science参考文献 70被引用 3

一句话总结

本文介绍了 IBM AIHWKit，这是一个基于 Python 的仿真工具包，可实现对模拟存内计算（AIMC）硬件的精确建模，用于深度神经网络（DNN）的训练与推理。它将器件非理想特性（如噪声和可变性）集成到类似 PyTorch 的工作流中，使研究人员能够在真实 AIMC 条件下训练和评估 DNN，主要贡献包括硬件感知训练、可定制的仿真组件，以及与托管云平台（Analog AI Cloud Composer）的集成。

ABSTRACT

Analog In-Memory Computing (AIMC) is a promising approach to reduce the latency and energy consumption of Deep Neural Network (DNN) inference and training. However, the noisy and non-linear device characteristics, and the non-ideal peripheral circuitry in AIMC chips, require adapting DNNs to be deployed on such hardware to achieve equivalent accuracy to digital computing. In this tutorial, we provide a deep dive into how such adaptations can be achieved and evaluated using the recently released IBM Analog Hardware Acceleration Kit (AIHWKit), freely available at https://github.com/IBM/aihwkit. The AIHWKit is a Python library that simulates inference and training of DNNs using AIMC. We present an in-depth description of the AIHWKit design, functionality, and best practices to properly perform inference and training. We also present an overview of the Analog AI Cloud Composer, a platform that provides the benefits of using the AIHWKit simulation in a fully managed cloud setting along with physical AIMC hardware access, freely available at https://aihw-composer.draco.res.ibm.com. Finally, we show examples on how users can expand and customize AIHWKit for their own needs. This tutorial is accompanied by comprehensive Jupyter Notebook code examples that can be run using AIHWKit, which can be downloaded from https://github.com/IBM/aihwkit/tree/master/notebooks/tutorial.

研究动机与目标

• 解决将 DNN 部署在模拟存内计算（AIMC）硬件上所面临的挑战，此类硬件存在器件噪声、非线性和电路非理想性等问题。
• 提供一个软件工具包，弥合高层深度学习框架与 AIMC 器件及外围电路精确物理模型之间的鸿沟。
• 使研究人员能够在无需访问物理芯片或低层次电路仿真器的情况下，模拟和优化 DNN 在 AIMC 硬件上的表现。
• 通过提供可定制的仿真组件和与基于云的物理硬件访问的集成，促进硬件感知训练算法的开发。
• 通过鼓励对维护良好且开源的工具包进行贡献，推动社区协作开发，减少 AIMC 研究领域的工具碎片化。

提出的方法

• 在 AIHWKit 内实现一个仿真框架，用于建模非易失性存储器（NVM）器件（如 PCM、ReRAM 和 MRAM）在交叉阵列中的物理行为。
• 将器件级非理想性（包括器件间差异、循环间差异、随机 Telegraph 噪声（RTN）以及电导非线性）集成到仿真流程中。
• 通过可配置的分辨率和非理想性，支持混合精度训练与推理，建模模拟到数字（ADC）和数字到模拟（DAC）转换过程。
• 通过模拟梯度累积（AGAD）和动态参考缩放支持硬件感知训练（HWA），并支持使用模拟权重更新的随机梯度下降（SGD）。
• 提供模块化软件架构，允许用户通过重写核心组件（如 AnalogTile 或 InferenceTile）来扩展功能，并可集成自定义外围逻辑。
• 与 Analog AI Cloud Composer（一个托管云平台）集成，提供对仿真和物理 AIMC 硬件的访问，支持端到端评估。

Figure 1: (a), Illustration of a potential AIMC chip. (b) AIMC devices implemented in AIHWKit and their properties.

实验结果

研究问题

• RQ1 如何在考虑器件和电路非理想性的前提下，有效实现 DNN 在模拟存内计算硬件上的训练与推理？
• RQ2 哪些关键的硬件感知训练策略能够在阻变存储器器件的噪声和非线性特性下维持 DNN 的准确性？
• RQ3 仿真工具包（如 AIHWKit）应如何架构，才能在无需 SPICE 级仿真的情况下，同时支持高层 DNN 框架集成和低层硬件建模？
• RQ4 在基于 AIMC 的推理中，不同精度和量化方案（例如位切片、混合精度）的性能与准确性权衡如何？
• RQ5 社区驱动的开源工具包开发如何减少碎片化，并加速模拟 AI 硬件研究的进展？

主要发现

• AIHWKit 能够在包含 RTN、电导非线性以及 ADC/DAC 量化等真实器件非理想性的条件下，精确仿真 DNN 在 AIMC 硬件上的推理与训练。
• 采用模拟梯度累积（AGAD）和动态参考缩放的硬件感知训练，即使在显著的器件可变性下，也能保持与数字训练相当的 DNN 准确性。
• 该工具包支持混合精度推理与训练，且 ADC/DAC 分辨率可配置，使研究人员能够评估低精度算术对模型性能的影响。
• Analog AI Cloud Composer 提供对仿真和物理 AIMC 硬件的无缝访问，支持在基于云的环境中对 DNN 进行端到端评估。
• 可通过拉取请求（pull requests）向 AIHWKit 贡献新外围操作或优化器实现等自定义扩展，确保长期可维护性和社区采纳。
• 该工具包已成功用于模拟和评估多种架构及器件参数下的 DNN，相关成果已发表于多篇研究（例如，[75, 76, 61, 22]）。

Figure 2: (a) Mapping of a neural network to an AIMC chip. (b) Implementation of in-memory SGD weight update. (c) Implementation of TTv2 weight update. (d) Implementation of mixed-precision weight update.