PrimeIntellect-ai/prime-diloco

해결하는 문제

Prime는 인터넷을 통해 효율적이고 글로벌하게 분산된 AI 모델을 훈련하기 위해 설계된 프레임워크입니다. 지리적으로 분산된 노드들 간에 대규모 모델을 훈련할 때 발생하는 문제들, 예를 들어 네트워크 불안정성, 높은 지연 시간, 체크포인트 및 복구 비용의 증가를 해결합니다.

작동 방식

Prime는 통신 오버헤드를 최소화하기 위해 DiLoCo(Distributed Low-Communication) 접근 방식을 사용합니다. 다음과 같은 핵심 최적화를 구현하고 있습니다:

• ElasticDeviceMesh: 동적 프로세스 그룹을 관리하는 장애 내성 추상화로, 노드가 훈련 실행 중에 참여하거나 탈퇴해도 콜드 리스타트 없이 가능합니다.
• 비동기 체크포인트: 메인 훈련 프로세스를 차단하지 않기 위해 체크포인트는 먼저 RAM 기반 파일 시스템(/dev/shm)에 저장된 후 비동기적으로 디스크나 원격 저장소로 업로드됩니다.
• 라이브 체크포인트 복구: 참여하는 노드는 사이드카 HTTP 서버를 통해 피어로부터 최신 모델 상태를 빠르게 가져와 훈련 중간에 참여할 수 있습니다.
• 사용자 정의 Int8 All-Reduce 커널: Pseudo-Gradient를 Int8로 양자화하는 C++ 구현의 릴링 리듀스 커널로, 네트워크 페이로드 크기를 4배 줄이되 손실 곡선에는 영향을 주지 않습니다.
• 대역폭 최적화: 여러 동시 연결을 활용한 샤딩된 Pseudo-Gradient와 VPN 기술을 통해 피어 투 피어 라우팅을 최적화하여 네트워크 활용률을 극대화합니다.
• 메모리 관리: PyTorch FSDP2/DTensor ZeRO-3를 사용해 모델 가중치, 기울기, 옵티마이저 상태를 노드 내 GPU 간에 샤딩하고, DiLoCo 옵티마이저 텐서를 CPU 메모리로 오프로드합니다.

대상 사용자

다수의 데이터 센터 또는 공공 인터넷을 통해 훈련을 분산하고자 하는 대규모 AI 모델을 훈련하는 연구자 및 엔지니어.

주요 특징

• ElasticDeviceMesh와 허트비트 메커니즘을 통한 장애 내성 훈련.
• 사용자 정의 Int8 All-Reduce 커널을 통해 통신 페이로드 4배 감소.
• RAM 기반 비동기 업로드를 통한 블로킹 없는 체크포인트.
• 고대역폭 인터넷 훈련으로 미국 내 데이터 센터 간 최대 4Gb/s 성능 달성.
• 효율적인 메모리 샤딩(ZeRO-3)을 위한 FSDP2 통합.