메모리 최적화로 SSD 오프로드 LLM 파인튜닝 한계 돌파

초록

MemAscend는 SSD‑오프로드 기반 LLM 파인튜닝에서 시스템 메모리 파편화와 핀드 메모리 비효율을 제거해 평균 55.7%의 메모리 사용량 감소와 I/O량 58% 절감을 달성한다. 이를 통해 128 GiB DRAM 환경에서도 컨텍스트 길이를 16 k에서 131 k 토큰으로, 배치 크기를 4에서 32까지 확장할 수 있다.

상세 분석

본 논문은 ZeRO‑Infinity와 같은 기존 SSD‑오프로드 기법이 GPU 메모리 부족을 해결하는 데는 성공했지만, 시스템 메모리 관리에서 심각한 비효율을 초래한다는 점을 지적한다. 특히 CUDA HostAlloc으로 확보되는 핀드 메모리는 64 Byte 혹은 256 Byte 단위의 과도한 정렬 요구로 전체 버퍼 풀의 70% 이상을 낭비하고, 동적 텐서 할당·해제 과정에서 파편화가 급격히 증가한다. 또한 CPU‑오버플로우 체크 로직이 중복 실행돼 피크 메모리 사용량을 1.25배까지 부풀린다. MemAscend는 네 가지 핵심 모듈로 이를 해결한다. 첫째, 어댑티브 버퍼 풀은 텐서 크기에 맞춰 가변적인 슬라이스를 재사용함으로써 메모리 파편화를 최소화한다. 둘째, 정렬‑프리 핀드 메모리 할당은 C++ 확장 연산을 통해 페이지 단위(4 KB)로 직접 매핑하고, 불필요한 256 Byte 정렬을 제거해 할당 오버헤드를 절반 이하로 낮춘다. 셋째, 오버플로우 체크를 기존 두 단계(전/후)에서 하나로 통합해 중복 메모리 복제를 방지하고, 체크 포인트 저장 시 발생하는 임시 버퍼를 재활용한다. 넷째, Direct NVMe 엔진은 파일 시스템을 우회해 사용자 공간에서 직접 NVMe 큐를 제어, I/O 레이턴시와 시스템 콜 오버헤드를 크게 감소시킨다. 추가적으로 반정밀(half‑precision) 옵티마이저를 도입해 파라미터·그라디언트·모멘텀을 16 bit로 전송·저장함으로써 SSD‑I/O 트래픽을 58% 절감하고, 전체 학습 스루풋을 최대 24.21% 향상시킨다. 실험 결과, 7B~70B 규모 모델에 대해 기존 ZeRO‑Infinity 대비 평균 55.7% 적은 시스템 메모리를 사용했으며, 동일 하드웨어에서 컨텍스트 길이와 배치 크기를 크게 확대할 수 있었다. 이러한 개선은 메모리 비용이 제한된 연구실이나 중소기업이 대형 LLM을 직접 파인튜닝할 수 있는 실질적인 문턱을 낮춘다.