제퍼넷 로고

AWS Trainium 및 Amazon EKS로 분산 교육 확장

시간

최근 딥 러닝의 발전으로 인해 GPT-3, BLOOM 및 OPT와 같은 점점 더 큰 모델이 생겨났으며 그 중 일부는 이미 100억 개의 매개변수를 초과했습니다. 더 큰 모델이 더 강력한 경향이 있지만 그러한 모델을 교육하려면 상당한 컴퓨팅 리소스가 필요합니다. FSDP 및 DeepSpeed와 같은 고급 분산 교육 라이브러리를 사용하더라도 교육 작업에는 한 번에 몇 주 또는 몇 달 동안 수백 개의 가속기 장치가 필요한 것이 일반적입니다.

2022년 후반에 AWS는 Amazon EC2 Trn1 인스턴스 에서 제공 AWS 트레이닝—클라우드에서 딥 러닝 모델을 교육하기 위해 비용 효율적이고 대규모로 확장 가능한 고성능 플랫폼을 제공하도록 최적화된 특수 목적 머신 러닝(ML) 가속기입니다. Trn1 인스턴스는 다양한 크기로 제공되며(다음 표 참조) 인스턴스당 최대 16개의 Trainium 액셀러레이터가 있습니다.

인스턴스 크기 Trainium 가속기 가속기 메모리(GB) vCPU 인스턴스 메모리(GiB) 네트워크 대역폭(Gbps)
trn1.2xlarge 1 32 8 32 최대 12.5까지
trn1.32xlarge 16 512 128 512 800
trn1n.32xlarge(출시 예정) 16 512 128 512 1600

Trn1 인스턴스는 소규모 교육 작업을 위한 독립 실행형 인스턴스로 배포하거나 수만 개의 Trainium 가속기에서 분산 교육을 지원하는 확장성이 뛰어난 울트라클러스터에 배포할 수 있습니다. 모든 Trn1 인스턴스는 독립 실행형 구성을 지원하는 반면 Trn1 울트라클러스터에는 trn1.32xlarge 또는 trn1n.32xlarge 인스턴스가 필요합니다. 울트라클러스터에서 여러 Trn1 인스턴스는 지정된 AWS 가용 영역에 함께 배치되며 집합적 컴퓨팅 작업을 위해 인스턴스당 800Gbps의 비차단 네트워크 대역폭을 제공하는 고속, 저지연, Elastic Fabric Adapter(EFA) 네트워킹에 연결됩니다. . 1년 초에 출시되는 trn32n.2023xlarge 인스턴스 유형은 이 대역폭을 인스턴스당 1600Gbps로 증가시킵니다.

많은 기업 고객은 클라우드에서 컨테이너 오케스트레이션을 위한 사실상의 표준인 Kubernetes를 사용하여 딥 러닝 워크로드를 배포하기로 선택합니다. AWS 고객은 종종 다음을 사용하여 이러한 워크로드를 배포합니다. Amazon Elastic Kubernetes 서비스 (아마존 EKS). Amazon EKS는 Kubernetes 클러스터의 생성, 구성, 수명 주기 및 모니터링을 간소화하는 동시에 업스트림 Kubernetes의 완전한 유연성을 제공하는 관리형 Kubernetes 서비스입니다.

오늘 Amazon EKS 및 EC2 Trn1 인스턴스를 사용하는 분산 교육 작업에 대한 공식 지원을 발표하게 되어 기쁩니다. 이번 발표로 이제 Trn1 인스턴스가 제공하는 가격 대비 성능, 확장성 및 사용 편의성을 최대한 활용하면서 Amazon EKS 내에서 대규모 컨테이너식 교육 작업을 쉽게 실행할 수 있습니다.

이 발표와 함께 Amazon EKS 및 Trn1 인스턴스를 사용하여 다중 인스턴스 분산 교육 작업(BERT 1단계 사전 교육)을 실행하는 데 필요한 단계를 안내하는 자세한 자습서도 게시하고 있습니다. 이 게시물에서는 솔루션 아키텍처에 대해 알아보고 자습서의 몇 가지 주요 단계를 검토합니다. 참조 공식 튜토리얼 저장소 완벽한 엔드투엔드 워크플로우를 위해.

따라가려면 다음과 같은 핵심 AWS 서비스에 대해 폭넓게 알고 있어야 합니다. 아마존 엘라스틱 컴퓨트 클라우드 (Amazon EC2) 및 Amazon EKS가 함축되어 있으며 딥 러닝 및 PyTorch에 대한 기본적인 지식이 있으면 도움이 될 것입니다.

솔루션 아키텍처

다음 다이어그램은 솔루션 아키텍처를 보여줍니다.

이 솔루션은 다음과 같은 주요 구성 요소로 구성됩니다.

  • EKS 클러스터
  • trn1.32xlarge 인스턴스로 구성된 EKS 노드 그룹
  • XNUMXD덴탈의 AWS 뉴런 SDK
  • Neuron 및 EFA용 EKS 플러그인
  • An Amazon Elastic Container Registry (Amazon ECR) 리포지토리
  • 학습 컨테이너 이미지
  • An Luster 용 Amazon FSx 파일 시스템
  • Volcano 배치 스케줄러 및 etcd 서버
  • TorchX 범용 작업 실행기
  • Trainium용 TorchX DDP 모듈

솔루션의 중심에는 EKS 서비스 엔드포인트를 통해 핵심 Kubernetes 관리 기능을 제공하는 EKS 클러스터가 있습니다. Amazon EKS의 이점 중 하나는 서비스가 부하를 기반으로 컨트롤 플레인을 능동적으로 모니터링하고 확장하여 분산 교육과 같은 대규모 워크로드에 대해 높은 성능을 보장한다는 것입니다. EKS 클러스터 내부에는 동일한 가용 영역에 상주하는 1.32개 이상의 trnXNUMXxlarge Trainium 기반 인스턴스로 구성된 노드 그룹이 있습니다.

Neuron SDK는 드라이버, 컴파일러, 런타임, 프레임워크 통합(예: PyTorch Neuron) 및 Trainium 가속기의 이점에 액세스할 수 있는 사용자 도구를 제공하는 소프트웨어 스택입니다. Neuron 장치 드라이버는 EKS 노드(Trn1 인스턴스)에서 직접 실행되며 노드에서 실행되는 교육 컨테이너 내에서 Trainium 칩에 대한 액세스를 제공합니다. Neuron 및 EFA 플러그인은 분산 교육에 필요한 Trainium 칩 및 EFA 네트워킹 장치에 대한 액세스를 제공하기 위해 EKS 클러스터 내에 설치됩니다.

교육 컨테이너 이미지를 저장하는 데 ECR 리포지토리가 사용됩니다. 이러한 이미지에는 Neuron SDK(Trn1 인스턴스에서 직접 실행되는 Neuron 드라이버 제외), PyTorch 훈련 스크립트 및 필수 종속성이 포함되어 있습니다. EKS 클러스터에서 교육 작업이 시작되면 먼저 컨테이너 이미지를 Amazon ECR에서 EKS 노드로 가져온 다음 이미지에서 PyTorch 작업자 컨테이너를 인스턴스화합니다.

공유 스토리지는 trn1.32xlarge 인스턴스와 동일한 가용 영역에 있는 고성능 FSx for Lustre 파일 시스템을 사용하여 제공됩니다. EKS 클러스터에 대한 FSx for Lustre 파일 시스템 생성 및 첨부는 다음을 통해 조정됩니다. Lustre CSI 드라이버용 Amazon FSx. 이 솔루션에서 공유 스토리지는 학습 데이터 세트와 학습 프로세스 중에 생성된 모든 로그 또는 아티팩트를 저장하는 데 사용됩니다.

솔루션은 TorchX 범용 작업 실행기 Amazon EKS 내에서 분산 교육 작업을 시작합니다. TorchX에는 Volcano 배치 스케줄러와 etcd 서버라는 두 가지 중요한 종속성이 있습니다. Volcano는 교육 작업의 예약 및 대기를 처리하는 반면 etcd 서버는 작업 시작 중 동기화 및 피어 검색을 위해 TorchElastic에서 사용하는 키-값 저장소입니다.

TorchX를 사용하여 교육 작업이 시작되면 시작 명령은 Trainium용으로 제공된 TorchX 분산 DDP 모듈을 사용하여 전체 교육 작업을 구성한 다음 각 PyTorch 작업자 포드에서 적절한 torchrun 명령을 실행합니다. 작업이 실행 중이면 표준 Kubernetes 도구(예: kubectl)를 사용하거나 TensorBoard와 같은 표준 ML 도구 세트를 통해 모니터링할 수 있습니다.

솔루션 개요

이 솔루션의 중요한 단계를 살펴보겠습니다. 이 개요 전체에서 다음을 참조합니다. TorchX 및 EKS를 사용하여 Trainium에서 다중 노드 PyTorch 뉴런 훈련 작업 시작 GitHub의 튜토리얼.

EKS 클러스터 생성

Trn1 인스턴스가 있는 Amazon EKS에서 분산 교육 작업을 시작하려면 먼저 에 설명된 대로 EKS 클러스터를 생성합니다. GitHub의 튜토리얼. 다음과 같은 표준 도구를 사용하여 클러스터를 생성할 수 있습니다. eksctlAWS 클라우드 포메이션.

EKS 노드 그룹 생성

다음으로 지원되는 리전에서 1.32개 이상의 trnXNUMXxlarge 인스턴스를 포함하는 EKS 노드 그룹을 생성해야 합니다. 튜토리얼에서 AWS CloudFormation은 Trainium 관련 EC2 시작 템플릿을 생성하는 데 사용됩니다., 분산 교육을 지원하는 데 필요한 적절한 Amazon 머신 이미지(AMI) 및 올바른 EFA 네트워크 구성으로 Trn1 인스턴스가 시작되도록 합니다. AMI에는 Trainium 가속기 칩을 지원하는 Neuron 장치 드라이버도 포함되어 있습니다. 이랑 eksctl Amazon EKS 관리 도구를 사용하면 새로 생성된 시작 템플릿을 참조하는 기본 YAML 매니페스트를 사용하여 Trainium 노드 그룹을 쉽게 생성할 수 있습니다. 예를 들어:

apiVersion: eksctl.io/v1alpha5
kind: ClusterConfig metadata: name: my-trn1-cluster region: us-west-2 version: "1.23" iam: withOIDC: true availabilityZones: ["us-west-xx","us-west-yy"] managedNodeGroups: - name: trn1-ng1 launchTemplate: id: TRN1_LAUNCH_TEMPLATE_ID minSize: 2 desiredCapacity: 2 maxSize: 2 availabilityZones: ["us-west-xx"] privateNetworking: true efaEnabled: true

이전 매니페스트에서 EKS 클러스터의 Trn1 인스턴스를 사용할 수 있도록 여러 속성이 구성되었습니다. 첫 번째, metadata.region Trn1 인스턴스를 지원하는 지역 중 하나로 설정됩니다(현재 us-east-1us-west-2). 다음으로, availabilityZones의 경우 Amazon EKS는 두 개의 가용 영역을 지정해야 합니다. 이러한 가용 영역 중 하나는 Trn1 인스턴스의 사용을 지원해야 하며 다른 하나는 임의로 선택할 수 있습니다. 자습서에서는 다음을 수행하는 방법을 보여줍니다. AWS 계정 내에서 Trn1 인스턴스를 허용할 가용 영역 결정. 동일한 Trn1 지원 가용 영역도 다음을 사용하여 지정해야 합니다. availabiltyZones EKS 노드 그룹과 연결된 속성입니다. efaEnabledtrue 분산 교육에 필요한 적절한 EFA 네트워크 구성으로 노드를 구성합니다. 마지막으로, launchTemplate.id 노드 그룹과 연결된 속성은 이전 단계에서 AWS CloudFormation을 통해 생성된 EC2 시작 템플릿을 가리킵니다.

이미 CloudFormation 템플릿을 적용하고 다음을 설치했다고 가정합니다. eksctl 관리 도구에서 다음 코드를 실행하여 Trainium 지원 EKS 노드 그룹을 생성할 수 있습니다.

> eksctl create nodegroup -f TEMPLATE.yaml

Trainium 및 EFA 장치용 Kubernetes 플러그인 설치

노드 그룹이 준비되면 다음 단계는 Trainium 가속기(Neuron 플러그인을 통해) 및 EFA 장치(EFA 플러그인을 통해)를 지원하는 Kubernetes 플러그인을 설치하는 것입니다. 이러한 플러그인은 표준을 사용하여 클러스터에 쉽게 설치할 수 있습니다. kubectl 튜토리얼에 표시된 관리 도구.

TorchX 범용 PyTorch 런처를 사용하여 분산 교육 작업을 시작하려면 Volcano 배치 스케줄러 및 etcd 서버라는 두 가지 전제 조건이 필요합니다. Neuron 및 EFA 플러그인과 마찬가지로 kubectl EKS 클러스터에 Volcano 및 etcd 서버를 설치하는 도구입니다.

EKS 클러스터에 공유 스토리지 연결

이 자습서에서는 FSx for Lustre를 사용하여 다양한 EKS 작업자 포드에서 액세스할 수 있는 고성능 공유 파일 시스템을 제공합니다. 이 공유 스토리지는 교육 데이터 세트와 교육 프로세스 중에 생성되는 모든 아티팩트 및 로그를 호스팅하는 데 사용됩니다. 이 튜토리얼에서는 다음을 사용하여 공유 스토리지를 생성하고 클러스터에 연결하는 방법을 설명합니다. Lustre CSI 드라이버용 Amazon FSx.

학습 컨테이너 이미지 만들기

다음으로 모든 종속성과 함께 PyTorch 교육 스크립트를 포함하는 교육 컨테이너 이미지를 생성해야 합니다. 예제 Dockerfile은 소프트웨어 종속성과 함께 BERT 사전 교육 스크립트를 통합하는 자습서에 포함되어 있습니다. Dockerfile은 교육 컨테이너 이미지를 빌드하는 데 사용되며 이미지는 클러스터에서 교육 작업이 시작될 때 PyTorch 작업자가 이미지를 가져올 수 있는 ECR 리포지토리로 푸시됩니다.

훈련 데이터 설정

교육 작업을 시작하기 전에 먼저 교육 데이터가 FSx for Lustre의 공유 스토리지 볼륨에 복사됩니다. 이 튜토리얼에서는 공유 스토리지 볼륨에 액세스할 수 있는 임시 Kubernetes 포드를 생성하는 방법을 설명하고 표준 Linux 셸 명령을 사용하여 훈련 데이터 세트를 다운로드하고 추출하기 위해 포드에 로그인하는 방법을 보여줍니다.

다양한 인프라와 소프트웨어 전제 조건이 갖춰졌으므로 이제 솔루션의 Trainium 측면에 집중할 수 있습니다.

모델 사전 컴파일

Neuron SDK는 통합 계층을 통해 PyTorch를 지원합니다. PyTorch 뉴런. 기본적으로 PyTorch Neuron은 JIT(Just-In-Time) 컴파일로 작동합니다. 여기서 교육 작업 내의 다양한 신경망 컴퓨팅 그래프는 교육 프로세스 중에 발생하는 대로 컴파일됩니다. 더 큰 모델의 경우 제공되는 것을 사용하는 것이 더 편리할 수 있습니다. neuron_parallel_compile 교육 시간에 그래프 컴파일을 방지하기 위해 사전에 다양한 컴퓨팅 그래프를 미리 컴파일하고 캐시하는 도구입니다. EKS 클러스터에서 교육 작업을 시작하기 전에 이 자습서에서는 다음을 사용하여 TorchX를 통해 사전 컴파일 작업을 먼저 시작하는 방법을 보여줍니다. neuron_parallel_compile 도구. 사전 컴파일 작업이 완료되면 Neuron 컴파일러는 모든 신경망 컴퓨팅 그래프를 식별 및 컴파일하고 나중에 실제 BERT 사전 교육 작업 중에 사용할 수 있도록 공유 스토리지 볼륨에 캐싱합니다.

분산 학습 작업 시작

사전 컴파일이 완료되면 TorchX는 두 개의 trn64xlarge 인스턴스에서 인스턴스당 1.32명의 작업자로 32명의 작업자 분산 교육 작업을 시작하는 데 사용됩니다. 각 trn32xlarge 인스턴스에는 1.32개의 Trainium 액셀러레이터가 포함되어 있고 각 액셀러레이터는 16개를 제공하기 때문에 인스턴스당 2개의 작업자를 사용합니다. 뉴런코어. 각 NeuronCore는 고유하게 액세스할 수 있습니다. 파이토치 XLA 장치 교육 스크립트에서. 자습서의 예제 TorchX 시작 명령은 다음 코드와 같습니다.

 torchx run -s kubernetes --workspace="file:///$PWD/docker" -cfg queue=test,image_repo=$ECR_REPO lib/trn1_dist_ddp.py:generateAppDef --name berttrain --script_args "--batch_size 16 --grad_accum_usteps 32 --data_dir /data/bert_pretrain_wikicorpus_tokenized_hdf5_seqlen128 --output_dir /data/output" --nnodes 2 --nproc_per_node 32 --image $ECR_REPO:bert_pretrain --script dp_bert_large_hf_pretrain_hdf5.py --bf16 True --cacheset bert-large

이전 TorchX 명령의 다양한 명령줄 인수는 자습서에서 자세히 설명합니다. 그러나 학습 작업을 구성할 때 다음 인수가 가장 중요합니다.

  • -cfg 대기열=테스트 – 학습 작업에 사용할 Volcano 대기열을 지정합니다.
  • -cfg 이미지_레포 – TorchX 컨테이너 이미지에 사용할 ECR 리포지토리를 지정합니다.
  • –script_args – PyTorch 훈련 스크립트에 전달되어야 하는 모든 인수를 지정합니다.
  • -nnodes 및 -nproc_per_node – 교육 작업에 사용할 인스턴스 및 인스턴스당 작업자 수
  • -스크립트 – 교육 컨테이너 내에서 시작할 PyTorch 교육 스크립트의 이름
  • -영상 – Amazon ECR의 교육 컨테이너 이미지 경로
  • -bf16 – BF16 데이터 유형 활성화 여부

학습 작업 모니터링

학습 작업이 시작된 후 다양한 방법으로 작업을 모니터링할 수 있습니다. 이 자습서에서는 다음을 사용하여 명령줄에서 기본 교육 스크립트 지표를 모니터링하는 방법을 보여줍니다. kubectl, TensorBoard에서 교육 스크립트 진행 상황을 시각적으로 모니터링하는 방법(다음 스크린샷 참조) 및 neuron-top Neuron SDK의 도구입니다.

환경 청소 또는 재사용

학습 작업이 완료되면 추가 학습 작업을 위해 클러스터를 재사용하거나 재구성할 수 있습니다. 예를 들어 EKS 노드 그룹은 eksctl 추가 Trn1 인스턴스가 필요한 학습 작업을 지원하기 위한 명령입니다. 마찬가지로 제공된 Dockerfile 및 TorchX 시작 명령을 쉽게 수정하여 추가 딥 러닝 모델을 지원하고 교육 토폴로지를 배포할 수 있습니다.

클러스터가 더 이상 필요하지 않은 경우 자습서에는 EKS 인프라 및 관련 리소스를 제거하는 데 필요한 모든 단계도 포함되어 있습니다.

결론

이 게시물에서는 Trn1 인스턴스와 Amazon EKS가 딥 러닝 모델의 고성능, 비용 효율적, 대규모 확장 가능한 분산 교육을 위한 관리형 플랫폼을 제공하는 방법을 살펴보았습니다. 또한 Trn1 인스턴스를 사용하여 Amazon EKS에서 실제 다중 인스턴스 분산 교육 작업을 실행하는 방법을 보여주는 포괄적인 자습서를 공유하고 솔루션의 몇 가지 주요 단계와 구성 요소를 강조했습니다. 이 자습서 콘텐츠는 다른 모델 및 워크로드에 맞게 쉽게 조정할 수 있으며 AWS에서 딥 러닝 모델의 분산 교육을 위한 기본 솔루션을 제공합니다.

Trainium 기반 Trn1 인스턴스를 시작하는 방법에 대한 자세한 내용은 다음을 참조하십시오. 뉴런 문서.


저자에 관하여

스콧 페리 AWS Annapurna ML 액셀러레이터 팀의 솔루션 아키텍트입니다. 캐나다에 거주하며 고객이 AWS Inferentia 및 AWS Trainium을 사용하여 딥 러닝 교육 및 추론 워크로드를 배포하고 최적화하도록 돕습니다. 그의 관심사는 대규모 언어 모델, 심층 강화 학습, IoT 및 유전체학입니다.

로레아 아리사발라가 그녀는 고객이 Amazon SageMaker를 사용하여 ML 솔루션을 설계하도록 돕는 영국 공공 부문의 솔루션 아키텍트입니다. 그녀는 또한 하드웨어 가속을 전담하는 Technical Field Community의 일원이며 AWS Inferentia 및 AWS Trainium 워크로드의 테스트 및 벤치마킹을 돕습니다.

spot_img

최신 인텔리전스

spot_img