LLM 효율적으로 훈련하기 - 양자화(Quantization)와 분산 훈련(Distributed Training)

※ 본 포스트는 Coursera 강의인 "Generative AI with Large Language Models"의 Week 1의 내용의 일부를 정리하고 필요한 내용을 추가하여 작성한 글입니다.

 

 

GPU를 사용하여 모델 훈련을 진행하다 보면 “CUDA out of memory”라는 메세지를 보셨을 것입니다. 언어 모델의 크기가 커질수록, 메모리 부족 문제를 많이 마주하게 됩니다. 아래는 현재 이 글을 작성하는 2023년 7월 기준으로 현재 올라온 모델의 사이즈를 나타낸 표인데, 21년도 이후에 등장한 언어 모델들은 10억 (1B) 파라미터는 기본이고, 1,000억 (100B) 이 훌쩍 넘는 초대형 언어 모델들이 등장합니다. 아무리 커도 3억 정도의 파라미터를 가졌던 BERT 모델 (340M) 과 비교하면 엄청난 차이입니다. 아래의 사이트에 들어가보시면, 최신 모델을 업데이트 해놓으니 참고해보기 좋을 것 같습니다. 

Source: https://informationisbeautiful.net/visualizations/the-rise-of-generative-ai-large-language-models-llms-like-chatgpt/

 

이 모델들을 기반으로 Pre-training/Fine-tuning을 하게되면 모델 파라미터를 GPU에 올려야하는데, 사이즈가 너무 크게 되면 아예 올라가지 않게 되어 훈련이 불가합니다. 이번 포스트에서는 모델 훈련을 진행하는 데 필요한 리소스를 계산하는 방법과 제한적인 컴퓨팅 리소스에서 훈련을 진행하는 방법에 대해서 정리해보겠습니다.

 

 

 

1B (10억) 파라미터를 저장하기 위해 필요한 GPU RAM은?

 

Source: https://en.wikipedia.org/wiki/Single-precision_floating-point_format

• 1 파라미터는 보통 32-bit float의 데이터 타입을 사용하여 저장합니다
  • 32-bit float란? 숫자를 32비트를 사용하여 0과 1으로 표현하는 것 (그림 참고)
• 1 파라미터 = 4 byte
• 1B (10억) 파라미터 = 4GB

 

 

1 파라미터를 저장하는데 4바이트가 필요하기 때문에, 1B 파라미터의 모델 가중치를 저장하는 데에만 4GB가 필요합니다. 하지만 1B 파라미터의 모델 훈련을 진행하는 경우에는,

• 두 개의 Adam 옵티마이저(8바이트)
• 그라디언트(4바이트)
• 활성화 함수(4바이트)
• 임시 변수(8바이트)

이들을 추가적으로 필요로 합니다. 즉, 모델 파라미터 당 4 바이트와 추가적으로 약 20 바이트가 필요하게 되어, 총 80 GB의 메모리가 필요하게 됩니다. 80GB는 한화 약 2천만원 정도 가격의 Nvidia A100 GPU의 메모리 용량이며, 소비자는 감당하기 힘들 뿐만 아니라 데이터 센터에서도 비용이 꽤 큰 리소스입니다.

 

 

모델의 크기가 커진다고 GPU를 무한정 늘릴 수 없을 뿐더러, 주어진 컴퓨팅 자원을 활용하여 효율적인 모델 훈련을 해야합니다. 어떻게 해야 메모리를 최소화하면서 거대한 모델을 훈련할 수 있을까요? 다음과 같은 방법을 사용할 수 있습니다:

 

1. 양자화 (Quantization)
2. 분산 훈련 (Distributed Training)

 

 

양자화(Quantization)

32비트 부동 소수점(32-bit float)이 4byte를 차지한다면, 메모리를 더 적게 차지하는 데이터 타입, 더 낮은 비트로 바꾸어 모델 가중치를 저장하면 필요한 전체 메모리도 줄어듭니다. 이것이 Quantization의 핵심 아이디어입니다.

 

딥러닝에서 Quantization은 두 가지로 나누어질 수 있습니다. 이 글에서 초점을 맞추는 것은 훈련 시에 양자화를 포함하는 QAT니다. 대체로 QAT가 PTQ보다 모델 성능이 좋으며, 낮은 정밀도를 사용하여 모델 배포 시에도 유용하다고 합니다. 

 

1. Post-training quantization (PTQ): 모델이 훈련된 뒤에 양자화를 하여 모델 크기를 줄이는 방법
2. Quantization-Aware Training (QAT): 훈련 과정에 scaling, clipping, rounding 등의 양자화 과정을 포함하여 훈련하는 방법

 

Source: https://cloud.google.com/blog/products/ai-machine-learning/bfloat16-the-secret-to-high-performance-on-cloud-tpus

 

그림에서 보이듯 FP32는 310^-38 to 310^38 범위의 숫자를 나타낼 수 있습니다. 더 낮은 부동 소수점 공간인 FP16, BFLOAT16, INT8으로 변환한다면, 정밀도에서의 손실이 생기지만 저장에 필요한 공간이 줄어듭니다. 예를 들어, 파이(π)를 각 데이터 타입으로 변환해보겠습니다.

• FP32 (full-precision float): 3.141592653589793 → 4 byte
• FP16 (half-precision float): 3.140625 → 2 byte
• BF16 (Bfloat16): 3.140625 → 2 byte
• INT8 (8-bit integer): 3 → 1 byte

 

FP32

모델 학습 시 기본값인 FP32는 정밀도가 높아서 모델의 정확도가 높아지지만, 메모리를 많이 차지하고 연산속도가 느립니다.

 

BF16 vs FP16

BF16과 FP16은 모두 FP32의 절반인 16비트에 인코딩을 하여 메모리를 반으로 줄일 수 있지만, 모델 정확도가 비교적 낮아질 수 있습니다. BF16은 구글에서 머신러닝 모델 훈련을 위해 개발한 데이터 타입인데, Google TPU에서는 작동을 하고, Nvidia GPU A100과 3090 RTX과 같은 최신 GPU들은 지원을 하지만, V100이나 T4와 같은 모델에서는 FP16밖에 지원하지 않는다고 합니다. BF16은 gradient scaling을 필요로 하지 않기 때문에, FP16에 비해 안정성이 높아서 훈련에 사용됩니다. 강의에서도 FLAN-T5의 훈련 방식과 같이 BF16이 가장 대중적으로 선택되는 훈련 시 데이터 타입이라고 이야기합니다. FP16은 훈련보다는 inference시에 빠른 연산을 위해 사용됩니다.

 

INT8

INT8 quantization은 다른 데이터 타입과 비교하면 정보 손실이 크지만, 최근에는 LLM의 크기가 커지면서 더 사이즈를 줄이는 연구가 많이 나오면서, 8비트 양자화를 진행하더라도 모델 성능이 크게 저하되지 않는다고 합니다. (읽어볼만한 글: LLM.int8()을 소개하는 허깅페이스 포스팅- 175B 파라미터의 16/32-bit 체크포인트를 불러와 트랜스포머의 피드포워드와 어텐션 레이어에서 int8 행렬 곱셈을 진행)

 

이렇게 Quantization을 포함하여 훈련을 하여도, 거대한 모델은 한 개의 GPU안에 다 들어가지 못하는 경우가 많습니다. 이를 해결하기 위해서 여러개의 GPU를 사용하여 훈련할 수 있습니다. 

 

 

 

분산 훈련 (Distributed Training)

분산 훈련에서는 거대한 딥러닝 모델을 훈련하는 동안, 훈련 작업을 여러 프로세서에 나눕니다. 이 프로세서는 worker node 또는 worker라고 불리는데, 이들은 훈련 과정에서 병렬로 훈련이 진행됩니다. 분산 훈련을 통해서 훈련 시간을 단축시켜 빠른 실험과 배포가 가능합니다. 

 

분산 훈련에는 크게 두 가지 종류가 있습니다. 

 

1. Model-parallel training: 한 모델이 쪼개져 worker에 할당되어 훈련
2. Data-parallel training: worker마다 모델을 복사하여 로드한 후, 각자 다른 훈련 데이터를 넣어 훈련 (모델을 모든GPU마다 복사하여 할당하기 때문에, 한 개의 GPU에 로드가 가능한 모델 사이즈여야 함)

Source: https://towardsdatascience.com/deep-learning-on-supercomputers-96319056c61f

 

Distributed Data Parallel (DDP)

Data-parallel training에서 흔히 쓰이는 기법은 Pytorch의 Distributed Data Parallel (DDP)입니다. DDP는 각 GPU에 모델을 복사하여 할당한 후, 데이터 배치를 병렬로 전송하여 처리합니다. Worker마다 각자 다른 결과와 그라디언트를 계산하여 synchronization step에서 이 결과들을 합쳐서 같은 값으로 각 GPU를 업데이트합니다. 

 

 

Fully Sharded Data Parallel (FSDP)

모델이 한 개의 GPU에 로드되지 않을 때, model parallelism 또는 model sharding을 통해서 여러개의 GPU에 한 개의 모델을 할당합니다. 가장 주목받는 기술은 Pytorch의 Fully Sharded Data Parallel(FSDP)인데, 이 기법은 2019년에 Microsoft가 소개한 ZeRO라는 기법에서 착안한 방법론입니다. ZeRO는 모델 파라미터, 그라디언트, 옵티마이저 상태를 여러개의 GPU에 걸쳐 분산하여 메모리를 최적화합니다. ZeRO는 세 가지 최적화 단계를 거치는데, 1단계(P_os)에서는 Optimizer State(os)만을, 2단계(P_os+g)에서는 Gradient(g)도 같이 분산하며, 3단계(P_os+g+p)에서는 Parameter까지 같이 분산하여 학습합니다. 이렇게 하면 사용되는 메모리량이 감소하는데, 아래 그림에서 보듯 처음에 120GB의 메모리를 쓰던 것이 1.9GB로 감소한 것을 볼 수 있습니다. 

Source: https://arxiv.org/pdf/1910.02054.pdf (ZeRO paper)

 

DDP에서는 각 배치를 처리하는데 필요한 모델 상태가 한 GPU안에 다 있었다면, FSDP에서는 여러 개의 GPU에 모델이 쪼개어져 있기 때문에 각 GPU마다 배치를 foward pass, backward pass를 진행할 때 매번 불러와야합니다. 이 작업을 진행할 때, 각 CPU는 다른 GPU에 데이터를 요청하여, 분산된(sharded) 데이터를 분산되지 않은(unsharded) 데이터로 만들고, 이후에는 다시 원래의 sharded 데이터로 만들어 각 GPU로 보냅니다. 

 

얼만큼 분산시킬 것인지 (Sharding)를 지정해줄 수 있는데, 

• 가능한 모든 GPU를 사용 → Full Sharding
• 1개의 GPU를 사용 → Full Replication (No sharding)

Full Sharding을 하게 되면, 사용하는 메모리는 줄지만, GPU간 통신량이 많아져서 성능의 문제로 이어지는 trade-off 가 발생합니다. Full Replication을 하게되면 큰 모델은 GPU에 올리지 못해 훈련이 불가한 상황이 발생할 수 있습니다. Hybrid Sharding은 이 두 가지를 적절하게 섞어 사용하는 방법입니다. 

 

 

이번 포스트에서는 사용하는 메모리를 최적화하여 거대한 모델을 훈련할 수 있는 방법으로 양자화와 분산 훈련을 알아보았습니다. 다음 포스트에서는 최적의 모델을 결정하는 방법에 대해서 알아보겠습니다. 

 

 

 

 

 

Reference

Quantization

• Quantization Aware Training
• A Study of BFLOAT16 for Deep Learning Training
• bf16, fp16, fp32의 차이점
• Reddit 답변: Mixed Precision Training: Difference between BF16 and FP16

Distributed Training

• https://neptune.ai/blog/distributed-training
• https://neptune.ai/blog/distributed-training-frameworks-and-tools