Unlocking the AMD Neural Processing Unit for ML Training on the Client Using Bare-Metal-Programming Tools

TL;DR

This study demonstrates the feasibility of training a GPT-2 model entirely on a client device by offloading the most compute-intensive GEMM operations to AMD's XDNA NPU using a bare-metal IRON tool-flow. By adopting a layer-by-layer approach and a carefully tiled GEMM implementation, the authors achieve substantial performance and energy-efficiency gains (e.g., average GEMM speedups ~2.8–3.1x and end-to-end improvements up to 1.7x throughput, 1.4x FLOPS/W) compared with a CPU baseline. They show that explicit data movement and minimal reconfiguration are critical to performance, while a pure-C, framework-free GPT-2 implementation enables tight optimization for edge hardware. The work highlights the practical potential and remaining bottlenecks of edge training, underlining the value of bare-metal toolchains for harnessing accelerator power without the overheads of high-level ML frameworks.

Abstract

There has been a growing interest in executing machine learning (ML) workloads on the client side for reasons of customizability, privacy, performance, and availability. In response, hardware manufacturers have begun to incorporate so-called Neural Processing Units (NPUs) into their processors for consumer devices. Such dedicated hardware optimizes both power efficiency and throughput for common machine learning tasks. AMD's NPU, part of their Ryzen AI processors, is one of the first such accelerators integrated into a chip with an x86 processor. AMD supports bare-metal programming of their NPU rather than limiting programmers to pre-configured libraries. In this paper, we explore the potential of using a bare-metal toolchain to accelerate the weight fine-tuning of a large language model, GPT-2, entirely on the client side using the AMD NPU. Fine-tuning on the edge allows for private customization of a model to a specific use case. To the best of our knowledge, this is the first time such an accelerator has been used to perform training on the client side. We offload time-intensive matrix multiplication operations from the CPU onto the NPU, achieving a speedup of over 2.8x for these operations. This improves end-to-end performance of the model in terms of throughput (1.7x and 1.2x speedup in FLOPS/s on mains and battery power, respectively) and energy efficiency (1.4x improvement in FLOPS/Ws on battery power). We detail our implementation approach and present an in-depth exploration of the NPU hardware and bare-metal tool-flow.

Figures (9)

• Figure 1: Overview of the XDNA architecture, showing the VLIW processor compute cores (L1), also called AI Engines, memory cores (L2), shim cores and the dedicated command processor. The small grey boxes between arrows are configurable interconnect switch boxes.
• Figure 2: Computation graph of GPT-2, including floating point operation count ($1 \text{MFLOP} = 1e6 \text{FLOP}$). The fine-tunable weights (as named in llm.c) are annotated in grey italics. The arrows on the left and right show the FLOP count in the forward and backward pass, respectively. The operations within the grey box are repeated twelve times.
• Figure 3: Available tools (blue) and the intermediate outputs (brown) they produce in the open-source IRON tool-flow
• Figure 4: Data movement and tiling of input and output matrices across the three memory levels. The data movement to the compute core (2, 3) is highlighted as an example: This core receives its sub-tile of the A matrix from the memory core in column 1, and its sub-tile of the B matrix from the memory core in column 2 (zero-indexed).
• Figure 5: Data layout transformations performed on the hardware DMAs for matrix sub-tiles for input matrix A. Matrices B and C are analogously tiled with their respective dimensions.