70B급 LLM을 단일 4GB GPU에서도 돌릴 수 있게 해주는 프레임워크, AirLLM

요즘 메모리 가격이 천정부지로 솟아오르고 있는데 이 메모리에 대한 비용 부담을 줄여줄 수 있는 프레임워크 airllm이 공개됐다.

https://github.com/lyogavin/airllm

GitHub - lyogavin/airllm: AirLLM 70B inference with single 4GB GPU

 

AirLLM은 70B급 LLM을 단일 4GB GPU에서도 돌릴 수 있게 해주는 메모리 최적화 추론 프레임워크이다. 사실상 모든 종류의 pc에서 70B급 대형 LLM 모델을 돌릴 수 있게 해준다고 할 수 있다.

 

메모리 사용량을 줄이는 방식에 대해서 간단하게 설명하자면, 모델을 여러 레이어 샤드(layer shard)로 쪼개서 필요할 때만 GPU에 올리고, 나머지는 CPU/RAM에 두는 식의 모델 샤딩으로 GPU 메모리 사용을 최소화했다고 한다. 기본 목표는 무손실” 추론(양자화·프루닝 없이도 4GB 단일 GPU에서 70B 모델 로딩 가능)이고, 선택적으로 블록 단위 양자화(block-wise quantization)로 추가 3배 속도 향상도 제공하고 있다.

 

또한 LLaMA2, ChatGLM, Qwen, Baichuan, InternLM, Mistral 등 Hugging Face 상의 여러 모델을 from_pretrained 인터페이스로 사용 가능하며 Hugging Face Transformers와 비슷한 AutoModel.from_pretrained, generate API를 제공해서 기존 코드 포팅이 쉽다.

 

참고로 레이어 샤딩은 대용량 모델 전체를 한 번에 GPU에 올리지 않고, 레이어를 여러 shard로 나눠서 순차적으로 GPU에 로드 후 계산, 그리고 다시 내리는 방식으로 GPU 메모리 피크를 줄이는 것을 말한다. 대부분의 파라미터를 CPU 메모리 또는 디스크에 두고, 필요한 부분만 옮겨 쓰기 때문에 GPU 메모리가 매우 작아도 동작하는 것이다.

 

다만 GPU 메모리를 아끼는 대신, CPU–GPU 데이터 이동이 많아져 지연 시간(latency) 은 커질 수 있다. 그렇기 때문에 CPU RAM·스토리지 용량과 I/O 대역폭 요구사항은 여전히 크다고 할 수 있다. 근데 뭐 가난한 개발자한테는 이 정도는 당연히 감수해야할 문제인 것 같긴하다.