galilai-group/stable-pretraining

해결하는 문제

"기초 모델" (CLIP, DINO, SimCLR와 같은 일반적인 시각적 특징을 학습하는 모델)을 위한 대규모 신경망 학습은 전례 없이 번거롭습니다. 연구자들은 데이터 로딩, 증강, 로깅, 평가, 클러스터 작업 관리 등을 동시에 다뤄야 하며, 작은 실수도 긴 실행을 무효화할 수 있습니다. 이 프로젝트는 PyTorch Lightning을 래핑하여 전체 워크플로우를 더 간단하고 유연하며 안정적으로 만드는 프레임워크입니다. 연구자들은 모델과 손실에 집중할 수 있고, 프레임워크가 하위 작업을 처리합니다.

작동 방식

이 프레임워크는 네 가지 구성 요소로 구성되며, 각각 Python 사전을 서로 전달합니다: DataModule (데이터), Module (모델 + 순전파), Callbacks (모니터링/평가 핸들러), Trainer/Manager (오케스트레이션). 핵심 아이디어는 순전파 함수가 {'loss': ..., 'embedding': ...}와 같은 텐서 사전을 반환하는 것이며, 로거, 평가 프로브, 체크포인트 등은 트레이닝 루프를 수정하지 않고도 그 사전에서 읽어들인다는 점입니다. SimCLR, DINO, MAE, BYOL, CLIP 등 많은 자기지도 학습 방법에 대한 사전 구축된 "순전파" 레시피가 제공되며, OnlineProbe (고정된 특징량 위에 작은 선형 분류기를 훈련하여 실시간 정확도 추적) 및 OnlineKNN (학습 없이 최근접 이웃 평가)와 같은 콜백도 포함됩니다. 또한 kornia를 통해 데이터 증강을 GPU로 이동하여 CPU 버퍼를 제거하고, Manager는 SLURM 클러스터 기능인 작업 재시도/재개, 원자적 체크포인트, 쿼리 가능한 실행 레지스트리 등을 제공합니다. 실험적인 JAX/Flax-NNX 백엔드도 동일한 설계를 반영합니다.

대상 사용자

기초 모델 또는 자기지도 학습 연구를 수행하는 연구자 및 엔지니어 — 특히 GPU 클러스터에서 대규모 시각 모델을 훈련하고 실시간 평가, 견고한 체크포인트, 최소한의 반복 코드를 원하는 사람. 또한 SimCLR, DINOv2, MAE 등 SSL 방법을 최소한의 코드로 빠르게 프로토타이핑하거나 벤치마킹하고 싶은 누구에게나 유용합니다.

주요 특징

• 사전 기반 설계: 순전파 함수가 상태 사전을 반환하므로, 중간 텐서는 자동으로 로깅 가능하고, 콜백은 트레이닝 루프를 건드리지 않고 첨부됩니다.
• 30개 이상의 내장 레시피: SSL, 감독 학습, 다중 모달 전처리를 아우르는 레시피 (SimCLR, DINO/DINOv2, MAE, BYOL, VICReg, Barlow Twins, LeJEPA, CLIP 등).
• 실시간 평가 콜백: OnlineProbe 및 OnlineKNN과 같은 실시간 표현 품질 모니터링 기능.
• GPU 기반 배치 증강 (kornia를 통해): 배치 전체에 걸쳐 증강을 벡터화하여 모델 크기와 정밀도에 따라 측정된 처리량 향상.
• SLURM 수준의 오케스트레이션: Manager는 선점/재시도, 원자적 체크포인트, 쿼리 가능한 실행 레지스트리 처리.
• 실험적 JAX/Flax-NNX 백엔드: torch 설계를 반영하며, 수치적 주장에 대한 동등성 회귀 테스트 제공.