Ray Train DLC와 EFA로 EKS에서 분산 학습 구성하기
Distributed training with Ray Train DLCs and EFA on EKS
AWS가 CUDA, NCCL, EFA 스택과 Ray를 미리 묶어 배포하는 Ray Train Deep Learning Container(DLC)를 사용해, 이미지를 직접 빌드하지 않고 EKS에서 PyTorch 분산 학습 클러스터를 띄우는 방법을 소개한다. 학습 코드와 클러스터 매니페스트만으로 두 인스턴스 간 노드 통신을 EFA로 처리할 수 있다는 점이 핵심이다. 기사는 Qwen2.5-1.5B를 PyTorch Fully Sharded Data Parallel(FSDP)으로 미세 조정하는 예제를 다룬다. 구성은 g6.12xlarge 인스턴스 2대에 NVIDIA L4 GPU 8개이며, 제어 노드(m7i.xlarge)에서 KubeRay 오퍼레이터와 Ray head가 돌고 GPU 워커 노드 2대가 같은 가용 영역에 놓인다. KubeRay는 하나의 RayCluster 매니페스트에서 head 그룹과 worker 그룹을 만들고, 각 GPU마다 워커 하나씩 총 8개의 워커를 배치한다. 사용 이미지는 public.ecr.aws/deep-learning-containers/ray:train-ml-cuda-v1.1 태그로, Ray 2.58.0, PyTorch 2.13.0, CUDA 13.0.2, NCCL 2.29.7, EFA 최신 버전과 AWS OFI NCCL 플러그인이 포함되어 있다. 인스턴스 간 집합 통신은 파드가 EFA를 요청하는 순간부터 EFA 경로로 흐른다. 기사는 같은 이미지가 KubeRay로 EKS에서 실행되든, EC2에서 직접 시작되든 동일하게 head와 worker 역할을 맡는다고 설명한다. 사용자가 추가하는 것은 인스턴스 종류, 파드당 GPU 수, NCCL이 쓸 네트워크 전송 방식, 작업에 필요한 Python 패키지 같은 설정뿐이며 모두 이미지 외부에서 지정한다. 샘플 저장소와 함께, 같은 KubeRay-on-EKS 기반을 쓰는 다중 노드 Ray Serve 예제도 링크되어 있다.
g6.12xlarge 2대와 L4 GPU 8개 구성에서 ray:train-ml-cuda-v1.1 태그와 EFA 요청만으로 인스턴스 간 NCCL 통신을 구성할 수 있어, CUDA·NCCL·aws-ofi-nccl 버전을 직접 맞추는 빌드 작업을 생략할 수 있다.
원문 출처
AWS Machine Learning Blog