Lichao Sun

We present MegaTrain, a CPU-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 and treats GPUs as transient compute engines. For each layer, we stream parameters in and compute gradients out, minimizing persistent GPU 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 256k token context on a single H100.

Residual connections are fundamental to deep transformer training, yet their theoretical relationship to more expressive alternatives, such as attention-based depth mixing (AttnRes), remains poorly understood. We show that the depth-axis mixing matrix in transformers is semiseparable, and its rank precisely controls the expressiveness of inter-layer information flow. Standard residual connections correspond to rank-1 semiseparable matrices; attention residuals to full rank. We introduce DEPTHSSD (Depth State Space Duality), which applies the structured state space duality framework from Mamba-2 to the depth axis, yielding a depth mixing matrix of controllable rank N with O(L · N) parameters per sub-layer. We evaluate DEPTHSSD at three model scales (125M, 455M, and 1B parameters) against four depth-mixing baselines: standard residuals, BLOCKATTNRES, DENSEFORMER, and Hyper-Connections. At 1B, DEPTHSSD-N=8 achieves the best validation loss (3.124, perplexity 22.7), a 2.4-point perplexity reduction over standard residuals (PPL 25.1) and outperforming DENSEFORMER (PPL 24.1), while running 1.72× faster (16.5 vs. 9.6 TFLOPS). At 455M, DEPTHSSD-N=8 again leads (PPL 29.4 vs. 30.1 for DENSEFORMER and 31.2 for standard) at 1.88× the throughput of DENSEFORMER. At 125M, DEPTHSSD matches the validation quality of all baselines while maintaining 1.3–1.6× higher throughput than the other depth-mixing methods.

Token-choice Mixture-of-Experts (TC-MoE) routes each token to a fixed number of experts, limiting dynamic computation allocation and requiring auxiliary losses to maintain load balance. We propose Expert Threshold (ET) routing, where each expert maintains an exponential moving average (EMA) threshold estimated from the global token distribution. At both training and inference, each token is independently routed to an expert if its score exceeds the expert's threshold, enabling dynamic computation allocation while achieving load balance without auxiliary losses. This fully causal mechanism eliminates dependence on other tokens in the batch, making it well-suited for autoregressive language modeling. In pretraining experiments scaling to 2.4B parameters on FineWeb-Edu, ET achieves 0.067 lower cross-entropy loss than TC-MoE, equivalent to reaching the same performance with 1.6$\times$ fewer tokens.