[LlamaFactory] LlamaFactory의 Triton 기반 Fused MoE 커널 도입: 40% 이상의 성능 향상

PR 링크: hiyouga/LlamaFactory#10481 상태: Merged | 변경: +856 / -10

들어가며

최근 대규모 언어 모델(LLM) 분야에서 Mixture of Experts(MoE) 아키텍처는 모델의 파라미터 수를 늘리면서도 연산 효율성을 유지하는 핵심 기술로 자리 잡았습니다. 하지만 기존의 HuggingFace 구현체는 각 전문가(Expert)별로 파이썬 루프를 돌며 연산을 수행하기 때문에, GPU의 병렬 처리 능력을 충분히 활용하지 못하는 병목 현상이 발생합니다. LlamaFactory의 이번 PR은 이러한 문제를 해결하기 위해 Triton을 사용하여 Fused MoE 커널을 구현했습니다.

코드 분석

1. Triton Fused MoE 커널 구현 (cuda_fused_moe.py)

핵심은 파이썬 루프를 제거하고 scatter와 gather 연산, 그리고 grouped GEMM을 하나의 Triton 파이프라인으로 통합한 것입니다.

# Before (Conceptual Python Loop)
for expert in experts:
    output = expert(hidden_states)

# After (Triton Fused Pipeline)
scatter_output = moe_scatter(hidden_states, scatter_index)
fc1_output = group_gemm_same_nk(scatter_output, fc1_weight, ...)
fc1_activation = torch.nn.functional.silu(fc1_1_output) * fc1_2_output
fc2_output = group_gemm_same_nk(fc1_weighted_output, fc2_weight, ...)
expert_output = moe_gather(fc2_output, scatter_index)

이 방식은 메모리 접근 패턴을 최적화하고 커널 실행 횟수를 줄여 GPU 활용도를 극대화합니다.

2. Trainer 최적화 (trainer.py)

여러 트레이너 클래스에서 create_optimizer 메서드의 시그니처를 수정하여 유연성을 확보했습니다.

# Before
def create_optimizer(self) -> "torch.optim.Optimizer":
    if self.optimizer is None:
        self.optimizer = create_custom_optimizer(self.model, self.args, self.finetuning_args)
    return super().create_optimizer()

# After
def create_optimizer(self, *args, **kwargs) -> "torch.optim.Optimizer":
    if self.optimizer is None:
        self.optimizer = create_custom_optimizer(self.model, self.args, self.finetuning_args)
    return super().create_optimizer(*args, **kwargs)

이는 다양한 학습 환경에서 커스텀 옵티마이저가 올바르게 주입되도록 보장하는 중요한 변경입니다.

왜 이게 좋은가

이번 최적화의 가장 큰 성과는 End-to-end 성능의 비약적인 향상입니다. 테스트 결과에 따르면, 기존 방식 대비 약 40%에서 최대 5배(Transformers v4.57.6 기준)까지 학습 속도가 빨라졌습니다.

• 성능 수치: train_samples_per_second가 2.595에서 3.599로 약 38.7% 증가했습니다.
• 교훈: MoE 모델의 성능 병목은 연산량 자체보다 '데이터의 이동'과 '커널 호출 오버헤드'에 있습니다. Triton과 같은 도구로 커널을 융합(Fusion)하면, 메모리 대역폭을 효율적으로 사용하고 파이썬 인터프리터의 오버헤드를 제거하여 GPU의 연산 성능을 온전히 끌어낼 수 있습니다.

리뷰어 피드백

리뷰 과정에서 transformers v5의 자체 MoE 최적화와 본 PR의 우선순위에 대한 논의가 있었습니다. frozenleaves는 Transformers v5 도입 시 성능은 3-5배 향상되지만, 약 5GB 정도의 추가 VRAM 사용량이 발생할 수 있음을 지적했습니다. 이는 성능과 메모리 사용량 사이의 트레이드오프를 고려해야 함을 시사합니다.

참고 자료

• https://triton-lang.org/main/index.html
• https://pytorch.org/docs/stable/generated/torch.compile.html

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.