트랜스포머에서 전문가 혼합 (MoEs)
Hugging Face는 transformers 라이브러리를 전면 개편하여 Mixture of Experts (MoE) 아키텍처에 대한 일등 지원을 제공했습니다. 이번 재설계는 MoE 모델이 체크포인트에 저장되는 방식과 하드웨어에서 실행되는 방식 사이의 근본적인 불일치를 해결하여 가중치 로딩을 크게 빠르게 하고, 플러그인 가능한 백엔드를 통한 보다 효율적인 추론을 가능하게 하며, 전문가 병렬성을 통한 확장 가능한 분산을 구현합니다.
MoE 기본 개념: 용량 vs. 활성 파라미터
Mixture of Experts 모델은 밀집 피드포워드 레이어를 전문가라고 불리는 학습 가능한 서브 네트워크 집합으로 교체합니다. 라우터는 각 토큰을 처리하기 위해 이 전문가 중 작은 부분집합을 선택하여 모델의 전체 용량과 추론 비용을 분리합니다.
- Total Parameters: 모델의 전체 용량과 품질을 결정합니다.
- Active Parameters: 토큰당 추론 속도와 연산 비용을 결정합니다.
예를 들어, gpt-oss-20b 모델은 총 210억 파라미터를 가지고 있지만 토큰당 약 36억 파라미터만 활성화합니다(32명 중 4명의 전문가 사용). 이를 통해 모델은 210억 파라미터 시스템의 품질을 유지하면서 36억 파라미터 모델의 속도로 동작하여 M3 Ultra Mac에서 초당 약 115 토큰을 달성합니다.
가중치 로딩 리팩터 및 WeightConverter
transformers에서 전통적인 가중치 로딩은 체크포인트 텐서와 런타임 파라미터 사이에 일대일 매핑이 있다고 가정했습니다. MoE 체크포인트는 일반적으로 전문가들을 독립적으로 직렬화합니다(예: 256개의 개별 텐서). 그러나 최적화된 런타임 커널은 그룹화된 GEMM 연산을 위해 전문가들을 하나의 연속 텐서로 패킹할 것을 요구합니다.
이를 해결하기 위해 Hugging Face는 WeightConverter 추상화를 통해 변환 파이프라인을 도입했습니다. 이 파이프라인은 단순한 키별 복사 방식에서 동적 변환 방식으로 로딩 프로세스를 전환합니다:
MergeModulelist: 여러 전문가 텐서를 하나의 연속 텐서로 쌓습니다.SplitModulelist: 패킹된 텐서를 개별 전문가 텐서로 분할합니다.- Lazy Materialization: 로더가 키를 한 번만 스캔하고, 종속성이 준비될 때 스레드 풀을 통해 텐서를 실제로 생성하여 메모리 피크를 감소시키고 반복 스캔을 방지합니다.
가중치 로딩 벤치마크
Qwen/Qwen1.5-110B-Chat을 단일 A100(80GB)에서 사용한 벤치마크는 v4에 비해 v5 파이프라인에서 상당한 속도 향상을 보여줍니다:
| 버전 | 전략 | 로드 모드 | 시간 |
|---|---|---|---|
| v4.57.6 | device_map="auto" |
Threadpool | 66.24s |
| v5 | device_map="auto" |
Async (default) | 20.71s |
| v5 | TP | Async | 10.1s |
플러그인 가능한 전문가 백엔드
전문가 연산을 모델 구현과 분리하기 위해 Hugging Face는 @use_experts_implementation 데코레이터를 사용한 Experts Backend 시스템을 도입했습니다. 이를 통해 모델은 런타임에 세 가지 실행 전략 사이를 전환할 수 있습니다:
eager: 선택된 전문가들을 순회합니다; 주로 디버깅 및 정확성 검증에 사용됩니다.batched_mm:torch.bmm을 사용해 토큰당 전문가 가중치를 복제합니다; 작은 배치와 GPU 중심 워크로드에 최적화되었습니다.grouped_mm:torch._grouped_mm을 사용해 토큰을 전문가 ID별로 정렬하고 단일 그룹화된 GEMM을 수행합니다; 큰 배치와 메모리 제한 환경에 최적화되었습니다.
전문가 병렬성 (EP)
전문가 병렬성은 수천억 개의 파라미터를 가진 모델이 전문가들을 장치에 분산시켜 여러 GPU에 걸쳐 확장할 수 있게 합니다. 표준 텐서 병렬성과 달리 각 장치는 할당된 전문가 부분집합(num_experts / num_devices)만 로드합니다.
이는 DistributedConfig(enable_expert_parallel=True)를 통해 구현되며 두 가지 핵심 구성 요소에 의존합니다:
GroupedGemmParallel: 전문가 차원(dim=0)을 따라 전문가 가중치를 샤딩합니다.RouterParallel: 전역 전문가 인덱스를 로컬 인덱스로 재매핑하고 all-reduce 연산을 사용해 장치 간 부분 출력을 결합합니다.
Unsloth와 함께 최적화된 MoE 훈련
Unsloth와의 협업을 통해 Hugging Face는 Expert Backend 추상화와 PyTorch의 torch._grouped_mm API를 활용하고, 맞춤형 Triton 그룹화 GEMM 및 LoRA 커널을 결합하여 더 빠른 MoE 훈련을 가능하게 했습니다. 이러한 최적화는 다음을 제공합니다:
- MoE 훈련 속도가 최대 12배 빨라집니다.
- VRAM 사용량이 35% 이상 감소합니다.
- 컨텍스트 윈도우가 약 6배 길어집니다.
- v4
transformers구현에 비해 전체 속도가 12~30배 향상됩니다.