[Huggingface] 모델 학습 시 GPU 메모리 사용 알아보기

원문 허깅페이스 - https://huggingface.co/docs/transformers/en/model_memory_anatomy

 

---

모델 학습 도중 GPU는 어떤 방식으로 활용되는가

모델 학습 속도와 메모리 활용의 효율성을 증대하기 위한 최적화 기법을 이해하기 위해,

①학습 도중 GPU가 어떤 식으로 활용되며 ②수행하는 작업에 따라 계산 강도가 어떻게 달라지는지 이해할 필요가 있다. 

 

Step-by-step 메모리 확인 가이드는 아래 huggingface 가이드에 있는 코드를 따라 해 보면 된다.

https://huggingface.co/docs/transformers/en/model_memory_anatomy

 

Pytorch를 통해 모델을 로드하고, 학습하는동안 GPU 사용률을 관찰해 보면 아래와 같이 변화한다:

 

(1) 최초의 GPU 상태

    : 아무것도 하지 않고 데이터만 가지고 있는 상황에서는 GPU 메모리가 사용되지 않는다.

• 하지만, 이때 모든 free memory를 사용자가 사용할 수 있는 것은 아니다.
• 모델을 GPU 커널에 올리기만 하더라도 1-2GB의 메모리를 사용하게 된다. 
  뿐만 아니라, 아주 작은 텐서라도 GPU에 올리는 순간, 꽤 많은 GPU 메모리를 사용하기 시작한다.

텐서 1개가 이미 1.3GB의 메모리를 사용하는 것을 관찰할 수 있다.

 

 

(2) 모델 로딩 후:

    : 모델 weight를 로딩하면, GPU 메모리를 사용하기 시작한다. 

• 예를 들어 Bert-large(355M) 규모의 모델을 로드하게 되면, 모델 파라미터만으로 2.6GB 정도의 메모리를 사용한다.
  정확한 숫자는 사용하는 GPU에 따라 달라질 수 있으며,
  최신 GPU에서는 모델 사용 속도를 높이도록 최적화된 방식으로 가중치를 로드하기 때문에 더 많은 공간이 사용될 수 있다.

(3)  기본적인 학습 도중 메모리 사용

   : huggingface Trainer만을 사용하여 별도의 GPU 퍼포먼스 최적화 작업 없이 배치사이즈 4로 학습을 수행해 보면,

     비교적 작은 배치 사이즈로도 GPU의 전체 메모리를 사용하는 것을 관찰할 수 있다. 

• 그러나 배치 사이즈가 클수록 모델 수렴이 빨라지거나 최종 성능이 좋아지는 경우가 많기에, 
  이상적으로 배치사이즈를 GPU 용량이 아닌 모델의 요구사항에 맞게 조절하는 것이 필요하다.
• 또한 흥미로운 점은, GPU가 모델 사이즈보다 훨씬 더 많은 메모리를 사용하고 있다는 것인데
  이를 이해하기 위해서는 모델의 작동과 메모리 요구사항에 대해 이해할 필요가 있다. 
  

 

모델 동작 방식에 대한 해부도

Transformer 아키텍처는 계산 강도에 따라 세 개의 동작으로 그룹 지을 수 있다.

이 내용은 performance bottleneck이 되는 부분을 분석하기 위해 유용하다.

1. Tensor Contraction

멀티헤드어텐션의 선형 레이어 맟 구성요소는 모두 배칭 된 행렬곱을 통해 이루어진다. 

이 연산은 transformer 학습에서 연산이 가장 많이 소요되는 부분이다.

 

2. Statistical Normalizations

소프트맥스와 레이어 정규화는 Tensor Contraction보다 계산 집약도가 낮고, 하나 이상의 reduction operation을 포함하며, 그 결과는 맵을 통해 적용된다.

 

3. Element-wise Operators

1과 2 이외의 나머지 연산 작용으로, biases, dropout, activations, and residual connections를 포함하며, 가장 계산 집약적이지 않은 부분이다. 

 

모델 메모리에 대한 해부도

위에서 실험을 통해 모델을 학습하는 과정이 모델을 GPU에 단순히 올려놓는 것보다 훨씬 더 많은 메모리를 사용한다는 것을 확인하였다. 학습 도중 GPU 메모리를 사용하는 많은 구성요소가 존재하기 때문이다. GPU 메모리를 구성하는 요소는 다음과 같다:

1. model weights - 모델 파라미터(웨이트)
2. optimizer states - 옵티마이저의 상태값
3. gradients - 그라디언트
4. forward activations saved for gradient computation - 그라디언트 계산을 위해 저장해 둔 forward activation
5. temporary buffers - 임시 버퍼
6. functionality-specific memory - 함수 특화적인 메모리

AdamW와 mixed-precision으로 일반적인 모델을 학습하기 위해서는 모델 파라미터당 18byte의 메모리와 더불어 activation 메모리가 필요하다. 추론 과정에서는 옵티마이저 상태값이나 그라디언트가 존재하지 않기 때문에 이를 제외하고 모델 파라미터당 6byte와 activation 메모리가 필요하게 된다. 

 

이를 상세히 정리하면 다음과 같다.

 

모델 파라미터

• fp32 학습의 경우 - 파라미터 개수 * 4byte
• mixed precision 학습 시 - 파라미터 개수 * 6byte 
  (모델을 fp32로 유지하면서 메모리에서는 fp16을 가지고 있기 때문)

옵티마이저 상태값

• 일반 AdamW - 파라미터 개수 * 8byte (2개의 state 보유)
• 8-bit AdamW 옵티마이저( bitsandbytes 등) - 파라미터 개수 * 2byte 
• SGD momentum 등의 옵티마이저 - 파라미터 개수 * 4byte (1개의 state 보유)

그라디언트

• fp32 / mixed precision 학습 모두 파라미터 개수 * 4byte (그라디언트는 항상 fp32로 저장함)

Forward Activation

• 시퀀스 길이, hidden size, batch size 등에 따라 상이
• forward 및 backward 함수에 의해 전달되고 반환되는 인풋/아웃풋, 그라디언트 계산을 위해 저장된 forward activation 등이 존재함!

Temporary Memory

• 추가적으로 계산이 끝나면 나오는 모든 종류의 임시 변수가 있는데, 순간적으로 추가 메모리가 필요하여 OOM이 될 수 있다. 따라서 코딩을 할 때 이러한 임시 메모리에 대해 전략적으로 생각하여 더 이상 필요하지 않은 시점에 즉각적으로 명시적으로 release 할 필요가 있다.

Functionality-specific memory

• 소프트웨어에 따라 특별 메모리가 필요할 수 있다
• 예를 들어 beam search를 수행할 때, 소프트웨어에서는 여러 벌의 입력-출력 복사본을 유지하고 있어야 한다. 

forward와 backward에서 실행 속도

• 컨벌루션과 선형 레이어의 경우, backward에서 forward pass보다 2배의 flop이 존재하고, 이는 일반적으로 속도가 2배 느린 것으로 생각할 수 있다. (backwatd에서 사이즈가 더 이상한 연산의 경우, 이 속도 저하가 더 커진다)
• 활성함수는 일반적으로 대역폭이 제한되고, forward보다 backward에서 일반적으로 더 많은 데이터를 읽어야 한다
  (예. activation forward에서는 1회 read + 1회 write vs activation backward에서는 2회 read + 1회 write)

---