Mixture of Experts (MoEs) in Transformers

Context

기존 Transformers 라이브러리는 조밀한(dense) 모델 기준으로 설계되었으나, MoE(Mixture of Experts) 아키텍처의 도입으로 무게 로딩, 장치 배치, 양자화, 백엔드 실행 등 핵심 파이프라인의 근본적 재설계가 필요했다. MoE 체크포인트는 각 전문가(expert)를 독립적으로 직렬화하지만, GPU 커널은 모든 전문가의 무게를 단일 연속 텐서로 패킹하는 최적화된 구현을 요구해 불일치가 발생했다.

Technical Solution

• 무게 로딩 재설계: WeightConverter 추상화 도입으로 체크포인트 레이아웃(256개 독립 텐서)을 런타임 레이아웃(1개 패킹 텐서)으로 자동 변환
• GroupedGemmParallel 구현: 전문가 무게를 전문가 차원(dim=0)으로 분할하여 각 장치가 num_experts / num_devices 만큼만 로드
• RouterParallel 메커니즘: 전역 전문가 인덱스를 로컬 인덱스로 재매핑하고, 각 장치에 할당되지 않은 전문가를 마스킹한 후 all-reduce로 부분 출력 결합
• Expert Backend 추상화: Unsloth와 협업으로 torch._grouped_mm API 표준화 및 커스텀 Triton grouped-GEMM + LoRA 커널 제공
• 훈련 최적화: LoRA 기반 커널을 활용한 MoE 훈련 가속화

Impact

• gpt-oss-20b(21B 총 파라미터)를 M3 Ultra Mac에서 초당 115 토큰 생성(예측값 111 tok/s와 일치)
• MoE 훈련 성능 Unsloth 최적화: 약 12배 빠른 훈련, 35% 이상 VRAM 감소, 약 6배 길어진 컨텍스트 길이
• v4 대비 전체 성능 12~30배 향상

Key Takeaway

MoE와 같은 새로운 아키텍처 도입 시 단순 모델 클래스 추가를 넘어 로딩, 직렬화, 분산 통신 전체 파이프라인을 재설계해야 하며, 체크포인트 형식과 런타임 실행 요구사항의 불일치를 변환 계층으로 투명하게 처리하는 것이 핵심이다.