Pipeline Gradient-based Model Training on Analog In-memory Accelerators

2024-10-19

arXiv.org

AI Summary

This work addresses two key challenges in large-model training on analog in-memory computing (AIMC) accelerators: (1) restricted data parallelism due to inefficient weight replication, and (2) stale weights and analog-domain deviations caused by asynchronous pipelined gradient updates. To this end, we propose Analog-SGD-AP—an asynchronous pipelined gradient descent algorithm tailored for AIMC. We establish the first convergence theory for AIMC architectures that jointly models physical non-idealities (e.g., device noise, nonlinear conductance response, weight update latency) and asynchronous timing behavior, rigorously deriving upper bounds on clock cycles and sample complexity. Analog-SGD-AP breaks the traditional data-parallelism bottleneck, enabling scalable multi-chip collaborative training. Evaluated on real datasets, it achieves convergence comparable to digital pipelined training while significantly improving training throughput and hardware efficiency.

Technology Category

• Machine Learning: Hardware-aware ML
• Cognitive Modeling & Cognitive Systems: Analogy
• Search and Optimization: Learning to Search

Application Category

• Search and Retrieval-Augmented AI: Agentic search
• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

Abstract

Aiming to accelerate the training of large deep neural models (DNN) in an energy-efficient way, an analog in-memory computing (AIMC) accelerator emerges as a solution with immense potential. In AIMC accelerators, trainable weights are kept in memory without the need to move from memory to processors during the training, reducing a bunch of overhead. However, although the in-memory feature enables efficient computation, it also constrains the use of data parallelism since copying weights from one AIMC to another is expensive. To enable parallel training using AIMC, we propose synchronous and asynchronous pipeline parallelism for AIMC accelerators inspired by the pipeline in digital domains. This paper provides a theoretical convergence guarantee for both synchronous and asynchronous pipelines in terms of both sampling and clock cycle complexity, which is non-trivial since the physical characteristic of AIMC accelerators leads to analog updates that suffer from asymmetric bias. The simulations of training DNN on real datasets verify the efficiency of pipeline training.

Problem

• Analyzes convergence of analog in-memory training with asynchronous pipelines
• Addresses hardware imperfections and stale weights in multi-layer DNN training
• Demonstrates efficient convergence despite analog and pipeline challenges

Innovation

• Analog in-memory computing for energy-efficient DNN training
• Asynchronous pipeline parallelism to utilize all accelerators
• Convergence theory for Analog-SGD-AP despite hardware imperfections