Las cargas de trabajo de entrenamiento de modelos de lenguaje grande (LLM) cada vez más se topan con los límites de memoria de la GPU antes de que el cómputo esté completamente utilizado. Los pesos del modelo, gradientes, estados del optimizador, búferes de comunicación y activaciones intermedias compiten por la limitada memoria de alto ancho de banda (HBM) de la GPU. A medida que crecen el tamaño del modelo, la longitud de la secuencia y el tamaño del lote, la capacidad de HBM a menudo se convierte en el principal cuello de botella para la escalabilidad.
Este artículo explica cómo la descarga al host puede utilizarse para mitigar estas restricciones en flujos de trabajo de entrenamiento basados en JAX.