El entrenamiento de modelos de lenguaje choca cada vez mas contra los limites de memoria de la GPU antes de aprovechar del todo el computo. Pesos del modelo, gradientes, estados del optimizador, buffers de comunicacion y activaciones intermedias compiten por la memoria de alto ancho de banda (HBM). A medida que crecen el tamano del modelo, la longitud de secuencia y el lote, la capacidad de HBM suele convertirse en el principal cuello de botella para escalar.
Una publicacion tecnica de NVIDIA explica como el host offloading en la libreria de codigo abierto JAX alivia esa presion. 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 atras. Es una alternativa a la rematerializacion de activaciones: en lugar de recalcularlas, el paso de entrenamiento las recarga desde la RAM del host.
¿Por que conviene en sistemas NVIDIA Grace Blackwell?
La tecnica 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 practica a la memoria del host como area de paso. La combinacion 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 util 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 politica de descarga mueve intermedios seleccionados de las proyecciones de atencion latente (MLA) y de la proyeccion 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 alcanzo 908,2 TFLOPs/s por dispositivo. Esto fue 57% mas rapido que la rematerializacion de activaciones con la misma configuracion de lote, y 67,7% mas rapido que la descarga sin LHS ni segmentacion.
A diferencia de la carga densa de Llama, donde el LHS por si 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.
¿Cuanto crece el lote que cabe en memoria?
El host offloading tambien cambia que configuraciones de lote son viables. Guardar activaciones seleccionadas en el dispositivo permitia un micro lote de 2 y un lote global de 256; con la descarga optimizada, se alcanzo un micro lote de 8 y un lote global de 1.024. Sin descarga ni rematerializacion, el dispositivo choca con un error de memoria agotada (OOM) al intentar esa configuracion 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 mas buffers de copia y activaciones precargadas mientras las transferencias se solapan con el computo: se cambia algo de capacidad de memoria por mejor solapamiento y mayor rendimiento.
Llama 3.1 405B
El experimento con Llama 3.1 405B corrio 10 pasos sobre datos sinteticos, con lote 2, longitud de secuencia 8.192, paralelismo de datos totalmente fragmentado (FSDP) en 128 y cuantizacion de pesos a 4 bits NVFP4. La descarga de activaciones QKV con LHS mejoro el rendimiento de 2.669 a 2.746 TFLOPs/s por dispositivo, un alza de 2,9% sobre la linea base. Al desactivar el LHS, el rendimiento cayo a 2.569 TFLOPs/s por dispositivo, lo que confirma que la tecnica 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 rematerializacion en la pasada hacia atras por transferencias que se solapan con el computo y la comunicacion.
¿Cuando conviene usar host offloading?
El host offloading es mas util cuando la memoria de la GPU limita el tamano del modelo, la longitud de secuencia o el lote, y cuando los tensores seleccionados son lo bastante grandes como para aliviar la presion de HBM. Rinde mejor cuando la carga tiene suficiente computo o comunicacion 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 aun mas ancho de banda, la tecnica se vuelve una palanca predecible para desacoplar el rendimiento del entrenamiento de los limites fisicos estrictos de memoria.
Es menos probable que ayude cuando los tensores son pequenos, cuando hay poco trabajo independiente para solapar las transferencias o cuando el cuello de botella esta en otro lado. Para equipos de LatAm que arriendan computo por hora, el atractivo es concreto: entrenar un modelo mas 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.
