mlc-ai/pith-train

해결하는 문제

실제 환경에서 사용되는 MoE(Mixture-of-Experts) 학습 프레임워크는 종종 지나치게 복잡하며, 무거운 컴파일 확장과 간접적인 설계 패턴에 의존하여 인간과 AI 코딩 에이전트가 수정하거나 최적화하기 어렵습니다. PithTrain은 에이전트나 인간이 끝까지 읽을 수 있을 정도로 작고, 동시에 실용적인 성능을 유지하는 고성능 학습 시스템을 제공합니다.

작동 방식

에이전트 네이티브 프레임워크로, 세 가지 계층으로 구성됩니다: 학습 루프를 위한 애플리케이션 계층, 모델 프로토콜과 분산 학습을 위한 엔진 계층, 저수준 커널을 위한 오퍼레이터 계층. 4D 병렬 전략(파이프라인, 데이터, 컨텍스트, 전문가 병렬)을 구현하고, 전방-후방 실행과 P2P 통신을 겹치는 데 DualPipeV 스케줄러를 사용합니다. FP8 학습을 지원하며, DeepGEMM 및 FlashAttention 같은 라이브러리와 통합됩니다.

대상 사용자

Qwen이나 DeepSeek 아키텍처와 같은 MoE 모델을 학습하고자 하며, AI 코딩 어시스턴트를 활용할 때 코드베이스를 쉽게 탐색하고 수정하며 확장하고 싶은 연구자 및 개발자.

주요 특징

• 에이전트 네이티브 설계: 약 11K 줄의 파이썬 코드로 구성된 컴팩트한 코드베이스이며, 최소한의 암묵적 간접 참조를 갖습니다.
• 실용적 성능: 4D 병렬화와 계산-통신 겹침을 지원합니다.
• FP8 학습: 효율성을 위해 FP8 선형 및 양자화를 내장 지원합니다.
• 유연한 체크포인팅: 리셰이딩 가능한 체크포인트와 Hugging Face 형식으로의 변환을 지원합니다.