LLM 분산 학습 : Accelerate와 FSDP

대규모 언어 모델(LLM) 학습에 사용되는 분산 학습 기술과 Accelerate 라이브러리에 대해 알아보려 한다.

 

LLM(Large Language Model)은 일반적으로 여러 GPU를 활용한 학습을 진행한다. 이는 하나의 모델을 여러 GPU에 나눠서 올리는 것을 의미한다. 이렇게 분산 학습을 효과적으로 수행하기 위해 Accelerate 라이브러리가 자주 사용되며, 그중 대표적인 기술로 FSDP(Fully Sharded Data Parallel)가 있다.

 

예를 들어, 모델 파라미터가 약 70B(70 billion)인 경우, 메모리 용량은 약 150GB에 달한다. 하지만 현재 사용하는 A100 GPU의 최대 메모리는 80GB로, 단일 GPU에 모델 전체를 올릴 수 없다. 이런 경우, FSDP를 활용하면 모델을 여러 GPU에 나누어 올릴 수 있다.

 

FSDP의 주요 기능

1. 파라미터 분산
FSDP는 모델 파라미터를 분할해 여러 GPU에 나눠 올릴 수 있게 한다. 이를 통해 큰 모델도 효율적으로 학습할 수 있다.

 

2. 메모리 최적화
FSDP는 단일 GPU에서도 사용할 수 있다. 메모리가 부족한 상황에서는 현재 사용하지 않는 파라미터를 GPU에서 내려놓고 필요할 때 다시 GPU로 불러오는 방식으로 메모리를 관리한다.

또한, Gradient Checkpointing을 통해 필요한 중간 결과물만 저장하고, 나머지는 필요할 때 다시 계산하는 방식으로 메모리 사용량을 줄일 수 있다.

 

Accelerate 라이브러리의 활용

이러한 장점들 때문에 Accelerate 라이브러리는 분산 학습뿐만 아니라 단일 GPU 환경에서도 자주 사용된다. 특히, LLM 학습이나 혼합 정밀도(Mixed Precision)와 같은 기술을 사용할 때 Accelerate를 활용하면 학습 속도를 높이고 메모리를 효율적으로 사용할 수 있다.