Picotron — GPU 특정 의존성을 제거한 경량 LLM 사전학습 프레임워크

Picotron은 Nanotron에서 발생하던 모듈 레벨의 하드웨어 특화 의존성 문제를 해결하기 위해 클린룸으로 재작성된 경량 LLM 사전학습 프레임워크이다. 이 프레임워크는 PyTorch가 지원되는 거의 모든 GPU에서 동작하도록 설계되었고 구형 카드에서는 FP16, 최신 카드에서는 BF16을 기본 정밀도로 사용하며 표준 PyTorch SDPA를 폴백으로 유지하는 한편 FlashAttention-2가 설치되어 있으면 런타임에 이를 후킹해 가속 경로를 활용한다. 작성자는 로컬에서 동작을 확인했고 FineWeb-Edu로 2M 모델을 학습해본 사례를 제시했으며 GitHub에 코드와 구성 파일을 공개했다.

프로젝트는 GQA/MLA, QK-Norm 및 logit soft-capping, Parallel FFN/Attn 실행, DDP에서의 ZeRO-1 래핑 같은 구성 옵션을 제공해 어텐션·정규화·분산 학습 관련 설정을 조정할 수 있는 환경을 마련했다. 런타임 감지와 폴백 설계는 환경별 수작업 설정을 줄이고 예산형 GPU에서도 실험을 시작할 수 있게 했으며 저장소에는 실제 설정 파일과 사용 가능한 옵션들이 포함되어 있어 재현이 가능한 근거를 제공한다. 로드맵으로는 MoE 준비(라우팅 용량 계수 및 부하 균형 로스)와 데이터셋 준비 간소화가 명시되어 있어 향후 확장과 사용성 개선이 예정되어 있다.

요약하면 Picotron은 CUDA 의존성으로 인한 진입 장벽을 낮추는 데 초점을 맞춘 도구로, 환경 감지 기반으로 가속 라이브러리를 선택적으로 활용하면서도 광범위한 GPU에서 동작하도록 설계되었다. 현재는 다양한 어텐션 및 분산 관련 구성 옵션을 제공해 연구 목적의 실험에 활용 가능하며 향후 MoE와 데이터 파이프라인 편의성 개선이 계획되어 있다. GitHub 링크를 통해 코드와 구성 파일을 확인할 수 있다.