LayUp: Layer-wise Parallelization for Energy-Efficient Edge LLM Training Exploiting Unified Memory Characteristics

Abstract

As demand for personalized AI grows, on-device large language model (LLM) training is increasingly pursued to preserve data privacy and enable real-time adaptation. However, existing training optimization techniques primarily target discrete memory systems, where CPU and GPU reside in separate memory spaces connected via PCIe. These approaches are not directly applicable to edge devices with unified memory systems, where CPU and GPU share the same physical memory. Applying offloading on such platforms introduces redundant data copies without any memory savings. In this paper, we propose LayUp, a layer-wise training optimization framework for edge devices with unified memory architectures. We analyze how memory allocation methods affect training performance across tensor types and apply allocation strategies suited to unified memory. By parallelizing GPU-side backward propagation with CPU-side layer-wise parameter updates via zero-copy shared memory, LayUp achieves up to 24.1% speedup and 6.1% energy reduction compared to baseline GPU-only training for GPT-2 models on the NVIDIA Jetson Orin NX, while conventional offloading increases latency by up to 81.7%. A double-buffered gradient management scheme reduces gradient memory by up to 86.4% and overall model memory by up to 10.8% for GPT-2 Large, enabling training where both baseline and offloading run out of memory. Our implementation is available at https://github.com/aica-ssu/LayUp.