PyTorch Monarch를 ROCm으로 확장하여 AMD Instinct에서 탄력적 장애 복구가 가능한 대규모 분산 학습 구현

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

TL;DR

대규모 LLM 훈련에서는 노드 실패가 빈번해 전통적 체크포인트 기반 복구로는 I/O 병목과 계산 손실, 클러스터 유휴 시간이 심각한 문제가 된다. PyTorch Monarch는 액터 기반 런타임과 프로세스 메쉬, 비동기 실행을 통해 단일 파이썬 프로그램으로 클러스터를 제어하면서 일부 노드가 실패해도 건강한 노드가 학습을 지속하고 실패 노드는 재접속해 상태를 재동기화하도록 하는 탄력적 복구 메커니즘을 구현한다. 이번 작업은 이러한 Monarch 런타임과 통신 스택을 ROCm으로 포팅해 AMD Instinct GPU에서 동일한 단일 컨트롤러 모델을 운영 가능하게 했으며, 이전 FP8 기반 대규모 실험에서 보고된 96.16% 확장 효율을 배경으로 신뢰성 문제 해결 측면에서 의미있는 진전으로 평가된다.

섹션별 상세

대규모 LLM 사전학습에서는 수백에서 수천 개의 GPU를 사용하는 분산 훈련이 필수이며 이 환경에서는 GPU 메모리 오류나 네트워크 분할, 노드 크래시가 빈번하게 발생해 전통적 체크포인트 기반 복구 방식이 한계에 봉착한다. 체크포인트 방식은 수백 기가바이트 규모의 모델 상태를 디스크에 기록하느라 I/O 병목을 유발하고 장애 발생 시 체크포인트 이후의 모든 계산이 유실되어 계산 자원이 낭비된다. 또한 실패 처리를 위해 클러스터 전체가 대기 상태에 들어가며 클러스터 규모가 커질수록 일정 기간 내 장애 발생 확률이 올라 확장 가능성에 제약을 준다.

PyTorch Monarch는 단일 파이썬 프로그램에서 전체 GPU 클러스터를 오케스트레이션하는 액터 기반 런타임과 프로세스 메쉬 추상화, 비동기 실행 모델을 결합해 분산 작업을 관리한다. 이 구조는 각 액터가 독립적으로 실행 흐름과 상태를 유지하게 하여 일부 노드가 실패하더라도 건강한 노드들이 학습을 계속하도록 허용하고 실패한 노드는 백그라운드에서 재접속·상태 재동기화를 통해 클러스터에 복귀한다. 글은 Monarch가 노드 장애 시 전체 작업을 중단하지 않고 동적으로 복구하는 동작을 시연했다고 기술해 탄력적 분산 학습을 가능하게 했음을 근거로 제시한다.

이번 작업에서는 Monarch의 GPU 런타임과 분산 통신 스택을 CUDA 환경을 넘어 ROCm으로 이식해 AMD Instinct GPU에서 단일 컨트롤러 모델을 운영 가능하게 했다. 포팅 작업은 런타임 호출, 메모리 관리 경로, 통신 계층의 ROCm 적응을 포함하며 그 결과 단일 프로그램 인터페이스로 AMD 하드웨어 생태계에서도 탄력적·장애 허용 분산 학습을 실행할 수 있게 되었다. 본문은 이전 연구에서 FP8 기반 1024-GPU MI325 클러스터에서 96.16% 확장 효율을 달성했으나 신뢰성은 여전히 과제였고, ROCm 포팅은 그 신뢰성 문제를 해결하는 중요한 진전임을 의미한다고 기술한다.

용어 해설

액터 기반 런타임(Actor-based runtime): — 액터 기반 런타임은 독립 실행 단위인 액터들이 메시지를 통해 비동기적으로 상호작용하며 작업을 수행하는 분산 실행 모델로, GPU 클러스터 전체를 하나의 파이썬 프로그램으로 제어하면서 노드 간 작업 분할과 실패 복구를 관리하는 데 핵심 역할을 한다.
프로세스 메쉬(Process Mesh): — 프로세스 메쉬는 분산 학습에서 프로세스들을 격자 형태로 추상화해 통신 경로와 연산 분할을 선언적으로 표현하는 구조로, 데이터·모델·파이프라인 병렬화와 장애 발생 시 재배치 전략을 단순화하는 데 사용된다.
탄력적 분산 학습(Elastic Distributed Training): — 탄력적 분산 학습은 일부 노드가 실패하거나 복구되는 동안에도 남은 노드가 학습을 계속하고 실패한 노드가 복귀하면 상태를 재조정해 훈련을 이어가는 방식으로, 체크포인트 재시작 방식보다 계산 낭비를 줄이고 GPU 활용도를 높이는 것이 목적이다.
FP8: — FP8은 8비트 부동소수점 숫자 형식으로, 대규모 모델의 메모리 사용량과 통신 비용을 줄이기 위해 사용되며 본문에서는 FP8 기반 대규모 사전학습에서 1024-GPU 클러스터 상의 확장 효율을 언급하는 맥락에서 중요하다.

기술

PyTorch Monarch
ROCm
AMD Instinct
FP8

활용 사례

대규모 LLM 사전학습에서의 탄력적 장애 복구
장기 실행 훈련 잡의 계산 자원 낭비 최소화
CUDA 이외의 하드웨어에서의 분산 학습 확장

언급된 리소스

문서Efficient MoE pre-training at scale with TorchTitan