PaperlayerView exact revision on arXiv

Source-linked AI summary

Fast and Memory Efficient Offload Training Framework with Hybrid XPU Computation

Zhiyi Yao, Zuning Liang, Yuedong Xu, Jin Zhao, Jessie Hui Wang, Tong Li

arXiv:2610.09657v1cs.DC

TL;DR

Offload training can leave GPUs idle and communication incompletely overlapped, while fixed offloading policies limit flexibility. MemFerry combines DHA and on-GPU computation, with a shadow model for parameter management, and its evaluations report faster training with lower GPU memory use.

  • Problem

    ZeRO-Offload has low GPU utilization and incomplete communication-computation overlap, while its fixed-proportion offloading policy is inelastic.

    Source: A. Limitations of Offload Training · p. 3

  • Method

    MemFerry uses DHA for selected layers while loading other parameters to GPU memory, and a shadow model provides unified management of parameter partitions.

    Source: §I · p. 2

  • Results

    MemFerry improves training speed by 1.32× on average and up to 1.68×, while reducing GPU memory usage by up to 61.7%.

    Source: §I · p. 2

  • Takeaways & Limitations

    GO-MemFerry enables training models up to 1.52× larger than those supported by ZeRO-Offload through gradient offloading.

    Source: §I · p. 2

    In ZeRO-Offload, communication and computation interference leaves 57.2% unhidden computation time, and fixed-proportion offloading is inelastic.

    Source: A. Limitations of Offload Training · p. 3

Abstract

from arXiv · show

With the ever-growing size of deep learning models, GPU memory is prone to being insufficient during training. A prominent approach is ZeRO-Offload, which moves the optimizer states to CPU memory and performs parameter update using CPU. However, the deficiencies of ZeRO-Offload include low GPU utilization, imperfect overlapping of communication and computation, and inflexible offloading. In this paper, we leverage Direct Host Access (DHA) on the GPU that can compute data in CPU memory, forming a novel hybrid on-GPU and DHA. We design and implement MemFerry consisting of an execution scheduler and a shadow model. The scheduler strategically chooses layers of parameters for DHA computation and transmits the remaining parameters to GPU memory simultaneously to shorten forward propagation time, and further loads DHA parameters to GPU memory to reduce backward propagation time. The shadow model presents a unified memory abstraction for the parameter partitions stored separately in GPU and CPU memories. To further reduce GPU memory usage, we present MemFerry along with its dynamic programming algorithm that offloads gradients to CPU memory via DHA. We further extend MemFerry to emerging scale-up domains with ScaleUp-MemFerry, which exploits otherwise underutilized accelerator interconnect bandwidth to assist host-to-accelerator data movement through adaptive multi-path transfer. Our experiments show that \system trains up to $1.68\times$ faster and MemFerry can train $1.52\times$ larger model compared to ZeRO-Offload on a single GPU, and increase training speed by at least $28.1\%$ when scaling to data parallelism on 8 GPUs. We further extend the design to a Huawei CloudMatrix384 scale-Up node with up to 8 NPUs, and our ScaleUp-MemFerry reduces the end-to-end iteration time by up to $20.7\%$ over DeepSpeed.

I. INTRODUCTION

Large models strain limited GPU memory, while ZeRO-Offload’s CPU-based optimizer offloading can constrain utilization and communication overlap. MemFerry combines DHA and GPU execution, with GO-MemFerry and ScaleUp-MemFerry extending gradient offloading and host-to-accelerator transfers.

I. INTRODUCTION

  • ZeRO-Offload’s fixed optimizer offloading leaves GPU utilization low and communication imperfectly overlapped, motivating faster training with lower GPU memory use.The paper reports 63% versus 99% GPU utilization for BERT-base on an A100.Source: §I · p. 1
  • DHA lets GPU SMs compute on CPU-resident data, and combining DHA with GPU computation can hide parameter-transfer time during forward propagation.The paper notes DHA performs better than on-GPU computation on certain large layers in forward propagation, while its backward-propagation disadvantage narrows when on-GPU gradient-transfer time is considered.Source: §I · pp. 1–2
  • MemFerry’s scheduler selects DHA-computed layers while loading other parameters to GPU memory, and its shadow model unifies parameter partitions across CPU and GPU memories.During the interval before gradient transmission, selected DHA parameters are loaded to GPU memory to reduce backward-propagation time.Source: §I · p. 2
  • GO-MemFerry writes gradients directly to CPU memory via DHA, while ScaleUp-MemFerry adaptively dispatches transfers over direct and assisted paths.The ScaleUp design uses otherwise underutilized accelerator interconnect bandwidth for host-to-accelerator data movement.Source: §I · p. 2
  • MemFerry improves training speed by 1.32× on average and up to 1.68×, while reducing GPU memory usage by up to 61.7%.GO-MemFerry reduces memory usage by 20.1% on average and up to 74.4%, enabling models up to 1.52× larger than ZeRO-Offload.Source: §I · p. 2

A. Memory Usage of DNN Training

C. Scale-up networks

Scale-up domains connect accelerators and pooled memory through direct accelerator-centric paths. The motivation and DHA measurements show why offloaded training needs better overlap and flexible ways to compute on CPU-resident data.

C. Scale-up networks

  • Scale-up fabrics connect dozens to hundreds of accelerators and can make CPU and pooled memory higher-performance tiers accessible through accelerator-centric paths.CloudMatrix384 is described as a 384-NPU, 192-CPU supernode connected by a peer-to-peer Unified Bus fabric.Source: C. Scale-up networks · p. 3

A. Limitations of Offload Training

  • GPT-2 ZeRO-Offload averages 63.1% GPU utilization, while its standard schedule retains 57.2% unhidden computation versus the best scenario and offloads fixed contents.The authors argue offloading should also be flexible enough to move gradients when GPU memory is limited.Source: A. Limitations of Offload Training · pp. 3–4

B. GPU Computation on CPU Memory

  • DHA lets GPU streaming multiprocessors compute directly from registered CPU memory while CPU-to-GPU transfers proceed in parallel; its forward-pass advantage depends on layer type.DHA is faster for Embedding and LayerNorm forward passes, while Linear and Convolution layers are faster on-GPU.Source: B. GPU Computation on CPU Memory · p. 4
  • DHA is slower for backward propagation, but writing gradients directly to CPU memory avoids a GPU HBM gradient buffer and significantly reduces GPU memory requirements.The reported DHA backward slowdown versus forward propagation is 5–10%; DHA-gradient backward is slower than in-HBM backward.Source: B. GPU Computation on CPU Memory · p. 4

C. An Ideal Execution Schedule

An ideal offload schedule combines DHA and on-GPU execution to hide parameter transfers during forward propagation and overlap gradient transfer with backward propagation and updates. The design still requires fine-grained scheduling, unified memory management, and flexible offloading.

C. An Ideal Execution Schedule

  • DHA computation on one layer can overlap GPU loading and execution of another, hiding parameter transmission and eliminating startup latency in the two-layer example.The proposed schedule transmits partial parameters in parallel with hybrid forward and backward execution.Source: C. An Ideal Execution Schedule · p. 5
  • Gradient transfer cannot start before gradients are generated, but partitioning it can overlap transmission with backward propagation and parameter updates.This differs from forward propagation, where parameters can be loaded while other layers compute via DHA.Source: C. An Ideal Execution Schedule · p. 5
  • The design challenges are layer-specific execution choices, unified CPU/GPU content management without costly blind copying, and flexible offloading beyond fixed contents.Further offloading should reduce peak GPU memory while sacrificing training speed only gently.Source: C. An Ideal Execution Schedule · p. 5

IV. SYSTEM DESIGN

MemFerry combines hybrid execution scheduling, shared logical storage, and flexible offloading. Its system overview also extends data movement to scale-up systems using accelerator interconnects as assisted paths.

IV. SYSTEM DESIGN

  • MemFerry’s overview covers hybrid training modes, memory management, flexible offloading, and scale-up-aware data movement.The DL Optimizer supports partial gradient transmission and parameter updates using the user’s optimizer.Source: §IV · p. 5, A. System Overview · p. 6

A. System Overview

  • MemFerry combines a layer scheduler, shadow model, and gradient offloading to overlap data movement with computation and allow a choice between speed and GPU memory use.Its system design also includes scale-up-aware data movement.Source: A. System Overview · p. 5
  • An adaptive dispatcher partitions each host-to-accelerator transfer across available paths according to their latency and bandwidth characteristics.Accelerator interconnects serve as assisted paths in scale-up domains.Source: A. System Overview · p. 5
  • The Profiler measures computation and transmission times, and the Planner uses those profiles to select an execution mode for each layer.The Training Engine manages memory and data transmission while enforcing the plan during training.Source: A. System Overview · p. 6

B. Hybrid Training Schedule

MemFerry assigns each layer a hybrid execution mode to overlap parameter transfer and DHA computation. Its schedule profiles layer costs, uses backtracking to expose more overlap when needed, and loads selected parameters for GPU backward computation.

B. Hybrid Training Schedule

  • GFGB transfers parameters to GPU before computation, DHA computes them directly from CPU memory, and DFGB transfers them after DHA forward execution for GPU backward computation.DFGB uses faster on-GPU backward computation to mitigate DHA’s backward-time cost.Source: B. Hybrid Training Schedule · p. 6
  • In a three-layer schedule, DHA and DFGB layers begin forward computation while the third layer’s parameters transfer; the DFGB layer is then loaded for GPU backward computation.Gradient transmission is hidden with backward computation and update, so the update need not wait for all gradients to transfer.Source: B. Hybrid Training Schedule · p. 6
  • MemFerry profiles per-layer parameter-transfer, forward, and backward times once to inform its scheduling decisions, with negligible overhead.The profile also includes DHA forward and backward times and inter-layer computation times.Source: B. Hybrid Training Schedule · p. 6

F P , TF P (A) Output: P

  • When reserved overlap time cannot cover a GFGB layer’s transfer, backtracking changes earlier GFGB layers to DHA to expose more overlap for loading that layer.The forward algorithm seeks to hide transfer time while minimizing total execution time; its comparison weighs additional overlap against any increase in total time.Source: B. Hybrid Training Schedule · p. 6, F P , TF P (A) Output: P · p. 7
  • MemFerry loads selected DHA parameters into GPU memory during the idle interval between parameter loading and gradient transmission, enabling their backward computation on GPU.A greedy algorithm decides which layers use DFGB mode.Source: F P , TF P (A) Output: P · p. 7

C. Memory Management

The shadow model presents a logical GPU model whose parameter storage can reside in GPU or CPU memory. Its dynamic mapping supports the different execution modes while avoiding redundant storage and copying.

C. Memory Management

  • The shadow model separates logical model parameters from their underlying storage, allowing GPU computation to access parameters held in either GPU or CPU memory.The design aims to provide shared CPU–GPU memory, minimal copying, and compatibility with ML frameworks.Source: C. Memory Management · p. 7
  • MemFerry maps GFGB parameters from CPU storage to GPU when loaded, maps DHA parameters directly from CPU memory, and switches DFGB parameter pointers before backward propagation.The logical GPU shadow model dynamically maps storage according to the execution plan.Source: C. Memory Management · p. 8
  • Sharing DHA parameter storage reduces redundant allocation and copying, while changing DFGB storage mappings between forward and backward propagation adds no burden or effect on learning.The CPU also hosts gradients and optimizer state; GPU runtime memory holds GFGB and DFGB parameter buffers.Source: C. Memory Management · p. 8

D. Gradient Offloading

GO-MemFerry offloads selected gradients to CPU memory through DHA and uses dynamic programming to meet a user-specified GPU-gradient reservation with minimal training-time impact. When gradients are fully offloaded, GPU memory is needed only for partial parameters and activations.

D. Gradient Offloading

  • GO-MemFerry maps CPU memory to gradient storage so backward propagation writes gradients directly to CPU memory, avoiding GPU-memory transfer and copying.It selects gradient offloading rather than further parameter offloading or activation offloading, which the paper describes as costly to training speed or PCIe traffic.Source: D. Gradient Offloading · p. 8
  • GO-MemFerry offloads gradients for only part of the parameters, selecting layers to meet users’ memory needs while maximizing training speed.The selected layers reflect differing gradient-offloading performance across model parameters.Source: D. Gradient Offloading · p. 8
  • A dynamic-programming algorithm chooses gradient layers to offload under the user-specified GPU gradient reserve, minimizing the effect on training speed.It uses layer and space dimensions, with the table recording execution-time gain for each plan.Source: D. Gradient Offloading · p. 8
  • When gradients are fully offloaded, GO-MemFerry requires GPU memory only for partial parameters and activations.The dynamic-programming plan meets the memory requirement with the least execution-time side effect.Source: D. Gradient Offloading · p. 8

E. Adapting to modern scale-up domains

ScaleUp-MemFerry extends MemFerry with assisted CPU-to-accelerator transfers that use otherwise idle interconnect capacity. An adaptive dispatcher splits transfers across direct and assisted paths according to their latency and bandwidth.

E. Adapting to modern scale-up domains

  • Scale-up assistance adds a CPU0→NPU1→NPU0 path alongside CPU0→NPU0, using another CPU egress link and otherwise idle NPU interconnect bandwidth.Multiple assisting NPUs can be used when available; added capacity remains subject to CPU-memory bandwidth and UB-fabric contention.Source: E. Adapting to modern scale-up domains · p. 9
  • As tensor size grows, more data shifts to the assisted path, exploiting its higher aggregate bandwidth while retaining the low-latency direct path for small transfers.The assisted route adds a forwarding hop, so it is less attractive for small, latency-dominated transfers.Source: E. Adapting to modern scale-up domains · p. 9
  • For each tensor, the dispatcher partitions bytes between direct and assisted paths and profiles their startup latencies and effective bandwidths.It balances concurrent path completion times, clips the allocation to [0, S], and sends the entire tensor directly when the assisted path is not beneficial.Source: E. Adapting to modern scale-up domains · p. 9

V. IMPLEMENTATION

MemFerry shares CPU memory between its CPU model and GPU shadow model, switches parameter storage at runtime, and schedules asynchronous transfers alongside computation. ScaleUp-MemFerry profiles and launches direct and assisted CPU-to-NPU transfers concurrently.

MemFerry Memory Management

  • A PyTorch memory allocator lets the CPU model and GPU shadow model access the same physical CPU-memory address for a DMA tensor.The allocator returns the CPU memory address corresponding to the tensor allocated in the GPU.Source: §V · p. 9
  • Parameter switching replaces each tensor’s storage field, avoiding data copies and computation-graph modification.The storage field holds memory pointers and access formats; switching it changes between DHA and CUDA tensors.Source: §V · pp. 9–10

MemFerry Tensor Scheduling

  • Nonblocking transfers let CPU computation begin after GPU transmission starts, while a synchronization barrier at the scheduling boundary preserves transfer consistency.The barrier is mandatory at that boundary.Source: §V · p. 10

ScaleUp-MemFerry Implementation

  • ScaleUp-MemFerry profiles direct and assisted paths, computes a data partition, then launches both transfers concurrently before synchronizing prior to tensor use.Profiling measures each path’s startup latency and effective bandwidth for the underlying UB topology.Source: §V · p. 10

A. Experimental Setup

Experiments compare MemFerry and GO-MemFerry with offload baselines across GPU configurations, models, and memory settings. MemFerry improves training speed, while GO-MemFerry supports larger models and offers configurable gradient offloading.

A. Experimental Setup

  • The evaluation uses V100 and A100 cloud instances, transformer models, and per-iteration time, GPU utilization, and GPU memory usage metrics.The instances have 32-GB V100 or 80-GB A100 GPUs; the scale-up evaluation additionally reports CPU-to-NPU transfer bandwidth.Source: A. Experimental Setup · p. 10

B. Training Efficiency of MemFerry

  • Across tested models, MemFerry improves training speed by 10% to 40% compared with DeepSpeed and Native-Offload.For BERT-base and RoBERTa, the reported improvement is 21.7% to 68.7%; for BERT-large and GPT-2, it is 7.1% to 32.1%.Source: B. Training Efficiency of MemFerry · pp. 10–11

C. Training Memory Usage

  • On average, MemFerry and GO-MemFerry reduce GPU memory usage by 25.4% and 41.6%, respectively.MemFerry reduces memory by 40.1% on average for smaller language models; GO-MemFerry reduces memory by 17.6% compared with MemFerry in that setting.Source: Figure 15 · p. 11, C. Training Memory Usage · p. 11

B. Training Efficiency of MemFerry

  • For GPT-2, MemFerry’s GPU utilization is 18.5% higher than DeepSpeed and 1.5% higher than Native-Offload; GO-MemFerry is 27% higher than DeepSpeed.The passage attributes MemFerry’s utilization improvement to reduced GPU idle time and parallel DHA computation and transmission.Source: B. Training Efficiency of MemFerry · p. 11

C. Training Memory Usage

  • 1.52× larger than DeepSpeed ZeRO-Offload, the largest model GO-MemFerry can train on the limited-memory V100 setup.Above 2.4B parameters, only GO-MemFerry trains normally; MemFerry is 1.12× faster than the baseline in the tested configurations, while fully offloading gradients makes GO-MemFerry 3.4% slower than the baseline.Source: C. Training Memory Usage · p. 11
  • GO-MemFerry’s gradient reserve factor lets users set GPU memory usage, with larger factors reducing iteration time as less gradient data is written through DHA.The reported iteration-time curves are convex, and the algorithm selects gradients with the smallest performance loss for offloading.Source: C. Training Memory Usage · p. 11

D. Scalability of MemFerry

MemFerry scales across model sizes and GPUs, while ScaleUp-MemFerry adds adaptive multi-path data movement for scale-up systems. Evaluations report faster training than offload baselines and higher transfer bandwidth from the adaptive paths.

D. Scalability of MemFerry

  • MemFerry is 1.32× faster than DeepSpeed on average across tested TransformerXL model sizes, and at batch size 128 it is 1.21× faster than DeepSpeed.At batch size 128, GO-MemFerry is 1.16× faster than DeepSpeed; MemFerry’s speed advantage decreases as model size increases.Source: D. Scalability of MemFerry · pp. 11–12

D. Scalability of MemFerry; E. Scaling to scale-up domain.

E. Scaling to scale-up domain.

  • For 16 MB transfers, adaptive routing achieves approximately 86 GB/s, versus approximately 53 GB/s for Direct and 43 GB/s for Assisted.Adaptive sustains more than 80 GB/s for transfer sizes from 16 MB to 256 MB; the evaluation uses four NPUs for training and four as assisting devices.Source: E. Scaling to scale-up domain · p. 12

VII. RELATED WORK

  • Prior work covers CPU/NVMe offloading, DHA or DMA for inference and specific training models, and scale-up interconnect use for inference or distributed training.MemFerry targets general offload training with DHA; ScaleUp-MemFerry uses adaptive multi-path transmission to assist CPU-to-accelerator movement.Source: §VII · pp. 12–13

VIII. CONCLUSION

  • The paper concludes that MemFerry uses hybrid GPU execution to accelerate offload training and reduce memory usage, with GO-MemFerry providing on-demand gradient offloading.The conclusion also identifies DHA, hybrid training patterns, and the shadow-model abstraction as contributions.Source: §VIII · p. 13
Loading 2610.09657v1…