POST
All
Product
Team
Tech
DocVLM: Make Your VLM an Efficient Reader
- 최윤진최
- Tech
Python 3.10 신규 문법 : Parenthesized context managers와 PEG Parser
- seunghoChoeS
- Tech
UReader: Universal OCR-free Visually-situated Language Understanding with Multimodal Large Language Model
- 최윤진최
- Tech
[팀 소개편] KPMG Lighthouse는 어떤 팀인가요?
- LighthouseL
- Team
[챕터 소개편] Backend Chapter를 소개합니다
- LighthouseL
- Team
[챕터 소개편]Frontend Chapter를 소개합니다
- LighthouseL
- Team
[챕터 소개편] AI Chapter를 소개합니다
- LighthouseL
- Team
ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
Created by
- Eunyoung LeeE
Created at
Category
- Tech
논문 개요
•
논문명: ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
•
출간일 : 2020.5
•
출간 학회 : -
•
저자 : Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, Yuxiong He
•
소속 : Microsoft
•
인용 수 : 1143
Parallelism
딥러닝에서 병렬 처리(Paralellism)란 한정된 하드웨어에 큰 모델을 끼워 넣거나 학습 속도를 빠르게 하는 작업
데이터 병렬 처리
•
데이터 병렬 처리(Data Parallelism)는 멀티 GPU로 학습을 할 때, 각 GPU가 모델의 전체 복제본을 갖고 있으며, 각 GPU가 데이터의 각기 다른 부분을 할당 받아 계산하고 결과를 합치는 과정
•
데이터셋을 동시적인 processing stream으로 분해하여 모두 같은 작업을 진행
•
각 GPU에서 계산한 결과를 모아 평균을 내고 다시 재분배하는 Synchronization 필요
•
모델 병렬 처리, 파이프라인 병렬 처리 등 다양한 멀티 GPU 학습 중 가장 널리 사용되는 방법
장점
•
확장성: 데이터셋의 사이즈나 계산의 복잡도가 커져도 GPU 추가를 통하여 쉽게 확장 가능
•
큰 데이터셋 사용 가능: 데이터를 나눠서 계산하기 때문에 큰 데이터를 다룰 수 있음
•
처리량 증가: 동시에 처리하기 때문에 계산에 필요한 시간이 줄어듦
단점
•
메모리 효율성: 복제본을 매번 메모리에 저장해야 하기 때문에 메모리 효율성이 떨어짐
모델 병렬 처리
•
모델 병렬 처리(Model Parallelism)는 모델의 사이즈가 너무 커서 모델을 하나의 GPU 메모리에 넣을 수 없을 때 사용하는 병렬 처리 방법
•
여러 GPU에 모델을 vertical하게 나누어 연산하여 각 GPU가 모델 연산의 일부 계산을 담당함
•
모델 병렬 처리의 종류로는 텐서 병렬 처리와 파이프라인 병렬 처리가 있으며, 둘 다 사용하는 방법은 2D Model Parallelism이라고 부름
◦
Tensor Parallelism
▪
가중치 행렬을 나누어 여러 GPU에서 연산을 하고 결과를 합치는 방식
▪
병렬 연산이 가능하기 때문에 학습 시간 단축 가능
◦
Pipeline Parallelism
▪
모델의 레이어를 여러 대의 GPU에 쪼개서 순차적으로 학습시키는 방식
▪
이전 GPU의 연산 결과를 다음 GPU에 넘겨 주는 순차적 방식이기 때문에 학습 시간 단축 불가
단점
•
Data parallelism에 비하여 synchronization 빈도가 훨씬 높아서 오버헤드 발생
•
데이터 병렬 처리와 다르게 실행순서를 지켜야 하기 때문에 까다로움
ZeRO
•
DP(Data Parallelism)의 계산/통신 효율성과 MP(Model Parallelism)의 메모리 효율성을 모두 달성할 수 있는 ZeRO-powered data parallelism
•
모델을 학습할 때 모델 사이즈를 8배 키우고, 성능을 10배 높일 수 있음
•
ZeRO 사용으로 MP 없이 13B까지의 모델 학습 가능 (32GB GPU 기준)
•
ZeRO-DP(Data Parallelism)과 ZeRO-R(Residual)두 가지의 최적화 방법이 통합됨
ZeRO-DP
(Zero Rudundancy Optimizer powered Data Parallelism)
•
ZeRO는 모델의 상태(옵티마이저 상태, 기울기, 파라미터)를 GPU 마다 복제하여 갖고 있는 대신 각 GPU에 값의 일부만 저장하는 DP 방법
◦
옵티마이저 상태: momentum and variances in Adam
•
학습 시 다른 GPU에 저장된 값이 필요할 때 모든 GPU가 synchronize하여 갖고 있지 않는 부분을 서로 가져다주고, 일시적으로 버퍼에 저장한 이후 계산을 마친 이후에는 다시 가져왔던 값을 버림
•
모델 상태에 따라 통신을 조정하는 동적 통신 스케줄을 사용하여 DP의 계산 세분성과 통신량을 유지함으로써 계산/통신 효율성을 유지
•
불필요한 중복을 제거하고 메모리 사용량을 최적화하여 효율적으로 분산학습이 가능
Mixed-Precision Training
•
mixed-precision training: 파라미터와 활성화 값을 fp16으로 저장하여 높은 throughput 가능, 순전파와 역전파 둘 다 fp16 사용
•
역전파 마지막에 업데이트를 효율적으로 계산하고 적용하기 위해 fp32 파라미터와 옵티마이저 상태 복사본을 갖고 있음
•
1.5B의 GPT-2로 계산하면 fp16로는 3GB만 필요하지만, 옵티마이저 상태까지 저장하기 위해서는 적어도 24GB의 메모리가 필요하며 residual states까지 포함하면 32GB GPU로도 학습 불가
ZeRO의 세 단계(ZeRO-1, ZeRO-2, ZeRo-3)
Ψ는 모델의 파라미터 개수, K는 optimizer specific constant term, $N_d$는 DP 단계
1.
Optimizer State Partitioning ($P_{os}$) – 메모리 4배 감소, DP와 같은 통신 볼륨
2.
Add Gradient Partitioning ($P_{os+g}$) – 메모리 8배 감소, DP와 같은 통신 볼륨
3.
Add Parameter Partitioning ($P_{os+g+p}$) – 메모리가 data parallelism degree와 선형적으로 감소(64개의 GPU에 분산하면 64배의 메모리 감소), 통신 볼륨 50% 증가
DP 정도에 따라 필요한 메모리
•
GPU 개수가 많아질수록 ZeRO로 학습할 수 있는 모델 크기가 커짐
ZeRO-R(Residual State Memory)
•
옵티마이저 상태, 기울기, 파라미터가 차지한 이후 나머지 부분인 학습 중 저장되는 중간 값
•
activation, temporary buffers, unusable fragmented memory
•
이차적으로 메모리 bottleneck이 될 수 있기 때문에 이 세가지 요소로 소모되는 메모리를 최적화하는 방법
◦
Activation(역전파를 하기 위해서 순전파 때 저장됨)
▪
activation 파티셔닝을 통해 MP 때 중복을 인식하고 제거하여 메모리 최적화
▪
필요할 때는 CPU에 오프로드 가능
▪
batch size 32, sequence length 1024인 100B 모델을 MP 16으로 학습할 경우, activation checkpoint를 저장하기 위해 GPU 별 33GB의 메모리가 필요하지만, ZeRO를 사용할 경우 2GB로 줄일 수 있음
◦
Temporary buffers
▪
메모리와 연산 효율의 균형을 위해 임시 버퍼의 적절한 크기 정의
▪
모델이 너무 커질 때 성능 효율적인 크기가 고정된 버퍼(fused buffer)를 사용하여 buffer size가 모델 사이즈에 의존하지 않도록 함
◦
Unusable fragmented memory
▪
텐서의 수명이 각기 다르기 때문에 연속적인(contiguous) 메모리가 부족해질 수 있음
▪
메모리를 사전에 관리하여 메모리 조각화를 방지함
Evaluation
40B 이상 모델은 MP 필요함
•
모델 사이즈
◦
MP와 결합하여 ZeRO-100B는 170B 모델 효율적으로 작동 가능
◦
SOTA 모델 사이즈의 8배
•
속도
◦
SOTA 속도의 10배
•
scalability
◦
GPU 수를 2배로 늘렸을 때 성능이 2배 이상 향상(super linear speedup)
◦
ZeRO-DP 덕분에 더 큰 배치 사이즈를 넣을 수 있음
Deepspeed
Deepspeed 라이브러리
•
한정된 컴퓨터 하드웨어로 좋은 병렬 처리를 사용하여 거대 모델을 학습하기 위해, 컴퓨팅 파워와 메모리 사용량을 줄이기 위한 라이브러리
•
Deepspeed 라이브러리는 ZeRO(Zero Redundancy Optimizer)를 구현하여 파이토치와 호환되는 가벼운 API를 통해 모델 학습을 지원함
•
딥러닝 모델과 학습 프레임워크(파이토치) 사이의 최적화 계층 역할
주요 특징
•
간편한 통합: 파이토치 코드의 약간의 수정만으로도 속도와 스케일을 향상시킬 수 있음
•
확장성: 기존 데이터 병렬 처리 방식은 1.3B 파라미터 이상의 모델에서는 메모리가 부족하지만, Deepspeed는 6B 파라미터까지 모델 병렬 처리 없이 데이터 병렬 처리만으로도 사용 가능(GPU 32GB 기준)
•
메모리 효율성: 모델 병렬 처리 없이도 메모리 효율적인 데이터 병렬 처리를 사용할 수 있음
간단한 코드 변경으로 가능
•
Deepspeed 초기화: 모델과 옵티마이저를 wrapping
•
순방향 및 역방향 패스, 파라미터 업데이트: 몇 줄의 코드 변경만으로 구현
Deepspeed 사용 방법
config.json 작성
{
"zero_optimization": {
"stage": 3, // stage 0, stage 1, stage 2, stage 3
"offload_optimizer": { // 모델 파라미터를 CPU에 오프로드
"device": "cpu",
"pin_memory": true
},
"offload_param": { // 옵티마이저 상태를 CPU에 오프로드
"device": "cpu",
"pin_memory": true
},
"overlap_comm": true, // 기울기 감소를 역전파 계산과 겹치도록 하여 속도 빨라짐
"contiguous_gradients": true, // 기울기가 생성되는 대로 연속적인 버퍼에 복사
"sub_group_size": 1e9, // 버킷으로 그룹화하여 파라미터 기울기 업데이트
"reduce_bucket_size": "auto", // 한 번에 줄어들 bucket 사이즈
"stage3_prefetch_bucket_size": "auto", // 파라미터 미리 불러올 때 고정된 버퍼 사이즈
"stage3_param_persistence_threshold": "auto", // 보다 작은 파라미터는 분할 x
"stage3_max_live_parameters": 1e9, // GPU에서 언제나 유지하려는 전체 파라미터 수의 상한선
"stage3_max_reuse_distance": 1e9, // 다시 사용될 파라미터 버리지 않기
"stage3_gather_16bit_weights_on_model_save": true // GPU에 퍼져 있는 가중치를 모아서 fp16으로 저장
},
"gradient_accumulation_steps": "auto",
"gradient_clipping": "auto",
"steps_per_print": 2000,
"train_batch_size": "auto",
"train_micro_batch_size_per_gpu": "auto",
}모델 분산 로드
from transformers import AutoConfig, AutoModelForForCausalLM
import deepspeed
with deepspeed.zero.Init():
config = AutoConfig.from_pretrained(base_model)
model = AutoModelForForCausalLM(config)Huggingface Trainer에 deepspeed config 전달
args = TrainingArguments(
.....,
deepspeed = 'path/to/deepspeed_config.json',
)실행
deepspeed --num_gpus=2 train.py <다른 arguments들...> --deepspeed ds_config.json 참고자료