더 크고, 더 빠른 모델을 향해, DeepSpeed.
최근에는 모델의 크기들이 점점 더 커지고 단일 GPU에서 학습하기 어려운 환경이 되어갑니다.
어디서는 몇백개의 GPU와 TPU를 이용해 자유롭고 빠르게 학습하는 반면 개인 딥러너들은 갈수록 따라잡기 어려워지는 현실.
하지만 하드웨어 및 소프트웨어의 성능은 날이 갈수록 빠르게 발전하고, 점점 더 큰 모델들을 학습하기 위한 SW들도 많이 공개 되고 있습니다. 그중의 하나가 DeepSpeed.
저의 경우 언어모델을 학습하는데 있어서, 더 GPT-3 175B 모델을 학습하는데, Nipa에서 지원 받은 GPU가 모델을 올리다가 다운 되는 현상이 있어서 여러가지 방법을 찾게 되었습니다.
최근에는 Deep Speed에 ZeRO-Infinity가 공개되어 GPT-3 보다 더 큰 매개변수로 학습이 가능하게 되었습니다. 저도 이번에 기존과 다르게 큰 모델을 학습시키는데 Deep Speed를 정리 및 학습해보고 활용해보고자 합니다.
DeepSpeed
deepspeed 깃헙에 들어가면 상단에 아래와 같은 문구가 보입니다. 최소한의 코드 변경으로 10배 큰모델을 10배 빠르게 학습시키는 것을 슬로건으로 합니다. DeepSpeed는 딥러닝 최적화 라이브러리로 학습을 쉽고, 편하고, 효율적으로 만들어 줍니다.
싱글 GPU를 사용하는 개인 개발자나 대규모의 슈퍼컴퓨터를 사용하는 데이터 사이언티스트들까지 초대형 모델을 학습할수 있도록 합니다.
DeepSpeed Features
대표적으로 DeepSpeed가 가지는 특징은 아래와 같습니다. 보다 자세한 문서들은 deepspeed.ai에 문서화로 특징들이 잘 정리되어 있습니다.
- Extreme Scale: 수백대의 GPU 클러스터를 사용하는 경우에도 DeepSpeed는 3D Parallelism을 이용해 수조개의 파라미터를 가진 모델을 효율적으로 학습할 수 있습니다.
- Extremely Memory Efficient: DeepSpeed의 ZeRO-Offload 를 이용하면 싱글 GPU를 사용해서 SOTA보다 10배 더큰 10B 파라미터의 모델을 학습할 수 있고, 수조개의 파라미터를 가진 모델을 학습해 더 큰 모델, 더 나은 모델을 찾기 위해 탐구할 수 있습니다.
- Extremely Long Sequence Length: DeepSpeed의
Sparse Attention을 활용하면 더 긴 입력 시퀀스에 대해 기존 트랜스포머 대비 6배 더 빠른 성능을 얻을 수 있습니다. - Extremely Communication Efficient: 3D Parallelism을 사용하면 클러스터 내에서 2~7 배 빠르게 학습시킬 수 있습니다.
DeepSpeed Installation
pip install deepspeedDeepSpeed Example
DeepSpeed를 사용한 훈련은 DeepSpeed 엔진을 사용해 학습됩니다. DeepSpeed 엔진을 사용하기 위해 torch.nn.module를 래핑하여 사용합니다.
초기화
deepspeed.initialize는 분산학습이나 혼합 정밀도 학습을 할 수 있도록 초기화 해줍니다.
model_engine, optimizer, _, _ = deepspeed.initialize(args=cmd_args,
model=model,
model_parameters=params)분산학습
아래 Torch 분산학습 코드를 사용하는 경우
torch.distributed.init_process_group(...)아래 코드로 교체해아합니다.
deepspeed.init_distributed()모델 학습
위와 같이 DeepSpeed engine을 초기화 하고나면, 3가지 API를 사용해 모델을 학습시킬수 있습니다.
for step, batch in enumerate(data_loader):
#forward() method
loss = model_engine(batch)
#runs backpropagation
model_engine.backward(loss)
#weight update
model_engine.step()DeepSpeed는 분산학습, 혼합 정밀도, 사전정의된 learning rate 스케쥴러들에 대해 자동으로 필요한 연산들을 수행합니다.
- Gradient Averaging: 데이터 병렬 학습에서
backward시에 학습데이터train_batch_size에 대해서 그래디언트가 평균임을 보장합니다. - Loss Scaling: Fp16/ mixed precision 학습시에 DeepSpeed 엔진은 precision loss를 피하기 위해 Loss를 자동으로 스케일링 합니다.
- Learning Rate Scheduler: DeepSpeed는 매 학습 Step 마다
step()을 호출
+ 스케쥴러를deepspeed.initialize에서 전달해서 DeepSpeed가 관리하게 할 수 있습니다.
+ 스케쥴러가 어떤 주기를 가지고 실행된다면 스케쥴러를 DeepSpeed에 전달하면 안되며, 독립적으로 처리될수 있도록 해야합니다.
모델 체크포인트 다루기
모델을 학습시키다보면 체크포인트 파일을 저장하고 로드하는 과정이 필요합니다. DeepSpeed는 유일한 Checkpoint를 확인하기 위한 두 매개 변수를 사용합니다.
- ckpt_dir: 체크포인트가 저장 될 경로
- ckpt_id: 체크포인트 폴더 안에서 구분되기 위한 유니크한 식별자 입니다.
#load checkpoint
_, client_sd = model_engine.load_checkpoint(args.load_dir, args.ckpt_id)
step = client_sd['step']
#advance data loader to ckpt step
dataloader_to_step(data_loader, step + 1)
for step, batch in enumerate(data_loader):
#forward() method
loss = model_engine(batch)
#runs backpropagation
model_engine.backward(loss)
#weight update
model_engine.step()
#save checkpoint
if step % args.save_interval:
client_sd['step'] = step
ckpt_id = loss.item()
model_engine.save_checkpoint(args.save_dir, ckpt_id, client_sd = client_sd)DeepSpeed 설정 다루기
DeepSpeed는 Json 형식을 통해서 설정을 다룹니다.
{
"train_batch_size": 8,
"gradient_accumulation_steps": 1,
"optimizer": {
"type": "Adam",
"params": {
"lr": 0.00015
}
},
"fp16": {
"enabled": true
},
"zero_optimization": true
}DeepSpeed Training 시작하기
DeepSpeed로 학습을 시작하기 전 아래 사항들을 준비해야합니다.
1. DeepSpeed 사용을 위한 모델 변경
2. client_entry.py: 모델 학습을 위한 스크립트
3. client args : argparse 명령어를 이용한 argument들
4. ds_config.json: DeepSpeed 학습을 위한 설정 파일
멀티 노드 학습
멀티 노드 학습 참고
싱글 노드 학습을 위한 설정
단일 GPU에서만 실행 중인 경우 DeepSpeed는 멀티노드학습에서 사용하는 호스트 파일을 필요로 하지 않습니다. 현재 노드 중 gpu1만 사용하고자 하는 경우 아래와 같이 사용할 수 있습니다.
deepspeed --include localhost:1
<client_entry.py> <client args> \
--deepspeed --deepspeed_config ds_config.json
마무리
위에서 DeepSpeed에 대해 알아보고, 단일 GPU에서 사용하는 방법에 대해 정리해보았습니다. 다음으로는 실제 대형 언어모델을 학습시키는데 적용해보고, 정리가 필요한 부분이 있으면 다시 정리하도록 하겠습니다.
