El entrenamiento de modelos de lenguaje choca cada vez más contra los límites de memoria de la GPU antes de aprovechar del todo el cómputo. Pesos del modelo, gradientes, estados del optimizador, buffers de comunicación y activaciones intermedias compiten por la memoria de alto ancho de banda (HBM). A medida que crecen el tamaño del modelo, la longitud de secuencia y el lote, la capacidad de HBM suele convertirse en el principal cuello de botella para escalar.
Una publicación técnica de NVIDIA explica cómo el host offloading en la librería de código abierto JAX alivia esa presión. El proceso mueve activaciones seleccionadas a memoria del sistema "anclada" (pinned host memory) durante la pasada hacia adelante y las trae de vuelta cuando se necesitan en la pasada hacia atrás. Es una alternativa a la rematerialización de activaciones: en lugar de recalcularlas, el paso de entrenamiento las recarga desde la RAM del host.
¿Por qué conviene en sistemas NVIDIA Grace Blackwell?
La técnica resulta especialmente ventajosa en los sistemas NVIDIA Grace Blackwell. La CPU NVIDIA Grace y la GPU NVIDIA Blackwell se conectan por NVLink-C2C con 900 GB/s de ancho de banda bidireccional, lo que vuelve práctica a la memoria del host como área de paso. La combinación de CPU Vera y GPU Rubin duplica esa cifra a 1,8 TB/s de ancho de banda coherente.
El ancho de banda alto entre CPU y GPU ayuda, pero no basta. Para mejorar el rendimiento, las transferencias de activaciones deben solaparse con trabajo útil de la GPU.
Resultados sobre cargas de MaxText
Los experimentos usan MaxText, un marco de entrenamiento de LLM en JAX que apoya el compilador XLA para escalar sobre GPU de NVIDIA. Todos los resultados se midieron en sistemas NVIDIA GB200 NVL72 con 128 GPU, con dos cargas de trabajo: Llama 3.1 405B, un modelo denso, y DeepSeek-V3 671B, un modelo disperso de mezcla de expertos (MoE).
DeepSeek-V3 671B
DeepSeek-V3 671B tiene 61 capas decodificadoras: las tres primeras usan bloques MLP densos y el resto usa bloques MoE, que dominan la pila. La política de descarga mueve intermedios seleccionados de las proyecciones de atención latente (MLA) y de la proyección ascendente de los expertos, activaciones lo bastante grandes como para decidir si un lote mayor cabe o no en memoria.
Con la descarga activada, junto al planificador que oculta latencia (LHS) y las transferencias segmentadas, la carga alcanzó 908,2 TFLOPs/s por dispositivo. Esto fue 57% más rápido que la rematerialización de activaciones con la misma configuración de lote, y 67,7% más rápido que la descarga sin LHS ni segmentación.
A diferencia de la carga densa de Llama, donde el LHS por sí solo bastaba para ocultar la latencia, la enorme huella de activaciones de las capas MoE y MLA de DeepSeek-V3 hace que activar las transferencias segmentadas aporte un impacto positivo y diferenciado sobre el rendimiento total.
¿Cuánto crece el lote que cabe en memoria?
El host offloading también cambia qué configuraciones de lote son viables. Guardar activaciones seleccionadas en el dispositivo permitía un micro lote de 2 y un lote global de 256; con la descarga optimizada, se alcanzó un micro lote de 8 y un lote global de 1.024. Sin descarga ni rematerialización, el dispositivo choca con un error de memoria agotada (OOM) al intentar esa configuración mayor.
Con LHS y transferencias segmentadas, la corrida con descarga usa 165,2 GiB de memoria de GPU, frente a 145,6 GiB sin esas optimizaciones. El aumento proviene de mantener más buffers de copia y activaciones precargadas mientras las transferencias se solapan con el cómputo: se cambia algo de capacidad de memoria por mejor solapamiento y mayor rendimiento.
Llama 3.1 405B
El experimento con Llama 3.1 405B corrió 10 pasos sobre datos sintéticos, con lote 2, longitud de secuencia 8.192, paralelismo de datos totalmente fragmentado (FSDP) en 128 y cuantización de pesos a 4 bits NVFP4. La descarga de activaciones QKV con LHS mejoró el rendimiento de 2.669 a 2.746 TFLOPs/s por dispositivo, un alza de 2,9% sobre la línea base. Al desactivar el LHS, el rendimiento cayó a 2.569 TFLOPs/s por dispositivo, lo que confirma que la técnica depende de un solapamiento efectivo con otro trabajo de la GPU.
La ganancia es menor que en DeepSeek-V3 671B, pero muestra el mismo mecanismo: la descarga dirigida de QKV reemplaza la rematerialización en la pasada hacia atrás por transferencias que se solapan con el cómputo y la comunicación.
¿Cuándo conviene usar host offloading?
El host offloading es más útil cuando la memoria de la GPU limita el tamaño del modelo, la longitud de secuencia o el lote, y cuando los tensores seleccionados son lo bastante grandes como para aliviar la presión de HBM. Rinde mejor cuando la carga tiene suficiente cómputo o comunicación independiente para ocultar la latencia de transferencia.
Resulta especialmente ventajoso en plataformas como NVIDIA Blackwell GB200 y Blackwell Ultra GB300, que usan el interconector dedicado NVLink-C2C para saltarse por completo los cuellos de botella del PCIe tradicional. Con la plataforma NVIDIA Vera Rubin ofreciendo aún más ancho de banda, la técnica se vuelve una palanca predecible para desacoplar el rendimiento del entrenamiento de los límites físicos estrictos de memoria.
Es menos probable que ayude cuando los tensores son pequeños, cuando hay poco trabajo independiente para solapar las transferencias o cuando el cuello de botella está en otro lado. Para equipos de LatAm que arriendan cómputo por hora, el atractivo es concreto: entrenar un modelo más grande sin sumar GPU adicionales, al costo de un ajuste fino de solapamiento que conviene validar con perfiladores como NVIDIA Nsight Systems antes de dar por buena la ganancia.
