ZeRO-2 & DeepSpeed: Shattering barriers of deep learning speed & scale

Microsoft Research Blog 글을 보고 번역하여 작성하였습니다.

출처 : https://www.microsoft.com/en-us/research/blog/zero-2-deepspeed-shattering-barriers-of-deep-learning-speed-scale/

 

2020년 2월에 Microsoft에서 오픈 소스 딥 러닝 훈련 최적화 라이브러리인 DeepSpeed와 라이브러리의 새로운 메모리 최적화 기술인 ZeRO(Zero Redundancy Optimizer)를 발표했습니다. DeepSpeed를 통해 MS연구원들은 170억 개의 매개변수로 공개적으로 알려진 가장 큰 언어 모델인 Turing Natural Language Generation(Turing-NLG)을 만들 수 있었습니다. 그 이후로 우리는 딥 러닝 training을 위한 속도와 규모의 한계를 뛰어 넘는 빠른 속도로 혁신을 계속해 왔습니다. Microsoft에서 개선된 ZeRO-2와 DeepSpeed를 통한 추가 개발을 소개하고 연구결과를 확인하고자 합니다.

 

• An order-of-magnitude larger and faster training with ZeRO-2 : ZeRO-2는 훈련 중 메모리 소비의 전체 ​​스펙트럼을 처리하여 원래 ZeRO의 메모리 최적화 범위를 확장합니다. 보다 구체적으로, ZeRO-2는 원래 ZeRO의 옵티마이저 상태 메모리 최적화 외에도 gradients, activation memory 및 fragmented memory의 memory footprint를 줄이는 새로운 기술을 도입합니다. 전체적으로 메모리 절약을 통해 DeepSpeed는 딥 러닝 training의 규모와 속도를 한 차원 높일 수 있습니다. 보다 구체적으로, ZeRO-2는 1,700억 개의 매개변수만큼 큰 모델을 최신 기술에 비해 최대 10배 더 빠르게 훈련할 수 있습니다.
• Fastest BERT training: ZeRO-2가 distributed training 중에 대형 모델을 최적화하는 동안 커널 최적화를 통해 단일 GPU 성능을 가속화하는 새로운 기술도 도입합니다. 이러한 최적화는 대규모 모델을 확장하기 위한 강력한 기반을 구축할 뿐만 아니라 BERT와 같이 고도로 조정되고 중간 크기의 모델의 단일 GPU 성능을 30% 이상 개선하여 하드웨어 피크의 50%인 V100 GPU당 64teraflops 의 놀라운 성능에 도달합니다. 이러한 최적화를 building block으로 사용하여 DeepSpeed는 가장 빠른 BERT training 속도 기록을 달성합니다. 1,024개의 NVIDIA V100 GPU에서 44분, 동일한 수 및 세대의 GPU에서 67분이라는 최고 발표 결과입니다.

ZeRO-2: Training models with 100 billion parameters up to 10x faster

Zero Redundancy Optimizer(Zero Redundancy Optimizer)는 대규모 분산 딥 러닝을 위한 새로운 메모리 최적화 기술입니다. data parallelism(효율적이지만 제한된 모델 크기만 지원할 수 있음) 또는 model parallelism(더 큰 모델 크기를 지원할 수 있지만 효율성을 제한하는 통신 오버헤드를 추가하면서 상당한 코드 리팩토링이 필요함)와 같은 기존 기술과 달리 ZeRO는 메모리에 더 큰 모델을 맞출 수 있습니다. 매우 효율적으로 유지하면서 코드 리팩토링이 필요하지 않습니다. ZeRO는 통신 오버헤드를 최소로 제한하면서 데이터 병렬 처리에 내재된 메모리 중복성을 제거함으로써 이를 수행합니다.

ZeRO는 데이터 병렬 프로세스에서 세 가지 모델 상태(optimizer states, gradients, and parameters)를 복제하는 대신 분할하여 데이터 병렬 프로세스에서 메모리 중복성을 제거합니다. 이렇게 하면 계산 단위와 통신 효율성을 유지하면서 기존의 데이터 병렬 처리에 비해 메모리 효율성을 높일 수 있습니다.

DeepSpeed의 2020년 2월 릴리스에는 최적화 프로그램 상태 메모리(ZeRO-1)를 줄이기 위한 최적화가 포함되었습니다. Microsoft는 gradient memory를 줄이기 위한 최적화를 포함하고 activation memory와 fragmented memory를 대상으로 하는 최적화를 추가하여 ZeRO-1을 확장한 ZeRO-2를 출시했습니다. ZeRO-1과 비교하여 ZeRO-2는 DeepSpeed로 훈련할 수 있는 모델 크기를 두 배로 늘리면서 훈련 효율성을 크게 향상시킵니다. ZeRO-2를 사용하면 모델 병렬 처리만 기반으로 하는 최첨단 기술보다 1000억 개의 매개변수 모델을 10배 빠르게 학습할 수 있습니다.

 

 

ZeRO-2 deep dive: Reducing gradients, activation, and fragmented memory

ZeRO-2는 model state(예: optimizer states and gradients), activation memory 및  fragmented memory를 포함하여 딥 러닝 트레이닝 중 메모리 소비의 전체 스펙트럼을 최적화합니다. 그림 1은 ZeRO-2의 핵심 기술을 보여주고 있으며, 자세한 내용은 아래와 같습니다.

Figure 1: ZeRO-2 optimization overview. Going beyond ZeRO-1, which optimizes partitioning optimizer states (Pos), ZeRO-2 introduces new technology to reduce the memory footprint of partitioning gradients (Pos+g), activation memory, and fragmented memory, tackling the full spectrum of memory optimizations.

Model state memory : ZeRO에는 model states를 최적화하기 위한 세 가지 누적 단계가 있습니다. 이러한 states는 각각 분할 최적화 states(Pos), 기울기(Pos+g) 및 매개변수(Pos+g+p)입니다. 2월에 공유한 ZeRO-1 구현은 모든 것을 복제하는 기존 데이터 병렬 처리를 사용할 때와 비교할 때 최대 4배의 메모리를 절약하는 첫 번째 단계인 파티셔닝 옵티마이저 states(Pos)를 지원합니다. ZeRO-2는 두 번째 단계인 파티셔닝 그라디언트(Pos+g)에 대한 지원을 추가하여 첫 번째 단계의 4배 감소 외에도 장치당 메모리 소비를 2배 더 줄입니다. 기본 데이터 병렬 처리와 비교하여 ZeRO-2는 동일한 통신 볼륨으로 model states에서 최대 8배의 메모리 절약을 얻습니다.

Activation memory : model states를 최적화한 후 activation(역방향 전달을 수행하기 위해 순방향 전달에서 저장)가 보조 메모리 병목 현상이 될 수 있음을 알았습니다. Activation checkpointing은 도움이 되지만 매우 큰 모델에는 충분하지 않습니다. ZeRO-2는 activation partitioning을 통해 기존 model parallelism 접근 방식에서 activation 복제를 제거하는 새로운 기술을 도입합니다. 또한 적절한 경우 Activation memory를 호스트 CPU로 오프로드합니다.

Fragmented memory : 다른 tensors의 수명 변화로 인해 훈련 중에 Fragmented memory를 관찰합니다. 사용 가능한 메모리가 충분하더라도 조각화로 인한 연속 메모리 부족으로 인해 메모리 할당 실패가 발생할 수 있습니다. ZeRO-2는 tensors의 다양한 lifetime을 기반으로 메모리를 사전에 관리하여 메모리 단편화를 방지합니다.

 

ZeRO-2 evaluation: Advancing size, speed, scalability, and usability

superlinear scalability과 improved usability을 통해 large model training을 문턱을 낮추기 위해 최대 10배 더 빠르고 더 큰 모델을 지원합니다. 이 네 가지 측면은 아래에 자세히 설명되어 있습니다.

Figure 2: ZeRO-2 scales to 170 billion parameters, has up to 10x higher throughput, obtains superlinear speedup, and improves usability by avoiding the need for code refactoring for models up to 13 billion parameters.

Model scale : Zero를 사용하지 않고 학습된 State-of-the-art large model들은 각각 OpenAI GPT-2는 1.5B, NVIDIA Megatron-LM은 8.3B, Google T5는 11B의 parameters를 가집니다.
 ZeRO-2는 170 Billion parameters의 모델을 효율적으로 실행할 수 있는 시스템 기능을 제공하며, 이는 이러한 large models보다 훨씬 더 큰 규모입니다(그림 2, 왼쪽 상단). 테스트는 400개의 NVIDIA V100 GPU를 사용하여 수행되었습니다. 더 많은 devices (예: 1,000개의 GPU)와 함께 ZeRO-2를 사용하면 200 Billion의 parameters로 확장할 수 있습니다.

Speed : 향상된 메모리 효율성은 더 높은 처리량과 더 빠른 Training을 제공합니다. 그림 2(왼쪽 하단)는 ZeRO-2, ZeRO-1 및 기본 model parallelism의 시스템 처리량을 보여줍니다. 여기서 우리는 최신 model parallelism 방식인 NVIDIA Megatron-LM을 기준 MP로 사용하는 반면, ZeRO-2 및 ZeRO-1은 모두 ZeRO 기반 data parallelism과 Megatron-LM model parallelism를 결합합니다. ZeRO-2는 400개의 NVIDIA V100 GPU가 있는 클러스터에서 GPU당 38테라플롭이 넘는 100 billion parameter models, hardware peak의 30%, 집계 성능이 15페타플롭 이상인 모델을 실행합니다. 동일한 크기의 모델의 경우 ZeRO-2는 model parallelism가 효율적이기 위해 높은 통신 대역폭이 필요하고 이러한 크기의 모델은 통신 대역폭이 제한된 노드 전체에서 model parallelism가 필요하기 때문에 기준선과 비교할 때 training 속도가 최대 10배 빠릅니다. ZeRO-2의 메모리 절약 덕분에 노드 간 model parallelism 없이 모델 병렬도를 줄이고 모델을 맞출 수 있어 통신 비용을 크게 절감할 수 있습니다. 또한 ZeRO-2는 추가 메모리 절약이 통신을 더욱 줄이고 더 큰 배치 크기를 지원하는 데 도움이 되기 때문에 ZeRO-1보다 최대 5배 빠릅니다.

Scalability : NVIDIA GPU 수가 두 배가 될 때 성능이 두 배 이상인 superlinear 속도 향상(그림 2, 오른쪽 상단)을 관찰합니다. ZeRO-2는 data parallelism 정도를 높일 때 모델 states의 메모리 공간을 줄여 GPU당 더 큰 배치 크기에 맞도록 하고 결과적으로 더 나은 성능을 제공합니다.

Democratizing large model training : ZeRO-2는 일반적으로 모델 리팩토링이 필요한 model parallelism 없이 모델 과학자가 최대 13 Billion 개의 parameters를 효율적으로 모델을 training할 수 있도록 합니다(그림 2, 오른쪽 하단). 13 Billion 개의 parameters는 대부분의 가장 큰 최신 모델(예: 110억 개의 parameters가 있는 Google T5)보다 큽니다. 처리량과 관련하여 우리는 2B ~ 13B 개의 parameters 범위의 모델 크기에 대해 V100 GPU당 37테라플롭(hardware peak 30%)의 평균 처리량을 관찰했습니다. 따라서 모델 과학자는 모델 병렬성에 대해 걱정하지 않고 대규모 모델을 자유롭게 실험할 수 있습니다. 이에 비해 classic data parallelism 방식(예: PyTorch Distributed Data Parallel)의 구현은 14억 개의 parameters 모델로 메모리가 부족하지만 ZeRO-1은 최대 60억 개의 parameters를 지원합니다.

ZeRO-2에 대한 자세한 내용은 DeepSpeed ​​GitHub 리포지토리 및 업데이트된 ZeRO 문서를 참조하십시오.

 

Achieving the fastest and most efficient BERT training with DeepSpeed

ZeRO는 주로 장치 클러스터에 걸쳐 distributed training 동안 large models에 이점을 제공하지만, 우리는 또한 새로운 기술, highly optimized transformer kernels, asynchronous I/O를 도입하여 각 개별 GPU에서 훈련의 컴퓨팅 및 I/O 속도를 향상시킵니다. 이 optimizations line은 large models을 확장할 때 견고한 기반을 구축할 뿐만 아니라 BERT와 같은 적당한 크기의 모델을 훈련하는 동안 이미 최적화된 성능의 마지막 비트를 짜냅니다(squeezes out the last bit of already optimized performance).

 

We achieve the fastest BERT training record: 44 minutes on 1,024 NVIDIA V100 GPUs. : 1,024 NVIDIA V100 GPU에서 44분. 또한, 향상된 training 시간은 과도한 하드웨어 리소스를 희생하는 것이 아니라 소프트웨어로 향상된 효율성에서 비롯됩니다. 동일한 수의 GPU 및 generation에서 최상의 결과와 비교할 때 training 처리량을 30% 이상 개선합니다. 단일 V100 GPU에서 64테라플롭의 처리량을 관찰하여 hardware peak의 50% 이상을 달성했습니다.

먼저 단일 GPU 성능부터 살펴보겠습니다. 그림 3은 잘 알려진 두 가지 구현인 NVIDIA BERT 및 HuggingFace BERT와 비교하여 BERT-Large training을 위해 DeepSpeed로 달성한 단일 V100 GPU 처리량을 보여줍니다. DeepSpeed는 각각 128 및 512의 시퀀스 길이에 대해 64 및 53테라플롭(초당 272 및 52샘플에 해당)의 높은 처리량에 도달하여 NVIDIA BERT에 비해 처리량이 최대 28%, HuggingFace BERT에 비해 최대 62% 향상되었습니다. 또한 메모리 부족 없이 최대 1.8배 더 큰 배치 크기를 지원합니다.

Figure 3: Performance evaluation of BERT-Large on a single V100 GPU, comparing DeepSpeed with NVIDIA and HuggingFace versions of BERT in mixed-sequence length training and with gradient accumulation step of 10. The labeled points show the highest throughput of each implementation in teraflops (Tflops). DeepSpeed boosts throughput and allows for higher batch sizes without running out of memory.

GPU 전반에 걸친 distributed training을 살펴보면, 표 1은 16~1,024개의 GPU를 사용한 end-to-end BERT-Large pretraining 시간(SQUAD의 경우 F1 점수 90.5)을 보여줍니다. 1,024개의 V100 GPU(64개의 NVIDIA DGX-2 nodes)를 사용하여 44분 만에 BERT pretraining 을 완료합니다. 1,472개의 V100 GPU를 사용하여 47분이 소요되는 NVIDIA의 가장 잘 알려진 결과와 비교하여 DeepSpeed는 리소스를 30% 적게 사용하면서 더 빠릅니다. 1,024개의 GPU를 사용하는 동안 NVIDIA BERT는 67분, DeepSpeed는 44분이 소요되어 training 시간이 30% 이상 단축됩니다. 마찬가지로 256개의 GPU에서 NVIDIA BERT는 236분, DeepSpeed는 144분이 걸립니다.

Table 1: BERT-Large training time using 1 to 64 DGX-2s with DeepSpeed.

V100 GPU의 2.5배 hardware peak 성능을 제공하는 NVIDIA A100 Tensor Core GPU와 같은 새로운 하드웨어와 소프트웨어 최적화를 결합하면 훨씬 더 높은 처리량을 얻을 수 있습니다. A100 GPU가 V100 GPU에서 얻은 것과 동일한 비율의 hardware peak 성능(50%)을 얻을 수 있다고 가정하면 1,024 A100 GPU 클러스터에서 BERT 교육 시간을 25분 미만으로 추가로 줄일 수 있습니다.

BERT pretraining에서 DeepSpeed의 시스템 성능 향상은 주로 두 가지 최적화 라인을 사용하는 highly optimized transformer kernels에서 비롯됩니다.

First, transformer-based networks가 producer-consumer 방식으로 작동하는 NVIDIA CUDA kernels의 많은 호출을 트리거하여 메모리에 대한 과도한 읽기/쓰기 요청을 생성한다는 것을 관찰했습니다. 이 패턴은 kernel fusion과 완벽하게 일치하며 우리는 데이터 이동을 줄이기 위해 fused transformer-kernels을 개발했습니다. 또한, fused kernel에서는 NVIDIA GPU 코어 간의 데이터 처리를 분산하여 동기화 및 데이터 통신을 줄이는 방식으로 처리합니다. DeepSpeed는 이러한 optimized fused transformer-kernels을 활용하여 데이터 재사용을 극대화하고 GPU에서 계산 효율성을 개선합니다.

Second, 메모리를 절약하고 큰 배치 크기를 지원하기 위해 가역 연산자(invertible operators)를 활용합니다. 가역 연산자는 backward pass가 입력과 독립적이고 출력에만 기반하여 공식화될 수 있는 연산자입니다. 우리는 Softmax 및 Layer Norm을 포함하여 여러 transformer 연산자가 그러한 동작을 하는 것을 보았습니다. 역전할 수 있는 속성을 활용하여 이러한 레이어에 대한 입력을 삭제하여 activation memory의 공간을 줄입니다. 원래의 가역적 Softmax 계산의 예가 그림 4에 나와 있습니다. 여기서 가역 버전(invertible version)은 operator 의 activation memory를 절반으로 줄입니다. activation memory를 줄이면 더 큰 배치 크기에 맞출 수 있어 종종(often) 효율성이 높아집니다.

Figure 4: Default PyTorch Softmax operation versus DeepSpeed invertible Softmax operation.

optimized kernels 외에도 asynchronous I/O를 사용하여 통신과 연산을 중첩하고, 더 높은 학습률과 더 빠른 수렴을 허용하는 pre-layer 표준을 활용하고, output layer에서 중복 계산을 줄입니다.

전체적으로 이러한 기술은 BERT training 시간과 효율성을 최적화합니다. 또한 이러한 최적화의 대부분은 BERT에 국한되지 않습니다. 대신 광범위한 워크로드에 적용할 수 있는 일반적인 최적화 전략입니다. 예를 들어 highly optimized kernels을 다른 transformer-based models에서 쉽게 재사용할 수 있는 custom transformer operators로 패키징합니다.

More updates of DeepSpeed, codes, tutorials and technical deep dive blogs are available in DeepSpeed.ai.

 

 

출처 : https://www.microsoft.com/en-us/research/blog/zero-2-deepspeed-shattering-barriers-of-deep-learning-speed-scale/

DeepSpeed & ZeRO-2: Shattering barriers of deep learning speed & scale