Skip to content
Book Open access

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

Aug 2026 · Proceedings of the ACM/IEEE International Symposium on Low Power Electronics and Design · 0 citations · 34 references

TL;DR

This paper proposes LayUp, a layer-wise training optimization framework for edge devices with unified memory architectures that achieves speedup and energy reduction compared to baseline GPU-only training for GPT-2 models on the NVIDIA Jetson Orin NX, while conventional offloading increases latency.

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.

Read PDF

Similar papers

Oct 2026

SynergyScale: Optimizing Offloading and Task Partitioning for Efficient Model Training

Deep neural networks (DNNs) with billions of parameters power many important applications, but their training is fundamentally constrained by the limited on-chip memory of GPUs. This memory wall forces training to rely on distributed execution or memory offloading, both of which introduce substantial inefficiencies. Existing offloading techniques can scale model size but often incur severe throughput degradation, while conventional distributed training suffers from poor hardware utilization due to limited cross-device bandwidth, leaving accelerator resources underused. We present SynergyScale, a software framework that addresses the GPU memory wall by jointly optimizing memory usage and computation efficiency. SynergyScale combines dynamic hierarchical memory offloading across GPUs, CPU, and secondary storage tiers to reduce GPU memory pressure with fine-grained multi-stream task partitioning that takes advantage of the concurrent execution capabilities of modern GPUs. The framework automatically selects offloading parameters that balance data movement overhead with parallel execution efficiency. Experimental results show that SynergyScale can train a 114.41B-parameter model on a single 80 GB NVIDIA A100 GPU and a 547.77B-parameter model on eight A100 GPUs, supporting models up to 29.5× larger than those enabled by existing approaches. By co-optimizing offloading and computation, SynergyScale achieves up to 4.5× higher training throughput than offloading-only methods, without altering training semantics, making billion-scale model training both efficient and cost-effective.

Xiaoyang Sun, Jie Xu, Zheng Wang · 0 citations
Preprint Jul 2026

Xema: Efficient Diffusion Serving through Fine-Grained Memory Management and Auto-Configuration

Xema is presented, a memory-efficient diffusion serving system that exploits predictable tensor lifetimes for trace-guided memory optimization and introduces an offline planner that jointly selects parallelism, concurrency, and memory control under GPU memory and SLO constraints.

Xueze Kang, Guangyu Xiang, Suyi Li et al. · 1 citation
Conference Jul 2026

Reducing Memory Requirements of LLM Inference Through Online rANS Decompression

The rapid growth of Large Language Models has outpaced GPU memory capacity, making state-of-the-art models unfit for most GPUs. This necessitates evicting model parameters to CPU memory or disk, further away from the processing elements. In these scenarios, inference throughput is restricted by interconnect bandwidth, and while quantization schemes like FP8 reduce memory usage and enable larger models to fit smaller memories, they compromise generation quality and rely on specialized hardware support for efficient inference. To this end, we propose an inference pipeline that utilizes lossless entropy encoding (rANS) alleviating the bottleneck while maintaining standard floating-point portability or output precision. We demonstrate the potential of this approach by integrating the pipeline into vLLM, with a hardware-aware parallelization strategy. Our method increases the effective GPU memory capacity up to 30% in inference workloads with evicted weights, with no loss of quality and at the same time achieving speed up of 2.5–4.4%.

Väinö-Waltteri Granat, Jakub Žádník, M. Mäkitalo et al. · 0 citations
Book Open access Jul 2026

Toward Low-Latency and Memory-Efficient Deployment of Irregular Sparse Deep Learning Workloads

This work introduces a tile-aware scheduling framework for efficient sparse Vision Transformer execution on GPUs and introduces a training-aware extension that reuses the inference tile schedule and augments it with backward computation and activation-memory strategies.

Changxin Li · 0 citations
Preprint Aug 2026

Rethinking Unified Memory for NPU-PIM Systems: Dual-View Memory for Dynamic Inference of LLM

PFM (PIM-as-Flexible-Memory), a dual-view memory system that decouples physical data layout from accessor-visible logical views, is presented, demonstrating its effectiveness and broad applicability as a unified memory management solution for NPU-PIM systems.

Shixin Zhao, Lian Liu, Tian Han et al. · 0 citations
Preprint Aug 2026

Enabling Memory-efficient Im2win Convolution with Multi-precision Support on GPU CUDA and Tensor Cores

Convolution is a principal computational bottleneck in deep neural networks, and its efficiency depends on tight integration between algorithms and GPU hardware. Existing GPU convolution methods suffer from large memory overhead, poor cache utilization, limited effectiveness across kernel sizes, or numerical instability. This work extends the im2win paradigm -- a universal, memory-efficient convolution method with contiguous memory access for all kernel sizes -- to run efficiently in full precision on CUDA cores and half precision on tensor cores. By introducing new kernel designs and optimizations such as zig-zag memory access and asynchronous data movement, im2win efficiently exploits hardware-accelerated half-precision matrix multiply-accumulate operations. Across twelve CNN benchmarks, im2win achieves up to 2.8x higher TFLOPS than its CUDA core implementation, 1.4x higher than cuDNN, and 6.4x higher than GEMM-based convolution with cuBLAS, while using as little as 53% and 35% of their memory, respectively. These results establish im2win as a unified, high-performance convolution framework for modern GPU architectures.

Xiang Fu, Jixiang Ma, Xinpeng Zhang et al. · 0 citations