meta-pytorch/torchft

해결하는 문제

대규모 AI 모델을 훈련할 때 단일 GPU나 노드가 다운되면 종종 훈련이 실패하며, 일반적으로 마지막 체크포인트에서 전체 훈련을 다시 시작해야 합니다. torchft는 스텝 단위의 장애 내성 기능을 제공하여 오류 발생 시 전체 작업을 중단하지 않고도 훈련을 계속할 수 있게 합니다.

작동 방식

이 라이브러리는 학습 스텝 단위의 멤버십 변경을 위한 기본 기능을 제공합니다. "라이트하우스" 서버를 사용하여 복제 그룹 간을 조정하고, 스텝 단위의 허트비트를 통해 워커의 건강 상태를 확인합니다. 또한 장애 내성 버전의 ProcessGroup과 체크포인트 전송 기능을 구현하여 스케일업 작업 중에 건강한 피어로부터 실시간 복구가 가능하게 합니다.

대상 사용자

PyTorch를 사용하여 대규모 모델(예: Llama 3 70B)을 훈련하는 머신러닝 엔지니어 및 연구자. 특히 DDP 또는 HSDP와 같은 복제된 가중치 전략을 사용하는 사람들.

주요 특징

• 장애 내성 알고리즘: 장애 내성 DDP, 장애 내성 HSDP, LocalSGD, DiLoCo를 즉시 사용 가능.
• 라이트하우스 조정: 워커의 건강 상태를 관리하고 복제 그룹을 조정하는 전용 서버로, "세계 정지"와 같은 훈련 일시 정지 현상을 방지합니다.
• 실시간 복구: 스케일업 작업 중에 특수한 체크포인트 전송을 사용해 건강한 피어로부터 상태를 복구할 수 있습니다.
• 통합성: 기존 PyTorch 훈련 스크립트에 최소한의 변경만으로 쉽게 통합 가능.