Beban kerja pelatihan model bahasa besar (LLM) semakin sering mencapai batas memori GPU sebelum komputasi sepenuhnya dimanfaatkan. Bobot model, gradien, keadaan pengoptimal, buffer komunikasi, dan aktivasi perantara semuanya bersaing untuk memori bandwidth tinggi (HBM) GPU yang terbatas. Seiring bertambahnya ukuran model, panjang urutan, dan ukuran batch, kapasitas HBM sering menjadi hambatan skalabilitas utama.
Artikel ini menjelaskan bagaimana offloading host dapat digunakan untuk mengurangi kendala ini dalam alur kerja pelatihan berbasis JAX.