PaperlayerView exact revision on arXiv

Source-linked AI summary

MegaTrain: Full Precision Training of 100B+ Parameter Large Language Models on a Single GPU

Zhengqing Yuan, Hanchi Sun, Lichao Sun, Yanfang Ye

arXiv:2604.05091v1cs.CLcs.DCcs.OS

TL;DR

MegaTrain addresses the challenge of full-precision training for 100B+ parameter language models when GPU memory is limited. It stores persistent state in host memory and streams layers through transient GPU templates using overlapped transfers and stateless execution. The system trains up to 120B parameters on one H200 and delivers 1.84× ZeRO-3 Offload throughput at 14B on GH200.

  • Problem

    Hundred-billion-parameter training remains inaccessible on limited hardware because full parameters and optimizer states must otherwise reside in memory, while layer streaming conflicts with persistent execution graphs.

  • Method

    MegaTrain makes host memory the authoritative store, streams parameters and gradients layer by layer, overlaps transfers with computation, and uses stateless templates instead of persistent autograd graphs.

  • Results

    MegaTrain trains models up to 120B parameters on one H200 with 1.5TB host memory and achieves 1.84× ZeRO-3 Offload training throughput at 14B on GH200.

  • Takeaways & Limitations

    The system decouples model scale from GPU memory capacity while sustaining efficient execution through pipelined streaming and bounded per-layer device memory.

  • Takeaways & Limitations

    The paper identifies extending MegaTrain to multiple GPUs and incorporating SSD-based tiered storage as future work.

Abstract

from arXiv · show

We present MegaTrain, a memory-centric system that efficiently trains 100B+ parameter large language models at full precision on a single GPU. Unlike traditional GPU-centric systems, MegaTrain stores parameters and optimizer states in host memory (CPU memory) and treats GPUs as transient compute engines. For each layer, we stream parameters in and compute gradients out, minimizing persistent device state. To battle the CPU-GPU bandwidth bottleneck, we adopt two key optimizations. 1) We introduce a pipelined double-buffered execution engine that overlaps parameter prefetching, computation, and gradient offloading across multiple CUDA streams, enabling continuous GPU execution. 2) We replace persistent autograd graphs with stateless layer templates, binding weights dynamically as they stream in, eliminating persistent graph metadata while providing flexibility in scheduling. On a single H200 GPU with 1.5TB host memory, MegaTrain reliably trains models up to 120B parameters. It also achieves 1.84$\times$ the training throughput of DeepSpeed ZeRO-3 with CPU offloading when training 14B models. MegaTrain also enables 7B model training with 512k token context on a single GH200.

1 Introduction

MegaTrain addresses the difficulty of training hundred-billion-parameter models on scarce GPU hardware by making host memory the authoritative store and GPUs transient compute engines. Its pipelined streaming and stateless execution designs support large-scale training while maintaining sustained performance.

  • Motivation and contribution: MegaTrain enables full-precision training of 100B+ parameter models on a single GPU.Parameters and optimizer states reside in host memory, while device memory is used transiently for computation.
  • System design: MegaTrain streams parameters layer by layer, offloads gradients, and keeps optimizer states in host memory.This design reduces persistent device state and decouples model scale from GPU memory size.
  • System design: A pipelined double-buffered engine overlaps parameter prefetching, computation, and gradient offloading across CUDA streams.The overlap targets the CPU-GPU bandwidth bottleneck while keeping memory footprints bounded.
  • System design: Stateless layer templates dynamically bind streamed weights, eliminating persistent autograd graph metadata and intermediate tensors.This supports layer-wise weight eviction and limits device memory to the footprint of a single layer.
  • Results: 120B parameters are trainable on a single H200 with 1.5TB host memory, while 14B training reaches 1.84× ZeRO-3 Offload throughput on a GH200.MegaTrain also supports 512K-token context training for 7B models on a single GH200.

2 Preliminaries

The preliminaries frame training as a memory-hierarchy and execution-scheduling problem. Existing systems treat GPU memory as the primary working store, while layer-wise streaming conflicts with persistent autograd and static CUDA Graph assumptions.

  • Training memory: Training memory comprises persistent state, activations, and operator workspaces.For mixed-precision Adam, persistent state requires at least 12P bytes for P parameters.
  • Training memory: A 70B model requires at least 840 GB of persistent state under the stated mixed-precision Adam accounting.The accounting includes BF16 weights, BF16 gradients, and FP32 optimizer moments.
  • Memory hierarchy: The memory hierarchy places frequently needed data in smaller, faster tiers and colder data in larger, cheaper tiers.Host memory offers greater capacity than device memory at lower cost per byte, while NVMe provides still larger capacity at lower bandwidth.
  • Memory hierarchy: H200 and GH200 systems expose different CPU-GPU boundaries, with GH200 using NVLink-C2C and H200 using PCIe.The stated interconnect and host-memory configurations differ substantially between the two evaluation platforms.
  • Execution challenges: MegaTrain treats host memory as a first-class parameter store and coordinates prefetch, computation, and gradient offload explicitly.The architecture uses transient layer templates on the GPU rather than persistent model state.
  • Execution challenges: Layer-wise streaming breaks global autograd assumptions because parameters are evicted and activations cannot remain on the GPU indefinitely.Dynamic buffer bindings and shifting synchronization points also prevent a single static CUDA Graph from capturing execution.

3 MegaTrain System

MegaTrain executes training by streaming layer state between host and device memory, recomputing activations as needed, and coordinating transfers with explicit asynchronous scheduling. Stateless templates and buffered CUDA streams keep device memory bounded while hiding data-movement latency.

  • Execution workflow: Persistent parameters, optimizer states, and accumulated gradients remain in host memory, while device memory holds lightweight reusable layer templates.Templates dynamically bind to streamed parameters during execution.
  • Execution workflow: During forward execution, MegaTrain prefetches each next layer’s parameters, binds them to a template, computes, and releases the layer buffer.The workflow therefore avoids retaining the full model on the GPU.
  • Execution workflow: During backward execution, block checkpoints support activation recomputation, while each layer’s gradients are immediately offloaded to host memory.Only one block’s activations need to remain available at a time, keeping device memory independent of total model depth.
  • Execution workflow: CPU-side Adam updates avoid repeatedly transferring optimizer states to and from the GPU.The optimizer phase is compute-light but I/O-intensive, making CPU execution suitable for reducing data movement.
  • Asynchronous execution: Three CUDA streams and double-buffered staging overlap compute, host-to-device weight transfer, and device-to-host gradient transfer.Background gradient evacuation prevents D2H latency from entering the critical path when overlap is maintained.
  • Runtime mechanisms: CUDA events coordinate weight readiness, backward completion, and buffer reuse without relying on a global training graph.Layer-contiguous tiling and pinned staging pools support large, efficient transfers with bounded host pinning.
  • Stateless execution: Stateless templates decouple mathematical structure from physical data and support dynamic weight binding as buffers change.MegaTrain also avoids CUDA graph capture because streamed execution has dynamic addresses and synchronization points.

4 Evaluation

The evaluation uses GH200 and H200 systems, with experiments involving model configurations, memory layouts, and training data. The supplied evaluation passages identify the platforms and benchmark setup but do not report outcome values.

  • Experimental systems: GH200 experiments use a single superchip with one GH200 GPU and 480 GB of host memory.Each superchip integrates a 72-core Grace ARM CPU and a GPU with 96 GB HBM3 memory.
  • Experimental systems: H200 experiments use one H200 SXM GPU with 141 GB HBM3e memory and 1.5 TB of host memory.The GPU connects to the host through PCIe Gen4.
  • Evaluation data: Accuracy is evaluated on MetaMathQA, comprising approximately 395,000 English math problem-answer pairs.The dataset is divided into approximately 70% training and 30% testing samples.
  • Model configurations: The evaluation reports base-model parameter counts and transformer depth among the model characteristics summarized in Table 2.The supplied passage does not include the table’s individual configurations.

W1 W2 W3

MegaTrain evaluates single-GPU training across model scales, architectures, depth, width, context length, correctness, and ablations. Across these settings, it sustains throughput and memory feasibility where offloading baselines degrade or fail.

  • Host Memory Scaling: 120B parameters remain within practical host-memory limits for MegaTrain, while existing offloading systems cross the feasibility boundary beyond 30B.The comparison attributes MegaTrain’s scaling to proportional memory growth without auxiliary duplication.
  • Compute Efficiency: 1.84× higher training throughput than ZeRO-3 Offload is achieved by MegaTrain at 14B, while throughput remains above 250 TFLOPS at 32B on GH200.On GH200, MegaTrain sustains 284 TFLOPS at 7B and 264 TFLOPS at 14B; on H200, it scales to 72B and 120B.
  • Depth Scalability: 20.1% throughput reduction occurs as depth grows from 28 to 180 layers, while both ZeRO-3 and FSDP encounter OOM beyond 84 layers.MegaTrain decreases from 284 to 227 TFLOPS as parameters increase from 10.9B to 43.0B.
  • Width Scalability: 35.0% throughput reduction occurs for MegaTrain from 1.0× to 3.0× width, versus 42.4% for ZeRO-3 and 43.9% for FSDP.At 3.5× width, MegaTrain sustains 193 TFLOPS while ZeRO-3 reaches 160 TFLOPS; beyond 4.0×, both baselines encounter OOM while MegaTrain reaches 5.0×.
  • Long-Context Training: Over 400 TFLOPS is sustained at 512K context on GH200, rising from 264.8 TFLOPS at 1K while memory usage remains stable.Layer-wise execution limits activation residency, and chunked MLP execution keeps the 512K setting within bounds.

5 Conclusion

MegaTrain enables full-precision training of 100B+ parameter models on a single GPU by making host memory authoritative and bounding GPU use to a per-layer footprint. The system suggests that memory and compute organization, rather than GPU capacity alone, is central to scaling single-node training.

  • MegaTrain enables full-precision training of 100B+ parameter models on a single GPU.
  • Host memory stores model state while GPUs act as transient compute engines, decoupling model scale from GPU memory capacity.
  • Pipelined parameter streaming and stateless layer templates bound GPU memory to a per-layer footprint while host memory scales linearly with model size.
  • Extending MegaTrain to multiple GPUs and tiered SSD storage is identified as a path toward larger-scale training.

A Implementation Details

MegaTrain combines a flexible PyTorch and CUDA runtime with C++ and CUDA extensions for critical-path data movement and optimization operations.

  • MegaTrain is implemented as a high-performance training runtime leveraging PyTorch and CUDA.
  • Critical-path operations such as batched parameter movement and SIMD-accelerated optimization run in C++ and CUDA extensions.

A.1 Host (CPU) Memory Management

MegaTrain organizes host memory as the authoritative parameter store and uses contiguous layouts, pinned staging slabs, and aliasing to support efficient layer-wise streaming.

  • The CPU parameter store is designed to maximize PCIe throughput and minimize host-side orchestration overhead.
  • Layer-Contiguous Tiling packs each Transformer layer’s weights, gradients, and Adam moments together in memory.
  • Pinned Slab Recycling stages only a small number of active slabs instead of pinning the entire model.
  • Flat-Tensor Layout uses a master store for capacity and two fixed-size pinned buffers matching the largest Transformer layer.
  • For tied embeddings, structural aliasing maps embedding and language-model head states to shared physical memory while tracking update readiness.
  • Weight tying avoids duplicate host-to-device transfers by remapping tied GPU pointers to the same device address.

A.2 Multi-Stream Pipeline and Scheduling

MegaTrain uses an event-driven, multi-stream pipeline with double buffering to overlap parameter transfers, computation, and gradient offloading, maintaining continuous GPU execution.

  • The scheduler manages concurrent ComputeStream, WeightStream, and GradStream operations.
  • Double buffering prefetches the next layer’s parameters while the compute stream executes the current layer.
  • The pipeline converts sequential execution into steady-state streaming so GPU compute units do not stall for parameters.
  • Weights-ready, backward-done, and buffer-free events explicitly coordinate transfers, computation, and buffer reuse.
  • Asynchronous gradient evacuation runs separately from computation so PCIe D2H latency does not enter the backward recomputation critical path.

A.3 GPU Buffer Management

MegaTrain manages GPU memory as a transient, explicitly controlled execution cache, using streaming buffers, pinned staging, gradient slabs, and preallocated workspaces. These mechanisms support a steady pipeline for large-model training while limiting allocation overhead and fragmentation.

  • GPU Buffer Management: The GPU acts as a transient execution cache with just-in-time parameter supply, high-bandwidth data movement, and a stateless device profile.Parameters are streamed into the GPU and gradients are evacuated while persistent device state is minimized.
  • GPU Buffer Management: Flat-buffer streaming packs each layer’s tensors into one contiguous pinned CPU buffer and transfers it asynchronously through a single H2D copy.GPU tensor views are created directly over the flat buffer to reduce transfer and unpacking overhead.
  • GPU Buffer Management: Batched Parameter Binding replaces hundreds of individual copy operations with one CUDA kernel and a metadata update for named layer parameters.The C++ extension targets Python dispatch overhead, which can exceed 10% of step time.
  • GPU Buffer Management: K=12 pinned host-memory slabs let GradStream offload completed layer gradients asynchronously while a background CPU thread accumulates them.This decouples GPU memory release from CPU optimization.
  • GPU Buffer Management: Preallocated, memory-mapped workspaces and reusable buffers eliminate allocation-induced pipeline bubbles and maintain a constant high-throughput steady state.The design targets training hundred-billion-parameter models near device capacity.
  • GPU Buffer Management: Two pinned staging buffers keep host pinning invariant to model depth, while recycled gradient slabs return only after CPU accumulation completes.Fixed-size host regions sustain the H2D and D2H streams without exhaustive pinning.

A.4 GPU Compute Dispatch

MegaTrain dispatches computation through stateless templates whose weights are bound dynamically as layers stream through the GPU. This avoids persistent autograd assumptions and overlaps binding with computation.

  • GPU Compute Dispatch: Stateless templates separate a layer’s mathematical structure from its physical weights, binding streamed parameters only immediately before execution.Templates contain Attention and MLP kernels but no persistent weight pointers.
  • GPU Compute Dispatch: Ping-pong binding executes F1 on Template A while W2 binds to Template B, removing weight preparation from the critical path.Alternating templates allow compute and parameter binding to proceed concurrently.
  • GPU Compute Dispatch: Layer-wise streaming makes persistent global autograd graphs unsuitable because parameters are evicted and activations cannot remain arbitrarily resident.MegaTrain therefore avoids CUDA graph capture and preserves explicit stream, ownership, and event dependencies.

B Ratel Reproduction on GH200

The GH200 reproduction of Ratel using its official codebase produced consistently low throughput across tested model sizes. The authors suspect SSD bottlenecks in the original implementation as the main cause.

  • Ratel Reproduction on GH200: Throughput remained consistently low across all tested Ratel model sizes on GH200 using the official codebase.The passage reports a reproduction result rather than a comparison with MegaTrain.
  • Ratel Reproduction on GH200: The authors suspect that SSD bottlenecks in Ratel’s original implementation mainly caused the low measured throughput.This is presented as a suspicion, not a confirmed causal diagnosis.
Loading 2604.05091v1…