초대용량 LLM 추론 효율화: CPU-GPU 통합 메모리로 KV 캐시 오프로드 구현

대규모 언어모델(LLM)의 추론 효율성을 높이기 위해 CPU와 GPU 간 메모리 공유 기술이 주목받고 있다. Llama 3 70B와 Llama 4 Scout 109B 같은 대규모 모델은 반정밀도(FP16)로 로드 시 각각 약 140GB, 218GB의 메모리가 필요하며, 특히 KV 캐시는 컨텍스트 길이와 배치 크기에 따라 선형 증가한다. 예를 들어 128K 토큰 길이의 캐시는 단일 사용자 기준 약 40GB를 차지해 GPU의 96GB 메모리로는 부족해 메모리 초과(OOM) 오류가 발생한다. 이 문제를 해결하기 위해 NVIDIA의 Grace Hopper 및 Grace Blackwell 아키텍처는 NVLink-C2C 기술을 통해 CPU와 GPU를 900GB/s의 고대역폭으로 연결하고, 메모리 일관성(메모리 공유)을 제공한다. 이로 인해 CPU의 480GB LPDDR 메모리와 GPU의 96GB 메모리가 하나의 통합 메모리 주소 공간으로 작동해, 모델 파라미터와 KV 캐시를 자동으로 공유할 수 있다. 실제로 GH200 슈퍼칩 환경에서 RAPIDS 메모리 매니저(RMM)를 활용해 관리 메모리 할당을 활성화하면, PyTorch가 GPU와 CPU 메모리를 통합해 사용할 수 있다. 이를 통해 Llama 3 70B 모델을 GPU에 로드할 때 OOM 오류 없이 성공적으로 실행 가능하다. 코드에서는 rmm.reinitialize(managed_memory=True)와 torch.cuda.memory.change_current_allocator(rmm_torch_allocator)를 설정해 메모리 관리를 통합하고, 텍스트 생성 파이프라인을 통해 질문에 대한 응답을 생성할 수 있다. 결론적으로, CPU-GPU 메모리 공유 기술은 모델 크기 증가에 따른 메모리 한계를 극복하고, 대규모 LLM의 추론과 훈련을 보다 효율적으로 가능하게 한다. 이는 미래의 AI 하드웨어 설계와 소프트웨어 최적화의 핵심 방향이다.