Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

SwiftRL: Towards Efficient Reinforcement Learning on Real Processing-In-Memory Systems

Kailash Gogineni, Sai Santosh Dayapule, Juan Gómez-Luna, Karthikeya Gogineni, Peng Wei, Tian Lan, Mohammad Sadrosadati, Onur Mutlu, Guru Venkataramani

TL;DR

SwiftRL addresses memory bottlenecks in reinforcement learning training by deploying Tabular Q-learning and SARSA on real Processing-In-Memory (PIM) hardware (UPMEM). The approach partitions experience data across thousands of PIM cores, uses fixed-point INT32 arithmetic to cope with hardware limitations, and explores single- and multi-agent configurations across two OpenAI Gym environments. Key contributions include a roofline-based analysis of RL memory behavior, a suite of six Q-learning and six SARSA variants with three sampling strategies, strong scaling up to 2,000 PIM cores, and comprehensive CPU/GPU comparisons, plus open-source releases. The results show near-linear performance scaling and substantial speedups over CPU and GPU baselines, highlighting PIM’s potential for memory-intensive RL workloads in practical settings.

Abstract

Reinforcement Learning (RL) trains agents to learn optimal behavior by maximizing reward signals from experience datasets. However, RL training often faces memory limitations, leading to execution latencies and prolonged training times. To overcome this, SwiftRL explores Processing-In-Memory (PIM) architectures to accelerate RL workloads. We achieve near-linear performance scaling by implementing RL algorithms like Tabular Q-learning and SARSA on UPMEM PIM systems and optimizing for hardware. Our experiments on OpenAI GYM environments using UPMEM hardware demonstrate superior performance compared to CPU and GPU implementations.

SwiftRL: Towards Efficient Reinforcement Learning on Real Processing-In-Memory Systems

TL;DR

SwiftRL addresses memory bottlenecks in reinforcement learning training by deploying Tabular Q-learning and SARSA on real Processing-In-Memory (PIM) hardware (UPMEM). The approach partitions experience data across thousands of PIM cores, uses fixed-point INT32 arithmetic to cope with hardware limitations, and explores single- and multi-agent configurations across two OpenAI Gym environments. Key contributions include a roofline-based analysis of RL memory behavior, a suite of six Q-learning and six SARSA variants with three sampling strategies, strong scaling up to 2,000 PIM cores, and comprehensive CPU/GPU comparisons, plus open-source releases. The results show near-linear performance scaling and substantial speedups over CPU and GPU baselines, highlighting PIM’s potential for memory-intensive RL workloads in practical settings.

Abstract

Reinforcement Learning (RL) trains agents to learn optimal behavior by maximizing reward signals from experience datasets. However, RL training often faces memory limitations, leading to execution latencies and prolonged training times. To overcome this, SwiftRL explores Processing-In-Memory (PIM) architectures to accelerate RL workloads. We achieve near-linear performance scaling by implementing RL algorithms like Tabular Q-learning and SARSA on UPMEM PIM systems and optimizing for hardware. Our experiments on OpenAI GYM environments using UPMEM hardware demonstrate superior performance compared to CPU and GPU implementations.
Paper Structure (25 sections, 1 equation, 7 figures, 1 table, 1 algorithm)

This paper contains 25 sections, 1 equation, 7 figures, 1 table, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background and Motivation
  3. Reinforcement Learning
  4. Processing-In-Memory
  5. SwiftRL Design and Implementation
  6. Memory Behavior of RL workloads
  7. Implementation of RL algorithms on PIM Architecture
  8. Q-learning
  9. Multi-agent Q-learning
  10. SARSA Learning
  11. Experimental Evaluation
  12. Experimental Setup
  13. RL Training Quality
  14. Performance Analysis of PIM Kernels: Scaling PIM Cores
  15. Comparison with CPU/GPU platforms
  16. ...and 10 more sections

Figures (7)

  • Figure 1: The Operational Workflow of an Offline RL System. An unknown behavior policy $\pi_\beta$ is used to gather the dataset, denoted as $\mathcal{D}$. Training occurs without any interaction with the environment, and the policy (i.e., final Q-table) is deployed only after it is fully trained. The architecture is adapted from levine2020offline.
  • Figure 2: The Roofline model depicts the performance characteristics of CPU versions in RL workloads, where "Q" refers to Q-learner sutton2018reinforcementwatkins1992qli2020sample, "S" signifies the SARSA learner sutton2018reinforcement, and "1M" and "20M" indicate the data size in millions of transitions. The workloads are tested on an Intel i7-9700K CPU.
  • Figure 3: Organization of a State-of-the-Art Processing-In-Memory (PIM) Architecture.Adapted from gomez2021benchmarkinggomez2023evaluatingchen2023simplepimUPMEMgiannoula2022towards.
  • Figure 4: SwiftRL execution on a real PIM system. The execution phase comprises four main steps: (1) loading the input dataset chunks into individual DRAM banks of PIM-enabled memory, (2) executing the RL workload (kernel) on PIM cores in parallel operating on different chunks of data, (3) retrieving partial results from DRAM banks to the host CPU, and (4) aggregating partial results on the host processor.
  • Figure 5: Execution time (measured in seconds) of RL workloads on 125, 250, 500, 1,000 and 2,000 PIM cores (x-axis) with each PIM core running with single thread for a frozen lake environment. In this illustration, SEQ-RAN-STR represent sequential, random and stride-based sampling techniques and both the Q-learning and SARSA learners are evaluated with FP32 and INT32 data types. The synchronization period $\tau$ and the stride value in this experiment is set to 50 and 4 respectively.
  • ...and 2 more figures