Amazon SageMaker HyperPod – 대규모 분산 학습을 위한 인프라 기능

오늘은 Amazon SageMaker HyperPod를 소개합니다. 이 기능은 대규모 분산 교육을 위해 특별히 구축된 인프라를 제공하여 파운데이션 모델(FM) 교육 시간을 감축하도록 지원합니다. 이제 SageMaker HyperPod를 사용하여 몇 주 또는 몇 달 동안 FM을 학습할 수 있고, SageMaker는 클러스터 상태를 능동적으로 모니터링하여 결함이 있는 노드를 교체하고 체크포인트에서 모델 학습을 재개하여 자동화된 노드 및 작업 복원력을 제공합니다.

클러스터는 SageMaker의 분산 학습 라이브러리로 사전 구성되어 제공되므로, 학습 데이터 및 모델을 모든 노드로 분할하여 병렬로 처리하고 클러스터의 컴퓨팅 및 네트워크 인프라를 완전히 활용할 수 있습니다. 추가 프레임워크, 디버깅 도구 및 최적화 라이브러리를 설치하여 교육 환경을 추가로 사용자 지정할 수 있습니다.

SageMaker HyperPod을 시작하는 방법을 보여드리겠습니다. 다음 데모에서는 SageMaker HyperPod를 생성하고 Lama 2 7B 모델을 AWS ML 교육 참조 아키텍처 GitHub 리포지토리에 공유된 예제를 사용하여 학습하는 방법을 보여줍니다.

클러스터 생성 및 관리
SageMaker HyperPod 관리자는 AWS Management Console 또는 AWS Command Line Interface(AWS CLI)를 사용하여 클러스터를 생성 및 관리할 수 있습니다. 콘솔에서 Amazon SageMaker로 이동하고, 클러스터 관리를 좌측 메뉴 내 HyperPod 클러스터 아래에서 선택한 다음, 클러스터 생성을 고릅니다.

다음 설정에서 클러스터 이름을 제공하고, 선택한 인스턴스 유형과 각 인스턴스 그룹에 할당할 인스턴스 수로 인스턴스 그룹을 구성합니다.

또한 클러스터 생성 중 각 인스턴스 그룹에서 실행되도록 하나 이상의 수명 주기 스크립트를 준비하고 Amazon Simple Storage Service(S3) 버킷에 업로드해야 합니다. 수명 주기 스크립트를 사용하여 클러스터 환경을 사용자 지정하고 필수 라이브러리 및 패키지를 설치할 수 있습니다. GitHub 리포지토리에서 SageMaker HyperPod용 예제 수명 주기 스크립트를 찾을 수 있습니다.

AWS CLI 사용
또한 AWS CLI를 사용하여 클러스터를 생성 및 관리할 수 있습니다. 데모에서는 클러스터 구성을 JSON 파일로 지정합니다. 두 개의 인스턴스 그룹을 생성하기로 선택합니다. 하나는 제가 “컨트롤러 그룹”이라고 부르는 클러스터 컨트롤러 노드용이고 다른 하나는 “worker-group”이라고 부르는 클러스터 워커 노드용입니다. 모델 교육을 수행할 워커 노드에는 Amazon EC2 Trn1 인스턴스(AWS Trainium 칩으로 구동)를 지정합니다.

// demo-cluster.json
{
   "InstanceGroups": [
        {
            "InstanceGroupName": "controller-group",
            "InstanceType": "ml.m5.xlarge",
            "InstanceCount": 1,
            "lifecycleConfig": {
                "SourceS3Uri": "s3://<your-s3-bucket>/<lifecycle-script-directory>/",
                "OnCreate": "on_create.sh"
            },
            "ExecutionRole": "arn:aws:iam::111122223333:role/my-role-for-cluster",
            "ThreadsPerCore": 1
        },
        {
            "InstanceGroupName": "worker-group",
            "InstanceType": "trn1.32xlarge",
            "InstanceCount": 4,
            "lifecycleConfig": {
                "SourceS3Uri": "s3://<your-s3-bucket>/<lifecycle-script-directory>/",
                "OnCreate": "on_create.sh"
            },
            "ExecutionRole": "arn:aws:iam::111122223333:role/my-role-for-cluster",
            "ThreadsPerCore": 1
        }
    ]
}

클러스터를 생성하려면 다음 AWS CLI 명령을 실행합니다.

aws sagemaker create-cluster \
--cluster-name antje-demo-cluster \
--instance-groups file://demo-cluster.json

생성 시 aws sagemaker describe-cluster 및 aws sagemaker list-cluster-nodes를 사용하여 클러스터 및 노드 세부 정보를 볼 수 있습니다. 컨트롤러 노드의 클러스터 ID와 인스턴스 ID를 기록해 둡니다. 클러스터에 연결하려면 해당 정보가 필요합니다.

또한 Amazon FSx for Lustre와 같은 공유 파일 시스템을 연결할 수도 있습니다. FSx for Lustre를 사용하려면 Amazon Virtual Private Cloud(VPC) 구성으로 클러스터를 설정해야 합니다. 다음 AWS CloudFormation 템플릿은 SageMaker VPC 생성 방법 및 Lustre용 FSx 배포 방법을 보여줍니다.

클러스터에 연결
클러스터 사용자는 클러스터 관리자가 프로비저닝한 클러스터에 액세스할 수 있어야 합니다. 액세스 권한이 있으면 SSH를 사용하여 클러스터에 연결하여 작업을 예약 및 실행할 수 있습니다. 사전 설치된 AWS Systems Manager용 AWS CLI 플러그인을 사용하여 클러스터의 컨트롤러 노드에 연결할 수 있습니다.

데모에서는 제어 노드의 클러스터 ID 및 인스턴스 ID를 대상으로 지정하여 다음 명령을 실행합니다.

aws ssm start-session \
--target sagemaker-cluster:ntg44z9os8pn_i-05a854e0d4358b59c \
--region us-west-2

Slurm을 사용하여 클러스터에서 작업을 스케줄링 및 실행합니다
출시 당시 SageMaker HyperPod는 워크로드 오케스트레이션을 위한 Slurm을 지원합니다. Slurm은 널리 사용되는 오픈 소스 클러스터 관리 및 작업 스케줄링 시스템입니다. 클러스터 생성의 일부로 수명 주기 스크립트를 통해 Slurm을 설치 및 설정할 수 있습니다. 예제 수명 주기 스크립트는 그 방법을 보여줍니다. 그런 다음 표준 Slurm 명령을 사용하여 작업을 예약 및 시작할 수 있습니다. 아키텍처 세부 정보 및 유용한 명령은 Slurm Quick Start 사용 설명서를 참조하세요.

이 데모에서는 Trn1 인스턴스로 Slurm에서 Lama 2 7B를 훈련하는 방법을 보여주는 AWS ML 교육 참조 아키텍처 GitHub 리포지토리 예제를 사용합니다. 클러스터는 이미 Slurm으로 설정되어 있고 FSx for Lustre 파일 시스템이 마운트되어 있습니다.

참고
Llama 2 모델은 Meta가 운영합니다. 메타 요청 액세스 페이지를 통해 액세스를 요청할 수 있습니다.

클러스터 환경 설정
SageMaker HyperPod은 Conda, venv, Docker 및 enroot와 같은 다양한 환경에서의 학습을 지원합니다. README의 지침에 따라 가상 환경 aws_neuron_venv_pytorch을 구축하고 torch_neuronx 및 neuronx-nemo-megatron 라이브러리를 Trn1 인스턴스 학습 모델용으로 설정합니다.

모델, 토크나이저, 데이터 세트 준비
지침에 따라 Llama 2 모델 및 토크나이저를 다운로드하고 모델을 Hugging Face 형식으로 변환합니다. 그런 다음 RedPajama 데이터 세트를 다운로드하여 토큰화합니다. 마지막 준비 단계로 모델 학습 속도를 높이기 위해 위해 AOT(Head-Of Time) 컴파일을 사용하여 Llama 2 모델을 사전 컴파일합니다.

클러스터에서 작업 시작
이제 sbatch 명령을 사용하여 모델 학습 작업을 시작할 준비가 되었습니다.

sbatch --nodes 4 --auto-resume=1 run.slurm ./llama_7b.sh

squeue 명령을 사용하여 작업 큐를 볼 수 있습니다. 교육 작업이 실행되면 SageMaker HyperPod 복원력 기능이 자동으로 활성화됩니다. SageMaker HyperPod는 이전 명령과 같이 자동 재개 파라미터가 설정된 경우 하드웨어 장애를 자동으로 감지하고, 필요에 따라 노드를 교체하며, 체크포인트에서 학습을 재개합니다.

다음 파일에서 모델 교육 작업의 결과를 볼 수 있습니다.

tail -f slurm-run.slurm-<JOB_ID>.out

모델 학습이 시작되었음을 나타내는 샘플 출력은 다음과 같습니다.

Epoch 0:  22%|██▏       | 4499/20101 [22:26:14<77:48:37, 17.95s/it, loss=2.43, v_num=5563, reduced_train_loss=2.470, gradient_norm=0.121, parameter_norm=1864.0, global_step=4512.0, consumed_samples=1.16e+6, iteration_time=16.40]
Epoch 0:  22%|██▏       | 4500/20101 [22:26:32<77:48:18, 17.95s/it, loss=2.43, v_num=5563, reduced_train_loss=2.470, gradient_norm=0.121, parameter_norm=1864.0, global_step=4512.0, consumed_samples=1.16e+6, iteration_time=16.40]
Epoch 0:  22%|██▏       | 4500/20101 [22:26:32<77:48:18, 17.95s/it, loss=2.44, v_num=5563, reduced_train_loss=2.450, gradient_norm=0.120, parameter_norm=1864.0, global_step=4512.0, consumed_samples=1.16e+6, iteration_time=16.50]

모델 교육 작업을 추가로 모니터링 및 프로파일링하려면 SageMaker에서 호스팅하는 TensorBoard 또는 원하는 다른 도구를 사용할 수 있습니다.

현재 사용 가능
SageMaker HyperPod는 미국 동부(오하이오), 미국 동부(버지니아 북부), 미국 서부(오레곤), 아시아 태평양(싱가포르), 아시아 태평양(시드니), 아시아 태평양(도쿄), 유럽(프랑크푸르트), 유럽(아일랜드) 및 유럽(스톡홀름) AWS 리전에서 정식 출시되었습니다.

자세히 알아보기:

요금 정보 및 지원되는 클러스터 인스턴스 유형 목록은 Amazon SageMaker HyperPod를 참조하세요.
개발자 가이드 확인하기
AWS Management Console을 방문하여 SageMaker HyperPod를 사용하여 FM 교육을 시작합니다.

— Antje

PS: 게시물 제목 아래에 이름이 하나만 표시되더라도 AWS 블로그 게시물은 항상 팀의 노력을 통해 작성됩니다. 이번 사례에서는 예제 코드로 아낌없이 도움을 주시고 대규모 모델 트레이닝 인프라, Slurm 및 SageMaker HyperPod 관리에 대한 전문 지식을 공유해 주신 Brad Doran, Justin Pirtle, Ben Snyder, Pierre-Yves Aquilanti, Keita Watanabe 및 Verdi March에게 감사의 말씀을 전합니다.

Amazon Web Services 한국 블로그

Amazon SageMaker HyperPod – 대규모 분산 학습을 위한 인프라 기능

주요 링크 모음

팔로우하기