기술 요약: DYNATRAIN

문제 정의

현대적인 대규모 언어 모델 (LLM) 학습은 모델이 방대한 GPU 클러스터에 적합하도록 데이터, 텐서, 파이프라인, 전문가, 옵티마이저 쉐이딩과 같은 병렬화 전략의 복잡한 조합에 의존합니다. 그러나 현재 학습 프레임워크는 시작 시 선택된 병렬화 레이아웃이 작업의 전체 수명 동안 최적이라고 가정하는 정적 실행 모델 하에서 작동합니다. 이 가정은 다음과 같은 이유로 점점 더 성립하기 어려워지고 있습니다:

1. 자원 변동: 생산 클러스터는 선점, 스팟 인스턴스 취소, 또는 비수시간에 GPU 를 방출하는 공존 서비스 워크로드 등으로 인해 자원을 동적으로 재할당합니다.
2. 동적 워크플로우: 인간 피드백을 통한 강화 학습 (RLHF) 과 같은 신흥 워크플로우는 추론 (롤아웃) 단계와 학습 단계 간에 자원을 전환해야 하므로, 런타임 시 병렬화 변경이 필요합니다.

기존 솔루션은 다음과 같은 중대한 한계에 직면해 있습니다:

• 체크포인트 및 재시작: 표준 접근 방식은 작업을 중단하고 테라바이트 규모의 상태를 디스크에 직렬화한 후 재시작하는 것입니다. 이는 저장소 I/O 병목 현상과 분산 통신자 재초기화의 높은 비용으로 인해 수십 초에서 수 분에 이르는 막대한 지연을 초래합니다.
• 기존 탄력적 시스템: 현재 메모리 내 "핫 스위칭" 접근 방식은 종종 사전 정의된 템플릿, 특정 병렬화 차원, 또는 정적 토폴로지로 제한되어 임의의 다차원 재구성을 지원하지 못합니다.

일반적인 핫 스위칭을 가능하게 하는 핵심 과제는 다음과 같이 세 가지입니다:

1. 상태 이질성: 현대적인 최적화 (예: ZeRO) 는 상태를 차원 전반에 걸쳐 분할하여 텐서 경계와 정렬되지 않은 불연속적이고 이질적인 쉐어드를 생성합니다.
2. 복잡한 라우팅: 두 개의 다차원 레이아웃 간 전환은 토폴로지 변경에 대한 전역적 관점 없이 데드락 없는 스케줄링이 필요한 M 대 N 피어 투 피어 통신 문제입니다.
3. 메모리 제약: 이전 상태와 새 상태는 통신 버퍼와 함께 일시적으로 공존해야 하므로, 용량이 제한된 GPU 에서 메모리 부족 (OOM) 실패가 발생할 위험이 있습니다.

방법론

DYNATRAIN 은 임의의 병렬화 전략에 걸쳐 초단위 온라인 재구성을 가능하게 하도록 설계된 분산 학습 시스템입니다. 이는 딥러닝 백엔드 (예: PyTorch) 와 학습 프레임워크 (예: Megatron-LM) 사이에 플러그 가능한 미들웨어 계층으로 작동합니다.

1. 가상 파라미터 공간 (VPS)

DYNATRAIN 의 핵심 추상화는 **가상 파라미터 공간 (VPS)**으로, 모든 파라미터, 그래디언트, 옵티마이저 상태가 완전하고 쉐어링되지 않은 형태로 존재하는 통합된 논리적 좌표 공간입니다.

• 논리적 표현: 모든 텐서는 논리적 다차원 인덱스 바운딩 박스로 표현됩니다. 이 공간에서 병렬 차원 (DP, TP, PP, EP, ZeRO) 은 크기가 1 로 축소됩니다.
• 결정적 매핑: 모든 병렬화 구성은 글로벌 VPS 를 랭크 로컬 서브-VPS 로 투영하는 수학적 매핑 함수 ($F_P$) 로 작용합니다.
• 기하학적 교차: 전략 간 재구성은 소스 및 목적지 서브-VPS 영역 간의 기하학적 교차를 계산하는 것으로 환원됩니다. 이는 복잡한 상태 이동을 어떤 논리적 영역을 전송, 수신 또는 유지할지 식별하는 결정적 문제로 변환하여, 쌍별 변환 규칙의 폭발을 방지합니다.

2. 시스템 구성 요소

VPS 위에 구축된 DYNATRAIN 은 세 가지 협력 구성 요소를 활용합니다:

• 상태 라우팅 플래너: 완전히 VPS 좌표 내에서 작동합니다. 각 랭크에 대한 소스 ($R^{(i)}_{src}$) 와 목적지 ($R^{(i)}_{dst}$) 논리적 영역을 비교하여 정확한 전송, 수신 및 유지 집합을 도출합니다. 옵티마이저 상태가 종종 시맨틱 텐서 경계가 아닌 평탄한 버퍼 오프셋으로 분할되는 불일치를 처리하기 위해 이를 VPS 좌표로 다시 매핑하여 표준 기하학적 교차로 취급합니다.
• 상태 전환 엔진: 플래너의 고수준 계획을 최적화된 실행 파이프라인으로 정제합니다.
  • 프라이미티브 최적화: 토폴로지가 허용하는 경우 포인트 투 포인트 전송을 집단적 프라이미티브 (브로드캐스트, 스캐터, 개더) 로 승격시키고, 인트라-노드 대역폭 (예: NVLink) 에 기반하여 최적의 소스 - 목적지 쌍을 선택합니다.
  • 메모리 인식 스테이징: OOM 을 방지하기 위해 쓸모없는 버퍼를 가비지 컬렉션하고, 사용 가능한 메모리에 기반하여 전송을 단계로 분할합니다. 연속적인 전송/수신 버퍼를 사용하여 분할된 전송을 통합하여 대역폭 활용도를 극대화합니다.
  • 데드락 없는 스케줄링: 대칭적인 통신 쌍을 보장하여 전역 조정 없이 데드락 없는 시퀀스를 보장하는 비동기 XOR 기반 스케줄링 알고리즘 ($Peer(i, s) = i \oplus s$) 을 사용합니다.
• 탄력적 장치 관리자 (EDM): 논리적 전략을 물리적 통신 그룹에서 분리합니다.
  • 통신 그룹의 수명 주기를 관리하고 즉시 재사용을 위해 캐싱합니다.
  • 비동기 오버랩을 지원하여, 새로운 "월드" 및 서브그룹이 백그라운드에서 초기화되는 동안 기존 토폴로지에서 학습 루프가 계속되도록 하여 토폴로지 변경의 지연을 마스킹합니다.

주요 기여

1. 가상 파라미터 공간 (VPS): 이질적인 분산 상태를 단일 논리적 좌표 체계 아래 통합하여 병렬화 전환을 기하학적 교차 문제로 축소하는 새로운 추상화.
2. 결합된 핫 스위칭 메커니즘: 상태 라우팅 플래너, 증명 가능한 데드락 없는 스케줄링을 갖춘 메모리 인식 상태 전환 엔진, 그리고 재구성과 학습을 오버랩시키는 탄력적 장치 관리자를 포함한 도구 세트.
3. 종단 간 시스템: DYNATRAIN 은 이러한 메커니즘을 통합하여 밀집형 및 MoE 모델에 대한 초단위 온라인 재구성을 제공하며, DP, TP, PP, EP 차원의 임의 변경을 지원합니다.

실험 결과

이 시스템은 16 노드 GPU 클러스터 (NVIDIA A800) 에서 13 억에서 2,350 억 파라미터에 이르는 모델 (LLaMA-2, GPT-3, Qwen-MoE 포함) 을 사용하여 평가되었습니다.

• 재구성 속도:
  • 밀집형 모델: 700 억 파라미터 밀집형 모델을 2 초 미만에 재구성합니다.
  • MoE 모델: 2,350 억 파라미터 MoE 모델을 4.36 초에 재구성합니다.
• 성능 향상:
  • **Megatron-LM 분산 체크포인팅 (MCP)**보다 밀집형 모델의 경우 최대 143 배, MoE 모델의 경우 최대 926 배 더 뛰어납니다.
  • Tenplex(체크포인트 기반 탄력적 시스템) 보다 단일 차원 변경의 경우 최대 33 배, 클러스터 간 작업 마이그레이션의 경우 최대 15 배 더 뛰어납니다.
  • HotSpa(메모리 내 베이스라인) 보다 하이브리드 재구성 시나리오에서 최대 2.11 배 더 뛰어납니다.
• 정확성: 1,300 억 파라미터 모델을 통한 손실 수렴 검증을 통해, DYNATRAIN 의 손실 궤적이 50 스텝마다 주기적인 자원 증가 및 감소가 있더라도 분산이나 스파이크 없이 정적 베이스라인을 완벽하게 추적함을 확인했습니다.
• 탄력적 확장: 동적 클러스터 확장 시나리오 (노드 추가/제거) 에서 시스템은 진행 중인 학습과 오버랩하여 초기화 오버헤드의 93% 이상을 마스킹하여 인지된 중단을 3 초 미만으로 줄였습니다.

의의

이 논문은 병렬화 핫 스위칭이 동적 자원과 다단계 워크플로우에 맞춰 학습을 적응시키기 위한 근본적인 시스템 프리미티브가 되어야 한다고 주장합니다. VPS 를 통해 분산 학습 상태를 물리적 레이아웃에서 분리함으로써 DYNATRAIN 은 온라인 재구성을 예외가 아닌 실용적인 것으로 만듭니다. 이 작업은 병렬화가 고정된 인프라 결정이 아니라 워크로드 및 클러스터 동역학에 지속적으로 대응할 수 있는 유연한 제어 차원이 되는 적응형 학습 런타임으로 나아가는 한 걸음을 의미합니다.