LLM 활용을 위한 GPU 메모리 계산

오픈소스 LLM중 가장 많이 활용되는 Meta의 Llama3 오픈소스 LLM을 통해 필요한 GPU 메모리량을 알아보자.
Llama 3의 8B 모델을 학습시키기 위한 메모리 요구사항은 가장 경제적으로 여유로운 설정과 가장 타이트한 최소 설정으로 나누어 계산하였다. 이는 학습 환경, 배치 크기, 시퀀스 길이, 최적화 기술에 따라 크게 달라진다.

 

---

1. Llama 3 8B 모델 학습 메모리 요구사항

경제적으로 여유로운 설정

가장 여유로운 설정에서는 FP32 데이터 형식과 기본 Adam 옵티마이저를 사용하며, Gradient Checkpointing이나 ZeRO와 같은 최적화 기법은 사용하지 않는다.

주요 가정

• 데이터 형식: FP32 (4바이트)
• 배치 크기: 8
• 시퀀스 길이: 2048
• 옵티마이저: Adam
• 최적화 기술: 사용하지 않음

메모리 계산

1. 모델 파라미터 메모리:
   8,000,000,000 (8B) × 4바이트 = 32GB
2. 활성화 메모리:
   8 (배치 크기) × 2048 (시퀀스 길이) × 4096 (히든 크기) × 4바이트 = 256MB
   • 32개의 레이어를 포함하면: 256MB × 32 = 8GB
3. 옵티마이저 메모리:
   32GB × 3 = 96GB
4. 그래디언트 메모리:
   32GB
5. 중간 계산 그래프 메모리:
   8GB × 2 = 16GB

총 메모리

32GB + 8GB + 96GB + 32GB + 16GB = 184GB

---

가장 타이트한 최소 설정

최소 설정에서는 FP16 데이터 형식을 사용하며, Gradient Checkpointing 및 ZeRO Stage 3와 같은 최적화 기술을 활용한다. 배치 크기와 시퀀스 길이도 축소된다.

주요 가정

• 데이터 형식: FP16 (2바이트)
• 배치 크기: 2
• 시퀀스 길이: 1024
• 옵티마이저: ZeRO Stage 3
• Gradient Checkpointing: 사용

메모리 계산

1. 모델 파라미터 메모리:
   8,000,000,000 (8B) × 2바이트 = 16GB
2. 활성화 메모리:
   2 (배치 크기) × 1024 (시퀀스 길이) × 4096 (히든 크기) × 2바이트 = 16MB
   • 32개의 레이어를 포함하면: 16MB × 32 = 0.5GB
3. 옵티마이저 메모리 (ZeRO Stage 3):
   약 16GB
4. 그래디언트 메모리:
   16GB
5. 중간 계산 그래프 메모리 (Gradient Checkpointing으로 대폭 절감):
   약 0GB

총 메모리

16GB + 0.5GB + 16GB + 16GB = 48.5GB

---

2. 70B 및 405B 모델 메모리 요구사항

Llama 3 70B 모델

1. 모델 파라미터 메모리 (FP16):
   70,000,000,000 (70B) × 2바이트 = 140GB
2. 활성화 메모리 (배치 크기 2, 시퀀스 길이 1024):
   약 20GB
3. 옵티마이저 메모리 (ZeRO Stage 3):
   약 140GB

총 메모리:

140GB + 20GB + 140GB = 300GB

---

Llama 3 405B 모델

1. 모델 파라미터 메모리 (FP16):
   405,000,000,000 (405B) × 2바이트 = 810GB
2. 활성화 메모리 (배치 크기 2, 시퀀스 길이 1024):
   약 60GB
3. 옵티마이저 메모리 (ZeRO Stage 3):
   약 810GB

총 메모리:

810GB + 60GB + 810GB = 1680GB

---

결론

• 8B 모델 학습:
  • 여유로운 설정: 184GB
  • 최소 설정: 48.5GB (NVIDIA L40S 한 장으로 가능)
• 70B 모델 학습:
  • 최소 설정: 약 300GB (NVIDIA L40S 7장 필요)
• 405B 모델 학습:
  • 최소 설정: 약 1680GB (NVIDIA H100 21장 필요)

Llama 모델의 학습 메모리는 최적화 기술에 크게 의존한다. 서버 환경과 모델 크기에 따라 적절한 최적화를 통해 메모리 사용량을 줄이는 것이 중요하다.