XBTorch: A Unified Framework for Modeling and Co-Design of Crossbar-Based Deep Learning Accelerators

TL;DR

XBTorch is a versatile simulation framework integrated with PyTorch that models crossbar-based memory technologies for deep learning accelerators, enabling efficient co-design, device modeling, and fault tolerance analysis.

Contribution

It introduces a unified, technology-agnostic simulation framework for crossbar-based accelerators, supporting detailed device and system-level modeling within PyTorch.

Findings

XBTorch accurately models crossbar devices like FeFET and ReRAM.

The framework facilitates hardware-aware training and inference with fault tolerance.

It enables comprehensive co-design of hardware and neural network models.

Abstract

Emerging memory technologies have gained significant attention as a promising pathway to overcome the limitations of conventional computing architectures in deep learning applications. By enabling computation directly within memory, these technologies - built on nanoscale devices with tunable and nonvolatile conductance - offer the potential to drastically reduce energy consumption and latency compared to traditional von Neumann systems. This paper introduces XBTorch (short for CrossBarTorch), a novel simulation framework that integrates seamlessly with PyTorch and provides specialized tools for accurately and efficiently modeling crossbar-based systems based on emerging memory technologies. Through detailed comparisons and case studies involving hardware-aware training and inference, we demonstrate how XBTorch offers a unified interface for key research areas such as device-level modeling, cross-layer co-design, and inference-time fault tolerance. While exemplar studies utilize ferroelectric field-effect transistor (FeFET) models, the framework remains technology-agnostic - supporting other emerging memories such as resistive RAM (ReRAM), as well as enabling user-defined custom device models. The code is publicly available at: https://github.com/ADAM-Lab-GW/xbtorch