MegaTrain: Обучение с полной точностью больших языковых моделей со 100+ миллиардами параметров на одном графическом процессоре

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

April 6, 2026

Авторы: Zhengqing Yuan, Hanchi Sun, Lichao Sun, Yanfang Ye

cs.AI

Аннотация

Мы представляем MegaTrain — систему с памятью в качестве центрального элемента, которая эффективно обучает большие языковые модели со 100+ миллиардами параметров с полной точностью на одном GPU. В отличие от традиционных систем, ориентированных на GPU, MegaTrain хранит параметры и состояния оптимизатора в памяти хоста (память CPU) и использует GPU как транзитные вычислительные модули. Для каждого слоя мы потоково загружаем параметры и вычисляем градиенты, сводя к минимуму постоянное состояние на устройстве. Для преодоления узкого места пропускной способности между CPU и GPU мы применяем две ключевые оптимизации. 1) Мы внедряем конвейерный исполняющий механизм с двойной буферизацией, который совмещает предварительную выборку параметров, вычисления и выгрузку градиентов в нескольких потоках CUDA, обеспечивая непрерывное выполнение на GPU. 2) Мы заменяем постоянные графы автоградиента на шаблоны слоев без состояния, динамически связывая веса по мере их потоковой передачи, что устраняет метаданные постоянного графа и обеспечивает гибкость планирования. На одном GPU H200 с 1,5 ТБ памяти хоста MegaTrain стабильно обучает модели до 120 млрд параметров. Она также демонстрирует в 1,84 раза большую пропускную способность обучения по сравнению с DeepSpeed ZeRO-3 с выгрузкой на CPU при обучении 14-миллиардных моделей. Кроме того, MegaTrain позволяет обучать 7-миллиардные модели с контекстом в 512 тыс. токенов на одном GH200.

English

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.84times 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.