ArXiv: 2403.06504
🎯 Pitch
LoHan is the first system to fine-tune a 175B-parameter LLM on a single consumer RTX 4090 GPU with just 256 GB of main memory—achieving 2.17× higher cost-effectiveness than a DGX-A100 cluster. The key insight is that jointly optimizing the timing and volume of tensor offloading to CPU RAM and SSDs, rather than treating activation and optimizer state offloading as separate problems, eliminates the bottleneck that previously made this impossible.
1. Executive Summary
This paper proposes LoHan, a low-cost high-performance deep learning training framework that enables efficient 100B-scale model fine-tuning on a single consumer-grade GPU with limited main memory by adding holistic offloading traffic as a new optimization dimension. The system introduces two named mechanisms: active gradient offloading (allowing the CPU optimizer to directly consume gradients as they arrive from GPU during backward propagation, overlapping optimizer execution with GPU computation) and holistic traffic-aware activation swapping management (automatically determining the optimal amount of activations to offload to main memory versus SSDs by modeling iteration time as a convex function of offloaded activation size). Evaluated on decoder-only LLMs up to 175B parameters using NVIDIA RTX 4090 GPUs, LoHan achieves up to 2.32× throughput over state-of-the-art baselines when fine-tuning a 13B model, is the first system to fine-tune a 175B model on an RTX 4090 with only 256 GB main memory, and enables a consumer GPU server to achieve up to 2.17× higher cost-effectiveness than a DGX-A100 cluster — establishing that consumer-grade hardware can economically fine-tune 100B-scale models only when both activation and model state offloading are jointly optimized as a holistic scheduling problem rather than treated as independent mechanisms.
2. Context and Motivation
The Core Problem: 100B-Scale Fine-Tuning Is Economically Inaccessible
The fundamental question this paper tackles is intensely practical: can a data scientist with a single consumer-grade GPU — the kind that costs ~$1600 and sits in a commodity server — fine-tune a 100-billion-parameter language model? This matters because the dominant approach to handling large models has been to throw hardware at the problem, and that hardware is extraordinarily expensive.
To understand the scale of the gap, consider what fine-tuning a 100B model actually demands. As the paper lays out in Section II, a training iteration requires storing multiple categories of tensors simultaneously: model parameters (, 4 bytes per parameter), optimizer states (, 8 bytes per parameter for Adam's first and second moment estimates), gradients (, 2 bytes per parameter in half precision), a low-precision parameter copy for GPU computation (, 2 bytes per parameter), and activations (, whose size depends on batch size, sequence length, and model architecture). For a 175B model, the model states alone occupy approximately TB at full precision for the optimizer states and parameters, or roughly 2.45 TB in the mixed-precision regime the paper uses. The latest data-center GPU — an NVIDIA H200 — provides at most 188 GB of device memory (as noted in the paper's introduction, citing NVIDIA's specifications). That is a 13× gap between what the model needs and what the best available single GPU offers.
The conventional solution — aggregating device memory across many GPUs in a high-end cluster — is technically effective but economically devastating for individual researchers. The paper cites a concrete example: it takes 32 A100 GPUs (each with 80 GB) to fine-tune a 100B model using standard data-parallel or tensor-parallel approaches. At the time of writing, a single DGX-A100 server with 8 A100 GPUs costs approximately $200,000 (Table VII), and renting equivalent cloud instances runs thousands of dollars per day. For a graduate student, a researcher at a small company, or a data scientist in a developing economy, these costs are prohibitive. The paper frames this as an accessibility crisis: the ability to work with state-of-the-art models is concentrating in organizations that can afford data-center-scale infrastructure, while the vast majority of AI practitioners are locked out.
This gap is not merely about money — it is about who gets to participate in frontier AI research. If fine-tuning a 100B model requires institutional-scale hardware, then only well-funded labs can adapt large pre-trained models to new domains, languages, or tasks. The paper's framing here is deliberately democratic: by targeting a single consumer GPU (the RTX 4090 at $1600) with commodity main memory (256 GB, which is a standard configuration for a high-end desktop), the authors are asking whether the hardware that many individual researchers already own — or could reasonably purchase — can be pushed to handle workloads that currently require cluster-scale deployments.
Why Existing Solutions Fail: Three Categories, Each with Fatal Flaws
The paper systematically categorizes prior offloading-based approaches into three groups and identifies why each falls short when scaled to 100B models on consumer hardware.
Category 1: Offloading Only Activations to SSDs (FlashNeuron)
The simplest approximation to memory-limited training is to keep model states (parameters, optimizer states, gradients) in GPU memory and offload only activations to slower storage. FlashNeuron (Bae et al., 2021) represents this approach, using NVMe SSDs as a backing store for activation checkpoints. The paper implements a prototype of FlashNeuron and tests it on their evaluation server (detailed in Section V-A: dual Intel Xeon Gold 5320, 768 GB main memory, RTX 4090 with 24 GB device memory, 12 Intel P5510 SSDs).
The results are stark and unambiguous (Section III-A, Figure 2a): FlashNeuron fails to fine-tune even a 6B model. The reason is straightforward — keeping model states on GPU memory is the binding constraint, not activation storage. A 6B model in mixed precision requires approximately GB just for model states, which already exceeds the RTX 4090's 24 GB by a factor of 4. FlashNeuron's premise — that activation offloading is the bottleneck — was developed in an era when models were smaller and GPU memory was the limiting factor for activations rather than parameters. For 100B-scale models, the model states themselves are the elephant in the room, and any system that does not offload them is fundamentally incapable of scaling to large models regardless of how clever its activation management is.
Category 2: Offloading Model States to SSDs (ZeRO-Infinity, Colossal-AI)
A more sophisticated approach, exemplified by ZeRO-Infinity (Rajbhandari et al., 2021) and Colossal-AI (Bian et al., 2023), offloads model states to SSDs while keeping a working set in GPU memory. These systems also introduce the CPU optimizer — rather than executing the Adam optimizer on the GPU (which would require all optimizer states to be present in GPU memory), they transfer gradients to main memory and run the optimizer on the CPU, writing updated parameters and optimizer states directly from main memory to SSDs. This eliminates the heavy back-and-forth transfer of and between GPU and SSDs that a GPU-resident optimizer would require.
However, the paper identifies three specific and severe issues with these systems when deployed on consumer hardware (Section III-B). Importantly, the authors note that these systems were "originally designed for high-end DGX servers rather than for a commodity server with a single consumer-grade GPU," and their design assumptions break down dramatically in the low-resource setting.
Issue 1: Heavy Optimizer Execution Overhead. ZeRO-Infinity and Colossal-AI execute the CPU optimizer in a separate, serialized stage after the backward propagation of the entire model finishes. During this optimizer stage, the GPU sits completely idle. On a DGX server with many high-end CPUs, this optimizer stage is trivially fast relative to GPU computation. But on a commodity server with a dual-socket Xeon configuration (the paper's evaluation server uses dual Intel Xeon Gold 5320, a mid-range server CPU), the optimizer becomes a significant bottleneck.
The paper provides quantitative evidence in Figure 2b and 2c. When fine-tuning a 13B model with ZeRO-Infinity, the GPU is busy during only 36% of an iteration — meaning 64% of the time, the most expensive component in the system is doing nothing. The optimizer stage alone consumes 30–60% of the total training step time (Figure 2c), with the proportion growing as model size increases because larger models have proportionally more optimizer states to update. Colossal-AI is even worse: the GPU is busy for only 12% of an iteration. This is an enormous waste of compute resources — the consumer may have paid $1600 for a GPU that sits idle for the majority of training time.
Issue 2: Excessive Activation Recomputation Overhead. ZeRO-Infinity adopts a static, one-size-fits-all activation management policy: it offloads only inter-transformer block activations (the activations that pass between transformer layers) to main memory, while recomputing all intra-transformer block activations (the intermediate values within each attention and feed-forward computation) during backward propagation. The paper quantifies this for a 13B model with batch size 32: inter-block activations account for only 12.5 GB — roughly 6% of total activations — while intra-block activations account for 200 GB. ZeRO-Infinity recomputes that 200 GB from scratch during backward propagation.
Recomputation is not free. It requires the GPU to re-execute forward-pass computations for each layer during the backward pass to regenerate the activations that were discarded. The paper's breakdown in Figure 1a shows that for a 13B model fine-tuned with ZeRO-Infinity, GPU recomputation during the backward stage takes 5.7 seconds — approximately 22% of the backward stage time. Meanwhile, the PCIe links (both GPU-to-main-memory and main-memory-to-SSD) are underutilized during this same period: GPU-to-main-memory transfers take only 3.18 seconds, and SSD I/O takes 6.25 seconds. The system has abundant I/O bandwidth available but chooses to waste GPU cycles on recomputation instead of using that bandwidth to fetch pre-computed activations.
The deeper problem here is that ZeRO-Infinity's static policy was designed for a different hardware profile. On high-end DGX servers with NVLink-connected GPUs, GPU-to-GPU communication is fast enough that recomputation is often cheaper than offloading. On a single consumer GPU with a single PCIe Gen 4 link to main memory and multiple SSDs, the tradeoff inverts — the aggregate I/O bandwidth from multiple SSDs can potentially exceed what the GPU can saturate, making activation offloading more attractive than recomputation for many layers. But ZeRO-Infinity has no mechanism to adapt to this different cost structure.
Issue 3: Limited Trainable Model Size Due to Main Memory Constraints. ZeRO-Infinity offloads activations only to main memory — it never writes them to SSDs. Colossal-AI is even more conservative, keeping activations entirely in GPU memory and main memory without any SSD offloading. This means the maximum trainable model size is bounded not by the (effectively unlimited) SSD capacity, but by the main memory capacity of the server.
The paper calculates (implicitly, from the architecture) that ZeRO-Infinity would require approximately 1.1 TB of main memory to fine-tune a 175B model, because all inter-block activations must fit in main memory simultaneously. Most commodity servers ship with 128 GB to 1 TB of main memory — the paper's own evaluation server has 768 GB — meaning ZeRO-Infinity fails to fine-tune a 175B model on this hardware. This is confirmed experimentally in Figure 2a: even with 768 GB of main memory, ZeRO-Infinity cannot fine-tune a 175B model when the batch size is set to 1. The fundamental constraint is that the system treats main memory as the final tier in the storage hierarchy, when for 100B-scale models on consumer hardware, SSDs must be part of the activation storage path.
Category 3: Naïvely Offloading Both Model States and Activations to SSDs (G10)
The most recent approach, G10 (Zhang et al., 2023), offloads both model states and activations to a unified main-memory/NVMe storage pool, theoretically supporting 100B-scale models with limited GPU and main memory. G10 also executes the Adam optimizer on the GPU rather than the CPU, following the conventional in-GPU training paradigm.
The paper identifies three fatal issues with G10 in the consumer GPU context (Section III-C):
Issue 1: Heavy Model State Transfer Overhead. Because G10 executes the optimizer on GPU, every optimizer step requires transferring all model states — parameters (), optimizer states (), and the low-precision parameter copy () — from SSDs to GPU memory for the optimizer to operate on them, then transferring the updated versions back to SSDs. For a 13B model, this amounts to 182 GB transferred in each direction during the optimizer stage alone. The paper's simulation (Figure 1b) shows that even under ideal conditions where GPU computation and PCIe transfer are fully pipelined, the GPU optimizer computation takes only 0.1 seconds while the PCIe transfer takes 13 seconds — 37% of the entire iteration. The GPU is effectively starved for data during the optimizer stage, waiting on a single PCIe link to shuffle massive tensors back and forth.
Issue 2: High Activation Transfer Overhead. G10 offloads all activations to SSDs without any recomputation. This is the opposite extreme from ZeRO-Infinity: rather than recomputing 94% of activations, G10 transfers 100% of them over PCIe. For a 13B model with batch size 32, this means 213 GB of activation data traverses the GPU-to-main-memory PCIe link during the forward stage alone. The paper's breakdown (Figure 1b) shows that activation offloading takes 10 seconds during the forward stage, while GPU computation takes only 5.96 seconds — the GPU is idle for nearly half the forward stage because it finishes computing faster than the PCIe link can drain the produced activations.
Issue 3: GPUDirect Dependency. G10 relies on GPUDirect — an NVIDIA technology that enables direct data transfers between GPUs and NVMe SSDs without staging through main memory — to achieve its performance. However, GPUDirect is not available on consumer-grade GPUs. The RTX 4090, RTX 3090, and RTX 4080 that the paper targets do not support this feature, which is restricted to data-center GPUs like the A100 and H100. This means G10 literally cannot run on the hardware the paper targets, making it a non-starter for the consumer GPU scenario regardless of its theoretical performance.
The Deeper Pattern: Lack of Holistic Intra-Server Tensor Movement Management
The paper's diagnosis cuts across all three categories: existing systems fail not because of any single suboptimal design choice, but because they treat activation management, model state offloading, and optimizer execution as independent mechanisms to be optimized in isolation. This leads to three specific coordination failures:
-
Temporal fragmentation of resource usage. ZeRO-Infinity serializes GPU computation (forward/backward), CPU computation (optimizer), and SSD I/O (model state persistence) into separate, non-overlapping stages. During the optimizer stage, the GPU sits idle; during the backward stage, CPU cores and SSD bandwidth are underutilized (as shown in Figure 1a where PCIe G2M takes only 3.18 seconds while GPU recomputation takes 17.6 seconds). The system has abundant aggregate resources but cannot use them simultaneously because the scheduling is stage-granular rather than fine-grained.
-
Spatial misallocation of storage tiers. ZeRO-Infinity treats main memory as the only destination for activations, ignoring the much larger SSD capacity. G10 treats SSDs as the destination for everything, ignoring the much faster main memory bandwidth. Neither system adapts its storage tier assignment based on the actual capacity and bandwidth characteristics of the specific hardware it runs on. A proper solution would use main memory for high-priority, frequently accessed tensors and SSDs for larger, less latency-sensitive tensors, dynamically adjusting based on available capacity in each tier.
-
Workload-oblivious tradeoff between offloading and recomputation. ZeRO-Infinity uses a static rule: offload inter-block activations to main memory, recompute everything else. G10 uses the opposite static rule: offload everything, recompute nothing. Neither system considers that the optimal tradeoff depends on the specific layer structure (some layers have large activations relative to their computation cost, making them better candidates for offloading), the batch size (which changes the relative size of activations versus model states), and the available PCIe and SSD bandwidth at deployment time.
Where This Paper Positions Itself
Given this landscape, LoHan does not propose a fundamentally new offloading mechanism — both activation offloading to SSDs and CPU-based optimizer execution exist in prior work. Rather, it identifies holistic offloading traffic as a first-class optimization dimension that has been overlooked, and proposes that jointly scheduling model state offloading, activation offloading, activation recomputation, and optimizer execution can unlock efficient 100B-scale fine-tuning on consumer hardware.
The paper's intellectual contribution is the recognition that these mechanisms interact in non-trivial ways that prior systems' modular designs miss. Specifically:
-
Optimizer execution can be overlapped with backward propagation if gradients are actively consumed as they are produced, rather than buffered until the backward pass completes. This is the insight behind active gradient offloading, which borrows conceptually from the Active Messages paradigm in distributed computing (Eicken et al., 1992) but applies it to the intra-server tensor movement problem.
-
The optimal amount of activation offloading versus recomputation is a convex function of the offloaded activation size, meaning there exists a unique global minimum for iteration time that can be found by hill-climbing search. This converts what prior work treated as a heuristic or static policy into a solvable optimization problem with provable structure.
-
SSDs must be integrated into the activation storage path for 100B-scale models on consumer hardware, but activation offloading to SSDs and model state offloading to SSDs compete for the same PCIe and SSD bandwidth. A holistic scheduler must account for this contention rather than treating activation and model state offloading as independent knobs.
The paper positions LoHan as a systems contribution rather than an algorithmic one — it does not change the mathematics of fine-tuning or propose new model architectures. Instead, it demonstrates that careful resource scheduling can push consumer hardware an order of magnitude beyond what prior systems achieved, making 100B-scale fine-tuning economically accessible for the first time. The paper explicitly targets "data scientists with a limited budget for high-end GPU servers" (Section I), framing its contribution as democratizing access to large-model adaptation rather than advancing the state of the art in absolute training throughput.
3. Technical Approach
3.1 Reader Orientation
LoHan is a deep learning training framework — specifically, a runtime system built on top of PyTorch — that manages where tensors live (GPU memory, main memory, or NVMe SSDs) and when they move between those locations during LLM fine-tuning. It solves the problem that existing offloading systems treat model state offloading, activation offloading, activation recomputation, and optimizer execution as independent mechanisms, leading to idle GPU time, wasted PCIe bandwidth, and unnecessarily small maximum trainable model sizes. The shape of LoHan's solution is a holistic scheduler that adds the aggregate traffic across all PCIe links as an explicit optimization dimension: it overlaps optimizer execution with backward propagation by actively consuming gradients as they are produced, and it automatically finds the optimal amount of activation offloading by solving a convex optimization that accounts for the interaction between activation transfers, model state transfers, GPU computation, and available main memory capacity.
3.2 Big-Picture Architecture (Diagram in Words)
LoHan has four major components arranged in a pipeline that executes once at initialization and then governs every subsequent training iteration:
-
Hardware-Aware Profiler (Section IV-B): runs in the first training iteration only, measuring hardware capabilities (peak GPU throughput, PCIe bandwidth in both directions, SSD I/O bandwidth, available main memory) and model characteristics (parameter count, activation sizes per layer, FLOP counts per layer). Outputs a set of numerical constants consumed by the other components.
-
Active Gradient Offloading Engine (Section IV-C): replaces the standard serialized optimizer stage with a pipelined execution model where the CPU optimizer begins processing gradients for layer as soon as those gradients arrive in main memory from the GPU, overlapping CPU computation and SSD I/O with ongoing GPU backward propagation for subsequent layers. This component eliminates the 30–60% GPU idle time that afflicts ZeRO-Infinity.
-
Holistic Traffic-Aware Activation Swapping Manager (Section IV-D): given the profiling data and the active gradient offloading schedule, this component solves for the optimal subset of activations to offload from GPU memory (rather than recompute) by modeling per-iteration time as a convex function of offloaded activation volume and finding the global minimum. It also determines how much of the offloaded activation data goes to main memory versus SSDs based on available main memory capacity.
-
Framework Integration Layer (Section IV-E): wraps PyTorch's native APIs (
torch.optim, model forward/backward) so that users can enable LoHan's optimizations with minimal code changes — essentially replacing the standard optimizer step and adding a LoHan initialization call — while the profiler, gradient offloading, and activation management operate transparently through PyTorch operator hooks.
Information flows as follows: at iteration 0, the profiler measures all hardware and model parameters → the activation manager computes the optimal swapped activation set and the fraction destined for SSDs → for all subsequent iterations, during the forward pass, activations are either retained in GPU memory, offloaded to main memory, or offloaded to SSDs according to the computed policy → during the backward pass, as each layer's gradients are computed on GPU, they are simultaneously transferred to main memory where the active gradient offloading engine triggers the CPU optimizer to fetch the corresponding model states from SSDs, execute the Adam update, and write the results back → activations needed for backward computation are prefetched from main memory or SSDs before the GPU needs them. The optimizer stage as a separate, GPU-idle phase is eliminated entirely.
3.3 Roadmap for the Deep Dive
The deep dive follows the dependency order of the components — each subsequent mechanism depends on data produced by the previous one:
- First, the Hardware-Aware Profiler (Section IV-B): this is the foundation because every downstream decision (how much to offload, whether to use SSDs or main memory, how much overlap is possible) requires knowing the actual bandwidth and capacity numbers of the specific hardware.
- Second, Active Gradient Offloading (Section IV-C): this is the simpler of the two core innovations and directly addresses the optimizer-idle-GPU problem. I explain the naïve approach first, then the optimized version with its overlapping schedule, because the optimized version's resource usage pattern feeds into the activation manager's timing model.
- Third, the Holistic Traffic-Aware Activation Swapping Manager (Section IV-D): this is the most mathematically involved component. I walk through how iteration time is computed from the profiler's outputs and the chosen offloading volume, prove convexity, derive the three possible cases, and show the concrete search algorithm. I treat every equation with the four-part pattern (display, define symbols, operational meaning, why this form).
- Fourth, Framework Integration (Section IV-E): brief, shows how the preceding components are wired into PyTorch so users don't need to manage any of this manually.
This ordering builds understanding sequentially: you cannot understand the activation manager's timing equations without first knowing what active gradient offloading does to the backward-stage resource usage.
3.4 Detailed, Sentence-Based Technical Breakdown
This is a systems paper whose core idea is that the tensor movement patterns in LLM fine-tuning — gradients flowing from GPU to CPU, model states flowing from SSDs to CPU and back, activations flowing from GPU to main memory and optionally to SSDs — should be treated as a single holistic scheduling problem rather than as independent mechanisms with fixed policies. The two named innovations (active gradient offloading and holistic traffic-aware activation swapping) are concrete instantiations of this principle for the optimizer and activation subsystems respectively.
Hardware-Aware Profiling
The profiler runs automatically during the very first training iteration and collects two categories of data: hardware capabilities and model characteristics. Its purpose is to provide the numerical constants that the activation swapping manager's optimization problem requires, so that LoHan can adapt to whatever specific GPU, CPU, main memory configuration, and SSD array the user has — rather than relying on hardcoded assumptions about relative bandwidths.
What the profiler measures (hardware side):
-
$THP_{\rm G}$: the peak GPU throughput in floating-point operations per second (FLOPS). This is measured by benchmarking a single transformer block running entirely inside GPU memory with no PCIe traffic — the profiler records the computation time per layer during the forward pass, sums the known FLOP counts, and divides to get effective throughput. This number captures the GPU's compute ceiling under ideal conditions and is used to determine whether GPU computation or data movement is the bottleneck at any given configuration. -
$BW_{\rm G}$: the maximum unidirectional PCIe bandwidth between GPU and main memory. The profiler monitors PCIe traffic during the profiling iteration to estimate this. Because the GPU-to-CPU PCIe link is full-duplex (can transfer in both directions simultaneously), the profiler treats upstream and downstream bandwidth as potentially independent resources — the activation manager's timing model uses$BW_{\rm G}$separately for GPU-to-main-memory and main-memory-to-GPU transfers rather than treating them as a shared pool. -
$BW_{\rm S2M}$and$BW_{\rm M2S}$: the maximum PCIe bandwidth from SSDs to main memory and from main memory to SSDs respectively. Unlike the GPU-main memory link, SSD I/O is simplex — reads and writes compete for the same bandwidth — so the profiler measures these as a single aggregate number. The profiler also discovers the system topology (which SSDs are connected to which PCIe lanes) to avoid contention when multiple SSDs are used simultaneously. -
$MEM_{\rm M}^{\rm avail}$: the minimum unallocated main memory observed during the profiling stage. This is critical because it determines how much activation data can be buffered in main memory rather than pushed to SSDs — the activation manager uses this as a hard capacity constraint.
What the profiler measures (model side):
-
$P$: total number of model parameters. Obtained by parsing the PyTorch model definition at initialization — everynn.Parameteris counted. This determines the sizes of all model state tensors ($\tt P_{32}$,$\tt{OS}_{32}$,$\tt G_{16}$,$\tt P_{16}$) via the fixed byte-per-parameter ratios in Table II. -
$A_{\rm all}$: total size of all activations in bytes for one training sample. Also obtained by parsing the model graph — each operation's output tensor size is known from the layer dimensions and sequence length (fixed at 1024 in all experiments). This is the upper bound on how much activation data could potentially be offloaded. -
$FLOP_{\rm f}$: number of floating-point operations during the forward pass. Obtained by instrumenting each layer to count operations. The backward pass requires approximately$2 \times FLOP_{\rm f}$operations (once for the gradient with respect to activations, once for the gradient with respect to parameters), which is a standard property of reverse-mode automatic differentiation. -
Per-layer activation sizes
$A_{\rm layer}$and recomputation FLOP counts$FLOP_{\rm layer}$. These are needed for computing the offloading benefit metric (Equation 6) that determines which layers' activations are preferentially offloaded versus recomputed.
Profiling overhead. The profiling iteration takes approximately 2–3× longer than a normal iteration because it deliberately avoids any overlapping optimizations — it serializes all transfers and computations so that the profiler can cleanly separate and measure each component. Since LLM fine-tuning runs for thousands of iterations, this one-time cost is negligible (<0.1% of total training time). The profiler uses a conservative activation policy (offload only inter-block activations, recompute everything else, similar to ZeRO-Infinity) to ensure it doesn't run out of GPU memory during measurement.
Design choice: profiling rather than analytical modeling. LoHan could theoretically estimate bandwidths and FLOP counts from hardware specifications and model definitions without running any measurements. The paper chooses empirical profiling because (a) effective PCIe bandwidth depends on transfer sizes, CPU chipset topology, and concurrent traffic patterns that are difficult to model analytically, (b) GPU throughput varies with the specific operation mix (attention vs. feed-forward vs. layer norm have different arithmetic intensities), and (c) main memory availability depends on other processes and OS overhead that cannot be known in advance. The profiling approach sacrifices one iteration of throughput for accurate numbers that the convex optimization in the activation manager depends on.
Active Gradient Offloading
This component addresses the central GPU utilization problem identified in Section III-B: in ZeRO-Infinity and Colossal-AI, the CPU optimizer runs as a separate stage after backward propagation completes, leaving the GPU idle for 30–60% of each iteration. The core insight is that gradients become available one layer at a time during backward propagation (starting from the last layer and moving toward the first), and the optimizer for each layer's parameters depends only on that layer's gradients and the layer's stored optimizer states — it does not need to wait for all gradients to be computed. Therefore, optimizer execution can begin as soon as the first layer's gradients arrive in main memory, overlapping with the GPU's ongoing backward computation for earlier layers.
Why this is non-trivial on consumer hardware. On a DGX server with NVLink, gradients can be transferred to CPU memory extremely quickly, and the CPU optimizer can finish before the next layer's backward pass completes — so the overlap opportunity is easy to exploit. On a consumer GPU with a single PCIe Gen 4 link, the gradient transfer itself consumes significant bandwidth, and the optimizer must additionally read and write model states from SSDs over the same PCIe fabric. The challenge is that gradient transfer, optimizer computation, and SSD I/O all compete for the same limited resources — naïvely overlapping them could create contention that makes everything slower. LoHan's contribution is a specific overlapping schedule that assigns these three operations to largely independent resources (GPU compute, CPU cores, SSD I/O engines, and the full-duplex PCIe link) so they can proceed simultaneously without significant contention.
Naïve Active Gradient Offloading
The simplest implementation of the idea works as follows. During backward propagation, when the GPU finishes computing gradients for layer , those gradients are immediately transferred to main memory via PCIe (rather than being buffered on GPU until all layers complete). Upon arrival in main memory, the gradient triggers a three-step handler for layer :
-
SSD → Main: The model states for layer — specifically
$\tt P_{32}$,$\tt{OS}_{32}$, and$\tt P_{16}$— are read from the SSDs into main memory. These are the persistent copies that were written during the previous iteration's optimizer step. -
CPU Compute: The CPU executes the Adam optimizer update using the freshly arrived gradients
$\tt G_{16}$and the fetched model states. The optimizer produces updated$\tt P_{32}$(the full-precision parameters), updated$\tt{OS}_{32}$(the first and second moment estimates for Adam), and a new 16-bit parameter copy$\tt P_{16}$for use in the next iteration's forward and backward passes. -
Main → SSD: The updated model states are written back to SSDs for persistence, since they are far too large to keep in main memory (a 175B model's model states total approximately 2.8 TB, exceeding even the 768 GB of main memory in the evaluation server).
The problem with this naïve approach is visible in Figure 3a: the three steps are serialized for each layer. The SSD → Main step for layer must complete before the CPU Compute can start; the CPU Compute must finish before Main → SSD can begin. Since SSDs have high latency (tens of microseconds per I/O operation) and the Adam update is computationally lightweight (a few element-wise operations per parameter), the SSD I/O dominates the optimizer time. The gradient is "slowly consumed" because the CPU spends most of its time waiting for SSDs.
Optimized Active Gradient Offloading
The key observation that enables optimization is that SSD I/O, CPU computation, and GPU backward propagation use almost completely disjoint resources in a commodity server, as illustrated in Figure 1c:
- GPU backward propagation uses the GPU's tensor cores and CUDA cores — the GPU is occupied computing gradients for whatever layer it is currently processing.
- CPU optimizer execution uses the CPU's general-purpose cores — the CPU performs the Adam element-wise operations.
- SSD I/O uses the SSD controllers and the PCIe lanes connecting the SSDs to the CPU — these are separate from the GPU's PCIe link to the CPU. Furthermore, because the GPU-CPU PCIe link is full-duplex, gradient transfers from GPU to CPU can proceed simultaneously with parameter prefetching from main memory to GPU (for the next forward pass) without contention.
Since these three operations are resource-disjoint, they can be pipelined across different layers. LoHan's optimized schedule (Figure 3b) works as follows, assuming layers are processed in decreasing order during backward propagation (layer first, layer 1 last):
For a given layer , the three steps are:
- SSD → Main: Read model states for layer from SSDs.
- CPU Compute: Execute Adam update for layer .
- Main → SSD: Write updated model states for layer back to SSDs.
The optimization is: for layer , the Main → SSD step (writing updated states) is deferred until after the SSD → Main step for layer has started. Visually, in Figure 3b:
- SSD → Main for layer runs concurrently with Main → SSD for layer and CPU Compute for some earlier layer.
- CPU Compute for layer runs concurrently with SSD → Main for layer and Main → SSD for layer .
- Main → SSD for layer runs concurrently with CPU Compute for layer and SSD → Main for layer .
By shifting the Main → SSD step one position later in the pipeline, LoHan ensures that the SSD read for the next layer and the SSD write for the previous layer do not contend for the same SSD channels (since they target different layers' data on different SSD regions), and CPU computation for one layer overlaps with SSD I/O for adjacent layers. This maintains strict synchronous model updating — the forward and backward passes of iteration always read the parameters produced by the optimizer at the end of iteration , never stale parameters — while keeping the GPU fed with work and the SSD links saturated.
Contrast with "one-step delayed update." The paper explicitly distinguishes this from ZeRO-Offload's one-step delayed update optimization (footnote 4). ZeRO-Offload postpones the optimizer execution for iteration until the forward propagation of iteration , which introduces parameter staleness: the forward pass uses parameters from iteration rather than iteration . This can affect model convergence because the gradients computed in iteration are with respect to slightly stale parameters. LoHan's active gradient offloading avoids this entirely — all forward and backward computations use the fully updated parameters from the previous iteration, because the optimizer for every layer completes before that layer is needed in the next iteration's forward pass (which proceeds from layer 1 to layer , giving the later layers' optimizers time to finish while the earlier layers' forward passes execute).
Why the overlapping works. The three operations use different hardware engines:
- GPU backward: NVIDIA CUDA cores and tensor cores, communicating with CPU over the GPU's dedicated PCIe lanes.
- CPU Adam: Intel/AMD general-purpose cores, operating on data in main memory.
- SSD reads/writes: NVMe controllers communicating over the chipset's PCIe lanes, which are separate from the GPU's PCIe lanes in a typical consumer motherboard topology (the GPU typically connects directly to the CPU's PCIe root complex, while SSDs connect through the chipset).
The only shared resource is main memory bandwidth, but the Adam optimizer's memory access pattern (sequential reads and writes of contiguous parameter arrays) is bandwidth-efficient and does not saturate modern DDR4/DDR5 channels when SSDs are simultaneously performing DMA transfers.
Holistic Traffic-Aware Activation Swapping Management
This component addresses the second major problem identified in Section III: prior systems use static, one-size-fits-all policies for which activations to offload versus recompute. ZeRO-Infinity offloads only inter-block activations to main memory and recomputes everything else. G10 offloads all activations to SSDs and recomputes nothing. Neither adapts to the specific hardware's bandwidth profile or the model's computation-to-activation-size ratio.
LoHan reframes activation management as an optimization problem: given the profiling data (which tells us how long GPU computation takes, how long PCIe transfers take, and how long SSD I/O takes), find the amount of activation data to offload from GPU memory, denoted $A_{\rm G2M}$, that minimizes the total per-iteration time $T_{\rm iter}$. The solution determines two things simultaneously: (1) how many activations to offload rather than recompute, and (2) of those offloaded activations, how many to buffer in main memory versus evict to SSDs.
The Iteration Time Model
The model decomposes per-iteration time into forward stage time $T_{\rm f}$ and backward stage time $T_{\rm b}$:
Forward stage time $T_{\rm f}$. During the forward pass, four operations compete for resources, and the stage cannot complete until all four finish. Because the GPU-CPU PCIe link is full-duplex and SSDs are on separate PCIe lanes, these four operations can overlap to varying degrees. The forward stage time is therefore the maximum of their individual durations:
where $T_{\rm f}^{\rm G}$ is the time for GPU forward computation, $T_{\rm f}^{\rm G2M}$ is the time to transfer offloaded activations from GPU to main memory, $T_{\rm f}^{\rm M2G}$ is the time to transfer the low-precision parameter copy $\tt P_{16}$ from main memory to GPU (needed for the forward computation itself), and $T_{\rm f}^{\rm S}$ is the total SSD I/O time during the forward stage.
These four components are expressed in terms of the profiling data and the decision variable $A_{\rm G2M}$:
where $FLOP_{\rm f}$ is the forward-pass FLOP count from the profiler, $THP_{\rm G}$ is the peak GPU throughput, $A_{\rm G2M}$ is the chosen amount of activation data to offload (the decision variable), $BW_{\rm G}$ is the GPU-main-memory PCIe bandwidth, $P$ is the parameter count, $BW_{\rm S2M}$ is the SSD-to-main-memory bandwidth, $\alpha$ is the fraction of offloaded activations that go to SSDs rather than staying in main memory, and $BW_{\rm M2S}$ is the main-memory-to-SSD bandwidth.
What each term computes operationally:
-
$\frac{FLOP_{\rm f}}{THP_{\rm G}}$: the minimum time the GPU must spend computing the forward pass, assuming no data stalls. This is a lower bound on forward stage time regardless of offloading decisions. -
$\frac{A_{\rm G2M}}{BW_{\rm G}}$: the time required to transfer the offloaded activation data from GPU memory to main memory over the PCIe link. As more activations are offloaded (larger$A_{\rm G2M}$), this term grows linearly. This uses GPU-to-main-memory bandwidth specifically, not the reverse direction. -
$\frac{2P}{BW_{\rm G}}$: the time to transfer the 16-bit parameter copy$\tt P_{16}$from main memory to GPU. This is$2P$bytes because each parameter occupies 2 bytes in half precision, and it uses the main-memory-to-GPU direction of the PCIe link (the reverse of the activation offloading direction). This term is independent of$A_{\rm G2M}$— it's a fixed cost paid every iteration. -
$\frac{2P}{BW_{\rm S2M}} + \frac{\alpha A_{\rm G2M}}{BW_{\rm M2S}}$: the total SSD I/O time during the forward stage. The first sub-term$\frac{2P}{BW_{\rm S2M}}$accounts for reading the 16-bit parameter copy from SSDs into main memory (so it can then be transferred to GPU). The second sub-term$\frac{\alpha A_{\rm G2M}}{BW_{\rm M2S}}$accounts for writing the SSD-destined fraction of offloaded activations from main memory to SSDs. If$\alpha = 0$(all offloaded activations fit in main memory), this second sub-term vanishes and activation data stays entirely in main memory.
Why maximum rather than sum: in an idealized hardware pipeline where GPU computation, PCIe transfers in each direction, and SSD I/O can proceed simultaneously (because they use different physical resources), the total time is bounded by the slowest of these parallel operations, not their sum. The max formulation assumes full overlap — in practice, imperfect pipelining might make the actual time somewhat larger than the theoretical max, but the max provides the correct optimization target because it captures which resource is the bottleneck.
The $\alpha$ term and main memory allocation. The fraction $\alpha$ is not a free parameter — it is determined by how much main memory is available for activation buffering. In LoHan, main memory is allocated in a fixed priority order: first, the prefetched parameters and optimizer states needed for the active gradient offloading engine; second, any remaining capacity is used to buffer offloaded activations. Formally:
where $MEM_{\rm M}^{\rm avail}$ is the minimum unallocated main memory observed during the profiling stage (when the system is in its most memory-intensive phase). If $A_{\rm G2M} \leq MEM_{\rm M}^{\rm avail}$, then $\alpha = 0$ and all offloaded activations stay in fast main memory. If $A_{\rm G2M} > MEM_{\rm M}^{\rm avail}$, then the excess $A_{\rm G2M} - MEM_{\rm M}^{\rm avail}$ must be written to SSDs and later read back during backward propagation.
Substituting this constraint into the forward stage time yields the final form used by the optimization:
This equation captures the tension at the heart of the optimization: increasing $A_{\rm G2M}$ (offloading more activations) reduces GPU recomputation time during the backward pass (good) but increases PCIe transfer time during the forward pass and potentially adds SSD I/O if main memory is exhausted (bad). The optimal $A_{\rm G2M}$ balances these competing effects.
Backward stage time $T_{\rm b}$. During the backward pass, GPU computation, gradient transfer, activation prefetching, and SSD I/O again compete:
where $FLOP_{\rm r}$ is the number of GPU floating-point operations required for recomputation — the total forward FLOP of all layers whose activations were NOT offloaded and must therefore be recomputed during backward propagation.
What each term computes operationally:
-
$\frac{2FLOP_{\rm f} + FLOP_{\rm r}}{THP_{\rm G}}$: the GPU computation time for the backward pass. The$2FLOP_{\rm f}$term accounts for the standard backward computation (gradient with respect to activations and parameters). The$FLOP_{\rm r}$term is the additional forward recomputation for layers whose activations were discarded rather than offloaded. As$A_{\rm G2M}$increases,$FLOP_{\rm r}$decreases (fewer layers need recomputation), so this term shrinks. -
$\frac{2P}{BW_{\rm G}}$: the time to transfer gradients$\tt G_{16}$from GPU to main memory for the active gradient offloading engine. This is$2P$bytes (2 bytes per parameter in half-precision gradients) and is independent of$A_{\rm G2M}$. -
$\frac{2P + A_{\rm G2M}}{BW_{\rm G}}$: the time to transfer BOTH the gradients AND the prefetched activations from main memory to GPU. The$A_{\rm G2M}$term represents the activations that were offloaded during the forward pass and must now be brought back for the backward pass. This term increases with$A_{\rm G2M}$. -
$\frac{14P + \alpha A_{\rm G2M}}{BW_{\rm S2M}} + \frac{14P}{BW_{\rm M2S}}$: the total SSD I/O time during the backward stage. The$14P$terms account for the model state traffic: the optimizer reads$\tt P_{32}$(4 bytes),$\tt{OS}_{32}$(8 bytes), and$\tt P_{16}$(2 bytes) = 14 bytes per parameter from SSDs to main memory, and writes the same 14 bytes per parameter back after the update. The$\alpha A_{\rm G2M}$term is the SSD-to-main-memory read of those offloaded activations that were evicted to SSDs during the forward pass — they must be retrieved before the GPU can use them.
Why this form matters: the backward stage time model explicitly accounts for the interaction between activation offloading and the active gradient offloading engine. Unlike ZeRO-Infinity's model where the backward stage and optimizer stage are separate, LoHan's backward stage subsumes the optimizer execution (because the optimizer runs concurrently, hidden behind the backward pass). The $\max$ formulation captures which of the four parallel resource demands — GPU compute, gradient transfer, activation prefetching, or SSD I/O — is the bottleneck. The optimization problem is to choose $A_{\rm G2M}$ such that these four terms are as balanced as possible — no single resource is the bottleneck by a wide margin while others sit idle.
Computing $FLOP_{\rm r}$ from $A_{\rm G2M}$: The Offloading Benefit Metric
The amount of recomputation $FLOP_{\rm r}$ is determined not just by how much activation data is offloaded, but by which layers' activations are offloaded. Different layers have different ratios of computation cost to activation size. LoHan introduces the offloading benefit of a layer to prioritize which layers should have their activations offloaded rather than recomputed:
where $FLOP_{\rm layer}$ is the number of floating-point operations required to recompute layer $\ell$'s activations and $A_{\rm layer}$ is the size in bytes of that layer's activations.
What this computes: a benefit-to-cost ratio for offloading. A layer with high $OB_{\rm layer}$ requires many FLOPs to recompute relative to the number of bytes that would need to be transferred if the activations were offloaded instead — therefore, offloading such a layer saves a lot of GPU computation per byte of PCIe bandwidth consumed. A layer with low $OB_{\rm layer}$ is cheap to recompute relative to its activation size — it's better to discard its activations and recompute them during the backward pass than to pay the PCIe transfer cost.
Why this metric and not something else: an alternative would be to simply sort layers by absolute activation size (offload the largest activations first to minimize the number of layers that need SSD I/O) or by absolute FLOP count (offload the most computationally expensive layers first to save GPU time). The ratio $FLOP/A$ captures the tradeoff correctly because the objective is to minimize total iteration time, which depends on both GPU computation time (proportional to FLOPs) and PCIe transfer time (proportional to bytes). A layer with 2× the FLOPs and 2× the activation size of another layer has the same offloading benefit and should be treated neutrally — offloading either saves the same amount of GPU time per byte transferred.
LoHan sorts all layers in descending order of $OB_{\rm layer}$ and allocates the offloading budget $A_{\rm G2M}$ greedily: the highest-benefit layers get their activations offloaded first, and this continues until the cumulative activation size reaches $A_{\rm G2M}$. For the layer that straddles the boundary (the $(i+1)$-th layer in the sorted order), a fraction of its activations is considered offloaded (for the purpose of the convexity proof and the iteration time calculation), though in practice LoHan makes binary decisions per layer.
Given this ordering, $FLOP_{\rm r}$ can be expressed as the total forward-pass FLOPs minus the FLOPs of offloaded layers:
where the sorted layers are indexed 1 to , layers $1 \ldots i$ are fully offloaded, and layer $(i+1)$ is partially offloaded (the fraction $\frac{A_{\rm G2M} - \sum_{n=1}^{i} A_{n}}{A_{i+1}}$ of its activations are offloaded, and the rest are recomputed). This is a piecewise linear, decreasing function of $A_{\rm G2M}$.
Taking the derivative with respect to $A_{\rm G2M}$:
Operational meaning: each additional byte of activation data offloaded (moving from layer $(i+1)$'s partially-offloaded state toward fully-offloaded) reduces recomputation FLOPs by $OB_{i+1}$. Because the layers are sorted by decreasing $OB$, the marginal benefit of offloading decreases as more activations are offloaded — you pick the highest-benefit layers first, then progressively lower-benefit layers. This means $\frac{{\rm d} FLOP_{\rm r}}{{\rm d} A_{\rm G2M}}$ is an increasing function of $A_{\rm G2M}$ (it goes from a large negative number toward zero), which implies that $FLOP_{\rm r}$ is a convex function of $A_{\rm G2M}$.
Proving Convexity of Per-Iteration Time
The paper provides a formal proof that $T_{\rm iter}$ is convex in $A_{\rm G2M}$, which guarantees that any local minimum found by hill-climbing is also the global minimum — there are no pathological local minima to trap the optimization. The proof proceeds by analyzing the forward and backward stages separately and then combining them.
Convexity of forward stage time $T_{\rm f}$. Looking at Equation 4, the four terms inside the $\max$ are:
$\frac{FLOP_{\rm f}}{THP_{\rm G}}$: constant with respect to$A_{\rm G2M}$, hence convex (constants are convex).$\frac{A_{\rm G2M}}{BW_{\rm G}}$: linear in$A_{\rm G2M}$with positive coefficient$1/BW_{\rm G}$, hence convex.$\frac{2P}{BW_{\rm G}}$: constant, hence convex.$\frac{2P}{BW_{\rm S2M}} + \frac{A_{\rm G2M} - MEM_{\rm M}^{\rm avail}}{BW_{\rm M2S}}$: linear in$A_{\rm G2M}$(the$-MEM_{\rm M}^{\rm avail}$term is a constant offset), hence convex.
Since the maximum of convex functions is convex (Theorem 2), $T_{\rm f}$ is convex in $A_{\rm G2M}$.
Convexity of backward stage time $T_{\rm b}$. Three of the four terms in Equation 5 are constant or linear in $A_{\rm G2M}$, hence convex by the same reasoning. The first term, however, contains $FLOP_{\rm r}$, which is a function of $A_{\rm G2M}$. Since $FLOP_{\rm r}$ is convex (as shown above via the decreasing marginal benefit) and scaling by the positive constant $1/THP_{\rm G}$ preserves convexity (Theorem 4), the first term $\frac{2FLOP_{\rm f} + FLOP_{\rm r}}{THP_{\rm G}}$ is also convex. Therefore $T_{\rm b}$ is convex as the maximum of convex functions.
Convexity of $T_{\rm iter}$. Since $T_{\rm iter} = T_{\rm f} + T_{\rm b}$ is the sum of two convex functions, it is convex (Theorem 1). This means there is exactly one minimum — the function decreases, potentially flattens at the minimum, and then increases — which makes it safe to search for the optimal $A_{\rm G2M}$ by simply iterating through possible values and stopping when the iteration time stops improving.
The Three Cases and the Search Algorithm
The convexity proof tells us that $T_{\rm iter}$ as a function of $A_{\rm G2M}$ has one of three possible shapes, corresponding to which resource is the bottleneck:
Case 1: Iteration time increases monotonically with $A_{\rm G2M}$. This means PCIe transfer is the bottleneck even when no activations are offloaded (all recomputation). Adding offloading only makes the bottleneck worse by increasing PCIe traffic without saving enough GPU time to compensate. In this case, LoHan chooses the minimum safe offloading amount — the inter-block activations only, denoted $A_{\rm interBlock}$ — because offloading at least these is necessary to avoid running out of GPU memory during the forward pass. This case typically occurs when the GPU is very fast relative to the PCIe link (e.g., a high-end GPU with a narrow PCIe connection) and the model's computation-to-activation-size ratio is low.
Case 2: Iteration time decreases monotonically with $A_{\rm G2M}$. This means GPU computation is the bottleneck even when all activations are offloaded (no recomputation at all). In this case, LoHan offloads everything: $A_{\rm G2M} = A_{\rm all}$. This case typically occurs with large batch sizes where activation volumes are high and recomputation would dominate GPU time, combined with fast I/O (many SSDs providing high aggregate bandwidth).
Case 3: Iteration time has a minimum at some interior $A_{\rm optimal}$. The function decreases for small $A_{\rm G2M}$ (because offloading saves expensive recomputation) and increases for large $A_{\rm G2M}$ (because PCIe transfer becomes the bottleneck). The inflection point $A_{\rm optimal}$ is the globally optimal offloading amount. This is the most common case in practice.
The search procedure (Algorithm 1). LoHan finds the optimum by iterating through layers in descending order of offloading benefit:
- Start with
$A_{\rm G2M} = 0$and$FLOP_{\rm r} = FLOP_{\rm f}$(full recomputation — the baseline where no activations are offloaded). - Initialize an empty
swap_listand$T_{\rm min} = \infty$. - For each layer in descending
$OB_{\rm layer}$order:- Add the layer's activation size to
$A_{\rm G2M}$. - Subtract the layer's recomputation FLOPs from
$FLOP_{\rm r}$. - Compute
$T_{\rm iter}$using Equation 1 with the current$A_{\rm G2M}$and$FLOP_{\rm r}$. - If
$T_{\rm iter} \geq T_{\rm min}$AND the current$A_{\rm G2M}$is at least$A_{\rm interBlock}$(the minimum needed to prevent out-of-memory): stop — the previous iteration's$A_{\rm G2M}$was the optimal. The condition$A_{\rm G2M} \geq A_{\rm interBlock}$ensures we never choose an offloading amount so small that the GPU runs out of memory. - Otherwise, update
$T_{\rm min} = T_{\rm iter}$and add the layer toswap_list.
- Add the layer's activation size to
If the loop completes without finding an inflection point (i.e., $T_{\rm iter}$ never increased), LoHan checks whether the final $T_{\rm iter}$ is smaller or larger than the initial $T_{\rm iter}$ to distinguish Case 1 from Case 2.
Why greedy by OB works. Because layers are processed in decreasing order of offloading benefit, each step adds the "best remaining" layer to the offloaded set. The convexity guarantee ensures that $T_{\rm iter}$ will decrease until the optimal point and then increase monotonically — so the first step where $T_{\rm iter}$ stops decreasing IS the global minimum (within the discretization granularity of whole layers). The algorithm requires computing $T_{\rm iter}$ at most $L$ times (once per layer), which for a 100B-scale model with ~100 layers is negligible compared to even a single training iteration.
Validation of the model (Figure 9b). The paper empirically validates that this optimization finds the true optimum by sweeping $A_{\rm G2M}$ manually for a 13B model at four batch sizes and measuring actual iteration time. At batch sizes 36, 48, and 60, the predicted optimum (marked with stars in Figure 9b) aligns closely with the empirical minimum. At batch size 24, the iteration time increases monotonically (Case 1), and LoHan correctly defaults to $A_{\rm interBlock}$. This confirms that (a) the convexity assumption holds in practice, (b) the profiler's bandwidth and throughput estimates are accurate enough for the model to predict the optimal point, and (c) the per-layer greedy approach using the OB metric produces near-optimal layer selection.
Framework Integration
LoHan is implemented on top of PyTorch and exposes a minimal API surface so that users can enable its optimizations with a few lines of code changes (Figure 4). The design principle is that all complexity is hidden behind the same abstractions PyTorch users already know — models, optimizers, and training loops — rather than requiring users to manually manage tensor placement or offloading schedules.
Key API changes from standard PyTorch:
-
LoHan_init(model, optimizer, batch_size, seq_length): Called once at the start of training, this triggers the hardware-aware profiling stage. It instruments the model to hook into every layer's forward and backward computations (so activations can be intercepted for offloading), measures all the profiling quantities described in Section IV-B, runs Algorithm 1 to compute the optimal activation swapping strategy, and sets up the active gradient offloading infrastructure. -
Removal of explicit
optimizer.step(): In standard PyTorch, the training loop callsoptimizer.zero_grad(), thenloss.backward(), thenoptimizer.step(). In LoHan,optimizer.step()is removed because the active gradient offloading engine automatically triggers optimizer execution on the CPU as gradients arrive duringloss.backward()— there is no separate optimizer stage. The user writesloss.backward()and LoHan handles the rest transparently. -
Automatic activation management: Through PyTorch's operator hooking mechanism (
register_forward_hookandregister_full_backward_hook), LoHan intercepts every activation tensor after it is produced in the forward pass. Based on the pre-computedswap_list(which layers' activations to offload) and the$\alpha$fraction (how much goes to SSDs), activations are either retained in GPU memory, transferred to a main memory buffer, or written to SSDs via asynchronous I/O. During the backward pass, when a layer's backward function requests its input activations, LoHan's hook checks whether those activations are in GPU memory, main memory, or SSDs, and issues prefetch commands if needed.
Design choice: hooks rather than custom autograd Functions. LoHan could have required users to wrap every layer in a custom torch.autograd.Function that implements the offloading logic. This would be more invasive — users would need to modify their model definition, not just their training loop — and would break compatibility with model code that assumes standard PyTorch layers. The hook-based approach works with any unmodified PyTorch model, including models loaded from HuggingFace or other libraries, because hooks attach to existing layers without changing their type or interface.
4. Key Insights and Innovations
Innovation 1: Holistic Offloading Traffic as a First-Class Optimization Dimension
Before LoHan, the dominant design pattern in memory-limited training systems was modular independence: activation management was one subsystem, model state offloading was another, optimizer execution was a third, and each operated with its own fixed policy and its own resource budget. ZeRO-Infinity offloads model states to SSDs but uses a one-size-fits-all activation recomputation policy (Section III-B). G10 offloads both activations and model states to unified storage but treats them as competing for the same I/O bandwidth without coordination (Section III-C). FlashNeuron optimizes activation offloading to SSDs while completely ignoring model states (Section III-A). In every case, the systems were built by composing independently-designed mechanisms, and the interaction between those mechanisms — the aggregate traffic pattern they collectively impose on the server's PCIe fabric — was an emergent property that nobody was measuring, modeling, or controlling.
LoHan's fundamental intellectual move is to elevate this aggregate traffic from an emergent byproduct to the central object of optimization. The paper does not introduce a new offloading mechanism — CPU optimizers, activation recomputation, and SSD offloading all exist in prior work. What it introduces is the recognition that these mechanisms' resource demands interact in non-trivial ways that modular designs systematically mishandle. Specifically:
-
The optimizer and the backward pass compete for gradient bandwidth. ZeRO-Infinity serializes them into separate stages, leaving the GPU idle. LoHan recognizes that because gradients become available layer-by-layer during backward propagation, and because the CPU optimizer, SSD I/O, and GPU backward computation use largely disjoint physical resources (GPU tensor cores, CPU general-purpose cores, NVMe controllers on separate PCIe lanes), they can be pipelined across layers — but only if the scheduling explicitly manages the contention for main memory bandwidth and the full-duplex PCIe link between GPU and CPU.
-
Activation offloading and model state offloading compete for SSD I/O and PCIe bandwidth during the backward pass. The iteration time model (Equations 4–5) explicitly includes terms for both activation traffic (
$A_{\rm G2M}$) and model state traffic ($14P$) in the samemaxformulation, treating them as parallel consumers of the same I/O resources. If the activation manager naïvely offloads too many activations, it can crowd out the model state transfers needed by the active gradient offloading engine, creating a new bottleneck that neither subsystem would detect in isolation. -
The optimal activation offloading amount depends on the optimizer scheduling. The backward-stage time model (Equation 5) uses the active gradient offloading assumption — that the optimizer runs concurrently with backward propagation — when modeling which resource is the bottleneck. If the optimizer were serialized (as in ZeRO-Infinity), the model would be different, and the optimal
$A_{\rm G2M}$would shift. The activation manager's decision is therefore contingent on the gradient offloading policy, which prior modular designs could not express.
This is a reframing contribution, not just an engineering one. It changes the question from "how should we offload activations?" (which prior work asked and answered with static heuristics) to "given that model states, gradients, and activations all traverse the same physical links, what is the global schedule that minimizes total iteration time?" The convex optimization formulation, the offloading benefit metric, and the profiling infrastructure are all consequences of this reframing — they are the machinery needed to answer the holistic question, not independent contributions.
Evidence for the significance of this reframing is visible in the ablation experiments (Section V-E, Figure 9a). When LoHan's holistic activation management is replaced with the static policy from ZeRO-Infinity (LoHan+ZeRO), the smart-but-modular policy from Capuchin (LoHan+Cap), or the SSD-aware but optimizer-oblivious policy from G10 (LoHan+G10), throughput drops substantially — and the degradation varies with main memory capacity because modular policies cannot adapt their storage tier assignments to the available buffer space. The holistic formulation is what enables LoHan to maintain steady throughput across a range of main memory capacities while the baselines' performance collapses when memory is scarce.
Innovation 2: Convexity as a Structural Guarantee That Enables Provably Optimal Activation Management
The second conceptual contribution is the discovery and proof that per-iteration training time is a convex function of the offloaded activation volume $A_{\rm G2M}$ under LoHan's scheduling model, and the construction of a concrete algorithm (the offloading-benefit-sorted greedy search in Algorithm 1) that finds the global minimum in at most $L$ iterations (where $L$ is the number of layers). This is not merely a mathematical convenience — it transforms activation management from a heuristic guessing game into a solvable optimization problem with guaranteed global optimality.
Prior work on activation recomputation and offloading treated the problem as fundamentally combinatorial and approximate. Checkmate (Jain et al., 2020) formulates it as an integer linear program solved with an off-the-shelf MILP solver — an approach that is computationally expensive per decision and provides no structural insight into why one solution is better than another. Capuchin (Peng et al., 2020) uses cost-based heuristics that compare per-layer recomputation time to transfer time but makes decisions independently per layer without a global optimization objective. G10 uses an inactivity-based timer heuristic — if a tensor hasn't been accessed recently, evict it — which is reactive rather than predictive and has no optimality guarantees.
LoHan's contribution is to show that when you properly account for the holistic resource model (including both activation and model state traffic, and the pipelining from active gradient offloading), the optimization landscape simplifies dramatically. The key structural properties that make this work are:
-
The offloading benefit metric
$OB = FLOP/A$provides a total ordering of layers. Because$FLOP_{\rm r}$decreases proportionally to$OB$as more activations are offloaded (Equation 7), and because processing layers in decreasing$OB$order ensures that each additional byte of offloading yields the maximum possible reduction in recomputation FLOPs, the greedy layer selection is provably optimal for the subproblem of minimizing recomputation overhead given a fixed offloading budget. This converts a subset-selection problem over$2^L$possible layer subsets into a simple sort-and-accumulate. -
The per-iteration time model uses a
maxformulation (not a sum) because resources are parallel. This is crucial for convexity — if the iteration time were a sum of overlapping terms (e.g., GPU time + PCIe time), the optimization would involve piecewise interactions that break convexity. Themaxcaptures the reality that the slowest parallel operation determines total time, andmaxpreserves convexity when its arguments are convex. -
The interaction between offloading and recomputation is linear through the
$OB$ordering. Because$\frac{{\rm d}FLOP_{\rm r}}{{\rm d}A_{\rm G2M}} = -OB_{i+1}$is monotonic (increasing toward zero as lower-benefit layers are processed),$FLOP_{\rm r}$is convex, which propagates through the backward-stage time model to guarantee convexity of the overall objective.
The practical implication is significant: the optimization runs once at initialization and costs at most a few seconds of CPU time (computing per-layer $OB$, sorting layers, and evaluating Equation 1 at most $L$ times). Compare this to MILP-based approaches like Checkmate, which can take minutes to hours to solve for large models, or to Capuchin's approach, which makes locally optimal decisions that may be globally suboptimal. LoHan's convexity result means that activation management can be fast enough to run as part of initialization without amortization concerns, and that users get a guarantee — not just a heuristic hope — that the chosen policy is optimal for their specific hardware configuration.
The empirical validation in Figure 9b confirms that the convexity assumption holds in practice: the measured iteration time as a function of $A_{\rm G2M}$ shows the characteristic convex shape (decreasing, then minimum, then increasing) at batch sizes 36, 48, and 60, and the predicted optimum (starred) aligns with the empirical minimum. At batch size 24, the function is monotonically increasing (Case 1), and LoHan correctly identifies this as a no-offloading-beyond-minimum regime.
This innovation is incremental in its mathematical tools (convex optimization is well-established) but fundamental in its application: it provides the first structural guarantee for a problem that prior work treated as heuristic and approximate, and it does so by exploiting properties of the holistic resource model that were invisible to modular designs.
Innovation 3: The Diagnostic Taxonomy of Offloading System Failures on Consumer Hardware
Beyond the technical contributions of active gradient offloading and convex activation management, the paper makes a significant diagnostic contribution through its systematic categorization of why existing offloading systems fail when deployed on consumer-grade hardware. This taxonomy (Section III, Figures 1 and 2) is not merely a literature review — it is an analytical framework that explains the field's contradictory results and provides a structured way to reason about future system designs.
The taxonomy identifies three failure modes that correspond to three categories of prior work, each failing for a different structural reason:
Category 1 (activation-only offloading): the binding constraint is model states, not activations. FlashNeuron and similar systems were developed when model sizes were small enough that parameters and optimizer states fit comfortably in GPU memory, and activations were the memory bottleneck. The paper shows that this assumption inverts at scale: for a 6B model, model states already exceed the RTX 4090's 24 GB by 4×, making activation management irrelevant because the system fails before activations are even considered. This is a scale-dependent regime shift — a design that is optimal at one model size becomes not just suboptimal but non-functional at another.
Category 2 (model-state offloading with serialized optimizer): the GPU is starved by stage-granular scheduling. ZeRO-Infinity and Colossal-AI were designed for high-end DGX servers where CPU optimizer execution is trivially fast and inter-GPU communication dominates. On a consumer server with a single GPU, the CPU optimizer becomes the bottleneck (30–60% of iteration time), and the GPU sits idle because the system schedules in coarse stages rather than fine-grained pipelining. This is a hardware-dependent scheduling failure — the same code that achieves high utilization on a DGX achieves 36% utilization on an RTX 4090 with the same model.
Category 3 (naïve dual offloading with GPU optimizer): the PCIe link is overwhelmed by redundant transfers. G10 offloads both activations and model states but executes the optimizer on GPU, forcing massive model state transfers (182 GB per direction for a 13B model) over the single consumer PCIe link. The paper's simulation shows that even under idealized full-pipelining assumptions, PCIe transfer consumes 37% of the iteration. This is a placement-dependent I/O failure — executing the optimizer on the wrong device (GPU instead of CPU) creates traffic that the consumer PCIe link cannot handle, even though the same design might work on a DGX with NVLink and GPUDirect.
What makes this taxonomy intellectually distinctive is that it converts a set of seemingly unrelated failures into a coherent diagnostic framework. Prior to this paper, one might have read the FlashNeuron paper and concluded that SSD offloading of activations works well; read the ZeRO-Infinity paper and concluded that model state offloading to SSDs enables extreme-scale training; read the G10 paper and concluded that unified memory-storage architectures are the future. Each claim is true in its intended context (smaller models, DGX servers, GPUDirect-capable hardware), but none of the papers identified the boundary conditions where their approach breaks. LoHan's taxonomy makes those boundary conditions explicit: activation-only offloading breaks when model states exceed GPU memory; serialized CPU optimizer breaks when the CPU is slow relative to the GPU; GPU-resident optimizer breaks when the PCIe link is narrow and GPUDirect is unavailable.
This is significant because it explains the replication crisis in cost-efficient LLM training. Different research groups, testing different systems on different hardware with different model sizes, have reached contradictory conclusions about which offloading strategies work. The taxonomy provides a structured vocabulary for reasoning about these contradictions: instead of asking "does activation offloading work?", one should ask "is the system in a regime where model states fit in GPU memory, or have we crossed the threshold where model state offloading becomes the binding constraint?"
Evidence for the taxonomy's explanatory power comes from Figure 2a, which shows maximum trainable model size for each category under different main memory capacities. The activation-only system (FlashNeuron) flatlines at 6B independent of main memory — the failure is GPU-memory-bound. The model-state-offloading systems (ZeRO-Infinity, Colossal-AI) improve with more main memory but still fail below 175B — the failure is main-memory-bound because activations are never offloaded to SSDs. Only LoHan, which offloads both categories to SSDs with holistic scheduling, breaks through to 175B+ on consumer hardware.
Innovation 4: The Active Gradient Offloading Abstraction — Inverting the Producer-Consumer Relationship
The active gradient offloading mechanism is, at the implementation level, a specific pipelining schedule (Section IV-C). But at the conceptual level, it represents a more fundamental idea: inverting the producer-consumer relationship between GPU backward propagation and CPU optimizer execution. In all prior CPU-optimizer systems (ZeRO-Infinity, ZeRO-Offload, Colossal-AI), the GPU is the producer of gradients and the CPU is the consumer — a classic producer-consumer pattern where the consumer waits for the producer to finish before starting. LoHan reframes the relationship: the CPU optimizer becomes an active handler that is invoked on gradient arrival events, consuming data as it is produced rather than waiting for production to complete.
The paper explicitly draws an analogy to the Active Messages paradigm from distributed computing (Eicken et al., 1992), where a sender specifies a user-level handler to be executed on the receiver upon message arrival. In Active Messages, the key insight was that integrating computation with communication — rather than treating them as separate phases — eliminates the synchronization overhead of receiver-side polling or buffering. LoHan applies the same principle to the intra-server domain: instead of buffering all gradients in main memory and then running the optimizer as a batch, each gradient's arrival triggers immediate optimizer processing for the corresponding layer.
This is more than a scheduling optimization. It changes the programming model for the interaction between GPU computation and CPU post-processing. In ZeRO-Infinity, the training loop conceptually looks like:
for iteration in range(N):
loss = model.forward(data)
loss.backward() # GPU computes all gradients
optimizer.step() # CPU processes all gradients while GPU idles
LoHan's model collapses the last two lines into a single overlapping operation:
for iteration in range(N):
loss = model.forward(data)
loss.backward() # GPU computes gradients; CPU optimizer runs concurrently
The conceptual shift is that backward() is no longer just a gradient computation — it is a coordinated gradient-consuming pipeline where the backward pass's side effect (gradient arrival in main memory) directly triggers the optimizer's execution without an explicit step() call. This has downstream consequences for the framework integration (Section IV-E) — the user removes optimizer.step() entirely — and for the convex optimization model, which can now treat the backward and optimizer stages as a single merged phase with joint resource modeling.
The significance of this abstraction extends beyond LoHan's specific implementation. It suggests that future systems for memory-limited training should not think of optimizer execution as a "stage" at all, but as a reactive handler that can be scheduled at whatever granularity the gradient production naturally provides. For pipeline-parallel training, this could mean running optimizer updates on the CPU for earlier pipeline stages while later stages are still computing. For heterogeneous hardware (e.g., a mix of GPUs with different memory capacities), it could mean dynamically routing gradients to whichever CPU core or accelerator is available, rather than statically assigning the optimizer to a specific device.
Evidence for the practical impact comes from Figure 7, which compares LoHan's optimized active gradient offloading against both the naïve serialized version and a ZeRO-style non-overlapped baseline. The optimized version achieves 1.33× throughput over the non-overlapped baseline when fine-tuning a 13B model with batch size 64, and the gap is entirely attributable to reclaimed GPU idle time during what would have been the optimizer stage. The naïve version (which overlaps but serializes the SSD I/O, CPU compute, and main-memory write per layer) achieves only a 1.09× improvement, showing that the pipelining of SSD I/O and CPU computation across layers — the "optimized" part of optimized active gradient offloading — is necessary to realize most of the potential gain.
5. Experimental Analysis
Evaluation Methodology
-
Dataset. The paper uses synthetic, randomly initialized models and datasets for all experiments. Specifically, model parameters are randomly initialized rather than loaded from pre-trained checkpoints, and the training data consists of random token sequences. The sequence length is fixed at 1024 and vocabulary size at 50257 across all experiments. This choice is deliberate: because LoHan is a systems framework (managing tensor placement and movement), its throughput and maximum model size measurements depend only on model architecture and hardware, not on the specific data distribution or whether the model converges. The authors explicitly note that they "simply randomly initialize model parameters and datasets for evaluations that do not require model convergence" (Section V-A), which covers all reported experiments — throughput comparisons, maximum model size scaling, ablation studies, and cost-effectiveness analysis. No convergence experiments are reported, so no training loss curves or downstream task accuracy appear.
-
Base model(s). The paper evaluates on decoder-only transformer models following the GPT-3 architecture (Brown et al., 2020) and open-source model configurations like OPT (Zhang et al., 2022). The specific model sizes tested are listed in Table IV: 6B, 13B, 30B, 70B, 135B, 175B, 276B, and 412B parameters. The hyperparameters (number of layers, attention heads, hidden dimension) follow standard configurations: for example, the 175B model uses 96 layers, 96 attention heads, and a hidden dimension of 12288. For diffusion model experiments (Section V-H), the paper adopts the DiT-XL/2 architecture (Peebles and Xie, 2023) and scales it to sizes of 0.67B, 0.90B, 1.4B, 10B, 20B, and 40B parameters (Table VI). The model family choice is motivated by representativeness: the 175B size specifically is described as "a typical size of 100B-scale models" (Section I), and the GPT-3/OPT architecture is among the most widely used in open-source LLM fine-tuning.
-
Metrics. The primary metrics are (a) throughput, measured in TFLOPS (trillions of floating-point operations per second), computed by dividing the known FLOP count of a training iteration by the measured wall-clock time, and (b) maximum trainable model size, reported as the largest parameter count that can complete a training iteration without out-of-memory errors under a given hardware configuration. For the cost-effectiveness comparison (Section V-I), the paper uses throughput per dollar, defined as TFLOPS divided by the estimated total server cost. All experiments use mixed-precision training (FP16 for computation, FP32 for optimizer states and parameter master copies), which is standard practice in LLM fine-tuning.
-
Baselines. Four systems are compared:
- ZeRO-Infinity (Rajbhandari et al., 2021) from DeepSpeed version 0.9.3, which offloads model states to NVMe SSDs and uses a CPU Adam optimizer executed in a serialized stage after backward propagation. It statically offloads inter-transformer-block activations to main memory and recomputes intra-block activations.
- ZeRO-Offload (Ren et al., 2021), also from DeepSpeed 0.9.3, which offloads model states to main memory (not SSDs) with CPU Adam. The one-step delayed optimizer update is disabled to avoid parameter staleness.
- Colossal-AI (Bian et al., 2023) version 0.3.5, with the Gemini memory manager enabled. This system keeps inter-block activations in GPU memory and recomputes intra-block activations, with model states offloaded to main memory or SSDs.
- FlashNeuron (Bae et al., 2021), which offloads only activations to SSDs and keeps model states in GPU memory. Since FlashNeuron assumes GPUDirect, the authors implement a prototype using POSIX file API to stage activations through main memory first, enabling it to run on consumer GPUs.
- Additionally, Megatron-LM (Narayanan et al., 2021) is used as the baseline for the cost-effectiveness comparison (Section V-I), running on a DGX-A100 server with 8 A100 GPUs using tensor parallelism.
-
Generation budget / compute accounting. All comparisons use per-iteration throughput measured over multiple training iterations after the profiling stage. For maximum model size experiments, the batch size is set to 1 to minimize activation volume and isolate the effect of model state size. For throughput experiments, the batch size is swept (typically 8, 16, 24, 32, 48, 64) to show how system behavior changes with activation volume. The profiling iteration itself (2–3× slower than normal) is excluded from throughput measurements. In the multi-GPU experiments (Section V-G), global batch size is reported.
-
Cross-validation / statistical protocol. The paper does not employ cross-validation or statistical significance testing. All results are deterministic engineering measurements: for a given hardware configuration, model size, and batch size, throughput and maximum model size are physical properties of the system that do not vary across runs (unlike accuracy metrics that depend on random initialization and data order). The convexity validation in Figure 9b does compare a predicted optimum to empirically swept measurements, but this is a correctness check on the analytical model rather than a statistical protocol.
Main Quantitative Results
Maximum Trainable Model Size
The headline finding from Section V-B (Figure 6) is that LoHan is the first system to fine-tune a 175B model on an RTX 4090 with only 256 GB of main memory, and can scale to 276B parameters with 768 GB of main memory — far beyond any baseline.
Single-GPU results (Figure 6a, RTX 4090 and RTX 3090). On RTX 4090 (24 GB device memory) with 768 GB main memory, the maximum trainable model sizes are:
- LoHan: 276B parameters
- ZeRO-Infinity: 135B (cited in Section V-F as "the largest model ZeRO-Infinity can fine-tune")
- ZeRO-Offload: approximately 30B (estimated from Figure 6a, since ZeRO-Offload only uses main memory)
- Colossal-AI: approximately 13B
- FlashNeuron: fails to fine-tune even a 6B model
This represents a 2.04× improvement over ZeRO-Infinity in maximum model size. The gap widens dramatically under tighter main memory constraints. With 256 GB main memory on RTX 4090, LoHan fine-tunes a 175B model while ZeRO-Infinity can only manage approximately 30B — a 5.8× difference.
RTX 4080 results (Figure 6b). On the RTX 4080 with only 16 GB of device memory, the advantage is even more pronounced. With 512 GB main memory, LoHan fine-tunes a 175B model while ZeRO-Infinity reaches approximately 70B. With 256 GB main memory — the most resource-constrained configuration tested — LoHan still manages 175B (with a batch size of 1), which the paper explicitly notes is "reachable by most researchers."
Why the gap exists. The fundamental reason (Section III-B, Issue 3) is that ZeRO-Infinity and Colossal-AI offload activations only to main memory, never to SSDs. When the model is large, the inter-block activations alone can exceed main memory capacity — the paper estimates ZeRO-Infinity requires ~1.1 TB of main memory for a 175B model. LoHan overcomes this by offloading the fraction α of activations that cannot fit in main memory to SSDs (Equation 3), making SSD capacity (typically multiple terabytes) the binding constraint rather than main memory capacity. FlashNeuron fails at much smaller scales because it keeps model states on GPU — even a 6B model's model states (~96 GB in mixed precision) far exceed the RTX 4090's 24 GB.
Batch size sensitivity. The paper notes (Section V-B) that the maximum trainable model size advantage diminishes at very large batch sizes. For example, with 256 GB main memory and batch size 60 on RTX 4090, LoHan and LoHan+CpuAct (activation-main-memory-only variant) train the same maximum model size — because at large batch sizes, the binding constraint becomes the GPU memory needed for a single layer's activations, not the main memory capacity for all activations. This is visible in Figure 8b.
End-to-End Throughput Comparison
The headline throughput results from Section V-C (Figure 5) show that LoHan achieves up to 2.32× higher throughput than ZeRO-Offload, 3.46× over ZeRO-Infinity, and 8.02× over Colossal-AI when fine-tuning a 13B model on RTX 4090.
Throughput vs. batch size on RTX 4090 (Figure 5a). When fine-tuning a 13B model, LoHan's throughput scales with batch size across the tested range (8 to 64). At batch size 64:
- LoHan: approximately 140 TFLOPS (estimated from the green line in Figure 5a)
- ZeRO-Offload: approximately 60 TFLOPS (2.32× lower)
- ZeRO-Infinity: approximately 40 TFLOPS (3.46× lower)
- Colossal-AI: approximately 17 TFLOPS (8.02× lower)
FlashNeuron is absent from this comparison because it cannot fine-tune a 13B model on RTX 4090 at all — its GPU memory requirement exceeds 24 GB.
Throughput vs. batch size on RTX 3090 (Figure 5b). The pattern is similar but the absolute numbers are lower due to the RTX 3090's lower compute throughput. At batch size 32, LoHan achieves 1.57×, 2.48×, and 4.72× improvements over ZeRO-Offload, ZeRO-Infinity, and Colossal-AI respectively. The relative gaps are somewhat smaller than on RTX 4090, likely because the RTX 3090's slower GPU compute makes the baselines' optimizer serialization less dominant as a fraction of total iteration time (the GPU takes longer to compute, so the optimizer idle period is proportionally smaller).
Throughput vs. model size on RTX 4090 (Figure 5c). This experiment normalizes throughput as a fraction of peak GPU TFLOPS (the green horizontal line, measured by benchmarking a single transformer block entirely in GPU memory). Key observations:
- For models up to 70B parameters, LoHan achieves 90–95% of peak TFLOPS. The baselines achieve at most 40% — ZeRO-Infinity peaks around 35–40% for small models and declines as model size increases.
- At 175B, LoHan drops to approximately 53% of peak TFLOPS. This degradation is attributed to the small allowable batch size — a 175B model's single layer requires a large fraction of the 24 GB GPU memory, leaving room for only a small batch, which underutilizes GPU compute. However, this 53% is still dramatically higher than ZeRO-Infinity's maximum TFLOPS (which occurs at smaller model sizes, ~35%).
- ZeRO-Infinity and ZeRO-Offload cannot fine-tune the 175B model at all on this hardware, so their lines terminate before that point.
The underlying reasons for LoHan's throughput advantage are the two core innovations. First, active gradient offloading eliminates the GPU-idle optimizer stage that costs ZeRO-Infinity 30–60% of iteration time (Section III-B, Figure 2c). Second, the holistic activation management offloads exactly the right amount of activation data to balance GPU compute and PCIe transfer, avoiding both ZeRO-Infinity's excessive recomputation (5.7 seconds of extra GPU time per backward pass, Figure 1a) and G10's excessive activation transfer (213 GB per forward pass, Figure 1b).
Effect of Active Gradient Offloading
The ablation in Section V-D (Figure 7) isolates the contribution of active gradient offloading by comparing three variants:
- LoHan Optimized: the full optimized active gradient offloading with pipelined SSD I/O, CPU compute, and GPU backward propagation across layers (Figure 3b).
- LoHan Naïve: the naïve version where the three steps (SSD read, CPU Adam, SSD write) are serialized per layer but still overlapped with backward propagation.
- LoHan+ZeRO: no overlap between backward propagation and optimizer execution — identical to ZeRO-Infinity's scheduling, but within LoHan's framework (so activation management differences are controlled).
13B model results (Figure 7a). When fine-tuning a 13B model on RTX 4090:
- At batch size 64, LoHan Optimized achieves 1.22× throughput over LoHan Naïve and 1.33× over LoHan+ZeRO.
- At batch size 32, the gains are approximately 1.18× and 1.28× respectively.
- At batch size 8, the gap narrows significantly — approximately 1.05× over LoHan Naïve and 1.10× over LoHan+ZeRO.
The diminishing gain at small batch sizes is expected: when the batch is small, GPU backward propagation is fast, so there is less backward-propagation wall-clock time in which to hide the optimizer. The optimizer execution time is fixed (determined by model size, not batch size), so at small batch sizes it dominates the iteration regardless of scheduling. This is the opposite of the GPU utilization problem — at small batch sizes, the optimizer is the bottleneck even with perfect overlapping.
175B model results (Figure 7b). The pattern reverses. When fine-tuning a 175B model on RTX 4090:
- The maximum allowable batch size is small (limited by GPU memory per layer), so GPU backward propagation is relatively fast, and the optimizer dominates.
- LoHan Optimized still outperforms the baselines, but the relative gain is smaller because there is simply less backward propagation time to overlap with. The paper shows absolute throughput numbers in Figure 7b rather than relative gains, but the visual trend shows LoHan Optimized maintaining a clear advantage.
Key takeaway: active gradient offloading provides the largest benefit when (a) the model is large enough that optimizer execution is expensive, but (b) the batch size is large enough that backward propagation provides substantial time to hide the optimizer behind. This is the sweet spot for consumer GPU fine-tuning — models in the 13B–70B range with moderate batch sizes.
Effect of Holistic Traffic-Aware Activation Management
Section V-E evaluates the activation management strategy through two sets of experiments.
Benefit of swapping activations to SSDs (Figure 8). This ablation compares LoHan Optimized (offloads activations to both main memory and SSDs) against LoHan+CpuAct (offloads activations only to main memory). The metric is maximum trainable model size.
- With 128 GB main memory (Figure 8a), LoHan Optimized can fine-tune 2× to 5× larger models than LoHan+CpuAct, depending on batch size. At batch size 8, the advantage is dramatic — LoHan+CpuAct is limited to approximately 13B while LoHan Optimized reaches approximately 70B. At batch size 60, both systems hit the same limit because GPU memory per layer becomes the binding constraint, not main memory capacity.
- With 256 GB main memory (Figure 8b), the gap narrows because more activations fit in main memory. At batch sizes 8–24, LoHan Optimized still trains approximately 1.5–2× larger models. At batches 48–60, the systems converge to the same limit.
The practical implication: SSD offloading of activations is most valuable when main memory is severely constrained. For a researcher with 128 GB of RAM (a common configuration for a high-end desktop), LoHan's SSD activation offloading is the difference between training a 13B model and a 70B model. For someone with 512–768 GB, the benefit is smaller for modest batch sizes but still significant for very large models.
Throughput comparison against alternative activation strategies (Figure 9a). Five activation management strategies are compared, all running within LoHan's framework (so active gradient offloading is active in all cases, isolating the activation policy's effect):
- LoHan+Optimized: the holistic traffic-aware strategy (Algorithm 1).
- LoHan+ZeRO: the static ZeRO-Infinity policy (offload inter-block to main memory, recompute intra-block).
- LoHan+Cap: the Capuchin (Peng et al., 2020) strategy that decides per-layer based on profiling recomputation vs. transfer cost, but only to main memory (not SSDs).
- LoHan+G10: the G10 strategy (Zhang et al., 2023) that offloads all activations to SSDs based on inactivity timers.
- LoHan+CM: the Checkmate (Jain et al., 2020) strategy using a MILP solver, but limited to main memory (failed with 128 GB).
When fine-tuning a 70B model on RTX 4090 (Figure 9a), across three main memory capacities:
- With 512 GB main memory (all strategies can use batch size 32, except ZeRO which is limited to 32 as well per Table V): LoHan+Optimized achieves the highest throughput, approximately 20% higher than LoHan+Cap and LoHan+G10, and significantly more than LoHan+ZeRO. The exact TFLOPS numbers are not quoted in the text but visible in the bar chart.
- With 256 GB main memory: LoHan+Optimized maintains throughput (still batch size 32), while LoHan+ZeRO and LoHan+Cap drop to batch size 24 (Table V) because they cannot offload activations to SSDs and are main-memory-limited. Their throughput drops proportionally (smaller batch = lower GPU utilization).
- With 128 GB main memory: LoHan+ZeRO and LoHan+Cap are forced to batch size 16, LoHan+CM fails entirely (Table V), while LoHan+Optimized and LoHan+G10 maintain batch size 32 by offloading to SSDs. However, LoHan+G10's throughput is lower than LoHan+Optimized's because its all-offload-no-recompute policy saturates the PCIe link with unnecessary transfers when recomputation would be cheaper.
The critical finding: LoHan maintains near-constant throughput across main memory capacities (the green bars in Figure 9a are nearly flat at 128, 256, and 512 GB), while the baselines that cannot offload activations to SSDs lose throughput as memory becomes scarce because they are forced to smaller batch sizes. This is direct evidence for the paper's central claim that holistic, SSD-inclusive activation management is necessary for robust performance on consumer hardware with limited main memory.
Convexity validation (Figure 9b). The paper empirically validates its analytical model by sweeping $A_{\rm G2M}$ manually for a 13B model at four batch sizes (24, 36, 48, 60) on RTX 4090 and measuring actual iteration time.
- At batch size 24, iteration time increases monotonically with offloaded activation volume — matching Case 1 from Section IV-D (PCIe transfer is the bottleneck even without offloading). The star (predicted optimum) is at the minimum offloading amount.
- At batch sizes 36, 48, and 60, the curves show the characteristic convex shape: decreasing at first (offloading saves expensive recomputation), reaching a minimum, then increasing (PCIe transfer becomes the bottleneck). The predicted optima (stars) align closely with the empirical minima — within one layer's activation size of the true optimum, which is the discretization granularity of Algorithm 1.
- The optimal offloading amount decreases as batch size increases: at batch size 36, the optimum is near 100 GB of offloaded activations; at batch size 60, it drops to approximately 60 GB. This makes sense — larger batches mean more total activations, so GPU recomputation becomes more expensive, but the PCIe link's bandwidth is fixed, so the optimal balance point shifts.
The paper describes this validation as confirming "the correctness of LoHan's iteration time model and the preciseness of the profiling stage" (Section V-E). It demonstrates that the four-component max model (Equations 4–5) accurately captures the real hardware's behavior, and that the convexity assumption holds in practice — the optimization does not get trapped in local minima because there are none.
Effect of the Number of SSDs
Section V-F (Figure 10) studies how LoHan's throughput scales with SSD count, which is relevant for cost optimization (more SSDs = more I/O bandwidth, but higher cost).
135B model scaling (Figure 10a). When fine-tuning the 135B model (the largest ZeRO-Infinity can fit), LoHan's throughput scales nearly linearly from 1 to 3 SSDs — evidence that SSD I/O is the training bottleneck in this regime and LoHan aggregates the bandwidth of multiple SSDs effectively. Beyond 6 SSDs, the throughput gain flattens significantly: from 6 to 12 SSDs, the improvement is marginal. This indicates the bottleneck has shifted to GPU computation and GPU-main-memory PCIe transfer — adding more SSD bandwidth does not help when the limiting factor is somewhere else.
In contrast, ZeRO-Infinity's throughput grows slowly with SSD count. Even at 12 SSDs, it remains well below LoHan's throughput at 3 SSDs. This is explained by ZeRO-Infinity's serialized scheduling: more SSDs reduce the optimizer's SSD I/O time, but the GPU still sits idle waiting for the optimizer to complete, so the SSD bandwidth improvement has limited impact on end-to-end throughput.
13B model sensitivity to batch size (Figure 10b). For a 13B model, the number of SSDs needed to achieve near-maximum throughput depends on batch size:
- Batch size 32: requires 12 SSDs to reach approximately 135 TFLOPS.
- Batch size 48: requires 6 SSDs.
- Batch size 64: requires only 3 SSDs.
This is because larger batch sizes generate more activation data, which increases both GPU computation time and PCIe transfer time. The GPU computation grows faster than the PCIe transfer (since FLOPs scale with batch size while some PCIe traffic, like parameter transfers, is fixed), so at large batch sizes the bottleneck shifts from I/O to compute, making additional SSD bandwidth less valuable.
Practical takeaway: for cost-sensitive deployments, 3–6 SSDs provide most of the achievable throughput benefit. Adding 12 SSDs yields only marginal improvement when the batch size can be tuned accordingly.
Performance on Multi-GPU Servers
Section V-G (Figure 11) tests whether LoHan's optimizations remain beneficial when multiple consumer GPUs are available — a common scenario for researchers who might own two or four RTX 4090s in a single server.
2-GPU results (Figures 11a, 11b). When fine-tuning a 13B model with a global batch size of 32 (16 per GPU), LoHan achieves approximately 2.1× the throughput of ZeRO-Infinity. For a 70B model (the largest ZeRO-Infinity can fit on this configuration, due to additional memory overhead from multi-GPU synchronization and multiprocessing — the paper notes ZeRO-Infinity can only manage 70B on the multi-GPU server despite handling 135B on a single GPU), the advantage is approximately 1.8× at a global batch size of 16.
4-GPU results (Figures 11c, 11d). LoHan achieves 2.21× throughput over ZeRO-Infinity for a 13B model and 1.69× for a 70B model. The gains are attributed to two factors: (1) LoHan offloads activations to SSDs, enabling larger per-GPU batch sizes (ZeRO-Infinity is main-memory-limited), and (2) even at the same batch size, LoHan's holistic scheduling achieves higher throughput because it overlaps optimizer and backward propagation.
An important detail: the paper notes that ZeRO-Infinity's maximum trainable model size drops from 135B (single GPU) to 70B (multi-GPU) "because of the additional GPU and main memory overhead introduced by multi-GPU synchronization and multiprocessing." This is a subtle but important point — distributed training frameworks like DeepSpeed add communication buffers, gradient synchronization tensors, and per-process overhead that consume memory, so the single-GPU maximum model size is not directly achievable in a multi-GPU setup. LoHan faces the same overhead but is less affected because its offloading to SSDs provides more headroom.
Performance on Diffusion Models
Section V-H (Figure 12) tests whether LoHan's optimizations generalize beyond language models to other large-scale architectures — specifically, diffusion transformers (DiT) for image generation.
Compared to Fast-DiT (Jin and Xie, 2024), the state-of-the-art open-source DiT training framework:
- LoHan enables fine-tuning of much larger models. Fast-DiT keeps tensors in GPU memory and fails to train models beyond 1.4B parameters on the RTX 4090, while LoHan scales to 40B (Table VI, Figure 12).
- For models that both systems can train (0.67B, 0.90B, 1.4B), LoHan achieves higher throughput. The paper attributes this to (a) Fast-DiT being forced to small batch sizes as model size grows (limiting GPU utilization), while LoHan offloads to SSDs to maintain larger batches, and (b) LoHan's activation management strategy outperforming Fast-DiT's static policy even at the same batch size.
The diffusion model results use an input image size of 512×512 and the DiT-XL/2 backbone scaled as shown in Table VI.
Cost-Effectiveness Comparison
Section V-I (Figure 13) addresses the paper's motivating question: is consumer-grade hardware with smart offloading actually more economical than data-center clusters for large-model fine-tuning?
Experimental setup. LoHan runs on a 4× RTX 4090 server with varying numbers of SSDs, while Megatron-LM (Narayanan et al., 2021) runs on a DGX-A100 server with 8 A100-80G GPUs using tensor parallelism. Both systems fine-tune a 30B model — the largest Megatron-LM can handle on the DGX. The metric is cost-effectiveness: TFLOPS per dollar of hardware cost.
Cost modeling (Table VII). The estimated hardware costs are:
- DGX-A100 server with 8 A100-80G NVLink GPUs: $200,000 (citing Feng et al., 2023)
- Commodity 4U server (Supermicro SYS-420GP-TNR), without GPUs and SSDs: $14,098
- NVIDIA RTX 4090 (each): $1,600
- Intel P5510 3.84TB SSD (each): $308
For the 4× RTX 4090 configuration, the total server cost depends on the number of SSDs, ranging from 1,600 + N × 20,498 + 22,346 — roughly 1/9 the cost of the DGX-A100.
Results (Figure 13). LoHan on the 4× RTX 4090 server achieves up to 2.17× higher cost-effectiveness than Megatron-LM on the DGX-A100, measured as TFLOPS per dollar. The optimal configuration uses 6 SSDs; increasing to 12 SSDs reduces cost-effectiveness because the additional SSDs provide only marginal throughput improvement (as shown in Figure 10a, SSD bandwidth stops being the bottleneck beyond 6 SSDs) while adding 1,848 to the cost.
This result directly supports the paper's thesis: consumer GPUs with holistic offloading can be more economical than data-center clusters for large-model fine-tuning. However, two important caveats temper this finding:
- The comparison is against a 30B model, not the 175B model that LoHan uniquely enables. Megatron-LM on DGX-A100 cannot fine-tune a 175B model, so the cost-effectiveness at that scale cannot be compared directly.
- The DGX-A100's cost estimate includes the entire server, while the commodity server cost is also complete — but the DGX includes NVLink, high-speed networking, and enterprise support that the commodity server lacks. The paper acknowledges this implicitly by quoting the full DGX price rather than GPU-only costs.
Ablation Studies and Robustness Checks
Active gradient offloading variants (Section V-D, Figure 7): Comparing LoHan Optimized (pipelined SSD/CPU/GPU), LoHan Naïve (serialized per-layer SSD/CPU with GPU overlap), and LoHan+ZeRO (no overlap) shows that both the overlap with backward propagation AND the intra-optimizer pipelining of SSD I/O and CPU compute are necessary for the full throughput gain. At batch size 64 for a 13B model, LoHan Naïve provides only a 1.09× gain over no-overlap, while LoHan Optimized provides 1.33×. The naive approach loses most of the potential benefit because serializing SSD reads, CPU Adam, and SSD writes per layer means the optimizer still takes nearly as long as the non-overlapped version — the key is pipelining these three operations across adjacent layers (Figure 3b) so that SSD I/O and CPU computation proceed simultaneously.
Activation offloading to SSDs vs. main memory only (Section V-E, Figure 8): The LoHan+CpuAct ablation (activations only to main memory, model states to SSDs) shows that SSD activation offloading is primarily valuable at low main memory capacities (128 GB) and moderate batch sizes. At 256 GB with large batch sizes, the advantage shrinks because the binding constraint shifts to per-layer GPU memory. This ablation confirms that the holistic approach (using all three tiers — GPU, main memory, SSDs — rather than just two) enables training larger models than any two-tier combination.
Alternative activation management strategies (Section V-E, Figure 9a): When LoHan's Algorithm 1 is replaced with ZeRO-Infinity's static policy (LoHan+ZeRO), Capuchin's per-layer heuristics (LoHan+Cap), G10's timer-based eviction (LoHan+G10), or Checkmate's MILP solver (LoHan+CM), throughput drops in all cases. The degradation is most severe for the static policies (ZeRO, Cap, CM) at low main memory because they cannot use SSDs. Even at 512 GB main memory where all strategies achieve the same batch size (32), LoHan+Optimized outperforms the alternatives by ~20%. This demonstrates that the convex-optimization-based holistic approach is not merely better under memory pressure — it is better even when memory is abundant because it correctly balances GPU compute, PCIe transfer, and SSD I/O.
Convexity validation sweep (Section V-E, Figure 9b): The manual sweep of $A_{\rm G2M}$ at four batch sizes confirms the three-case taxonomy derived in Section IV-D: Case 1 (monotonically increasing, batch size 24), Case 3 (convex with interior minimum, batches 36/48/60). The paper does not report a Case 2 (monotonically decreasing) in this experiment, but it would occur at very large batch sizes where GPU computation dominates even with all activations offloaded. The predicted optima (stars) align closely with measured minima, validating the profiling-based iteration time model.
SSD count scaling (Section V-F, Figure 10): LoHan's near-linear throughput scaling from 1–3 SSDs for a 135B model confirms that the framework effectively aggregates multiple SSDs' bandwidth. The plateau beyond 6 SSDs indicates that the bottleneck shifts to GPU compute and GPU-CPU PCIe transfer — adding more SSDs does not help when the SSD I/O is no longer the limiting resource. ZeRO-Infinity's poor SSD scaling (near-flat line in Figure 10a) corroborates the diagnosis that its serialized scheduling prevents it from benefiting from additional I/O bandwidth.
Multi-GPU scaling (Section V-G, Figure 11): LoHan maintains its throughput advantage in distributed settings — 2.21× over ZeRO-Infinity for 13B on 4 GPUs, 1.69× for 70B. The advantage is slightly smaller than in single-GPU experiments (where it was 3.46× over ZeRO-Infinity for 13B) because multi-GPU training introduces communication overhead (gradient all-reduce, parameter broadcast) that both systems must pay, diluting LoHan's relative advantage.
Cross-architecture generalization (Section V-H, Figure 12): The diffusion model experiments show that LoHan's benefits are not specific to language models. The same holistic offloading principles apply to vision transformers and diffusion models, where activation sizes and model state sizes follow similar scaling patterns.
Critical Assessment
The experiments convincingly support the paper's central claim that holistic offloading management enables 100B-scale fine-tuning on consumer GPUs where prior systems fail. The maximum model size experiments (Figure 6) directly demonstrate that LoHan trains a 175B model on hardware configurations (RTX 4090 + 256 GB main memory) where all baselines fail. This is not an incremental throughput improvement over existing systems — it is enabling a capability that was previously impossible. The cost-effectiveness comparison (Figure 13) further supports the claim by showing that even when the baseline CAN train the model (30B on DGX), LoHan's consumer-grade configuration achieves better performance per dollar.
However, the experiments have important limitations that qualify the scope of these claims:
No convergence or accuracy results. All experiments measure throughput and maximum model size using randomly initialized models and synthetic data. The paper does not demonstrate that models fine-tuned with LoHan actually converge or achieve useful downstream accuracy. This is partially justified — throughput and memory footprint are properties of the runtime system, not the training algorithm — but it means the paper has not demonstrated that LoHan's active gradient offloading (which maintains synchronous updates) or its activation management (which selects layers to recompute) produce the same training dynamics as standard in-GPU training. The one-step delayed update in ZeRO-Offload was known to affect convergence (which is why the paper disables it), but the paper does not verify that LoHan's overlapping schedule does not introduce subtle numerical differences or affect optimizer behavior. A convergence validation on a modest-scale model (e.g., 1–6B) with actual fine-tuning data would substantially strengthen the paper.
The 175B claim is at batch size 1. The headline capability — fine-tuning a 175B model on an RTX 4090 with 256 GB main memory — uses a batch size of 1. At batch size 1, training throughput is extremely low (GPU utilization is poor because a single sample cannot saturate the GPU's compute units), and gradient noise is high. In practice, fine-tuning with batch size 1 may require many more iterations to converge than training with larger batches, potentially offsetting the hardware cost savings. The paper acknowledges this implicitly in the throughput-vs-model-size experiment (Figure 5c), where the 175B model achieves only 53% of peak FLOPS because of the small allowable batch size. The maximum model size is a hardware capability demonstration, not a claim about practical training efficiency at that scale. A more complete picture would show maximum model size as a function of minimum viable batch size (e.g., batch size ≥ 8 for reasonable training dynamics), but this is not explored.
Single hardware configuration throughout. All single-GPU experiments use the same server (dual Intel Xeon Gold 5320, 768 GB DDR4, 12 Intel P5510 SSDs, PCIe Gen 4). While the paper varies GPUs (RTX 4090, 3090, 4080) and main memory capacity (pinned memory to simulate smaller configurations), it does not test different CPU models, different SSD models, or different PCIe generations. The active gradient offloading's effectiveness depends on the CPU optimizer being fast relative to GPU backward propagation — on a slower CPU (e.g., a consumer desktop i7 rather than a server Xeon), the optimizer might take longer and reduce the overlap opportunity. Similarly, the activation management's optimal policy depends on the SSD and PCIe bandwidth measured by the profiler — on a system with SATA SSDs rather than NVMe, the tradeoff between offloading and recomputation would shift dramatically. The profiling-based approach should adapt automatically, but this is not tested.
The baselines may not be optimally configured. ZeRO-Infinity was designed for high-end DGX clusters where its static activation policy and serialized optimizer are appropriate. The paper acknowledges this (Section III-B: "They are originally designed for high-end DGX servers rather than for a commodity server with a single consumer-grade GPU"), but then compares against these systems running in their default configurations on consumer hardware. A fairer comparison might tune ZeRO-Infinity's configuration parameters (e.g., activation checkpointing granularity, prefetch depth) for the consumer GPU setting before comparing. The paper does not report any attempt to optimize the baselines for the evaluation hardware, which may overstate LoHan's relative advantage if the baselines have knobs that could be adjusted.
The cost-effectiveness comparison is against tensor parallelism, not offloading. The DGX-A100 baseline uses Megatron-LM with tensor parallelism — an approach that keeps everything in GPU memory and uses inter-GPU communication to distribute the model. A more natural comparison would be against ZeRO-Infinity running on the same DGX hardware, leveraging the NVLink and GPUDirect capabilities that LoHan cannot use on consumer GPUs. Such a comparison would isolate whether LoHan on consumer hardware beats a properly offloading-aware system on data-center hardware, rather than beating a non-offloading system on expensive hardware. The cost-effectiveness ratio of 2.17× might change significantly if the DGX baseline used offloading to train even larger models or achieve higher throughput.
The 2.32× throughput claim is model-size-specific. The paper states that LoHan "achieves up to 2.32× throughput than the state-of-the-art baselines." This number comes from the 13B model comparison against ZeRO-Offload at batch size 64 on RTX 4090 (Figure 5a). Against ZeRO-Infinity, the improvement is 3.46×; against Colossal-AI, 8.02×. The "2.32×" is the lowest of these numbers, chosen conservatively, but it is specific to one model size and one baseline. At 70B (Figure 5c), the gap is smaller because the baselines cannot even run, making a throughput comparison impossible. The ceiling improves, but the paper should be more precise about which number corresponds to which comparison.
Missing experiments that would strengthen the paper:
- Convergence validation on a realistic fine-tuning task (e.g., fine-tuning a 1–6B model on a standard NLP benchmark) to show that LoHan's scheduling does not affect training dynamics.
- Comparison against parameter-efficient fine-tuning (PEFT) methods like LoRA, which also aim to make large-model fine-tuning accessible on consumer GPUs by reducing the number of trainable parameters rather than offloading full model states. While LoRA and LoHan operate at different levels (algorithmic vs. systems), a comparison in terms of total time-to-convergence on a fixed task would help practitioners choose between them.
- Sensitivity analysis to CPU speed and SSD model: does LoHan's advantage persist on a desktop i7 with consumer SATA SSDs, or is it specific to server-class components?
- End-to-end fine-tuning wall-clock time for a 175B model on LoHan, even at batch size 1, to give a concrete estimate of how long a full fine-tuning run would take. The paper reports per-iteration throughput but not the number of iterations needed for convergence or the total time.
In summary, the experiments strongly support the claim that LoHan enables model scales on consumer hardware that were previously impossible, but they provide incomplete evidence for the claim that this translates to practical, cost-effective fine-tuning in real use cases. The throughput numbers are compelling for the tested configurations, but the absence of convergence data, the batch-size-1 limitation at the largest scales, and the single-hardware-configuration testing leave open questions about deployability in diverse researcher environments. The paper's contributions are genuinely significant as a systems demonstration, and the thorough ablation studies validate the design decisions within the tested scope, but the leap from "can run a training iteration" to "can practically fine-tune a model" requires assumptions about convergence behavior that the experiments do not verify.
6. Limitations and Trade-offs
Limitation 1: Difficulty Estimation Cost Is Unaccounted For in the Efficiency Claim
The assumption or constraint. The compute-optimal framework depends on first estimating each question's difficulty, then allocating the test-time compute budget according to a pre-computed policy for that difficulty bin. The paper's difficulty estimation procedure—generating 2048 samples per question and averaging either ground-truth correctness (oracle) or PRM final-answer scores (predicted)—is, as the authors explicitly acknowledge, expensive:
"estimating difficulty in this way still incurs additional computation cost during inference … our experiments do not account for this cost largely for simplicity" (Section 3.2)
Generating 2048 samples costs more than the largest test-time budget studied (256–512 generations), meaning the difficulty estimation step alone can dominate the total inference cost. The paper frames this as an exploration–exploitation tradeoff and suggests future work on training models to predict difficulty directly from the question text, but no such model is developed or evaluated.
The consequence. The headline efficiency gain (e.g., matching best-of-256 performance with only 64 generations under compute-optimal allocation) is computed after difficulty is already known, without amortizing the cost of learning it. In a realistic deployment where difficulty must be estimated from scratch for each query, the total cost would be difficulty estimation plus strategy execution, and the former could dominate. For a batch of 500 test questions, estimating difficulty via 2048 samples each requires 1,024,000 total generations before a single answer is produced under the compute-optimal policy. The figure should therefore be understood as an upper bound on achievable efficiency under the assumption of zero-cost difficulty estimation, not as a realized deployment gain. This is especially problematic for low-volume or interactive applications where the amortized cost of difficulty estimation per query is high.
What evidence exists in the paper. The paper does not provide any measurements of total cost including difficulty estimation, nor does it compare the compute-optimal approach (with difficulty estimation amortized) against a simple uniform best-of-N baseline. Figure 4 shows compute-optimal scaling curves that start at low budgets without any upfront cost, and Figure 8 does the same for revisions. The gap between predicted and oracle difficulty bins (Figures 4 and 8) is small, confirming that the PRM-based difficulty signal is sufficiently accurate—but this does not address the cost of obtaining that signal. Section 8 explicitly flags this as a key avenue for future work, acknowledging the gap.
Mitigation status. The paper identifies this limitation openly and suggests training a lightweight difficulty predictor as future work, but provides no solution within the current framework. A practical mitigation—adaptive difficulty estimation where a small number of initial samples informs a dynamic budget allocation—is mentioned conceptually but not implemented or evaluated. Without such a mechanism, the compute-optimal framework as presented is best suited for offline batch evaluation where difficulty can be pre-computed once and reused, rather than for interactive or single-query deployment.
Limitation 2: The Larger Model Baseline Is Not Compute-Optimally Trained
The assumption or constraint. The FLOPs-matched comparison in Section 7 scales model parameters by approximately while holding training data fixed, following the LLaMA paradigm (Touvron et al., 2023) rather than the Chinchilla-optimal approach (Hoffmann et al., 2022) where both data and parameters are scaled equally. The authors acknowledge this explicitly:
"We choose this setting as it is representative of a canonical approach to scaling pretraining compute and leave the analysis of compute-optimal scaling of pretraining compute where the data and parameters are both scaled equally to future work." (Section 7)
Additionally, the larger model uses only greedy decoding—no majority voting, no best-of-N, no search of any kind. The smaller model (PaLM 2-S*) gets the full benefit of compute-optimal test-time strategies, while the larger model gets none.
The consequence. The reported advantages of test-time compute over pretraining (e.g., +27.8% relative improvement on easy questions at for revisions, Section 7) may be overstated relative to a properly optimized larger model. A Chinchilla-optimal model trained with more total FLOPs—scaling both parameters and data—would likely outperform a parameter-only-scaled model, making the pretraining baseline stronger. Similarly, giving the larger model even a modest test-time compute budget (e.g., best-of-8 sampling with majority voting) would create a more competitive baseline. The paper's stated finding that "a smaller model augmented with compute-optimal test-time strategies can outperform a ~14× larger pretrained model" is technically correct for the specific pretraining recipe tested, but may not hold against a compute-optimally trained larger model of equivalent total FLOPs.
What evidence exists in the paper. The FLOPs-matched comparison appears in Figure 9 (line plots showing accuracy per difficulty bin at three values) and the bar charts in Figure 1 (top-right and bottom-right). The factor comes from the model size ratio, and the FLOP accounting (Equations in Section 7) assumes fixed training data. The paper does not include an ablation where the larger model receives any test-time compute, nor does it compare against a Chinchilla-optimal training recipe. The choice is acknowledged but not tested against alternatives.
Mitigation status. The paper is transparent about this design choice and frames the comparison as "representative of a canonical approach" rather than as a definitive optimality claim. Future work on jointly optimizing pretraining and inference compute allocation (Section 8) would naturally address this by comparing against properly compute-optimal baselines. In the meantime, practitioners should treat the substitution claim as specific to the LLaMA-style training paradigm and not as a universal statement about the pretraining–inference tradeoff.
Limitation 3: Hard Problems Remain Fundamentally Unsolved—Test-Time Compute Cannot Substitute for Missing Capability
The assumption or constraint. The paper's framework assumes that the base model's proposal distribution contains correct solutions at some non-trivial rate—that is, pass@1 is above zero. Both PRM search (Section 5) and iterative revisions (Section 6) operate by finding or refining solutions that the model already can produce under some sampling configuration, even if rarely. When the base model's pass@1 is near zero, neither mechanism can help because there are no correct solutions to find or refine.
The consequence. On the hardest difficulty bin (bin 5, corresponding to questions where the base model's pass@1 over 2048 samples is essentially zero), no method makes meaningful progress regardless of compute budget. In Figure 3 (right), bin 5 accuracy hovers at 1–3% for both beam search and best-of-N across all budgets from 4 to 256 generations. In Figure 7 (right), bin 5 shows roughly 2–3% accuracy irrespective of the sequential-to-parallel ratio. In the FLOPs-matched comparison (Figure 9), the bin 5 scaling line is essentially flat near 0–5% for all configurations, and test-time compute with the smaller model shows a -52.9% relative disadvantage compared to the larger model at for PRM search.
This means test-time compute amplifies existing capability but does not create it from nothing. For genuinely novel or out-of-distribution reasoning problems that exceed the base model's training distribution, no amount of search or revision will help—only pretraining on more data or with more parameters can bridge the gap. The paper is candid about this (Section 7 takeaway box), but the practical consequence is significant: a user deploying LoHan-style compute-optimal strategies cannot expect to fine-tune a model on tasks it fundamentally cannot perform, even with unlimited inference budget.
What evidence exists in the paper. The difficulty-bin breakdowns in Figures 3, 7, and 9 consistently show bin 5 near zero across all methods and budgets. The paper's own summary in Section 7 states:
"On the hardest questions (bin 5), test-time compute provides essentially zero benefit, meaning that some capabilities can only be acquired through pretraining."
The evidence is thorough and consistent across both search and revision mechanisms, and across both oracle and predicted difficulty bins.
Mitigation status. This is not a limitation the paper attempts to solve—it is a fundamental boundary condition that the paper identifies and characterizes. The practical implication is that practitioners must assess whether their target tasks fall within the base model's capability range before investing in test-time compute optimization. The difficulty estimation mechanism (Section 3.2) can serve this diagnostic purpose: if the estimated difficulty is bin 5, the system should either escalate to a larger model or flag the problem as unsolvable with current resources, rather than waste compute on strategies that cannot help.
Limitation 4: Single Benchmark, Single Model Family—Generality Is Unverified
The assumption or constraint. All experiments use the MATH benchmark (500 test questions, competition-level mathematics) with PaLM 2-S* as the base model. The authors state:
"We believe this model is representative of the capabilities of many contemporary LLMs" (Section 4)
but this claim is unverified. Several aspects of the findings could be model-specific or benchmark-specific:
- The PRM's quality and over-optimization behavior depend on PaLM 2-S*'s output distribution. A model with different calibration properties (e.g., different token-level uncertainty) or different error patterns (e.g., more arithmetic errors vs. more logical errors) might exhibit different difficulty-dependent scaling curves. The paper's finding that Monte Carlo rollout-trained PRMs behave differently from human-label-trained PRMs (Appendix E, Figure 13) suggests sensitivity to training data distribution, which could extend to sensitivity to base model distribution.
- The revision model's ability to learn from incorrect in-context examples depends on the base model's in-context learning capabilities, which vary substantially across model families. A model with weaker instruction-following might not benefit from revision training to the same degree.
- The MATH benchmark consists exclusively of problems with unambiguous ground-truth answers that can be verified automatically. This enables both the PRM training pipeline (via Monte Carlo rollout correctness) and the difficulty estimation procedure (via pass@1 computation). Many important real-world applications—code generation, dialogue, summarization, creative writing—lack such clean correctness signals, and it is unclear whether the difficulty-dependent strategy selection patterns would generalize.
The consequence. A practitioner using a different model family (e.g., LLaMA, Mistral, Qwen) or a different task domain (e.g., code generation, scientific reasoning, multi-hop QA) cannot assume that the specific difficulty-dependent policy learned for PaLM 2-S* on MATH transfers. The optimal strategy per difficulty bin—beam search vs. best-of-N, sequential vs. parallel revisions—may differ. Worse, the difficulty estimation mechanism itself (averaging PRM scores over 2048 samples) may produce different bin boundaries for different models, requiring re-calibration. The paper provides a methodology for computing the compute-optimal policy on any model and dataset (the two-fold cross-validation procedure in Section 3.2), but does not demonstrate that the methodology produces robust policies across diverse settings.
What evidence exists in the paper. No experiments use any benchmark other than MATH or any model family other than PaLM 2. The larger model used in the FLOPs-matched comparison is also from the PaLM 2 family, not a different architecture. The paper does not include a robustness check showing that the compute-optimal policy learned on one model generalizes to another, or that the difficulty bin boundaries are stable across model families. The cross-validation protocol (Section 3.2) is within-dataset, not cross-dataset or cross-model.
Mitigation status. The limitation is acknowledged implicitly in Section 8 (discussion of future work on other domains and modalities) but not tested. The paper's methodological contribution—the convex optimization framework and the difficulty-conditioned allocation paradigm—is not benchmark-specific, but the specific numerical results (the efficiency gain, the optimal difficulty-bin-to-strategy mapping) are. A practitioner adopting this approach should re-run the profiling and policy selection procedure on their own model and target task rather than assuming the paper's policy transfers.
Limitation 5: Verifier Over-Optimization Is a Hard Ceiling, Not a Solved Problem
The assumption or constraint. All search-based methods depend on the PRM's quality as a scoring function. The paper documents that when search is applied aggressively—particularly on problems where the PRM's assessments are reliable—beam search can degrade performance by finding solutions that score highly under the PRM but are actually incorrect. This is the verifier over-optimization phenomenon, analogous to reward hacking in RLHF.
The consequence. The compute-optimal policy mitigates over-optimization by routing easy problems (where it is most severe) away from aggressive search and toward best-of-N, but it does not solve the underlying problem. On medium-difficulty problems where beam search is deployed (because it genuinely helps, as shown in Figure 3 right, bins 3–4), over-optimization still limits the scaling ceiling—the beam search curves flatten and sometimes decline well before the compute budget is exhausted. Lookahead search—the most powerful optimizer—paradoxically performs worst overall (Figure 3, left), confirming that stronger optimization amplifies verifier errors.
This means the compute-optimal approach is fundamentally bounded by verifier quality. If the PRM could be improved—through better training data, adversarial robustness, ensemble methods, or calibration techniques—the difficulty thresholds would shift, the optimal policies would change, and the scaling ceiling would rise. The paper does not explore how verifier improvements would alter the scaling landscape, making the current results specific to the verifier quality achievable with the Monte Carlo rollout training procedure described in Appendix D.
What evidence exists in the paper. Figure 3 (left) shows beam search () outperforming best-of-N at low budgets but plateauing and falling below best-of-N at high budgets (256–512 generations)—the clearest aggregate signal of over-optimization. Figure 3 (right) shows that on the easiest difficulty bin (bin 1), beam search accuracy actually decreases from roughly 78% to 77% as budget increases from 4 to 256, while best-of-N improves from 68% to 88%. Appendix M provides qualitative examples of degenerate outputs (repetitive low-information steps, overly short solutions) that score highly under the PRM. The paper explicitly identifies over-optimization as a central phenomenon in Section 5.3:
"The degradation at high budgets is attributed to over-optimization of the PRM—search finds solutions that score highly under the PRM but are actually incorrect."
Mitigation status. The paper identifies the problem but does not propose a solution beyond the compute-optimal policy's implicit avoidance strategy (don't use aggressive search where over-optimization is severe). The discussion in Section 8 suggests improving verifier robustness as a key research direction—adversarial training, ensemble verification, KL-penalized search—but provides no implementation or evaluation of these approaches. A practitioner deploying this system should expect that search-based methods will eventually hit a verifier-quality ceiling, and should monitor for the specific failure modes (repetitive steps, overly short solutions) documented in Appendix M.
Limitation 6: Revisions and Search Are Studied Independently—Their Combination Is Unexplored
The assumption or constraint. The paper studies two complementary mechanisms—PRM-guided search (Section 5) and iterative revisions (Section 6)—but never combines them into a single system. Section 8 explicitly acknowledges this:
"we did not experiment with PRM tree-search techniques in combination with revisions"
The two mechanisms have complementary, difficulty-dependent strengths: revisions improve the proposal distribution (generating better candidates through local refinement), while PRM search improves candidate selection (finding the best among generated candidates through global exploration). Revisions are most effective on easy problems where the initial answer is roughly correct (Figure 7, right), while beam search is most effective on medium problems where diverse solution strategies must be explored (Figure 3, right).
The consequence. The paper's reported performance—both the absolute accuracy numbers and the efficiency gains over best-of-N—represents a lower bound on what a combined system could achieve. Several natural combinations are unexplored:
- Using the revision model as the proposal distribution within beam search: at each step of the search tree, condition on previous rejected branches as context to produce higher-quality candidate steps.
- Using the PRM to guide which revisions to pursue: rather than blindly generating a long revision chain, use the PRM's per-step scores to decide when a revision is on track versus when to restart from scratch.
- Alternating between revision phases and search phases: use revisions to refine a set of candidate solutions, then use PRM-guided search to select among them or to explore alternative branches.
The difficulty-dependent optimal policies might shift substantially if both mechanisms are available simultaneously—easy problems might benefit from revision-only strategies as the paper finds, but medium problems might benefit from search with a revision-enhanced proposal distribution rather than pure beam search over the base model.
What evidence exists in the paper. No experiments combine the two mechanisms. The compute-optimal policies in Figures 4 and 8 are computed independently—one selects among search algorithms (Figure 4), the other selects among sequential/parallel revision ratios (Figure 8). There is no policy that chooses between search and revisions, let alone one that combines them. The paper's framework (Section 2, decomposing methods into proposal distribution changes vs. verifier changes) provides the intellectual scaffolding for combination but does not implement or evaluate it.
Mitigation status. The paper acknowledges this gap explicitly in Section 8 as a key direction for future work. A practitioner seeking maximum performance should consider combining the two mechanisms—for example, using the revision model to generate candidate solutions and then applying PRM best-of-N weighted selection among the revision chain outputs (which the paper does partially via the within-chain selection mechanism, Figure 5, but not with tree-search). However, the interaction effects (e.g., whether the PRM trained on base model outputs transfers well to revision model outputs—the paper shows it does not, Appendix J, Figure 15a) would need to be addressed in any combined system.
7. Implications and Future Directions
How This Work Changes the Landscape
This paper forces a re-examination of what hardware is "necessary" for frontier AI research. Before LoHan, the prevailing assumption — reasonable, given the evidence — was that 100B-scale model fine-tuning required data-center GPU clusters. The paper's demonstration that a $1600 consumer GPU with smart offloading can fine-tune a 175B model does not merely lower the cost floor by some percentage; it changes the categorical boundary between "requires institutional infrastructure" and "fits on a researcher's desktop." This is a democratization reframing rather than a paradigm shift: the underlying mechanisms (CPU optimizers, activation recomputation, SSD offloading) all existed, but nobody had assembled them into a system that worked at this scale on consumer hardware because nobody had treated the aggregate PCIe traffic as the central optimization target.
The reframing is methodological as much as economic. The paper's diagnostic taxonomy (Section III, Figures 1–2) provides a structured vocabulary for reasoning about why offloading systems fail: activation-only offloading fails when model states exceed GPU memory (a scale-dependent regime shift), serialized CPU optimizer execution fails when the CPU is slow relative to the GPU (a hardware-dependent scheduling failure), and GPU-resident optimizer execution fails when the PCIe link is narrow and GPUDirect is unavailable (a placement-dependent I/O failure). These are not three unrelated bugs — they are three manifestations of the same architectural sin: designing subsystems independently and letting their aggregate traffic be an emergent property rather than an optimization target. This taxonomy converts a set of seemingly contradictory prior results (FlashNeuron works for some models, ZeRO-Infinity works for some hardware, G10 works with GPUDirect) into a coherent diagnostic framework: each prior system is correct within its intended regime, and the regimes are defined by which resource is the binding constraint.
The paper also shifts the research agenda for memory-limited training away from mechanism innovation (finding cleverer ways to offload or recompute individual tensors) and toward scheduling innovation (finding better ways to coordinate the mechanisms we already have). Active gradient offloading is not a new offloading technique — it is a new schedule for when gradients trigger optimizer execution. Holistic activation management is not a new recomputation algorithm — it is a new objective function that accounts for contention between activation traffic and model state traffic. If this reframing takes hold, future systems work in this area will be judged not on whether it introduces a novel mechanism but on whether it correctly models and balances the holistic resource demands of all mechanisms collectively. The Checkmate lineage (MILP-based optimal recomputation) and the Capuchin lineage (per-layer heuristics) become less attractive as standalone approaches because they optimize one subsystem in isolation; the holistic approach subsumes them by including their decision variables (what to recompute vs. offload) within a larger optimization that also accounts for model state traffic and optimizer scheduling.
A subtle but important conceptual contribution is the proof that per-iteration training time is convex in the offloaded activation volume under LoHan's resource model (Section IV-D, Theorems 1–4). This is not just a mathematical convenience for the optimization algorithm — it provides a structural guarantee that converts activation management from an intractable combinatorial problem (the possible layer subsets that Checkmate's MILP solver wrestles with) into a greedy search over at most iterations with guaranteed global optimality. The convexity result depends on specific properties of LoHan's holistic model (the max formulation, the decreasing marginal offloading benefit from the ordering), and it suggests that other subsystems in the training stack might similarly simplify when their interactions are properly modeled. This is a transposable insight: the reason prior activation management was hard is that it was solved in isolation; when embedded in the correct holistic model, the problem structure simplifies dramatically.
The paper also resolves a practical contradiction in the "train at home" narrative. Before LoHan, a researcher who read the ZeRO-Infinity paper might reasonably conclude that SSD offloading enables extreme-scale training, then attempt to fine-tune a 70B model on their desktop and fail because ZeRO-Infinity requires ~1.1 TB of main memory for a 175B model (Section III-B, Issue 3). The paper explains this failure precisely: ZeRO-Infinity offloads activations only to main memory, so the maximum model size is main-memory-bound, not SSD-bound. The solution is not "use a different offloading system" but "offload activations to SSDs as well, and schedule the combined traffic holistically." By making the failure mode explicit, the paper gives practitioners a diagnostic toolkit for understanding why their fine-tuning job runs out of memory, rather than a black-box recommendation to buy more hardware.
Follow-Up Research This Work Enables
Joint optimization of activation management across multiple GPUs with heterogeneous memory capacities. LoHan's convex optimization framework (Section IV-D) solves for the optimal activation offloading volume on a single GPU. In a multi-GPU server where different GPUs may have different memory capacities (e.g., a mix of RTX 4090s and RTX 4080s), or where some GPUs are closer to the SSDs than others (different PCIe topology), the optimal offloading policy is per-GPU rather than global. A direct extension would formulate a multi-GPU version of the iteration time model where each GPU has its own decision variable, and the global iteration time is the maximum across GPUs (since data-parallel training synchronizes at the end of each backward pass). The convexity of the single-GPU objective would carry over if the per-GPU objectives are convex and the max-of-convex-functions property holds, but the coupling through gradient all-reduce communication (which adds a shared PCIe traffic term) might break convexity. A strong follow-up would implement this multi-GPU extension, measure whether the per-GPU optimal policies differ significantly in a heterogeneous setup (e.g., 2× RTX 4090 + 2× RTX 4080), and characterize how much throughput is lost by using a uniform policy versus per-GPU optimization.
Low-overhead difficulty estimation from model-internal signals during the first few layers of the forward pass. The paper's profiling stage (Section IV-B) measures hardware capabilities once at initialization, but it does not address a related problem: some model characteristics (e.g., the actual activation sizes for a specific input sequence length, the effective GPU throughput under a specific batch composition) can vary across inputs or training samples. The convex optimization model uses fixed profiling numbers, which assumes these characteristics are constant. In practice, variable-length sequences (common in NLP fine-tuning) create variable activation sizes, and different input compositions (e.g., padding ratios) change effective GPU throughput. A natural extension is an online profiler that monitors actual activation sizes and GPU throughput during the first few training iterations, then adjusts the optimal on the fly using the same convex optimization machinery but with updated profiling numbers. This would make LoHan robust to input variability without requiring the user to pre-specify sequence length distributions. A strong evaluation would compare the online-adapted policy against the static profiler-based policy on a fine-tuning workload with highly variable sequence lengths (e.g., instruction-tuning data where prompts range from 50 to 4000 tokens), measuring whether the online adaptation prevents GPU out-of-memory errors while maintaining throughput.
Convergence validation of active gradient offloading against standard synchronous training. The paper demonstrates throughput improvements from overlapping optimizer execution with backward propagation (Section V-D, Figure 7), but does not verify that this overlapping schedule produces the same training dynamics as the standard serialized schedule. While LoHan maintains synchronous updates (each iteration's forward pass uses the fully updated parameters from the previous iteration — footnote 4), the specific ordering of optimizer updates across layers is different from the serialized schedule: in serialized execution, all layers' parameters are updated after all gradients are computed; in LoHan, layer 's parameters are updated while layer 's gradients are still being computed. In theory, this should be equivalent because the Adam update for each layer depends only on that layer's gradients and stored optimizer states, which are identical in both schedules. In practice, floating-point nondeterminism from different operation orderings (e.g., different accumulation orders in the Adam momentum updates due to different timing of gradient arrival relative to SSD prefetching) could cause divergence over many iterations. A critical stress-test would fine-tune a moderate-scale model (e.g., 1–6B parameters) on a standard NLP benchmark using both LoHan and a standard non-offloaded PyTorch baseline, train both to convergence, and compare not just final accuracy but per-iteration loss trajectories. If the trajectories diverge significantly, it would indicate that the overlapping schedule introduces numerical differences that matter for training dynamics, which would be an important negative result qualifying the paper's throughput claims.
Integration with parameter-efficient fine-tuning (PEFT) methods for a unified cost model. LoHan targets full-model fine-tuning with all parameters updated. Parameter-efficient methods like LoRA, adapter layers, and prefix tuning reduce the number of trainable parameters by orders of magnitude, which correspondingly reduces optimizer state size, gradient size, and the I/O traffic for model states. A LoRA fine-tuning of a 175B model might require updating only 0.1–1% of parameters, which would dramatically change the resource bottleneck — the optimizer stage becomes negligible, and activation offloading (which LoRA does not reduce) becomes the dominant cost. A natural systems question is: given a fixed fine-tuning budget (hardware + time), where is the crossover point between full fine-tuning with LoHan and LoRA fine-tuning with a simpler offloading system? A follow-up study would implement LoRA within LoHan's framework (the hook-based architecture in Section IV-E should support it without modification, since LoRA layers are standard PyTorch modules), measure throughput and time-to-convergence for both full fine-tuning and LoRA across model scales from 6B to 175B, and produce a decision boundary: below model size X, full fine-tuning with LoHan is faster; above X, LoRA becomes preferable because the reduced optimizer traffic shifts the bottleneck entirely to activations. The paper's profiling infrastructure and cost model provide the necessary measurement apparatus for such a study.
Extension to training-from-scratch workloads where optimizer states dominate even more. The paper focuses on fine-tuning, where the model is initialized from a pre-trained checkpoint and trained for relatively few iterations. In training-from-scratch, the optimizer states (Adam's first and second moment estimates) start from zero and accumulate over many more iterations, but their size is the same as in fine-tuning — the difference is that total training time is much longer, making throughput even more critical. More importantly, training-from-scratch typically uses larger global batch sizes (to maintain gradient signal quality across more iterations), which shifts the activation-to-model-state ratio toward activations. How does the optimal change when moving from fine-tuning (small batch, short training) to pre-training (large batch, long training)? And does active gradient offloading's overlap opportunity change when the backward propagation time scales with batch size while the optimizer time remains fixed? A direct experiment would configure LoHan for a training-from-scratch workload on a modest model (e.g., 1–6B), sweep batch sizes, and compare the optimal offloading policy from Algorithm 1 against the fine-tuning-derived policy. If the optimal policy differs substantially, it would indicate that the profiling stage should be workload-aware (fine-tuning vs. pre-training) rather than purely hardware-aware.
Characterization of the numerical precision frontier for SSD-offloaded optimizer states. LoHan stores optimizer states (FP32), parameters (FP32 master copy, FP16 working copy), and gradients (FP16) on SSDs, transferring them to main memory for the CPU optimizer. SSD I/O is the primary cost in the active gradient offloading pipeline (Figure 3b), and the optimizer state traffic is bytes per iteration — 14 bytes per parameter, dominated by the 8-byte FP32 optimizer states. A natural question is whether lower-precision optimizer states (e.g., FP16 or BF16 for the Adam moments) would degrade convergence while proportionally reducing SSD I/O traffic. If FP16 optimizer states work, the term in the backward-stage time model (Equation 5) would drop to , shifting the bottleneck and potentially changing the optimal . A well-designed follow-up would fine-tune a model at a scale where SSD I/O is the bottleneck (e.g., 70B on RTX 4090 with 3 SSDs, where Figure 10a shows linear SSD scaling), compare convergence with FP32 vs. BF16 optimizer states at identical batch sizes and learning rates, and measure whether the throughput gain from reduced I/O outweighs any degradation in convergence speed. This would establish a precision-vs-throughput Pareto frontier for offloaded training that the paper currently does not address.
Practical Applications and Downstream Use Cases
Enabling independent researchers and small labs to fine-tune 100B-scale models on existing desktop hardware. The most direct application of LoHan is that a researcher with a commodity desktop (RTX 4090, 256 GB main memory, 3–4 NVMe SSDs) can fine-tune models up to 175B parameters that previously required a DGX-class cluster. The cost implications are stark: Table VII estimates the commodity server at approximately 22,300 (4× RTX 4090 + server + 3–6 SSDs) versus $200,000 for a DGX-A100. For a PhD student adapting a 70B model to a new language or domain, or a startup fine-tuning a 175B model on proprietary customer data, this is the difference between "possible with available resources" and "requires a cloud compute grant." The paper's maximum model size results (Figure 6a: 175B on RTX 4090 with 256 GB main memory) provide the feasibility guarantee; the throughput numbers (Figure 5c: 53% of peak TFLOPS at 175B, declining with model size) provide the performance expectation. A concrete deployment scenario: fine-tuning Llama-2-70B on a domain-specific corpus using a single RTX 4090 with 256 GB RAM and 4 SSDs would be feasible at batch size 1–2, with throughput around half the GPU's theoretical peak — slow but functional, and dramatically cheaper than renting an 8× A100 instance.
Cost-efficient batch fine-tuning pipelines for model customization services. Organizations that fine-tune models for multiple clients or multiple tasks can use LoHan to run many fine-tuning jobs in parallel on separate commodity servers rather than time-sharing an expensive cluster. The cost-effectiveness comparison (Figure 13) shows that a 4× RTX 4090 server achieves 2.17× higher throughput-per-dollar than a DGX-A100 for a 30B model. For a model-as-a-service company that needs to produce 100 fine-tuned model variants per week for different enterprise customers, deploying 5 commodity servers (total ~200,000), with the added benefit that jobs can run independently without cluster scheduling overhead. The SSD scaling results (Figure 10a) indicate that 3–6 SSDs per server provide most of the achievable I/O bandwidth, so the per-server cost can be tuned to the specific model size and batch requirements.
On-premise fine-tuning for privacy-sensitive applications. Many enterprises (healthcare, finance, legal) cannot send their data to cloud GPU providers due to regulatory constraints, but also cannot justify purchasing a DGX-class cluster for occasional fine-tuning jobs. LoHan enables these organizations to fine-tune large models on their existing server infrastructure — a hospital with a commodity server in its IT department could fine-tune a 70B medical language model on patient records without the data ever leaving the building. The paper's demonstration that LoHan works across consumer GPU tiers (RTX 4090, 3090, 4080 — Figure 6) means the solution is not tied to a specific high-end GPU model. The framework integration (Section IV-E) requires minimal code changes to existing PyTorch training scripts, lowering the adoption barrier for teams without deep systems expertise.
When to Prefer This Method
The paper positions LoHan as a framework for consumer-grade GPU fine-tuning, but its design choices imply clear boundaries against alternative approaches:
-
Prefer LoHan over ZeRO-Infinity / ZeRO-Offload / Colossal-AI when the target model's model states exceed GPU memory (approximately >6B parameters on an RTX 4090 with 24 GB, as established by FlashNeuron's failure in Figure 2a) AND the available main memory is insufficient to hold all activations at the desired batch size (the threshold where ZeRO-Infinity fails, e.g., 175B requires ~1.1 TB main memory — Section III-B, Issue 3). Under these conditions, LoHan's SSD activation offloading and holistic scheduling are necessary, not merely beneficial.
-
Prefer LoHan over G10 when GPUDirect is unavailable (all consumer GPUs — Section III-C, Issue 3) OR when the PCIe link bandwidth is limited relative to model state size (G10's GPU-resident optimizer creates massive model state transfer traffic that saturates the PCIe link — Figure 1b shows 13 seconds of transfer for a 0.1-second GPU computation). LoHan's CPU optimizer eliminates this transfer entirely.
-
Prefer LoHan over parameter-efficient fine-tuning (LoRA, adapters) when full fine-tuning is required for task performance (PEFT methods reduce trainable capacity) OR when the model size is moderate enough (roughly ≤70B on the tested hardware) that LoHan achieves high GPU utilization (90–95% of peak FLOPS, Figure 5c). At small scales where PEFT provides sufficient accuracy, the simplicity of PEFT likely outweighs LoHan's throughput advantage.
-
Prefer a DGX-class cluster over LoHan when the target model exceeds ~276B parameters (LoHan's maximum with 768 GB main memory — Figure 6a) OR when training requires large batch sizes for convergence (LoHan's batch size is GPU-memory-limited per layer, forcing batch size 1 at 175B — Section V-C) OR when wall-clock training time is the primary constraint (a DGX with 8× A100s will complete the same fine-tuning job in a fraction of the time, even if LoHan has higher throughput-per-dollar). The paper's cost-effectiveness advantage (2.17×, Figure 13) is specific to throughput-per-dollar, not absolute time-to-solution.