새로운 기능 – Amazon SageMaker 관리형 데이터 병렬 처리를 통해 대규모 데이터 세트 훈련 간소화

오늘 Amazon SageMaker가 수백 또는 수천 기가바이트의 데이터 세트에 대한 모델을 더 쉽게 훈련할 수 있는 새로운 데이터 병렬 처리 라이브러리를 지원한다는 것을 발표하게 되어 기쁩니다.

데이터 세트와 모델이 점점 커지고 정교해짐에 따라 대규모 분산 훈련 작업을 수행하는 기계 학습(ML) 실무자는 Amazon Elastic Compute Cloud (EC2) p3 및 p4 인스턴스와 같은 강력한 인스턴스를 사용하는 경우에도 점점 더 긴 훈련 시간에 직면해야 합니다. 예를 들어 8개의 NVIDIA V100 GPU가 장착된 ml.p3dn.24xlarge 인스턴스를 사용하면 공개적으로 사용 가능한 COCO 데이터 세트에 대한 마스크 RCNN 및 빠른 RCNN과 같은 고급 물체 감지 모델을 훈련하는 데 6시간이 걸립니다. 마찬가지로, 최첨단 자연어 처리 모델인 BERT 훈련은 동일한 인스턴스에서 100시간 이상 걸립니다. 자율 주행 차량 회사와 같은 일부 고객은 대규모 GPU 클러스터에서 며칠 동안 실행되는 대량의 훈련 작업을 일상적으로 처리합니다.

상상할 수 있듯이, 이러한 긴 훈련 시간은 ML 프로젝트의 심각한 병목 현상으로 생산성이 저하되고 혁신이 느려집니다. 고객들이 도움을 요청했으므로 저희가 도와드려야 했습니다.

Amazon SageMaker에서 데이터 병렬 처리 소개
이제 Amazon SageMaker는 SageMaker SDP(데이터 병렬 처리) 라이브러리를 통해 ML 팀이 분산된 훈련 시간과 비용을 줄일 수 있도록 지원합니다. TensorFlow와 PyTorch에서 사용할 수 있는 SDP는 보다 효율적인 계산 분산을 구현하고 네트워크 통신을 최적화하며 가장 빠른 p3 및 p4 GPU 인스턴스를 완벽하게 활용합니다.

이제 GPU 리소스의 최대 90%를 데이터 전송이 아닌 훈련에 사용할 수 있습니다. 분산된 훈련 작업은 관련된 GPU 수에 관계없이 선형에 가까운 확장 효율성을 달성할 수 있습니다. 즉, 훈련 작업이 단일 인스턴스에서 8시간 동안 실행되는 대신 8개의 인스턴스에서 약 1시간만 소요되며 비용은 최소화됩니다. SageMaker는 훈련 비용과 훈련 시간 간의 절충을 효과적으로 제거하여 ML 팀이 더 빨리 결과를 얻고, 더 빠르게 반복하고, 혁신을 가속화할 수 있도록 합니다.

Swami Sivasubramanian는 AWS re:Invent 2020에서 기조 연설을 하는 중에 T5-3B 및 마스크-RCNN에 대한 가장 빠른 훈련 시간을 보여주었습니다.

• T5-3B 모델은 30억 개의 파라미터 가지고 있으며 자연어 처리 벤치마크에 대한 최첨단 정확성을 달성하며 일반적으로 성능을 훈련하고 조정하는 데 몇 주 간의 작업이 필요합니다. 우리는 이 모델을 6일 동안 256개의 ml.p4d.24xlarge 인스턴스에서 훈련했습니다.
• MASK-RCNN은 고객이 사용하는 인기 있는 인스턴스 세분화 모델입니다. 작년 re:Invent에서 우리는 PyTorch에서 26분 동안 마스크-RCNN을 훈련했고 TensorFlow에서 27분 동안 훈련했습니다. 올해 우리는 Mask-RCNN에 대해 TensorFlow에서 6분 12초 및 PyTorch에서 6분 45초로 지금까지 가장 빠른 훈련 시간을 기록했습니다.

Amazon SageMaker가 어떻게 이러한 속도 향상을 달성할 수 있는지 설명하기 전에 먼저 데이터 병렬화가 작동하는 방식과 확장이 어려운 이유를 설명해 보겠습니다.

데이터 병렬 처리에 대한 입문서
단일 GPU에서 모델을 훈련하는 경우 모델 파라미터, 옵티마이저 파라미터, 기울기(역전파에 의해 계산된 파라미터 업데이트) 등 전체 내부 상태를 로컬로 사용할 수 있습니다. 그러나 GPU 클러스터에 훈련 작업을 배포할 때는 상황이 다릅니다.

“데이터 병렬 처리”라는 기술을 사용하여 훈련 집합은 GPU에 균등하게 분산된 미니 배치로 분할됩니다. 따라서 각 GPU는 전체 데이터 세트의 일부에서 모델을 훈련합니다. 분명히 이것은 모델 상태가 다른 배치를 처리하기 때문에 각 GPU에서 약간 다를 것임을 의미합니다. 훈련 컨버전스를 보장하려면 모든 노드에서 모델 상태를 정기적으로 업데이트해야 합니다. 이 작업은 동기식으로 또는 비동기식으로 수행할 수 있습니다.

• 동기식 훈련: 모든 GPU는 기울기 업데이트를 다른 모든 GPU(다대다 통신) 또는 이를(다대일에서 일대다 순서) 재배포하는 중앙 파라미터 서버에 보고합니다. 모든 업데이트가 동시에 적용되므로 모델 상태가 모든 GPU에서 동기화되고 다음 미니 배치를 처리할 수 있습니다.
• 비동기식 훈련: 기울기 업데이트는 다른 모든 노드 또는 중앙 서버로 전송됩니다. 그러나 즉시 적용되므로 모델 상태가 GPU마다 다를 수 있습니다.

불행히도 이러한 기술은 잘 확장되지 않습니다. GPU 수가 증가함에 따라 파라미터 서버에 필연적으로 병목이 발생합니다. 파라미터 서버가 없어도 n개의 GPU가 n*(n-1)*모델 크기 바이트의 총 양에 대해 매번 반복 후에 n*(n-1) 메시지를 교환해야하므로 네트워크 혼잡이 곧 문제가 됩니다. 예를 들어, Resnet-50은 컴퓨터 비전 애플리케이션에 널리 사용되는 모델입니다. 2,600만 개의 파라미터를 사용하여 각 32비트 기울기 업데이트에는 약 100MB가 소요됩니다. 8개의 GPU를 사용할 경우 각 반복마다 56개의 업데이트를 보내고 수신해야 하며 총 5.6GB를 사용합니다. 빠른 네트워크에서도 약간의 오버헤드가 발생하고 훈련 속도가 느려집니다.

Horovod 프로젝트 덕분에 2017년에 중요한 발걸음을 내딛었습니다. Horovod는 “ring-allreduce”라는 분산 훈련에 최적화된 통신 알고리즘을 구현했습니다. 이 알고리즘은 곧 인기 있는 딥 러닝 라이브러리와 통합되었습니다.

nutshell에서 ring-allreduce는 분산된 비동기식 알고리즘입니다. 파라미터 서버는 없습니다. 노드는 방향 순환 그래프로 구성되며 간단히 말하면 단방향 링입니다. 각 반복에 대해 노드는 이전 노드로부터 기울기 업데이트를 수신합니다. 노드가 자체 배치를 처리하면 자체 업데이트 및 수신한 업데이트를 모두 적용하고 결과를 인접 노드에 전송합니다. n개의 GPU를 사용하면 모든 GPU가 업데이트되기 전에 각 GPU가 2* (n-1) 메시지를 처리합니다. 따라서 GPU당 교환되는 총 데이터 양은 2*(n-1)*model size이며 n*(n-1)*model size보다 훨씬 많습니다.

여전히 데이터 세트가 계속 증가함에 따라 네트워크 병목 현상이 다시 발생하는 경우가 많습니다. SageMaker와 새로운 AllReduce 알고리즘을 입력합니다.

Amazon SageMaker의 새로운 데이터 병렬 처리 알고리즘
AllReduce 알고리즘을 사용하면 GPU가 더 이상 서로 통신하지 않습니다. 각 GPU는 기울기 업데이트를 GPU 메모리에 저장합니다. 특정 임계값을 초과하면 이러한 업데이트가 공유되어 GPU 인스턴스의 CPU에서 실행되는 파라미터 서버로 전송됩니다. 이렇게 하면 전용 파라미터 서버가 필요하지 않습니다.

각 CPU는 모델 파라미터의 하위 집합을 담당하며 모든 GPU에서 오는 업데이트를 수신합니다. 예를 들어, 1개의 GPU가 장착된 3개의 훈련 인스턴스를 사용하면 훈련 클러스터의 각 GPU는 세 개의 CPU 각각에 기울기 업데이트의 1/3을 전송합니다.

그런 다음 각 CPU는 수신한 모든 기울기 업데이트를 적용하고 통합 결과를 모든 GPU에 다시 배포합니다.

이제 이 알고리즘의 작동 방식을 이해했으므로 인프라를 관리하지 않고도 자체 코드와 함께 사용할 수 있는 방법을 살펴 보겠습니다.

Amazon SageMaker에서 데이터 병렬 처리를 사용한 훈련
SageMaker 데이터 병렬 처리 API는 사용하기 쉽도록 설계되었으며 기존 분산 훈련 툴킷과 원활하게 통합되어야 합니다. 대부분의 경우 훈련 코드에서 Horovod(TensorFlow) 또는 분산 데이터 병렬 처리(PyTorch)에 대한 import 명령문만 수정하면 됩니다.

PyTorch의 경우 다음과 같습니다.

import smdistributed.dataparallel.torch.parallel.distributed as dist
dist.init_process_group()

그런 다음 각 GPU를 단일 SDP 프로세스에 고정해야 합니다.

torch.cuda.set_device(dist.get_local_rank())

그런 다음 평소와 같이 모델을 정의합니다. 예를 들면 다음과 같습니다.

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout2d(0.25)
        self.dropout2 = nn.Dropout2d(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)
...

마지막으로 모델을 인스턴스화하고 이를 사용하여 다음과 같이 DistributedDataParallel 객체를 다음과 같이 만듭니다.

import torch
from smdistributed.dataparallel.torch.parallel.distributed import DistributedDataParallel as DDP
device = torch.device("cuda")
model = DDP(Net().to(device))

코드의 나머지 부분은 바닐라 PyTorch이며, SageMaker SDK에서 사용할 수 있는 PyTorch 추정기를 사용하여 훈련할 수 있습니다. 여기서는 8개의 NVIDIA V100 GPU가 있는 ml.p3.16xlarge 인스턴스를 사용합니다.

from sagemaker.pytorch import PyTorch
estimator = PyTorch(
    entry_point='train_pytorch.py',
    role=sagemaker.get_execution_role(),
    framework_version='1.6.0',
    instance_count=1,
    instance_type='ml.p3.16xlarge',
    distribution={'smdistributed':{'dataparallel':{enabled': True}}}
)
estimator.fit()

이후 SageMaker가 모든 필요한 인프라를 맡아서 프로비저닝합니다. 여러분은 훈련 작업이 실행되는 동안 다른 작업에 집중할 수 있습니다.

시작하기
훈련 작업이 여러 GPU에서 몇 시간 또는 며칠 동안 지속될 경우 SageMaker Data Parallelism 라이브러리를 사용하면 시간과 비용을 절약하고 실험과 혁신을 더 빠르게 할 수 있습니다. 현재 SageMaker를 사용할 수 있는 모든 리전에서 추가 비용 없이 사용할 수 있습니다.

신속하게 시작하려면 예제를 따라 하세요. 사용해보시고 의견을 알려주세요. 언제든지 AWS 지원 담당자나 SageMaker에 대한 AWS 포럼을 통해 피드백을 보내 주시기 바랍니다.

– Julien