Amazon SageMaker에서 분할된 데이터 병렬 처리를 사용하여 선형에 가까운 확장으로 거대한 모델 훈련

뛰어난 정확도를 추구하기 위해 자연어 처리 및 컴퓨터 비전과 같은 분야의 딥 러닝 모델은 지난 몇 년 동안 크기가 크게 증가했으며 종종 수십억에서 수천억 개의 매개 변수로 계산됩니다. 이러한 거대한 모델을 훈련하는 것은 어렵고 복잡한 배포 전략이 필요합니다. 데이터 과학자와 기계 학습 엔지니어는 교육 컴퓨팅을 최적화하는 가장 좋은 방법을 끊임없이 찾고 있지만 전체 클러스터 크기와 함께 증가할 수 있는 통신 오버헤드로 어려움을 겪고 있습니다.

이것이 우리가 최근에 출시한 이유입니다. s하드 데이터 병렬 처리 on 아마존 세이지 메이커, 새로운 메모리 절약 분산 교육 기법 SageMaker 모델 병렬(SMP) 라이브러리. 샤딩된 데이터 병렬 처리는 대규모 모델을 위해 특별히 제작되었으며 Amazon 내부에서 사용합니다. 마이크 내부 기술, 매개변수 수집 및 기울기 동기화에 기반을 둔 값비싼 통신 오버헤드를 줄여 통신 규모를 최소화하려는 과학적인 노력입니다. 시퀀스 길이가 30인 2B 매개변수 GPT-2048 모델에서 이 새로운 기능은 DeepSpeed ​​ZeRO-141에 비해 39.7% 향상된 3 TFLOP를 달성했습니다. 시퀀스 길이가 10인 2B GPT-512 모델의 경우 이 새로운 기능은 PyTorch의 FSDP(Fully Sharded Data Parallel)에 비해 564% 향상된 초당 13.9개의 샘플을 달성했습니다. 거대한 모델 교육에서 속도 향상의 모든 비율은 팀의 비용 절감과 생산성 향상으로 이어집니다.

이 블로그 게시물에서는 먼저 샤딩된 데이터 병렬 처리의 주요 차별화 요소와 사용 시기에 대해 자세히 살펴보겠습니다. 그런 다음 이 새로운 기능을 사용하여 SageMaker에서 30B 매개변수 GPT-2 모델을 쉽게 훈련하는 방법을 배우게 됩니다. 마지막으로 39.7 GPU에서 DeepSpeed ​​ZeRO를 최대 256% 능가하는 다른 오픈 소스 옵션과 성능을 비교할 것입니다.

샤딩된 데이터 병렬 처리 방식 및 사용 시기

분할 데이터 병렬 처리를 소개하기 전에 더 광범위한 기술 제품군을 살펴보겠습니다. 대형 모델에 대한 최근의 분산 교육 접근 방식은 모델 매개변수, 기울기 및 최적화기 상태가 데이터 병렬 노드에서 공유되는 패러다임으로 이동했습니다. 특히 프레임워크가 지원하지 않는 경우 장치 간에 분할할 계층을 선택하는 고유한 복잡성이 있는 파이프라인 병렬 처리와 달리 자동화된 모델 분할, 이 패러다임은 데이터 병렬화의 단순성을 우아하게 유지하면서 모델이 단일 GPU에 맞아야 하는 데이터 병렬화의 제약을 제거합니다.

이 패러다임에 해당하는 기존 프레임워크, 특히 DeepSpeed ​​ZeRO-3 및 FairScale에서 업스트림된 PyTorch의 FSDP에서 모델 상태는 모든 GPU는 클러스터 크기에 따라 증가하는 큰 통신 오버헤드를 발생시키는 대가로 각 GPU의 메모리 소비를 낮추고 그에 따라 확장성을 규모에 따라 크게 떨어뜨리는 전략입니다. 대조적으로, SMP 라이브러리 파티션 모델의 샤딩된 데이터 병렬 처리는 규모 인식 내에서만 모델 상태의 각 복제본을 분할하여 하위 집합 GPU의.

자세히 살펴보자 스케일 인식 모델 파티셔닝 샤딩된 데이터 병렬의 핵심 기술인 MiCS에서. 이 설계 이면의 직관은 전체 데이터 병렬 그룹에 걸쳐 훈련 상태를 분할하는 것이 수백억 개의 매개변수가 있는 모델을 훈련하는 데 필요하지 않을 수 있다는 것입니다. 예를 들어, 8개의 V100 GPU(각각 32GB)는 혼합 정밀도를 사용하여 Adam 옵티마이저로 훈련할 때 약 10GB의 메모리가 필요한 200B 매개변수 모델의 모델 상태 복제본을 보유하기에 충분합니다. 모델 상태의 완전한 복제본을 제한함으로써 가장 작은 GPU의 하위 집합을 사용하면 DeepSpeed ​​및 PyTorch FSDP에 비해 통신 오버헤드 규모를 효과적으로 줄일 수 있습니다. 샤딩된 데이터 병렬은 또한 계층적 통신 및 2-홉 경사 동기화와 같은 MiCS의 다른 기술을 활용합니다. 자세한 내용은 다음을 확인하세요. AWS에서 거대한 모델 교육의 선형에 가까운 확장 or MiCS: 퍼블릭 클라우드에서 거대한 모델 훈련을 위한 선형에 가까운 확장.

이제 다른 분산 교육 기술보다 병렬로 분할된 데이터를 선택해야 하는 시점을 어떻게 알 수 있습니까? 일반적인 규칙은 모델에 1억 개 미만의 매개변수가 있고 GPU 메모리에 들어갈 수 있는 경우 SageMaker 데이터 병렬 라이브러리 or SageMaker 교육 컴파일러 당신에게 충분할 수 있습니다. 더 큰 언어 또는 컴퓨터 비전 모델이 있는 경우 다음과 결합된 분할된 데이터 병렬 처리 기술을 사용하여 이를 훈련하는 것이 좋습니다. 활성화 체크포인트 과 활성화 오프로딩 SageMaker 모델 병렬 라이브러리에서 먼저 다음과 같은 다른 기술보다 텐서 병렬 처리 또는 파이프라인 병렬 처리.

분할 데이터 병렬 처리를 사용하여 Amazon SageMaker에서 GPT-2 교육

이제 복잡성을 캡슐화하는 SMP를 사용하여 분할된 데이터를 병렬로 사용하여 GPT-2 모델을 훈련하는 방법을 알아보겠습니다. 이것 완전한 튜토리얼 노트북 데이터 처리, 교육 작업 정의 및 제출, 교육 로그 모니터링에 이르기까지 전체 프로세스를 안내합니다. 다음은 이 기능을 사용하기 위한 주요 단계를 강조한 간략한 개요입니다.

1. 시작하기

분할 데이터 병렬 처리는 PyTorch v1.12.0+에서 사용할 수 있으며 FP16 및 BF16에서 모두 작동합니다. SMP 라이브러리를 사용하는 가장 쉬운 방법은 사전 구축된 PyTorch용 AWS Deep Learning Container를 사용하는 것입니다. 그러나 자신의 Docker 컨테이너를 가져오려는 경우 다음을 참조할 수 있습니다. SageMaker 분산 모델 병렬 라이브러리를 사용하여 고유한 Docker 컨테이너를 생성합니다. 시작하려면 다음을 따르십시오. PyTorch 교육 스크립트 수정 훈련 스크립트에서 SMP의 API를 조정합니다. 이 섹션에서는 바로 사용할 수 있는 교육 스크립트의 코드 조각으로 몇 가지 주요 단계만 설명합니다. train_gpt_simple.py. 스크립트의 주석을 따를 수 있으며 API 문서 SMP API가 사용되는 위치에 대해 자세히 알아보세요.

먼저 다음을 호출하여 라이브러리를 가져오고 초기화합니다. smdistributed.modelparallel.torch.init() 훈련 스크립트 시작 부분:

import smdistributed.modelparallel.torch as smp

smp.init(smp_config)

둘째, 분할할 모델을 래핑합니다. smdistributed.modelparallel.torch.DistributedModel 반환 된 사용 DistributedModel 앞으로 개체:

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_config(model_config)
model = smp.DistributedModel(model, trace_device="gpu", backward_passes_per_step=args.gradient_accumulation)

옵티마이저를 다음으로 래핑합니다. smdistributed.modelparallel.torch.DistributedOptimizer 옵티마이저 상태를 저장하고 로드합니다.

from torch import optim

optimizer = optim.Adam(
    param_groups, betas=(args.beta1, args.beta2), lr=args.lr, weight_decay=args.weight_decay
)

optimizer = smp.DistributedOptimizer(
        optimizer, 
        static_loss_scale=None, 
        dynamic_loss_scale=True,
        dynamic_loss_args={"scale_window": 1000, "min_scale": 1, "delayed_shift": 2},
        )

전진 및 후진 논리를 단계 함수에 넣고 다음으로 장식하십시오. smdistributed.modelparallel.torch.step.  내부에 정의된 모든 계산 smp.step-decorated 기능은 분산 방식으로 실행됩니다.

@smp.step
def train_step(model, optimizer, input_ids, attention_mask, args):
    loss = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)["loss"]
    model.backward(loss)

    return loss

@smp.step
def test_step(model, input_ids, attention_mask):
    loss = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)["loss"]
    
    return loss

2. 데이터세트 준비

우리는을 사용하여 오픈웹텍스트 이 예제에서 사용하는 데이터 세트입니다. 노트북은 스크립트를 사용합니다 data_prep_512.py 데이터 세트를 다운로드하고 사전 처리합니다. 다음을 수정하여 다른 데이터세트로 훈련할 수도 있습니다. data_pipeline.py. 대규모 데이터세트 및 모델을 처리할 때 저장된 데이터를 사용하여 교육 작업의 속도를 높일 수 있습니다. Luster 용 Amazon FSx, 기본적으로 통합된 고성능 파일 시스템을 제공합니다. 아마존 단순 스토리지 서비스 (S3). 의 지침을 참조하십시오. Amazon FSx for Lustre를 사용하도록 데이터 입력 채널 구성 FSx Lustre 파일 시스템을 데이터 입력 채널로 설정하는 방법에 대한 지침.

3. 교육 작업 시작

이 단계에서는 이미 훈련 스크립트를 수정했습니다 이전 섹션에서 언급한 대로 데이터 세트를 준비했습니다. 에게 분할된 데이터 병렬 처리를 활성화하려면 간단히 설정하십시오. sharded_data_parallel_degree FBI 증오 범죄 보고서 파이토치 추정기. 이 튜토리얼에서는 다음을 설정합니다. sharded_data_parallel_degree=128 과 instace_count=32 p4d.24xlarge 노드의 경우 모델 상태가 총 128 GPU 중 256 GPU에 걸쳐 분할됨을 나타냅니다. 이 선택된 값을 기반으로 SMP는 자동으로 데이터 병렬 정도를 2로 설정합니다(256/128=2이기 때문에). 즉, 데이터 병렬 처리를 위한 두 개의 복제본이 있습니다. 이상적인 값을 선택하는 일반적인 규칙 sharded_data_parallel_degree 모델 매개변수의 3B마다 공유 그룹에 노드를 하나 더 추가하는 것입니다. 이 튜토리얼에서는 모델 크기가 30B이므로 샤딩에 최소 10개의 노드를 사용해야 합니다. 그리고 16개의 노드(128개의 GPU)는 임계값 위의 가장 작은 2의 거듭제곱이므로 다음을 설정합니다. sharded_data_parallel_degree=128.

체크포인트를 위해 우리는 또한 체크포인트 유틸리티 세트를 제공합니다. sharded_data_parallel_checkpoint.py , 전체를 재구성하는 유틸리티 포함 state_dict 고급 사용 사례의 경우. 마지막으로 Estimator에서 fit()을 호출하여 분산 교육 작업을 시작할 수 있습니다.

smp_estimator = PyTorch(
    entry_point="train_gpt_simple.py",
    instance_type="ml.p4d.24xlarge",
    source_dir=os.getcwd(),
    volume_size=500,
    instance_count=32,
    distribution={
        "mpi": {
            "enabled": True,
            "processes_per_host": processes_per_host,
            "custom_mpi_options": mpioptions,
        },
        "smdistributed": {
            "modelparallel": {
                "enabled": True,
                "parameters": {
                    "ddp": True,
                    "skip_tracing": True,
                    "delayed_parameter_initialization": True,
                    "offload_activations": True,
                    "activation_loading_horizon": 4,
                    # To enable sharded data parallelism.
                    # Here we shard model states across 128 GPUs. 
                    "sharded_data_parallel_degree": 128, 
                    "fp16": False,
                    "bf16": True,
                    # This is to disable pipeline parallelism.
                    "partitions": 1,
                },
            }
        },
    },
    framework_version="1.12",
    py_version="py38",
    hyperparameters=hyperparameters,
    checkpoint_s3_uri=checkpoint_s3_uri if not use_fsx else None,
    checkpoint_local_path=hyperparameters["checkpoint-dir"] if use_fsx else None,
    ...
)

smp_estimator.fit(inputs=data_channels)

4. 교육 작업 모니터링

교육 로그에 액세스하고 GPU 및 메모리 사용률을 추적할 수 있습니다. 아마존 클라우드 워치. "algo-1"의 로그를 확인하십시오. 출력 스트림에 모든 인스턴스의 교육 작업 로그가 있는 기본 노드이기 때문입니다.

벤치마킹 성능

우리는 각각 시퀀스 길이 16 및 32에 대해 4 및 24 p512d.2048xlarge 노드에서 SMP 라이브러리의 분할 데이터 병렬 처리를 벤치마킹했습니다. 30B-파라미터 GPT2 모델은 7168, 48 레이어 및 64 헤드의 숨겨진 너비를 사용하도록 구성됩니다. 다음을 설정하여 시퀀스 길이가 2048인 동일한 구성을 채택할 수 있습니다. model_config = "gpt2-30b" 튜토리얼 노트북에서. 이 설정으로 SMP는 초당 73.52개의 샘플을 달성했으며 DeepSpeed ​​ZeRO-39.7에 비해 3% 빨라졌습니다. 토큰 크기가 500억이면 이 속도 향상은 p367d.4xlarge 노드에서 거의 24시간의 비용 절감을 의미하며, 이는 교육당 $12,000 이상의 예산 절감에 해당합니다! 다음 표에는 벤치마크 결과가 요약되어 있습니다.

1/ 각 모델 구성에 대해 DeepSpeed ​​ZeRO의 다양한 기능, 단계 및 구성을 테스트하고 DeepSpeed ​​기준으로 최고의 처리량을 제공하는 것을 선택했습니다. 벤치마크는 다음에서 실행되었습니다. 아마존 엘라스틱 컴퓨트 클라우드 (아마존 EC2). 2/ 이러한 결과는 곧 제공될 AWS에 최적화된 개선된 통신 집합체에 의존합니다. 3/ 훈련 시간은 처리된 토큰 수를 기반으로 속도에서 예상됩니다.

요약하면, 우리는 다양한 모델 및 구성에서 DeepSpeed와 비교할 때 SMP에서 샤딩된 데이터 병렬 처리로 일관되게 더 높은 처리량을 관찰했습니다. 이 새로운 기능은 또한 DeepSpeed에 비해 더 나은 메모리 효율성을 보여주므로 SMP가 더 큰 배치 크기에 적합하고 특정 전역 배치 크기에 맞추는 데 필요한 그래디언트 누적 수준을 줄일 수 있습니다.

결론

이 게시물에서는 분할된 데이터 병렬 처리라는 새로운 분산 교육 기술과 Amazon SageMaker에서 선형에 가까운 확장으로 거대한 모델 교육의 속도를 높이는 방법을 소개했습니다. 우리는 또한 다음과 같은 새로운 기술로 GPT-2 모델을 훈련시키는 방법을 살펴보았습니다. 완전한 예. 당신은 따라갈 수 있습니다 Amazon SageMaker 예제 GitHub 리포지토리 모든 SageMaker 모델 병렬 예제를 추적하거나 다음 회의에 참석하려면 분산 교육 워크샵. 샤딩된 데이터 병렬 처리에 대해 자세히 알아보려면 다음을 참조하십시오. 선적 서류 비치.

저자 소개

에밀리 웹버 SageMaker가 출시된 직후 AWS에 합류했으며 그 이후로 전 세계에 알리기 위해 노력해 왔습니다! Emily는 고객을 위한 새로운 ML 경험을 구축하는 것 외에도 티베트 불교 명상과 공부를 즐깁니다.

캔 카라쿠스 AWS에서 대규모 분산 딥 러닝을 최적화하는 AWS의 수석 응용 과학자입니다. 그의 연구 관심 분야는 딥 러닝, 분산 최적화, 분산 시스템 및 정보 이론입니다. 일 외에는 자전거 타기, 여행, 독서 및 학습을 즐깁니다.

라훌 후일골 AWS의 수석 소프트웨어 엔지니어입니다. 그는 클라우드에서 대규모 딥 러닝 모델을 쉽고 효과적으로 훈련할 수 있도록 분산 딥 러닝 시스템에서 일하고 있습니다. 여가 시간에는 사진, 자전거 타기, 정원 가꾸기를 즐깁니다.

수힛 코굴 딥 러닝 프레임워크를 작업하는 AWS 인공 지능 그룹의 소프트웨어 개발 엔지니어입니다. 여가 시간에는 하이킹, 여행, 요리를 즐깁니다.

에린 호 AWS Deep Learning의 제품 관리자입니다. 그녀는 고객이 AWS에서 딥 러닝 모델을 더 쉽게 교육할 수 있는 제품을 개발하고 있습니다. 그녀는 즐거운 야외 활동을 위해 하이킹과 스키를 즐깁니다.