LazyTrain: Limited-resource Allocation toward Zero-waste Yield Optimization in Large Language Model Training

2026-08-12

AI Summary

This work addresses performance bottlenecks arising from inefficient communication and computation scheduling when training large language models under constrained hardware resources. The authors propose a unified mixed-integer programming framework that jointly optimizes checkpoint selection, activation placement, recomputation strategies, and overlapping of CPU–GPU–NVMe communications, formulating multidimensional resource scheduling as a mixed-integer problem for the first time. Additionally, they introduce a Hybrid 8-bit operator that integrates 8-bit optimizer state compression with fast gradient clipping to reduce memory usage while preserving computational efficiency. Experiments on training the Qwen3.6-27B model on H800 GPUs demonstrate a peak throughput of 1,361 tokens/s (a 1.24× improvement), requiring only 68.84 GB of GPU memory, achieving 95.42% accuracy, and delivering an effective compute performance of 219.95 TFLOPS.

Abstract

Training large language models on limited hardware is increasingly a scheduling problem across GPU compute, host memory, PCIe transfer, and storage bandwidth. Existing offloading systems reduce GPU residency, and MegaTrain shows that a CPU-master layer-streaming executor can train large models on a single GPU, but fixed checkpointing and placement heuristics still leave communication exposed on the critical path. We propose LazyTrain, an optimization layer over a layer-streaming executor. LazyTrain formulates checkpoint selection, activation placement, recomputation, and CPU-GPU-NVMe communication overlap as a mixed-integer scheduling problem, then executes the solved policy during training. It further couples 8-bit optimizer states with fast gradient clipping as a single Hybrid 8-bit operator: state compression reduces optimizer-state memory, while fast clipping counteracts the additional CPU-side update overhead. Across H800 experiments from Qwen2.5-3B to Qwen3.6-27B, LazyTrain improves sustained TFLOPS over matched baselines runs by approximately 1.24$\times$; RTX 3090 experiments likewise increase the maximum feasible batch size by one at each model scale. In the primary Qwen3.6-27B H800 MetaMathQA run, LazyTrain reaches 219.95 TFLOPS and 1361 tokens/s at batch size 72, peaks at 68.84\,GB of GPU memory, and obtains 95.42\% exact-match accuracy on the full evaluation split. The source code is available at https://github.com/DataArcTech/LazyTrain.

Problem

• large language model training
• resource allocation
• checkpointing
• communication overlap
• memory optimization

Innovation

• layer-streaming
• mixed-integer scheduling
• checkpoint optimization
• Hybrid 8-bit operator
• resource-constrained training