MegaTrain: Entrenamiento de Precisión Completa de Modelos de Lenguaje Grandes de 100B+ Parámetros en una Única GPU

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

April 6, 2026

Autores: Zhengqing Yuan, Hanchi Sun, Lichao Sun, Yanfang Ye

cs.AI

Resumen

Presentamos MegaTrain, un sistema centrado en la memoria que entrena eficientemente modelos de lenguaje grandes de más de 100B parámetros en precisión completa en una sola GPU. A diferencia de los sistemas tradicionales centrados en la GPU, MegaTrain almacena los parámetros y los estados del optimizador en la memoria del host (memoria de la CPU) y trata las GPUs como motores de cálculo transitorios. Para cada capa, transmitimos los parámetros hacia la GPU y calculamos los gradientes hacia fuera, minimizando el estado persistente en el dispositivo. Para combatir el cuello de botella del ancho de banda entre la CPU y la GPU, adoptamos dos optimizaciones clave. 1) Introducimos un motor de ejecución segmentado con doble búfer que superpone la precarga de parámetros, el cálculo y la descarga de gradientes a través de múltiples flujos CUDA, permitiendo una ejecución continua en la GPU. 2) Reemplazamos los grafos de autograd persistentes con plantillas de capas sin estado, vinculando los pesos dinámicamente a medida que se transmiten, eliminando así los metadatos persistentes del grafo mientras proporcionamos flexibilidad en la planificación. En una sola GPU H200 con 1.5TB de memoria del host, MegaTrain entrena de forma fiable modelos de hasta 120B parámetros. También logra un rendimiento de entrenamiento 1.84 veces superior al de DeepSpeed ZeRO-3 con descarga a CPU al entrenar modelos de 14B. MegaTrain también permite el entrenamiento de modelos de 7B con un contexto de 512k tokens en un solo 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.