# Evaluating Expansion Memory for Optimizer State Offloading for Large Transformer Models

Moiz Arif, Avinash Maurya, Sudharshan Vazhkudai, Bogdan Nicolae

► To cite this version:

Moiz Arif, Avinash Maurya, Sudharshan Vazhkudai, Bogdan Nicolae. Evaluating Expansion Memory for Optimizer State Offloading for Large Transformer Models. HPAI4S’25: HPC for AI Foundation Models & LLMs for Science (co-located with IPDPS’25), Jun 2025, Milan, Italy. hal-04996269

HAL Id: hal-04996269

https://hal.science/hal-04996269v1

Submitted on 18 Mar 2025

HAL is a multi-disciplinary open access
archive for the deposit and dissemination of sci-
entific research documents, whether they are pub-
lished or not. The documents may come from
teaching and research institutions in France or
abroad, or from public or private research centers.

L'archive ouverte pluridisciplinaire HAL, est
destinée au dépôt et à la diffusion de documents
scientifiques de niveau recherche, publiés ou non,
émanant des établissements d'enseignement et de
recherche français ou étrangers, des laboratoires
publics ou privés.

<figure><img src="image_1.png" /></figure>

Distributed under a Creative Commons Attribution 4.0 International License
# Evaluating Expansion Memory for Optimizer State Offloading for Large Transformer Models

Moiz Arif*, Avinash Maurya†, Sudharshan Vazhkudai*, Bogdan Nicolae†

*Micron Technology Inc., Austin, TX, USA

†Argonne National Laboratory, Lemont, IL, USA

Email: *{marifa, svazhkudai}@micron.com; †{amaurya, bnicolae}@anl.gov

**Abstract**—Transformers and large language models (LLMs) have seen rapid adoption across various domains, with their sizes exploding to hundreds of billions of parameters and have revolutionized natural language processing and other artificial intelligence (AI) applications. However, the increasing complexity of such models poses significant challenges for the available high-bandwidth, and main memory resources. One critical issue is the growth of optimizer states, which often exceed the memory capacity of GPUs. To address this, several strategies have been explored, including offloading optimizer states to main memory and NVMe storage. In this paper, we investigate the use of compute express link (CXL) as an expansion memory solution for offloading large optimizer states. We compare the performance and efficiency of CXL against next-generation NVMe devices, highlighting the potential benefits and limitations of each approach. Our findings show that CXL-based offloading reduces training time by up to 40% compared to NVMe offloading, while also enabling the training of significantly larger models than DRAM only offloading. This research provides valuable insights into optimizing memory management for large-scale transformer models, paving the way for more efficient and scalable AI systems.

**Index Terms**—LLM Training, Scalable AI, Expansion Memory, Heterogeneous Memory Systems, CXL

## I. INTRODUCTION

Large Language Models (LLMs) have revolutionized the field of natural language processing (NLP) by introducing powerful AI systems capable of understanding and generating human-like text with remarkable fluency and coherence. These models, trained on vast amounts of data, are capable of performing a wide range of tasks: language translation, text summarizing, knowledge distillation, etc. In a quest to achieve emergent abilities and reasoning, state-of-the-art LLMs are exploding to large sizes (billions of parameters). GPT [1], LLaMA [2], and DeepSeek [3], [4], primarily built on the transformer architecture [5], are just a few examples. As illustrated in Figure 1 [6], GPU memory capacity has not kept pace with model growth, creating a widening gap. As a consequence, it is increasingly more difficult to pre-train and fine-tune LLMs without aggregating the memory of a large number of GPUs. This trend has resulted in a so called *memory wall* [6], [7] that becomes the main limitation of scalability, as opposed to compute capability.

To democratize LLM training, several state-of-the-art LLM training frameworks (e.g. DeepSpeed [8] and Megatron [9]) introduce a family of techniques that include quantization,

<figure><img src="image_2.png" /><figcaption>Fig. 1. Growth of Models vs. Accelerator Hardware.</figcaption></figure>

redundancy elimination (often referred to as ZeRO), mixed-precision training, and offloading. These techniques offer a cost-effective solution compared to adding multiple GPUs to satisfy the memory demands of LLM training. In this paper, we focus on offloading techniques, particularly optimizer state offloading, which reduces GPU memory usage by transferring the optimizer state to main memory or storage. In a nutshell, the idea is to store the model parameters on the GPU memory in FP16, while a larger FP32 version is stored on the slower tiers such as host memory and disk. The optimizer state (8× larger than the FP16 model parameters) is stored in FP32 on the host memory only. During the forward pass, FP16 activations are computed on the GPUs based on the FP16 model parameters. Similarly, during the backward pass, FP16 gradients are computed on the GPUs based on the FP16 activations. Then, the update phase is executed on the CPU: the FP16 gradients are moved to the host memory and converted to FP32, the new FP32 model parameters are computed based on the FP32 gradients and the FP32 optimizer state, then the FP32 model parameters are converted to FP16 and copied to the GPU memory for the subsequent iteration. By employing this approach, it becomes possible to increase GPU utilization and train large LLMs on a relatively small number of GPUs, albeit with a performance penalty due to the slower data transfers between host and GPU memory.

**Limitations of State-of-the-art Offloading:** Augmenting GPU memory with host memory helps mitigate the memory wall, though it only provides a partial solution. For example, an Amazon AWS p4de.24xlarge instance provides 640 GB of aggregate GPU memory with only 1 TB of host memory, while the ALCF Polaris cluster (Argonne National
Laboratories) offers 320 GB of GPU memory and 512 GB of host memory. Both configurations exhibit a GPU-to-host memory ratio of only 1:1.6, which is not enough to keep up with the model parameter growth trend mentioned earlier. As a consequence, offloading techniques have adopted multi-tiering strategies: in addition to the host memory, tertiary storage such as node-local NVMe devices can be used for additional offloading. Although using node-local NVMe devices significantly lowers the cost of increasing memory capacity, their I/O throughput still lags considerably behind that of host memory. Furthermore, there are additional overheads in managing a tertiary storage tier at all levels of the software stack: (1) the OS kernel exposes NVMe devices as block devices whose I/O involve system calls that trigger expensive context switches; (2) multi-threaded read/write operations for peak I/O throughput of NVMe limit the CPUs' computational capability; (3) multi-GPU systems perform parallel I/O on NVMe creating competition and bottlenecks; and (4) the AI runtime needs to parallelize and coordinate the access to multiple tiers, which involves additional metadata and synchronization overheads. As a consequence, the effective I/O throughput that AI runtimes can achieve for NVMe offloading is even less than the already limited theoretical peak I/O throughput that NVMe devices are capable of. In turn, this leads to a large overall end-to-end training and fine-tuning performance penalty that needs to be mitigated.

**Key Insights and Contributions:** In this paper, we explore ways to alleviate the memory wall and address the limitations of NVMe devices by leveraging emerging memory technologies. Specifically, we explore CXL which is a next-generation, open industry standard interconnect designed to provide high-bandwidth, low-latency connectivity between the host processor and various memory subsystems, including main memory (node-local DRAM), and GPU high-bandwidth memory (HBM). Compared to NVMe, CXL offers higher throughput, reduced latency, and the ability to expand main memory, allowing workloads to access faster memory resources before resorting to slower storage options. Our main contribution is to analyze the challenges and opportunities of adopting CXL offloading vs. NVMe offloading. We summarize our contributions as follows:

1. We enable the use of CXL offloading for DeepSpeed and devise a framework to collect low-level I/O performance metrics at runtime for both our CXL offloading approach and the existing NVMe offloading implementation (§ III).

2. We run extensive experiments in various LLM pre-training scenarios that make use of DeepSpeed. We collect and analyze both end-to-end performance metrics and the low-level I/O performance metrics in order to explain the differences in emerging I/O patterns and their impact on end-to-end runtime between CXL offloading and NVMe offloading. This forms the basis of a discussion that emphasizes the challenges and opportunities of adopting CXL offloading and future possible optimizations (§ IV).

## II. BACKGROUND AND RELATED WORK

**3D parallelism:** Data, pipeline, and tensor parallelism techniques are widely used to accelerate training of large models, including CNNs [10], DLRMs [11], LLMs [12], and LVMs [13]. Data parallelism speeds up training by running multiple model replicas on different mini-batches in parallel, synchronizing gradients during the update step. Pipeline parallelism [14] partitions models vertically across GPUs, assigning different layers to different devices, while tensor parallelism [15] splits individual model tensors horizontally across GPUs. However, as LLMs continue to grow in size, further optimizations beyond 3D parallelism is necessary to enable training in resource-constrained environments.

**Mixed Precision Training:** Mixed precision training, introduced by Baidu and Nvidia research [16], is a widely used technique to improve throughput and reduce memory footprint in LLM training [17], [18]. As shown in Figure 2(a) (two left dotted blocks), this method maintains two copies of model parameters: a low-precision FP16 or BF16 copy used in the forward and backward passes and a high-precision FP32 master copy used by the optimizer for better stability [19]. Activations and gradients are stored in FP16/BF16, enabling faster communication and lower memory usage. During the update step, gradients are upscaled to FP32 for use by the optimizer. Several large-scale LLMs, including BLOOM-176B [18], OPT-175B [17], GPT-3 [1], and GLM-130B [20], have successfully adopted mixed precision training, demonstrating its efficiency, scalability, and stability.

**Data Parallel Redundancy Elimination:** Modern LLM training runtimes such as PyTorch's Fully Sharded Data-Parallel (FSDP) [21], DeepSpeed's ZeRO Redundancy (ZeRO) [22], and Colossal-AI [23] employ redundancy elimination techniques to reduce GPU memory usage by removing duplicate model and optimizer states across data-parallel ranks. DeepSpeed, widely used in large-scale training [18], [24], introduces a three-stage redundancy elimination approach: ZeRO-1 partitions optimizer states, ZeRO-2 extends partitioning to include gradients, and ZeRO-3 further divides model parameters [22]. As shown in Figure 2(a), training a model with $\Phi$ parameters on $K$ GPUs using ZeRO-3 reduces memory requirements to approximately $[O(\Phi)/K]$, but at the cost of a $1.5 \times$ increase in communication overhead [22].

In this paper, we focus on ZeRO due to its high redundancy, making it ideal for training large models on limited resources. Since ZeRO-3 distributes model states across GPUs, it requires frequent scatter-gather operations to reconstruct layers on demand, leading to higher communication costs. Consequently, although it is memory efficient, it is not easily compatible with pipeline parallelism, which relies on effective inter-stage communication. Instead, ZeRO-3 employs a hybrid model and data parallelism, where training on $K$ GPUs effectively creates $K$ "virtual" data-parallel replicas. These replicas are not fully independent, as model states are dynamically synchronized across GPUs rather than being fully duplicated [25], [26].

**Partitioning Model and Optimizer States Into Subgroups:** To further reduce intermediary memory usage during
<figure><img src="image_3.png" /><figcaption>Fig. 2. Illustration of DeepSpeed ZeRO-3 runtime showing (a) hybrid model and data parallelism; wherein the model parameters, activations, gradients, and optimizer states are further sharded into subgroups; (b) memory and storage components for optimizer offloading; and (c) compute node architecture.</figcaption></figure>

ing computation, DeepSpeed’s ZeRO-3 shards large models and optimizer states on each GPU into smaller subgroups, as illustrated in Figure 2(a) (zoom on ZeRO-3 subgroup sharding). This subgroup sharding technique is exclusive to DeepSpeed and is not available in other runtimes such as FSDP or Colossal-AI. The subgroup size, denoted as $τ$, is user-defined and determines the number of parameters per subgroup. Thus, for a model with $Φ$ parameters running on $K$ GPUs under ZeRO-3, each GPU holds approximately $[Φ/K]$ parameters, which are further divided into $[[Φ/K]/τ]$ subgroups, each containing $τ$ parameters. This subgroup-based sharding approach enables efficient piecewise computation and overlapping communication, significantly reducing memory overhead, as detailed in ZeRO-Infinity [25].

**Multi-tier Offloading:** Beyond redundancy elimination, offloading techniques further reduce GPU memory usage by transferring model parameters and optimizer states to larger but slower memory tiers, such as host memory or storage subsystem. Since optimizer states– comprising FP32 parameters, momentum, variance, and gradients– occupy substantial memory but are only needed during the update step, offloading them is an effective strategy in environments with limited GPU memory. In contrast, FP16 model parameters are actively used during forward and backward passes, and since multiple forward-backward steps can be accumulated (gradient accumulation) before an update, offloading model parameters to the slower tiers incurs more significant I/O overheads as compared to offloading the optimizer states. As illustrated in Figure 2(b), ZeRO-3 enables optimizer offloading to main memory, memory expansion subsystems, and disk, keeping FP16 model parameters, activations, and FP16 gradients for a single subgroup on the GPU while storing FP32 optimizer states and gradients on the slower tiers. When offloaded to NVMe, a part of the main memory is reserved as a buffer to facilitate asynchronous read/write to minimize I/O overhead.

**CXL Memory Expansion:** Numerous studies [27]–[29] have demonstrated the advantages of incremental memory capacity expansion over slower disk-based swap storage [28], significantly enhancing workload performance. For bandwidth-intensive workloads, such as AI inference, CXL memory offers higher aggregate memory bandwidth [27], [30],

achieved through software page-level interleaving strategies, which distributes memory allocation between local DRAM and CXL-attached memory in user-specified ratios, thereby optimizing performance. Additionally, various tiered memory management strategies [29], [31], [32] have been proposed to improve access latency for frequently accessed pages based on page temperatures (hotness/coldness). Similarly, CXL memory pooling [33] and sharing enables workloads to dynamically access a shared pool of memory, improving resource utilization and reducing latency, resulting in increased workload performance, especially in data-intensive applications.

### III. ADOPTING CXL FOR LLM PRE-TRAINING

#### A. Hardware Architecture

1) **CXL Protocols and Memory Expansion:** CXL is an industry-standard high-bandwidth, low-latency interconnect designed to extend memory and heterogeneous compute capabilities beyond conventional DRAM-attached architectures. Built atop PCIe infrastructure, CXL offers cache-coherent memory sharing through three key protocols: CXL.io, which manages device discovery, I/O virtualization, and DMA operations; CXL.cache, which allows accelerators to cache host memory, reducing latency and redundant data movement; and CXL.mem, which provides direct load/store access to attached memory devices, treating them as part of the main memory hierarchy. CXL architecture supports three types: Type 1, which is primarily used for caching and memory expansion, providing a bridge between the CPU and other memory devices; Type 2, which combines the main memory and the accelerator memory, allowing for more complex data processing and computation tasks; and Type 3, which focuses on pure memory expansion, offering large-scale memory capacity to support data-intensive workloads. For offloading large LLM training states, we utilize CXL Type 3 as high-bandwidth memory tier alongside DRAM. The CXL.mem protocol is employed to allow direct memory access without software-driven staging, ensuring minimal latency penalties compared to traditional multi-tier offloading solutions like NVMe.

2) **Enabling CXL across System, Firmware, Kernel, and DeepSpeed:** Our system architecture, illustrated in Figure 2(c), consists of four H100 GPUs evenly distributed across two
sockets. Each socket is equipped with dedicated DDR5 memory and CXL memory expansion modules connected over PCIe Gen5. Additionally, two NVMe Gen5 drives are RAID-stripped on socket-0, requiring inter-socket NVMe accesses from socket-1 via the UPI links. Detailed bandwidth and capacity specifications are provided in Section IV.

To enable CXL, the motherboard must have a dedicated PCIe x16 slot with a direct connection to the host processor. The CXL modules use x8 PCIe lanes and connect to the PCIe x16 slot via an EDSFF adapter card. Additionally, modern CPUs such as Intel Sapphire Rapids (SPR) and AMD EPYC Genoa support CXL's memory coherence protocols. Several BIOS changes are necessary to enable CXL on the host, and Linux kernel 6.9+ is required to expose CXL memory as a NUMA node. Our setup includes E3.S form-factor CXL modules and NVMe drives, along with Intel SPR processors running Linux kernel 6.9.1.

At the DeepSpeed level, we enforce NUMA-aware process binding to optimize memory locality. Each training process is explicitly pinned to its corresponding CPU cores, local DRAM, and CXL memory, ensuring that GPU-CPU memory accesses remain intra-socket. While this approach is effective for single-node configurations, in the future, we plan to extend it for multi-node disaggregated memory pools, similar to CXL memory pooling architectures [33], to further enhance I/O performance in distributed settings.

## B. Profiling Application and System-Level Characteristics

To analyze the application-level memory and I/O characteristics of DeepSpeed, we instrument the training engine to capture macro events and collect global timestamps– e.g., forward, backward, and update phases of each GPU/process. Despite being macro-level traces, such low-overhead profiling techniques provide windows of interest from the application perspective that can be further analyzed through system-level characterization to analyze the combined application and system-level behavior.

Subsequently, we obtain in-depth system-level traces corresponding to multi-tier memory capacity utilization, read/write throughput, and GPU/CPU PCIe link utilization to understand at a finer granularity the impact on other resources due to the introduction of the CXL. More specifically, we utilize Nvidia's Management Library, nvml, to evaluate GPU core utilization; and Intel's Performance Counter Monitor (PCM) to obtain bandwidth utilization of the main memory and monitor inter-socket communication. We then correlate and overlap the macro-level time windows from application-level traces to understand the system-level behavior which we detail in the next section.

# IV. PERFORMANCE EVALUATION

## A. Experimental Setup

We conduct our experiments on a single system consisting of 4xH100 PCIe GPUs with 80 GB HBM each, totaling to 320 GB of HBM capacity, 2x Intel Xeon Platinum 8480+ SPR processors with 56 CPUs each (total 112 cores, 224 threads).

The total system memory ranges from 1 TB to 2 TB utilizing configurations of 64 GB, and 128 GB DDR5 modules evenly split across 2 NUMA domains, and shared by consecutive GPU IDs, i.e., GPU0 and GPU1 are mapped to NUMA0, and GPU2 and GPU3 are mapped to NUMA1. Similarly, a total of 1 TB of expansion memory is evenly split across both sockets using 4x 256 GB CXL expansion cards (each appearing as a CPU-less NUMA node). The GPUs are independently connected to the host with PCIe Gen5 interface, providing 55 GB/s unidirectional D2H and H2D throughput for pinned host memory. For pageable host memory, the peak unidirectional D2H and H2D throughput are 16 GB/s and 9 GB/s, respectively. We use 2x NVMe Gen5 drives of 7.68 TB each capable of achieving sequential read of 14GB/s and write of 10GB/s. A striped volume is provisioned on top of these drives to provide higher aggregated read/write performance.

## B. Compared Approaches

To evaluate the trade-offs and benefits of different offloading strategies and optimize memory management for large-scale transformer models, we have designed three offloading backends for the GPU.

**Main Memory (DRAM-Only)**: Evaluate the performance of transferring optimizer states to DRAM memory, excluding the use of expansion memory or storage. The complete memory footprint for offloading remains within the main memory.

**Main and Expansion Memory (DRAM+CXL)**: Evaluate the performance of DRAM and CXL memory where CXL memory is used as in memory expansion capacity without interleaving. The expansion memory approach provides large memory capacity at slightly higher latency compared to the DRAM memory. CXL memory offers lower latency and higher bandwidth compared to NVMe.

**NVMe Storage (NVMe)**: Evaluate the performance of Gen5 NVMe drives for optimizer state offloading. These drives offer higher capacity but incur high latency overheads compared to main and expansion memory.

## C. Methodology and Performance Metrics

**Models:** The configurations of the models used in our evaluations, reflecting real-world LLM training scenarios: 20B [34], 30B [2], 40B [34], 52B [35], and 70B [2]. We consider the case of moderately large models that are large enough to spill over the system memory but small enough to be practically trained on GPU-constrained systems in real-world scenarios.

**Dataset:** For our evaluations, we utilize a subset of the OSCAR-en dataset, comprising 79K records, as included in the Bloom model repository [18]. We employ the default LLaMA2 tokenizer to pre-process the dataset into tokens. Following the OPT training chronicles [17], we use a default sequence length of 2048 and set the micro-batch size to 1 to prevent out-of-memory (OOM) errors in any configuration.

**Key Performance Metrics:** We use the following metrics for evaluating the aforementioned approaches: (1) the average time for computing the forward, backward, and update
<figure><img src="image_4.png" /><figcaption>Fig. 3. Average iteration time breakdown of different models.</figcaption></figure>

<figure><img src="image_5.png" /><figcaption>Fig. 4. Iteration time breakdown for increasing batch sizes.</figcaption></figure>

phases of 10 training iterations for various model sizes; (2) up-
date throughput (expressed as millions of parameters updated
per second); (3) average CPU core utilization observed per
iteration. (4) effective read/write throughput observed while
fetching and flushing disk-offloaded subgroups; and (5) several
system-level statistics for each training phase per iteration
(such as. DRAM read/write throughout, GPU utilization, GPU
memory/bandwidth utilization, PCIe/UPI link utilization, etc.)

D. *Experimental Results*

1) **Breakdown of Iteration Time:** In the first set of experiments, we evaluate iteration time, encompassing the forward pass, backward pass, and update phase, across different model sizes. We test models ranging from 20B to 70B parameters. This analysis focuses on the effects of offloading progressively larger models, which introduce substantial I/O demands– requiring up to 1.6 TB of optimizer states to be loaded and written back to local memory and storage tiers during each iteration. All four GPUs in the node are fully utilized, synchronizing via a barrier at the end of the update step. Figure 3 illustrates the average iteration time breakdown into forward, backward, and update phases for the studied models. As the number of parameters in the model increase, the average iteration time also increases, with the update phase consuming the majority of the time, followed by the backward phase. The forward phase takes a negligible amount of time for all model sizes compared to the overall iteration time because it only needs to operate with the GPU resident data structures and does not offload activations to slower tiers. For models with 20B and 30B parameters, the DRAM+CXL offload shows identical iteration times compared to DRAM only, as the memory footprint of the optimizer states remains within DRAM and does not spill over to CXL memory. For models with 40B parameters and above, the optimizer states start overflowing to CXL memory, and the higher access latency of CXL memory causes a slight slowdown of 7% for small models compared to DRAM only, with the slowdown continuing to increase up to 40% for large models. However, offloading with DRAM+CXL memory expansion results in iteration times that are 22% to 87% lower compared to NVMe offloading, depending on the model size. Overall, CXL expansion memory enables the training of larger parameter models

whose optimizer states would otherwise exceed the available
DRAM capacity, potentially causing training failures, and
significantly outperforming NVMe-based offloading strategies.

2) **Update Throughput:** In the next experiment, we evaluate the update throughput, measured in millions of parameters updated per second, as shown in Figure 5. The highest update throughput is observed when the optimizer state resides entirely in DRAM, which is comparable to the DRAM+CXL configuration for smaller models whose optimizer states can fit entirely in the available system DRAM. As the optimizer states for larger models exceed the available DRAM memory, we utilize higher-capacity DRAMs (128GB instead of 64GB DIMMs) to accommodate the optimizer states of all models, allowing us to capture baseline numbers for DRAM-only configurations. For larger models, as DRAM memory becomes exhausted and more states are offloaded to CXL memory, we observe a lower overall throughput by up to 24%. However, as the optimizer states exceed the available DRAM memory, the throughput begins to drop, depending on the volume of states being offloaded to CXL memory. This drop occurs because of the slightly higher access latency of CXL memory affecting the update speed. Despite this, the throughput of the DRAM+CXL configuration remains 3.6× to 5.3× faster compared to the NVMe offload scenario across all evaluated models. This demonstrates the efficiency of CXL memory in handling larger model sizes, where NVMe offloading would otherwise introduce substantial performance bottlenecks. Lastly, we observe an average throughput degradation of 82% for all studied models when using NVMe offloading compared to the DRAM-only baseline. Overall, CXL-based expansion memory based offloading provides higher update throughput of up to 5.3× compared to NVMe based offloading while reducing the iteration time and total training time.

3) **Scaling Batch Sizes:** Next, we examine the impact of increasing batch sizes on the average iteration time. Due to memory constraints, the mini-batch size is set to one, meaning each mini-batch runs a single forward-backward pass, accumulating gradients that are applied during the update phase once all mini-batches have been processed. In our experiments, we use the 70B model running on 4xH100 GPUs to evaluate the impact on iteration time for global batch sizes ranging from 4 to 256. Figure 4 illustrates that the average
<figure><img src="image_6.png" /><figcaption>Fig. 5. Average update throughput of different models.</figcaption></figure>

<figure><img src="image_7.png" /><figcaption>Fig. 7. Average DRAM bandwidth during training.</figcaption></figure>

iteration time increases as batch sizes grow for DRAM-only, DRAM+CXL, and NVMe offload scenarios. We observe that for all batch sizes, the DRAM-only configuration consistently achieves the lowest iteration time, while the NVMe offload configuration results in the longest iteration time per iteration. The DRAM+CXL configuration, offering higher capacity than DRAM-only, avoids the need for offloading to NVMe but incurs slightly higher latency. The average iteration time doubles with each doubling of the batch size, starting from a batch size of 4 and continuing this trend up to a batch size of 256. This analysis highlights the trade-offs between different memory configurations, emphasizing the efficiency of DRAM-only for smaller batch sizes and the benefits of DRAM+CXL for larger models, where it mitigates the performance penalties associated with NVMe offloading.

4) **End-to-end Training Time:** Next, we evaluate the end-to-end training performance to understand the influence of asynchronous offloading during the update step on subsequent iterations. We measure the total execution time of 10 iterations and report the results in Figure 6. The end-to-end performance trends are similar to that of a single iteration (observed in Figure 3); thereby demonstrating that despite asynchronous data movement, the high throughput of DRAM+CXL approach significantly outperforms the NVMe-based approach. For DRAM only offloading, the training time is the lowest due to its high bandwidth and low latency, ensuring efficient execution of all three steps compared to DRAM+CXL and NVMe offloading. With DRAM+CXL, the forward and backward passes remain relatively efficient, but the update step incurs additional overhead due to the CPU's role in

<figure><img src="image_8.png" /><figcaption>Fig. 6. Total training time for 10 iteration.</figcaption></figure>

<figure><img src="image_9.png" /><figcaption>Fig. 8. Average GPU utilization during training.</figcaption></figure>

managing CXL memory access, leading to increased training time of up to 13% for larger models compared to DRAM only offloading. Compared to NVMe, DRAM+CXL reduces overall training time by up to 40%, while still accommodating the optimizer states of larger models within the system memory. In the case of NVMe offload, the training time is significantly higher across all steps due to the high I/O wait times and latency, which cause prolonged idle periods for the CPU, resulting in the slowest overall training time among the three configurations with up to 95% slowdown compared to DRAM only offloading.

5) **Analysis of System Statistics during Training:** To study the behavior of various system components involved in the training process, we monitor several system performance counters for CPU, GPU, memory, and storage. We breakdown our analysis of performance counters for each forward, backward, and update phase to capture the system behavior and identify bottlenecks.

**Average DRAM Bandwidth During Training:** When analyzing average DRAM bandwidth utilization for optimizer state offloading during training, distinct patterns emerge across different strategies. We observe the DRAM bandwidth throughout the training process and report the results in Figure 7. For DRAM offloading, the average bandwidth utilization remains high and consistent due to direct and efficient memory access, leveraging DRAM's high bandwidth and low latency. As larger parameter models are offloaded, the effective bandwidth increases, with peak bandwidth rising from 300 GB/s for 20B models to 350 GB/s for 70B models. DRAM offloading achieves up to 20% higher average bandwidth compared
<figure><img src="image_10.png" /><figcaption>Fig. 9. Memory footprint for studied models during training.</figcaption></figure>

to other offloading strategies. In the case of DRAM+CXL offloading, while the bandwidth utilization is initially high, it decreases for larger models as more states are offloaded to CXL memory. This results in a drop in average DRAM bandwidth utilization, as some optimizer states are read from CXL memory, which has a lower peak bandwidth compared to DRAM. For NVMe offloading, a portion of the optimizer states are cached in DRAM, depicting similar behaviour to DRAM only with a peak bandwidth of 250 GB/s. However, the overall efficiency is impacted by higher I/O wait times and latency associated with NVMe storage.

**Average GPU Utilization During Training:** When training 20B to 70B models on GPUs with the optimizer state offloaded to the CPU, the average and peak GPU utilization can vary significantly based on the model size and offloading strategy. The results are shown in Figure 8. For smaller models, such as 20B, the peak GPU utilization reaches 100% during the forward and backward phase, as the GPU can quickly access the necessary data and perform computations with minimal delay. The utilization drops during the update phase and depending on the offloading strategy the average utilization varies due to the bottlenecks observed with each strategy. As the model sizes increase to 70B, we observe sustained peak bandwidth of 100% for longer durations as the time for forward and backward phases increases. Larger models also maintain higher average GPU utilization since GPUs must frequently communicate with the CPU to access the optimizer states. Overall, the impact on GPU utilization depends on the efficiency of the data transfer between the GPU and CPU, as well as the ability of the CPU to manage and provide the required optimizer states in a timely manner.

**DRAM Memory Footprint for Offloading:** As model sizes increase from 20B to 70B parameters, the memory required to store the offloaded optimizer states in DRAM also increases. The results are shown in Figure 9. As the model size grows to 70B parameters, the optimizer states require substantially more memory, leading to a larger DRAM footprint. For DRAM+CXL we observe a slightly higher capacity utilization due to redundant data copies and inefficient management of main and expansion memory. For NVMe offload, we observe that a portion of DRAM is used for all studied models because even though the primary storage

<figure><img src="image_11.png" /><figcaption>Fig. 10. Average CPU utilization during the update phase.</figcaption></figure>

for optimizer states is NVMe, some data still needs to be temporarily stored in DRAM to facilitate asynchronous data movement and update computations. We observe that 226GB, 420GB, 720GB, 660GB and 1020GB of NVMe storage are used for offloading optimizer states for 20B, 30B, 40B, 52B, and 70B respectively in addition to the DRAM usage shown in the figure.

**CPU Utilization:** With offloading, the CPU becomes the core coordinator responsible for the orchestration of (a) GPU to CPU gradient transfer; (b) up-scaling FP16 gradients to FP32 such that the optimizer can use them; (c) prefetching the optimizer parameter, momentum, variance, and gradient in subgroup-by-subgroup fashion (if offloaded to NVMe tier); (d) performing gradient clipping to prevent overflow; (e) running update step on a single subgroup; (f) transferring the updated parameters from host memory to the GPU memory for use in subsequent iterations; and (g) lazily flushing out the subgroup of the optimizer state back to slower tier (if offloaded to NVMe tier). Therefore, in the last set of experiments, we study the CPU utilization to analyze how efficiently the CPU core cycles are used by different approaches. As shown in Figure 10, given the low read/write throughput of DRAM and DRAM+CXL approaches, we observe significantly higher CPU utilization as compared to the case of NVMe offloading, which incurs higher wait times. For models larger than 40B in size, we observe that the DRAM+CXL approach utilizes CPU more as compared to the DRAM-only approach. This is due to the progressively larger number of states being offloaded to CXL memory in proportion to the model size, as well as the CPU's role in coordinating and managing data transfers between the CXL memory and the GPU.

## V. CONCLUSION AND FUTURE WORK

The rapid adoption of transformers and LLMs across various domains has led to models with hundreds of billions of parameters, revolutionizing NLP and other AI applications. However, this increasing complexity poses significant challenges for high-bandwidth and main memory resources, particularly due to the growth of optimizer states that exceeds GPU memory capacity. To address this, strategies such as offloading optimizer states to main memory and NVMe storage have been explored. In this paper, we investigate the use of CXL-based expansion
memory for offloading optimizer states and compare it to NVMe-based offload and achieve up to 40% lower training time on average with CXL memory expansion compared to NVMe. Similarly, we can train higher parameter models with CXL memory expansion over DRAM to avoid training failures. In the future, we aim to optimize the management of CXL memory through intelligent prefetching and caching techniques for optimizer states to mask the overhead of accessing CXL memory. Additionally, we plan to develop sophisticated algorithms that dynamically balance the distribution of optimizer states between DRAM, CXL memory, and NVMe storage, optimizing memory bandwidth and minimizing access latency. These advancements will facilitate the efficient scaling of large language models, ensuring robust performance and resource utilization in high-demand AI applications.

## ACKNOWLEDGMENTS

This material is based upon work supported by the U.S. Department of Energy (DOE), Office of Science, Office of Advanced Scientific Computing Research, under Contracts DE-AC02-06CH11357. Additionally, this work is partially supported by the National Science Foundation (NSF), Office of Advanced Cyberinfrastructure, under Grant CSSI-2411386.

## REFERENCES

[1] T. B. Brown, B. Mann, N. Ryder, M. Subbiah, J. Kaplan et al., "Language models are few-shot learners," in Proc. of the 34th International Conference on Neural Information Processing Systems, 2020.

[2] H. Touvron, L. Martin, K. Stone, P. Albert, A. Almahairi et al., "Llama 2: Open foundation and fine-tuned chat models," arXiv preprint arXiv:2307.09288, 2023.

[3] D. Dai, C. Deng, C. Zhao, R. Xu, H. Gao et al., "Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models," arXiv preprint arXiv:2401.06066, 2024.

[4] DeepSeek-AI, D. Guo, D. Yang, H. Zhang, J. Song et al., "Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning," arXiv preprint arXiv:2501.12948, 2025.

[5] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones et al., "Attention is all you need," in Proc. of the 31st International Conference on Neural Information Processing Systems, 2017, p. 6000-6010.

[6] A. Thompson, "Inside large language models," https://lifearchitect.ai/models, 2024.

[7] A. Gholami, Z. Yao, S. Kim, C. Hooper, M. W. Mahoney, and K. Keutzer, "AI and Memory Wall," IEEE Micro, pp. 33-39, 2024.

[8] J. Rasley, S. Rajbhandari, O. Ruwase, and Y. He, "Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters," in Proc. of the SIGKDD International Conference on Knowledge Discovery & Data Mining, 2020, pp. 3505-3506.

[9] M. Shoeybi, M. Patwary, R. Puri, P. LeGresley, J. Casper, and B. Catan-zaro, "Megatron-lm: Training multi-billion parameter language models using model parallelism," arXiv preprint arXiv:1909.08053, 2019.

[10] X. Zhao, L. Wang, Y. Zhang, X. Han, M. Deveci, and M. Parmar, "A review of convolutional neural networks in computer vision," Artificial Intelligence Review, vol. 57, no. 4, p. 99, 2024.

[11] S. Zhang, L. Yao, A. Sun, and Y. Tay, "Deep learning based recommender system: A survey and new perspectives," ACM computing surveys (CSUR), vol. 52, no. 1, pp. 1-38, 2019.

[12] W. X. Zhao, K. Zhou, J. Li, T. Tang, X. Wang, Y. Hou, Y. Min, B. Zhang, J. Zhang, Z. Dong et al., "A survey of large language models," arXiv preprint arXiv:2303.18223, 2023.

[13] Y. Jiang, X. Yan, G.-P. Ji, K. Fu, M. Sun, H. Xiong, D.-P. Fan, and F. S. Khan, "Effectiveness assessment of recent large vision-language models," Visual Intelligence, vol. 2, no. 1, p. 17, 2024.

[14] Y. Huang, Y. Cheng, A. Bapna, O. Firat, M. X. Chen et al., GPipe: efficient training of giant neural networks using pipeline parallelism. Curran Associates Inc., 2019.

[15] Z. Cai, X. Yan, K. Ma, Y. Wu, Y. Huang et al., "Tensoropt: Exploring the tradeoffs in distributed dnn training with auto-parallelism," IEEE Transactions on Parallel and Distributed Systems, vol. 33, no. 8, pp. 1967-1981, 2021.

[16] P. Mickevicius, S. Narang, J. Alben, G. Diamos, E. Elsen et al., "Mixed precision training," 2018. [Online]. Available: https://arxiv.org/abs/1710.03740

[17] S. Zhang, S. Roller, N. Goyal, M. Artetxe, M. Chen, S. Chen, C. Dewan, M. Diab, X. Li, X. V. Lin et al., "Opt: Open pre-trained transformer language models," arXiv preprint arXiv:2205.01068, 2022.

[18] B. Workshop, ;, T. L. Scao, A. Fan, C. Akiki, E. Pavlick et al., "BLOOM: A 176B-Parameter Open-Access Multilingual Language Model," 2023. [Online]. Available: https://arxiv.org/abs/2211.05100

[19] J. Lee, J. Bae, B. Kim, S. J. Kwon, and D. Lee, "To fp8 and back again: Quantifying the effects of reducing precision on llm training stability," 2024. [Online]. Available: https://arxiv.org/abs/2405.18710

[20] A. Zeng, X. Liu, Z. Du, Z. Wang, H. Lai et al., "GLM-130B: An Open Bilingual Pre-trained Model," 2023. [Online]. Available: https://arxiv.org/abs/2210.02414

[21] Y. Zhao et al., "Pytorch fsdp: Experiences on scaling fully sharded data parallel," VLDB Endow., vol. 16, no. 12, p. 3848-3860, Aug. 2023. [Online]. Available: https://doi.org/10.14778/3611540.3611569

[22] S. Rajbhandari, J. Rasley, O. Ruwase, and Y. He, "Zero: Memory optimizations toward training trillion parameter models," in Proc. of the International Conference for High Performance Computing, Networking, Storage and Analysis (SC), 2020.

[23] S. Li, H. Liu, Z. Bian, J. Fang, H. Huang et al., "Colossal-ai: A unified deep learning system for large-scale parallel training," 2023. [Online]. Available: https://arxiv.org/abs/2110.14883

[24] S. L. Song, B. Kruft, M. Zhang, C. Li, S. Chen et al., "DeepSpeed4Science Initiative: Enabling Large-Scale Scientific Discovery through Sophisticated AI System Technologies," 2023. [Online]. Available: https://arxiv.org/abs/2310.04610

[25] S. Rajbhandari, O. Ruwase, J. Rasley, S. Smith, and Y. He, "Zero-infinity: Breaking the gpu memory wall for extreme scale deep learning," in Proc. of the International Conference for High Performance Computing, Networking, Storage and Analysis (SC), 2021.

[26] J. Ren, S. Rajbhandari, R. Y. Aminabadi, O. Ruwase, S. Yang, M. Zhang, D. Li, and Y. He, "Zero-offload: Democratizing billion-scale model training," in Proc. USENIX Annual Technical Conference (ATC), 2021.

[27] M. Ahn, A. Chang, D. Lee, J. Gim, J. Kim et al., "Enabling cxl memory expansion for in-memory database management systems," in Proc. of the 18th International Workshop on Data Management on New Hardware (DaMoN), 2022.

[28] M. Arif, K. Assogba, M. M. Rafique, and S. Vazhkudai, "Exploiting cxl-based memory for distributed deep learning," in Proc. of the 51st International Conference on Parallel Processing (ICPP), 2023.

[29] M. Arif, A. Maurya, M. M. Rafique, D. S. Nikolopoulos, and A. R. Butt, "Application-attuned memory management for containerized hpc workflows," in Proc. IEEE International Parallel and Distributed Processing Symposium (IPDPS), 2024.

[30] J. Wahlgren, M. Gokhale, and I. B. Peng, "Evaluating emerging cxl-enabled memory pooling for hpc systems," in Proc. Workshop on Memory Centric High Performance Computing (MCHPC), 2022.

[31] H. A. Maruf, H. Wang, A. Dhanotia, J. Weiner, N. Agarwal et al., "Tpp: Transparent page placement for cxl-enabled tiered-memory," in Proc. of the Intl. Conf. on Architectural Support for Programming Languages and Operating Systems, 2023.

[32] S. Kumar et al., "Radiant: Efficient page table management for tiered memory systems," in ACM SIGPLAN International Symposium on Memory Management, 2021.

[33] H. Li, D. S. Berger, L. Hsu, D. Ernst, P. Zardoshti et al., "Pond: Cxl-based memory pooling systems for cloud platforms," in Proc. of the International Conference on Architectural Support for Programming Languages and Operating Systems, 2023.

[34] Z. Wang, Z. Jia, S. Zheng, Z. Zhang, X. Fu, T. E. Ng, and Y. Wang, "Gemini: Fast failure recovery in distributed training with in-memory checkpoints," in Proc. of the 29th Symposium on Operating Systems Principles (SOSP), 2023.

[35] X. Li, Y. Yao, X. Jiang, X. Fang, C. Wang, X. Liu, Z. Wang, Y. Zhao, X. Wang, Y. Huang, S. Song, Y. Li, Z. Zhang, B. Zhao, A. Sun, Y. Wang, Z. He, Z. Wang, X. Li, and T. Huang, "Tele-flm technical report," 2024.