AWS SageMaker HyperPod (w. EKS) 도입기

도입 배경

딥러닝 모델을 학습하는 과정을 온프레미스 환경에서 수행하고 있었다. (온프레미스 환경의 GPU 워크로드는 쿠버네티스로 연동되어 있는 상태)
모델의 크기 및 데이터 볼륨이 증가함에 따라 FSDP, DDP와 같은 분산 학습 기법을 활용했음.
멀티 노드 분산 학습과정을 위해 Kubeflow의 Training Operator를 활용했지만, 인프라 관리 비용과 네트워크 병목이 발생했다.
- Kustomize, Istio와 같은 스택들을 관리하는 데에서 오는 부담감
- InfiniBand가 아닌 일반 이더넷 기반에서 NCCL 백엔드의 TCP 소켓 통신에 대한 대역폭 병목
적합한 서비스를 찾던 와중, Slurm과 같은 HPC 플랫폼과 Kubernetes를 오케스트레이터로 지원하는 SageMaker의 HyperPod를 발견했다.

Overview

우선 HyperPod를 이해하기 이전에 AWS에서 정의한 생성형 인공지능 스택에 대해 3단계로 나눌 수 있음을 인지하자.
- 1st Layer : 인프라 영역으로 / SageMaker AI, 훈련 및 추론용 인스턴스 유형(Trainium, Inferencetia)가 있다.
- 2st Layer : 애플리케이션 영역 / 작년 2024 Summit부터 AWS가 엄청 밀고 있는 Bedrock이 있다.
  - Knowledge Base나 Flow와 같은 기능들을 통해 벡터 데이터베이스와의 통합 및 LLM 워크플로우를 쉽게 정의할 수 있다. (개인적으로 Flow의 버전에 대한 CI/CD 기능이 나오면 좋을 것 같기도)
- 3rd Layer : 프로덕트 영역 / B2B, B2C 버전으로 추상화된 제품인 Amazon Q 세트가 있다.
  - Q Developer의 경우 개발자들이 Amazon SageMaker Unified Studio(ML 라이프사이클을 통합된 인터페이스에서 수행하는 서비스)에서 도움을 받을 수 있는 에이전트의 형태로도 기능이 제공된다.

백문이 불여일견, 위 아키텍처는 위에서 설명한 AWS 생성형 인공지능 스택의 1 Layer, HyperPod 클러스터를 도식화한 것이다.
기본적으로 HyperPod는 Slurm과 같은 HPC 플랫폼과 통합되지만, 2024년 12월 9일부로 EKS도 지원하기 시작했다. (https://aws.amazon.com/ko/blogs/korea/amazon-sagemaker-hyperpod-introduces-amazon-eks-support/)
HyperPod 클러스터를 제어하기 위해서는 오케스트레이터(Orchestrator)가 필요하다.
- 오케스트레이터의 옵션으로는 Slurm과 EKS가 있다. EKS가 익숙하다면 훨씬 유리할 것이다.
- EKS 자체는 사용자의 관리형 VPC에서 동작하지만, HyperPod 클러스터는 AWS에서 자체적으로 관리하는 VPC에 프로비저닝된다.
- EKS에서는 노드 그룹(ASG)단위로 워커 노드가 동작하지만, HyperPod 클러스터에서는 인스턴스 그룹을 생성하여 EKS 클러스터에 조인하는 방식으로 동작한다.
- 아무래도 머신러닝을 목적으로 하다보니 FSx for Lustre를 PVC로 Dynamic Provisioning을 하게 되지만, 데이터가 적다면 S3 PV를 Static Provisioning하여 사용해도 괜찮은 것 같다.
HyperPod의 특성
- Deep Health Check : 클러스터의 상태가 변경(creation, update, node replacement)될 때 GPU와 AWS Trainium instance에 대한 스트레스 테스트 기능 (EFA 동작 여부 확인)
- Automated Node Recovery : HyperPod 모니터링 에이전트가 지속적으로 노드의 상태를 모니터링 하고, 이슈 발생시 대체하거나 재부팅
- Job Auto Resume : 제출되는 PytorchJob(학습 요청 CR)은 Kubeflow Training Operator에 의해 복구 가능하고 실패하거나 인터럽션이 발생하더라도 지속적인 수행이 가능하다.
  - 참고) 제출된 Job들은 Queue에 적재되며, 사용가능 할때 폴링하여 수행되는 듯 하다. 이때, 타팀 컴퓨팅 리소스가 부족하면 우리팀의 Queue에서 선점형 스케줄링 방식으로 컴퓨팅 리소스를 사용할 수 있다. 반대로 타팀에서 수행중이더라도 우리팀의 Priority가 더 높으면 선점이 가능하기도 하다.

도입 과정

처음에는 AWS Data & AI Roadshow 2025에 참석해서 HyperPod의 컨셉을 듣고 근거 없는 자신감이 생겨 레포에서 찾은 매뉴얼 스텝을 보고 구축을 시작했다.

Terraform으로 EKS 클러스터 프로비저닝
HyperPod 관련 Helm Chart를 설치
1. Affinity 때문인지 training-operator가 GPU 인스턴스가 아닌 노드에서는 스케줄링되지 않았다.
Life Cycle 스크립트 S3 버킷에 업로드
1. 이 과정이 가장 헷갈렸는데, EKS의 경우에는 별도로 설정할 필요없이 AWS에서 기본제공되는 on_create.sh만 아무 버킷에 주입하면 된다.
2. Life Cycle 스크립트는 인스턴스 그룹에서 프로비저닝되는 인스턴스의 환경을 구성하는 역할을 한다.

위 과정대로만 하면 쉽게 될줄 알았는데, HyperPod에서 EKS로 연동하는 과정에서 계속 실패해서 한동안 골머리를 앓았다.. 결국 AWS 본사에서 다시 HyperPod 전담팀과 미팅을 가지면서 워크숍에서 제공되는 방법으로 쉽게 구축할 수 있다고 답변을 주셨다.
살펴보니 AWS CLI 설치부터 네트워크 설정, CloudFormation 스택으로 프로비저닝까지 전 과정이 한번에 되는 헬퍼 스크립트가 있었다. (이 좋은 것이.. -_-)
도입 후에는 g5 인스턴스 그룹에서 분산학습을 수행할 때 큰 문제없이 동작했다. 모델은 간단하게 MLP를 토이 모델로 하여 DDP 코드를 작성한 도커 이미지를 ECR에 올렸다.
이제 HyperPod 클러스터에 Job을 제출할 차례인데, 수동으로 PytorchJob CRD 리소스를 프로비저닝하는 방법도 있지만, hyperpod-cli를 설치해서 클러스터를 제어하는 방법도 있다.
- 필자의 경우에는 hyperpod 명령어를 쓰는 것 조차 귀찮아서 아예 배시스크립트로 만들었다. (전담팀분들은 alias에서 명령어들을 한글자 단위로 정의해서 사용하시기도 함)
- 이슈) 헬퍼 스크립트로 프로비저닝된 상태에서 hyperpod start-job 커맨드를 실행시킬 때, --scheduler-type 옵션을 Kueue로 수정했다. (validate_scheduler_related_fields 함수를 살펴봤다.)
  - scheduler-type 옵션 기본값이 "SageMaker"인데, 이 경우 sagemaker에서 관리되는 유효한 네임스페이스가 있어야 한다. 유효한 네임스페이스란 (sagemaker.amazonaws.com/sagemaker-managed-queue=true 레이블이 있는 네임스페이스)
  - priority 옵션이 주어졌을 때, 유효한 workloadpriorityclasses CR이 없어서 생기는 문제로 보인다.
  - 이러한 레이블링이나 CR을 자체적으로 만드는 건 아닌것 같아서 임시방편으로 스케줄러 유형을 Kueue로 수정했다. (헬퍼 스크립트에서 뭔가 누락했거나 아직 인지하지 못하고 넘어간 부분이 있는 것 같다..)

echo "Connecting to Hyperpod cluster $HYPERPOD_CLUSTER_NAME in region $REGION and namespace $NAMESPACE ..."
hyperpod connect-cluster --cluster-name $HYPERPOD_CLUSTER_NAME \
--region $REGION \
--namespace $NAMESPACE

echo "Creating Hyperpod job $JOB_NAME with image $IMAGE_SHA ..."
hyperpod start-job --job-name $JOB_NAME \
--job-kind kubeflow/PyTorchJob \
--scheduler-type Kueue \
--image $IMAGE_SHA \
--entry-script "$ENTRY_SCRIPT" \
--pull-policy IfNotPresent \
--instance-type $INSTANCE_TYPE \
--node-count $NODE_COUNT \
--tasks-per-node $TASKS_PER_NODE \
--results-dir $RESULTS_DIR \
--persistent-volume-claims "$PERSISTENT_VOLUME_CLAIM:$MOUNT_PATH_IN_CONTAINER" \

최종적으로 위와 같은 배시 스크립트를 작성함으로써 문제없이 S3 버킷에 학습된 체크포인트가 저장되는 것을 확인했다.
만약 p5 계열과 같은 인스턴스를 프로비저닝 하려면 아래와 같은 사항들을 고려해야 한다.
- 예시) AWS 콘솔에서 Service Quotas - Amazon SageMaker - ml.p5.48xlarge를 검색한 뒤에, 필요한 quotas들을 상향 요청해야 한다. (ml.p5.48xlarge for cluster usage)
- 훈련 플랜(Training Plan)을 구매하는 것을 권장한다. 온디맨드로 p5 인스턴스의 경쟁이 치열해서 쉽게 스케줄링도 안되고(배시 스크립트로 while문 돌려도 안됨), 훈련 플랜으로 예약시 온디맨드 대비 68% 비용 절감이 가능하다. (24시간 기준 $850)
- 훈련 플랜에서 p5 인스턴스 예약시, us-east-1 리전 기준 use1-az5 AZ에 주로 스케줄링 되는 것 같다. 고로, HyperPod 클러스터 구축 과정(헬퍼 스크립트)에서 AZ ID를 입력해야 하는 란이 있는데 기본값이 use1-az1이므로, use1-az5로 입력해야 한다. (높은 계열의 인스턴스 사용시, 훈련 플랜부터 먼저 구매해놓고 HyperPod 클러스터 구축이 유리해보임)

(2025.04.22 추가)

멀티 노드 분산학습시, EFA(Elastic Facbric Adapter) 활성화

싱글 노드만 학습할 때는, 메인보드의 PCI Express를 통해 프로세스간 통신이 이루어져 인지하지 못했는데, 훈련 플랜으로 p5.48xlarge 2EA를 구매 후, 분산 학습을 돌리니 훨씬 느려졌다.
이는 ENI(Elastic Network Interface, 일반 이더넷)를 통해 NCCL 백엔드의 소켓 통신이 이루어져 느린 것이다. 즉, Infiniband 활성화가 안되었기 때문..
EFA를 활성화하기 위해 다음 링크(https://docs.aws.amazon.com/ko_kr/sagemaker/latest/dg/your-algorithms-training-efa.html)를 참조했다. 하지만, 예시에 있는 내용 그대로 수행하면 안되니 아래와 같이 수정하자.

(base)  ✘ shawn 🧁   ~  k exec -it hyperpod-cli-test-eeldi-worker-1 -n default /bin/bash                                
kubectl exec [POD] [COMMAND] is DEPRECATED and will be removed in a future version. Use kubectl exec [POD] -- [COMMAND] instead.
root@hyperpod-cli-test-eeldi-worker-1:/app# python
Python 3.11.9 | packaged by conda-forge | (main, Apr 19 2024, 18:36:13) [GCC 12.3.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> torch.cuda.nccl.version()
(2, 20, 5)

우선 위 커맨드 처럼 분산학습을 수행할 컨테이너에서 NCCL 백엔드 버전을 확인한다. (이미 알고 있다면 상관없다.)

FROM pytorch/pytorch:2.4.1-cuda12.1-cudnn9-devel

ARG DEBIAN_FRONTEND=noninteractive
ENV NVIDIA_VISIBLE_DEVICES=all
ENV NVIDIA_DRIVER_CAPABILITIES=compute,utility

RUN apt-get update && apt-get upgrade -y && apt-get install curl -y
RUN apt-get install -y --no-install-recommends gcc \
  git \
  build-essential
RUN pip install --upgrade pip && pip install poetry

ARG OPEN_MPI_PATH=/opt/amazon/openmpi
ENV NCCL_VERSION=2.20.5
ENV EFA_VERSION=1.39.0
ENV BRANCH_OFI=1.13.x-aws
ENV FI_PROVIDER=efa
ENV FI_EFA_USE_DEVICE_RDMA=1
ENV NCCL_PROTO=simple

#################################################
## EFA and MPI SETUP
RUN cd $HOME \
  && curl -O https://s3-us-west-2.amazonaws.com/aws-efa-installer/aws-efa-installer-${EFA_VERSION}.tar.gz \
  && tar -xf aws-efa-installer-${EFA_VERSION}.tar.gz \
  && cd aws-efa-installer \
  && ./efa_installer.sh -y --skip-kmod -g \

ENV PATH="$OPEN_MPI_PATH/bin:$PATH"
ENV LD_LIBRARY_PATH="$OPEN_MPI_PATH/lib/:$LD_LIBRARY_PATH"

#################################################
## NCCL, OFI, NCCL-TEST SETUP
RUN cd $HOME \
  && git clone https://github.com/NVIDIA/nccl.git \
  && cd nccl \
  && git switch -c v${NCCL_VERSION}-1 \
  && make -j64 src.build CUDA_HOME=/usr/local/cuda BUILDDIR=/usr/local

RUN apt-get update && apt-get install -y autoconf libtool
RUN apt-get install -y libhwloc-dev 

RUN cd $HOME \
  && git clone https://github.com/aws/aws-ofi-nccl.git \
  && cd aws-ofi-nccl \
  && git switch -c v${BRANCH_OFI} \
  && ./autogen.sh \
  && ./configure --with-libfabric=/opt/amazon/efa \
       --with-mpi=/opt/amazon/openmpi \
       --with-cuda=/usr/local/cuda \
       --with-nccl=/usr/local --prefix=/usr/local \
  && make && make install
  
RUN cd $HOME \
  && git clone https://github.com/NVIDIA/nccl-tests \
  && cd nccl-tests \
  && make MPI=1 MPI_HOME=/opt/amazon/openmpi CUDA_HOME=/usr/local/cuda NCCL_HOME=/usr/local

WORKDIR /app

COPY pyproject.toml /app/
RUN poetry config virtualenvs.create false && poetry install --only main --no-root

COPY ./ /app/

CMD [ "/bin/sh" ]

베이스 이미지는 -runtime이 아닌 -devel로 정해주어야 EFA 관련 빌드할 때 필요한 도구들이 포함되는 듯 하다.
- 추가로, 밑에서 git 커맨드와 make 커맨드를 사용하기 때문에 apt-get을 통해 git, build-essential 의존성이 요구된다.
NCCL_VERSION 환경변수에 위에서 확인한 NCCL 백엔드 버전을 주입한다.
- NCCL 버전과 호환되는 EFS_VERSION과 BRANCH_OFI를 주입한다. BRANCH_OFI의 경우에는 해당 레포지토리(https://github.com/aws/aws-ofi-nccl)에서 유효한 브랜치로 수정해야 한다.
공식 문서에는 git clone과 동시에 -b 옵션으로 브랜치를 명시하는데, 이 과정에서 detach HEAD가 발생해서 직접 클론 후에 브랜치 전환을 하도록 수정했다.
NCCL 빌드 과정에서 CUDA_HOME을 명시적으로 수정해주었다.

2025-04-22 09:03:59,210 | Trainer 0 | WARNING | step:1800 | Gradient contains NaN/inf. Skipping update...
2025-04-22 09:03:59,213 | Trainer 0 | INFO | Train epoch: 1 | step: 1800/200000, elapsed=18.0898s, tokens/s=31897, train loss=2.2058, train ppl=9.08, lr=0.000030, grad norm=4.3306
2025-04-22 09:04:16,888 | Trainer 0 | INFO | Train epoch: 1 | step: 1900/200000, elapsed=17.6076s, tokens/s=30241, train loss=2.2167, train ppl=9.18, lr=0.000032, grad norm=4.4011
2025-04-22 09:04:34,686 | Trainer 0 | INFO | Train epoch: 1 | step: 2000/200000, elapsed=17.7321s, tokens/s=31871, train loss=2.1512, train ppl=8.59, lr=0.000033, grad norm=4.3355
2025-04-22 09:04:39,953 | Trainer 0 | INFO | Eval epoch: 1 | step: 2000, elapsed=5.2646s, tokens/s=303844, val loss=1.8955, val ppl=6.66, 
2025-04-22 09:04:42,619 | Trainer 0 | INFO | Saving checkpoint to /app/pv/inst_pretrain/pv-7x3wzi/20250422-0857/step_2000-tl_2.1512-vl_1.8955.pt
2025-04-22 09:04:51,937 | Trainer 0 | INFO | Saving checkpoint to /app/pv/inst_pretrain/pv-7x3wzi/20250422-0857/best-step_2000-tl_2.1512-vl_1.8955.pt

싱글 노드에서 DDP 수행시, Step 당 14초
멀티 노드에서 EFA 미적용시, Step 당 270~280초로 느렸지만, 적용하고나서 FSDP는 17~18초, DDP는 10~11초까지 단축
- 위 로그가 Llama 3 1B 모델에 대한 EFA + FSDP
- 온프레미스 A100 4EA 싱글노드로 학습할 때보다 대략 2~3배 NLL 수렴 속도가 빨라졌다.

결론

인프라 레이어를 추상화해주고, 분산 학습을 Job 단위로 동작할 수 있는 점은 매우 편하다고 생각한다.
- Kubeflow를 운영해보면 알겠지만, Istio, Kustomize를 처음 써보는 입장으로써 상당히 번거롭다.
- 필자같은 경우에는 Training Operator를 GUI에서 활용할 수 있도록 직접 풀스택으로 코딩하고 Virtual Service로 엮어서 Kubeflow내에서 커스텀 메뉴를 만드는 엔지니어링 공수가 들었다. (...)

아쉬운 점은 아직 한국 리전(ap-northeast-2)는 지원하지 않고, 그에 따라 번역 문제가 많다. 예를 들면,,
- HyperPod 클러스터를 삭제하고 싶어서 "확인"을 입력해도 삭제 버튼이 활성화되지 않았다. (confirm 입력해야함..)
- 인스턴스 그룹 구성 메뉴에서 수명 주기 스크립트 입력 UI가 상당히 번역퀄이 문제였던것 같다..
  - 왜 둘다 경로라고 하는지 모르겠다. 실제로는 "수명 주기 스크립트 파일의 S3 경로"에는 (별도의 경로 설정 없이) 버킷만 선택했고, "생성 시 수명 주기 스크립트에 대한 디렉터리 경로"에는 on_create.sh만 적었다.
- 헬퍼 스크립트를 알았기에 망정이지만, 몰랐더라면 어마무시한 삽질을 했을 것 같다... 아직 초기 서비스라 그런지 레퍼런스를 찾기가 어려운 점이 있다.

긴 글 읽어주셔서 감사드리고, 누군가에게 이 글이 도움되었으면 합니다.

도입 배경

Overview

도입 과정

결론

티스토리툴바