Transformers에서의 Mixture of Experts (MoEs) 지원 및 최적화

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

핵심 요약

대규모 언어 모델의 성능 향상을 위해 파라미터 수를 늘리는 dense 모델 방식은 비용과 지연 시간 측면에서 한계에 직면했다. 이에 대한 대안으로 부상한 Mixture of Experts (MoE) 아키텍처는 토큰당 활성 파라미터 수를 제한하여 효율성을 확보한다. Hugging Face는 Transformers 라이브러리를 개편하여 MoE 모델의 가중치 로딩 속도를 3배 높이고, 최적화된 실행 백엔드와 전문가 병렬성(EP)을 지원한다. 또한 Unsloth와의 협업을 통해 학습 속도를 12배 향상시키고 메모리 사용량을 35% 절감하여 실무적인 sparse 모델 활용을 가능하게 한다.

배경

Transformer 아키텍처에 대한 이해, PyTorch 및 분산 학습(Distributed Training) 기본 지식, Hugging Face Transformers 라이브러리 사용 경험

대상 독자

대규모 MoE 모델을 프로덕션 환경에서 배포하거나 효율적으로 학습시키고자 하는 AI 엔지니어 및 연구자

의미 / 영향

이 기술적 업데이트는 DeepSeek R1과 같은 대규모 sparse 모델의 접근성을 획기적으로 높인다. 특히 로딩 속도 개선과 학습 효율화는 인프라 비용 절감으로 이어져 중소규모 기업에서도 고성능 MoE 모델을 직접 운영할 수 있는 기반을 마련한다.

섹션별 상세

MoE는 Transformer 백본을 유지하면서 특정 dense 레이어를 전문가(Expert) 네트워크 세트로 대체한다. 각 토큰은 라우터에 의해 소수의 전문가에게만 할당되어 연산되므로, 전체 파라미터 대비 추론 시 활성 파라미터 수를 획기적으로 줄여 효율성을 확보한다.

Transformers v5에서는 WeightConverter를 통해 체크포인트 로딩 파이프라인을 재설계했다. 체크포인트에 분산 저장된 전문가 가중치를 런타임에 필요한 단일 연속 텐서로 병합하거나 비동기적으로 로드하여, Qwen 110B 모델 기준 로딩 시간을 v4 대비 약 3배 단축했다.

python

WeightConverter(
    ["block_sparse_moe.experts.*.w1.weight", "block_sparse_moe.experts.*.w3.weight"],
    "mlp.experts.gate_up_proj",
    operations=[
        MergeModulelist(dim=0),
        Concatenate(dim=1),
    ],
)

개별 전문가 가중치를 단일 패킹 텐서로 병합하는 WeightConverter 설정 예시

Experts Backend 아키텍처는 모델 구현과 전문가 연산을 분리하여 grouped_mm과 같은 최적화된 커널을 유연하게 적용한다. 특히 대규모 배치 처리 시 토큰을 전문가 ID별로 정렬하여 한 번에 연산하는 grouped_mm 방식은 메모리 제약이 있는 환경에서 높은 성능을 발휘한다.

Expert Parallelism (EP)은 모델의 전문가들을 여러 GPU 장치에 분산 배치하여 단일 장치 메모리 한계를 극복한다. enable_expert_parallel 플래그 하나로 복잡한 샤딩 전략과 라우팅 인덱스 매핑, 결과 집계 과정을 자동화하여 개발자 경험을 개선했다.

python

from transformers import AutoModelForCausalLM
from transformers.distributed.configuration_utils import DistributedConfig

distributed_config = DistributedConfig(enable_expert_parallel=True)
model = AutoModelForCausalLM.from_pretrained(
    "openai/gpt-oss-120b",
    dtype="auto",
    distributed_config=distributed_config,
)

DistributedConfig를 사용하여 전문가 병렬성(EP)을 활성화하는 방법

Unsloth와의 기술 협력을 통해 Triton 기반의 전용 커널을 도입하여 MoE 학습 성능을 극대화했다. 이를 통해 기존 방식 대비 학습 속도는 최대 12배 빨라졌으며, VRAM 사용량은 35% 감소하여 더 긴 컨텍스트를 학습할 수 있는 환경을 제공한다.

실무 Takeaway

대규모 MoE 모델 도입 시 Transformers v5의 비동기 로딩과 WeightConverter를 활용하면 로딩 병목을 해결하고 메모리 효율을 높일 수 있다.
enable_expert_parallel 설정을 통해 수천억 파라미터 규모의 sparse 모델을 추가적인 복잡한 코드 작성 없이 여러 GPU에 분산 배치하여 실행 가능하다.
Unsloth 최적화가 적용된 Transformers 환경을 구축하면 MoE 모델 학습 시 VRAM을 35% 절감하면서도 학습 속도를 12배까지 향상시킬 수 있다.

언급된 리소스

튜토리얼Unsloth's Official Guide to MoE Training

GitHubTransformers GitHub Repository