대규모 언어 모델(LLM) 학습 워크로드에서는 컴퓨트가 완전히 활용되기 전에 GPU 메모리 한계에 부딪히는 경우가 점점 더 많습니다. 모델 가중치, 그래디언트, 옵티마이저 상태, 통신 버퍼 및 중간 활성화는 모두 제한된 GPU 고대역폭 메모리(HBM)를 경쟁합니다. 모델 크기, 시퀀스 길이 및 배치 크기가 증가함에 따라 HBM 용량이 주요 확장 병목 현상이 되는 경우가 많습니다.

이 기사에서는 JAX 기반 학습 워크플로우에서 이러한 제약을 완화하기 위해 호스트 오프로딩을 어떻게 사용할 수 있는지 설명합니다.