As cargas de trabalho de treinamento de modelos de linguagem grande (LLM) estão cada vez mais atingindo os limites de memória da GPU antes que o cálculo seja totalmente utilizado. Os pesos do modelo, gradientes, estados do otimizador, buffers de comunicação e ativações intermediárias competem pela limitada memória de alta largura de banda (HBM) da GPU. À medida que o tamanho do modelo, o comprimento da sequência e o tamanho do lote aumentam, a capacidade da HBM frequentemente se torna o principal gargalo de escalabilidade.
Este artigo explica como o offloading para o host pode ser usado para mitigar essas restrições em fluxos de trabalho de treinamento baseados em JAX.