DeepSpeed - 쉽고 효율적인 딥러닝 학습 최적화 라이브러리

DeepSpeed?

 

DeepSpeed is a deep learning optimization library that makes distributed training easy, efficient, and effective.

 

10x Larger Models | 10x Faster Training | Minimal Code Change

 

microsoft에서 공개한 딥러닝 모델 학습 최적화 라이브러리입니다. 특히 Pytorch 기반 모델에만 적용 가능합니다.

 

 

장점

 

Speed

gpt-2 모델 학습 시 DeepSpeed 적용 전 대비 3.5배 더 빠르다고 합니다. 

 

Memory efficiency

메모리 효율적인 데이터 병렬화를 제공합니다. Deepspeed를 사용하면 13 billion 파라미터 모델을 v100에서 학습 가능합니다. 하지만 PyTorch’s Distributed Data Parallel를 사용하면 1.4billion 파라미터 모델 학습 시에 out of memory가 발생한다고 합니다.

 

Scalability

바로 위에 설명한 대로 큰 크기의 파라미터 모델 학습이 가능합니다. 

 

Fast convergence for effectiveness

더 개선된 하이퍼 파라미터 튜닝과 lamb같은 batch size optimizer를 지원합니다..

 

Good Usability

Pytorch 모델에 DeepSpeed를 적용하기 위해선 단지 몇 줄의 코드만 바꾸면 됩니다. 다른 모델 병렬화 라이브러리와 비교했을 경우, DeepSpeed는 코드를 다시 디자인하거나 모델 재구조화가 필요하지 않습니다.

 

 

Features

다음은 DeepSpeed의 기능들입니다. 자세한 설명은 링크를 통해 보시길 바랍니다.

• Distributed training with mixed precision
  • 16-bit mixed precision
  • Single-GPU/Multi-GPU/Multi-Node
• Model parallelism
  • Support for Custom Model Parallelism
  • Integration with Megatron-LM
• The Zero Redundancy Optimizer (ZeRO)
  • Optimizer State and Gradient Partitioning
  • Activation Partitioning
  • Constant Buffer Optimization
  • Contiguous Memory Optimization
• Additional memory and bandwidth optimizations
  • Smart Gradient Accumulation
  • Communication/Computation Overlap
• Training features
  • Simplified training API
  • Activation Checkpointing API
  • Gradient Clipping
  • Automatic loss scaling with mixed precision
• Training optimizers
  • Fused Adam optimizer and arbitrary torch.optim.Optimizer
  • Memory bandwidth optimized FP16 Optimizer
  • Large Batch Training with LAMB Optimizer
  • Memory efficient Training with ZeRO Optimizer
• Training agnostic checkpointing
• Advanced parameter search
  • Learning Rate Range Test
  • 1Cycle Learning Rate Schedule
• Simplified data loader
• Performance analysis and debugging

 

 

 

이제 튜토리얼에 있는 예시들을 통해 어떻게 사용하는지 살펴보겠습니다.

 

To initialize the DeepSpeed engine

model_engine, optimizer, _, _ = deepspeed.initialize(args=cmd_args,
                                                     model=model,
                                                     model_parameters=params)

여기서 torch.nn.module 형식의 모델은 전부 적용 가능합니다.

 

 

distributed environment setup

# need to replace :
torch.distributed.init_process_group(...)

# to :
deepspeed.init_distributed()

 분산 학습의 경우 torch.distributed가 아닌 deepspeed.init_distributed()로 변경하여 사용합니다.

 

Training

for step, batch in enumerate(data_loader):
    #forward() method
    loss = model_engine(batch)

    #runs backpropagation
    model_engine.backward(loss)

    #weight update
    model_engine.step()

학습 시 바뀌는 부분은 없어 보입니다.

 

 

 

Model Checkpointing

_, 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)

학습 중 checkpoint마다 모델을 손쉽게 저장할 수 있습니다.

 

 

이상으로 DeepSpeed 라이브러리에 대해 알아보았습니다.

 

---

https://github.com/microsoft/DeepSpeed

GitHub - microsoft/DeepSpeed: DeepSpeed is a deep learning optimization library that makes distributed training easy, efficient,

https://www.microsoft.com/en-us/research/project/deepspeed/

DeepSpeed - Microsoft Research

https://www.deepspeed.ai/

DeepSpeed