초대형 언어 모델을 위한 그래프 기반 계층 메모리 관리 프레임워크

초록

본 논문은 슈퍼노드 아키텍처의 대용량 공유 메모리를 활용해, 원격 메모리 접근을 컴퓨테이션 그래프의 일급 연산으로 승격시키는 컴파일러 기반 프레임워크 HyperOffload를 제안한다. 그래프‑드리븐 캐시 연산과 전역 실행 순서 최적화를 통해 메모리 사용량을 최대 26% 절감하면서도 통신 지연을 연산에 가려 성능 저하를 최소화한다.

상세 분석

HyperOffload는 기존 런타임 중심의 오프로드·스와핑 방식이 갖는 “지역적 시야”와 “반응형 스케줄링”의 한계를 극복하기 위해, 원격 메모리 접근을 컴파일 타임에 명시적인 IR 연산으로 삽입한다. 이를 위해 MindSpore의 MindIR에 ‘cache‑load’·‘cache‑store’ 연산자를 추가하고, 새로운 컴파일 패스에서 텐서의 라이프타임과 의존 관계를 전역적으로 분석한다. 분석 결과를 바탕으로 전역 실행‑순서 정제 알고리즘을 적용해, 데이터 전송을 연산이 풍부한 구간에 미리 배치함으로써 통신‑연산 겹침을 최적화한다. 특히, KV‑cache와 파라미터 같은 장기 저장 데이터는 사전 프리패치 전략으로 원격 메모리 풀에 위치시키고, 필요 시 즉시 로드하도록 설계했다. 이러한 설계는 (1) 피크 메모리 사용량 감소, (2) 메모리 압축·재배치에 따른 런타임 오버헤드 최소화, (3) 대규모 모델·긴 시퀀스 상황에서도 일정한 대역폭 활용을 가능하게 한다. 실험에서는 LLaMA‑3‑8B와 같은 8 B 파라미터 모델을 384 NPUs 슈퍼노드 환경에서 실행했을 때, 기존 런타임 프리패치가 2.7배 지연을 초래하는 반면, HyperOffload는 5.5 s 내에 완료하며 피크 메모리를 26 % 절감했다. 또한, 학습 워크로드에서도 활성화 체크포인트와 옵티마이저 상태를 원격 메모리 풀에 효율적으로 배치해 OOM을 방지하고, 전체 학습 속도에 미치는 영향을 5 % 이하로 억제했다. 이와 같이 컴파일러 레벨에서 메모리 관리와 통신 스케줄링을 통합함으로써, 슈퍼노드와 같은 차세대 대용량 메모리 하드웨어의 잠재력을 최대한 활용할 수 있음을 입증한다.