Amazon SageMaker HyperPod를 활용한 비전 파운데이션 모델 'ZERO' 대규모 분산 학습

주요 처리 흐름

• 슈퍼브에이아이(Superb AI)는 비전 파운데이션 모델 'ZERO'의 효율적인 대규모 분산 학습을 위해 Amazon SageMaker HyperPod를 도입했습니다.
• 개발 프로세스는 '데이터 정제 및 학습 알고리즘/파라미터 개선 → 학습 → 결과 분석'의 사이클로 구성되며, SageMaker Training Plan을 활용하여 고성능 GPU 클러스터 사용을 예약합니다.
• 저렴한 GPU 인스턴스를 활용한 트레이닝 리허설을 통해 데이터 무결성 검증, 트레이닝 스크립트 버그 파악 및 FSx for Lustre에 데이터 사전 로딩을 수행합니다.
• 리허설 완료 후, 예약된 고성능 GPU 클러스터로 인스턴스를 교체하여 실제 학습을 진행합니다.

기술적 상세 및 구현 세부사항

• Amazon SageMaker HyperPod: 유연한 인스턴스 크기 조절, EFA(Elastic Fabric Adapter) 탑재, 최신 드라이버 및 OS 지원으로 인프라 구성 및 종속성 관리에 유리합니다.
• Data Preparation: 약 400만 장의 고품질 데이터를 선별 및 큐레이션했으며, Hugging Face Datasets 라이브러리와 Arrow 파일 기반 포맷을 사용하여 1G 샤드 약 1200개로 분할 후 Amazon S3에 업로드했습니다.
• FSx for Lustre 연동: 데이터셋 관리를 위해 Amazon S3 버킷을 FSx의 Data Repository로 등록하며, 'file system path'와 'data repository path' 간의 관계를 신중하게 설정합니다.
• Training Plan: 대규모 학습 기간을 확보하기 위해 사용하며, 수요가 높아 개발 일정에 맞춰 조정이 필요합니다.
• Instance Replacement: 리허설 인스턴스 제거 후, provisioning_parameters.json 파일을 편집하여 Training Plan으로 예약된 고성능 인스턴스 그룹을 추가하고 Target Instance Count를 업데이트합니다.
• Slurm Configuration: 새로 추가된 인스턴스가 Slurm에 정상 등록되도록 sinfo 및 scontrol reconfigure 명령을 사용하며, 필요 시 Slurm 파티션 설정을 수동으로 변경합니다.

문제 해결 및 비용 효율성

• 트레이닝 리허설: WandB API key 오류, 체크포인트 저장 디렉토리 소유권 문제 등 트레이닝 중단 가능성이 있는 버그를 사전에 발견하고 수정했습니다.
• 데이터 사전 로딩: 리허설 과정에서 FSx로 데이터를 미리 로드하여 실제 학습 시 수 시간의 데이터 로딩 시간을 절약했습니다.
• 비용 절감: 저렴한 인스턴스로 리허설을 진행함으로써 고성능 인스턴스 예약 시간을 효율적으로 활용하고 불필요한 비용 낭비를 방지했습니다.

결론 및 제언

• 본 워크플로우는 대규모 모델 학습 시 인프라 운영 부담을 최소화하고 모델 성능 고도화에 집중할 수 있도록 지원합니다.
• SageMaker HyperPod의 유연성, 고성능 네트워크, 간편한 데이터 통합 기능은 성공적인 분산 학습에 기여했습니다.
• ZERO 모델은 AWS Marketplace에서 구독 및 배포가 가능하며, SageMaker 엔드포인트로 수 분 내 추론 환경 구성이 가능합니다.