AI 基礎設施與訓練系統
LazyTrain Overlaps Memory Transfers With Mixed-Integer Scheduling, Speeding Up 27B Model Training on a Single H800 by 1.24×
Before training, LazyTrain jointly determines where activations should be saved, offloaded, or recomputed, fitting PCIe and NVMe transfers into GPU compute windows. In Qwen3.6-27B experiments, performance increased from 176.90 to 219.95 TFLOPS, although the system has so far been validated only on a single GPU with offline scheduling.
When memory is insufficient, single-GPU training of large models typically relies on fixed rules to choose among activation checkpointing, CPU offload, and recomputation. LazyTrain reformulates these choices as a mixed-integer linear program. Given GPU, system memory, and NVMe capacities, along with PCIe and SSD bandwidth and per-layer compute times, it jointly determines which boundaries require checkpoints, which storage tier should hold them, which blocks should be recomputed, and which compute windows should accommodate the transfers.
The system constructs path-flow constraints spanning the beginning to the end of backpropagation, while also enforcing HBM, DRAM, and NVMe capacity limits and bidirectional traffic constraints for parameters, gradients, and activations. The SCIP solver generates the schedule before training, so the runtime does not perform an online search. The main Qwen3.6-27B configuration has 64 Transformer blocks, a sequence length of 1,024, and a batch size of 72. Of the 52 materialized checkpoints, 30 remain on the GPU, 21 are placed on the CPU, and one is stored on NVMe, while parameters and optimizer state reside primarily on the CPU.
On the same H800, with identical data partitioning and batch size, LazyTrain increased sustained compute performance from MegaTrain’s 176.90 TFLOPS to 219.95 TFLOPS, an improvement of approximately 1.24×. Throughput reached 1,361 tokens per second, with peak GPU memory usage of about 68.84 GB. Removing the mixed-integer scheduling reduced performance to 193.17 TFLOPS, indicating that most of the gain comes from jointly configuring transfers and recomputation. By contrast, removing the so-called Hybrid 8-bit path lowered performance only to 219.29 TFLOPS, and quantized states were used for only about 2% of the parameters.
The open-source alpha is built on MegaTrain’s CPU-master, layer-by-layer streaming architecture and includes configurations for Qwen3.6-27B, MetaMathQA, and NVMe scheduling. However, the paper reports experiments only on single H800 and RTX 3090 GPUs, with most table entries based on a single run. The schedule is not reconfigured in response to bandwidth fluctuations during execution, and the claimed zero additional exposed transfer time is a solver estimate rather than a result validated through step-by-step stall instrumentation. The more important question is whether the gains persist across multiple GPUs, multiple nodes, and different Mixture-of-Experts (MoE) architectures—not the peak result from a single 27B case.