구형 GPU에서 충돌 없이 실제로 실행되는 LLM 학습 프레임워크를 구축했습니다 [P]
요약
구형 GPU(T4, V100 등)에서도 충돌 없이 실행 가능한 LLM 학습 프레임워크인 Picotron을 소개합니다. FlashAttention-2 등 하드웨어 특화 의존성을 제거하여 범용성을 높였으며, PyTorch SDPA를 통해 다양한 GPU 환경을 지원합니다.
핵심 포인트
- 구형 GPU에서도 작동하는 클린룸 재작성 프레임워크
- FlashAttention-2 등 무거운 의존성 제거로 호환성 확보
- GQA, MLA, QK-Norm 등 최신 모델 아키텍처 지원
- PyTorch SDPA 폴백을 통한 유연한 연산 지원
여러분 안녕하세요,
최근 Nanotron을 가지고 이것저것 시도해보고 있었는데, 모듈 레벨에서 너무 많은 무겁고 하드웨어에 특화된 의존성(flash-attn, triton, functorch 등)을 가져오는 방식 때문에 매우 좌절했습니다. T4나 V100 같은 구형 또는 저가형 GPU에서 실행하려고 하면, 임포트(import) 단계에서 바로 충돌이 발생합니다.
그래서 이를 해결하기 위해 Picotron (https://github.com/Syntropy-AI-Labs/picotron)을 작성했습니다. 이는 모든 필수적인 GPU 특화 의존성을 제거한 클린룸 재작성(clean-room rewrite) 버전입니다.
이 프레임워크는 PyTorch를 지원하는 거의 모든 GPU에서 실행됩니다 (연산 능력(compute capability) 8.0 미만의 구형 카드는 FP16을 기본값으로 사용하고, 최신 카드는 BF16을 사용합니다). 기본적으로 표준 PyTorch SDPA로 폴백(fallback)되지만, 설치된 것을 감지하면 런타임에 FlashAttention-2에 연결됩니다.
많은 보일러플레이트(boilerplate) 및 코드 모듈을 작성하는 데 AI 어시스턴트를 사용했지만, 로컬에서 작동하도록 만들었으며 방금 FineWeb-Edu 데이터셋으로 아주 작은 2M 모델을 학습시켰습니다.
또한 다음 기능들을 위한 설정(configs)을 추가했습니다:
• GQA / MLA (Multi-head Latent Attention)
• QK-Norm 및 logit soft-capping (Gemma 2 스타일)
• 병렬 FFN/Attn 실행
• DDP 상의 ZeRO-1 래핑(wrapping)
현재 로드맵은 꽤 짧습니다:
• MoE 준비 (라우팅 용량 계수(routing capacity factors) 및 부하 분산 손실(load balancing loss))
• 데이터셋 준비를 수동 스트리밍보다 더 쉽게 만들기
CUDA 의존성 지옥(dependency hell)과 싸우고 계셨다면 확인해 보세요: https://github.com/Syntropy-AI-Labs/picotron
submitted by /u/Capital_Savings_9942 to r/MachineLearning
[link] [comments]
AI 자동 생성 콘텐츠
본 콘텐츠는 r/OpenAI Codex (search)의 원문을 AI가 자동으로 요약·번역·분석한 것입니다. 원 저작권은 원저작자에게 있으며, 정확한 내용은 반드시 원문을 확인해 주세요.
원문 바로가기