NVIDIA menjelaskan bagaimana host offloading di JAX dapat mengurangi tekanan memori bandwidth tinggi GPU selama pelatihan model bahasa yang besar. Teknik ini menggeser aktivasi terpilih ke memori host yang dipin dan mengembalikannya selama backpropagation alih-alih menghitung ulang.

Perusahaan tersebut mengatakan pendekatan ini sangat efektif pada sistem Grace Blackwell, di mana NVLink-C2C memungkinkan komunikasi CPU-GPU bandwidth tinggi. Dikombinasikan dengan penjadwalan penyembunyian latensi dan transfer pipelined, host offloading dapat tumpang tindih antara pergerakan data dengan komputasi untuk meningkatkan efisiensi.

Uji kinerja menggunakan MaxText dengan Llama 3.1 405B dan DeepSeek-V3 671B menunjukkan throughput pelatihan yang lebih tinggi dan dukungan untuk konfigurasi batch yang lebih besar. NVIDIA mengatakan pengembang dapat mengadopsi metode ini melalui JAX Toolbox dan contoh MaxText yang dioptimalkan untuk meningkatkan pemanfaatan memori dan skalabilitas pelatihan.