AWS가 NVIDIA Blackwell GPU(B200)를 탑재한 P6-B200 인스턴스에서 Amazon SageMaker AI 훈련 작업을 최적화하는 방법을 공개했다. P6-B200 인스턴스는 GPU 8개를 장착하고 있으며 Flexible Training Plan을 통해 예약 접근이 가능하다. Blackwell은 이전 세대 GPU 대비 확장된 메모리 용량(B200 180GB, B300 268GB)과 NVLink 5 인터커넥트를 통한 최대 1.8TB/s의 GPU 간 대역폭을 제공해, 대규모 모델 훈련의 주요 병목이던 메모리 압박과 통신 오버헤드를 줄인다.
최적화의 핵심은 세 가지 영역이다. 배치 크기의 경우, Blackwell의 넉넉한 메모리를 활용하면 그래디언트 동기화 횟수를 줄여 전반적인 처리량을 높일 수 있다. 모델 샤딩 측면에서는 GPU당 메모리가 늘어나 일부 모델에서 병렬화 수준을 낮추거나 아예 없앨 수 있어, 노드 간 통신 오버헤드가 감소한다. 정밀도 형식은 FP8, MXFP8, NVFP4 등 낮은 정밀도 형식이 Blackwell 5세대 텐서 코어에서 하드웨어 가속을 받는다. 14B 파라미터 이하 소형 모델에서는 배치 크기 조정이 가장 유효하고, 14B 이상 대형 모델에서는 활성화 체크포인팅을 통한 메모리-연산 교환이 필수적이다. 1B 파라미터 LLM을 MXFP8 정밀도와 활성화 체크포인팅을 함께 적용하고 배치 크기를 키우면 기준 대비 처리량이 약 8배 증가한다는 실험 결과가 제시됐다.
실행 환경 구성 측면에서 AWS는 단계별 절차를 제시한다. 먼저 FSDP(Fully Sharded Data Parallel) 훈련 스크립트를 작성하고, NVIDIA TransformerEngine이 포함된 커스텀 Docker 컨테이너를 Amazon ECR에 등록한 후, Flexible Training Plan으로 P6-B200 용량을 예약한다. 안정적 생산 훈련에는 Flexible Training Plan을, 비용 절감이 우선일 때는 체크포인트 자동 재개가 지원되는 Spot 인스턴스를 선택할 수 있다. 훈련 작업은 Amazon CloudWatch 로그를 통해 실시간 모니터링된다.
AWS는 Blackwell 기반 훈련 작업을 올바르게 구성하면 1B~64B 파라미터 트랜스포머 모델 전반에서 일관된 성능 향상을 얻을 수 있다고 밝혔다. 핵심은 자신의 워크로드가 연산 바운드인지 메모리 바운드인지를 파악한 뒤, 각 병목에 맞는 최적화를 단계적으로 적용해 실측하는 방식이다. 구체적인 코드와 Dockerfile, IAM 정책은 Amazon SageMaker AI 공식 문서에서 확인할 수 있다.
