PyTorch조회 2회

TorchAO의 MXFP8 MoE 학습 프리미티브를 활용한 Llama4 Scout 학습 가속화

TorchAO의 MXFP8 양자화 기법을 MoE 모델 학습에 적용하여 bfloat16 대비 수렴 성능 저하 없이 학습 속도를 30.2% 개선했다.

섹션별 상세

Llama4 Scout 모델 학습에서 MXFP8을 적용해 bfloat16 대비 30.2%의 속도 향상을 달성했으며, 3,000 스텝 이상의 실험에서 두 방식의 손실 곡선이 거의 일치함을 확인했다.

bfloat16과 mxfp8의 학습 손실 곡선 비교 그래프 — Chart3,000 스텝 이상의 학습 과정에서 mxfp8(주황색)과 bfloat16(하늘색)의 손실 값이 거의 동일하게 수렴하는 것을 보여준다. 이는 저정밀도 양자화 학습이 모델의 정확도에 부정적인 영향을 미치지 않음을 입증하는 핵심 증거이다.

핵심 기술인 _to_mxfp8_then_scaled_grouped_mm API는 입력 활성화 함수와 가중치를 MXFP8로 동적 양자화한 후 스케일링된 Grouped GEMM을 수행하며, 이는 기존 bfloat16 대비 최대 1.8배 빠른 속도를 제공한다.

BF16 Grouped GEMM과 MXFP8 Grouped GEMM의 실행 시간 비교 다이어그램 — Diagram입력을 동적으로 양자화하는 오버헤드가 발생함에도 불구하고, MXFP8 기반의 연산 속도가 훨씬 빨라 전체적인 순 가속(Net speedup)이 발생함을 시각적으로 나타낸다. 이는 커널 최적화의 중요성을 강조한다.

Blackwell GPU의 5세대 Tensor Core(tcgen05)를 활용하기 위해 스케일 팩터를 특수한 블록 레이아웃(128x4 타일)으로 변환하는 커널을 구현하여 하드웨어 가속 성능을 최적화했다.

MXFP8 양자화 데이터와 스케일 팩터의 메모리 레이아웃 구조 — Diagram데이터는 e4m3 형식을 사용하고 스케일 팩터는 e8m0 형식을 사용하며, 각 1x32 블록이 하나의 스케일을 공유하는 구조를 보여준다. 하드웨어 가속을 위한 데이터 정렬 방식을 설명한다.

128x4 타일 메모리 레이아웃에서의 스케일 팩터 배치 상세표 — InfographicBlackwell GPU의 Tensor Core 요구사항에 맞춘 복잡한 블록 레이아웃 변환 방식을 상세히 보여준다. 단순 행 우선 방식에서 특수한 타일 기반 구조로 변환되는 과정을 수치적으로 나타낸다.

MoE의 순전파(Forward)와 역전파(Backward) 과정에서 동적 양자화가 적용되며, 특히 가중치 그래디언트 계산 시에는 수축 차원(Contracting Dimension)에 따른 복잡한 레이아웃 변환을 처리한다.

MoE 순전파 과정에서의 입력 활성화 함수와 전문가 가중치 행렬 곱셈 구조 — Diagram2D 입력 텐서와 3D 전문가 가중치 텐서가 어떻게 오프셋을 기반으로 그룹화되어 연산되는지 설명한다. MXFP8이 적용되는 구체적인 연산 지점을 명시한다.

실험 결과, 출력 임베딩 투영이나 특정 Attention 레이어(Wk, Wv)는 저정밀도에 민감하거나 연산 규모가 작아 MXFP8 적용 대상에서 제외하는 것이 성능과 수렴 안정성 측면에서 유리하다.

기술

TorchAO
TorchTitan
MXFP8
PyTorch
Blackwell Tensor Cores

활용 사례

Llama4 Scout 학습 가속화
대규모 MoE 모델의 효율적 학습
저정밀도 동적 양자화 적용

언급된 리소스

GitHubTorchAO MoE Training Documentation

문서TorchTitan MXFP8 MoE Configs

AI 분석 전체 내용 보기

AI 요약 · 북마크 · 개인 피드 설정 — 무료

출처 · 인용 안내

원문 발행 2026. 03. 13.수집 2026. 03. 13.출처 타입 RSS

인용 시 "요약 출처: AI Trends (aitrends.kr)"를 표기하고, 사실 확인은 원문 보기 기준으로 진행해 주세요. 자세한 기준은 운영 정책을 참고해 주세요.