기술 요약: MXFP4 를 활용한 대규모 언어 모델 사전 학습

문제 제기

대규모 언어 모델 (LLM) 의 급속한 확장으로 하드웨어 효율적인 학습 알고리즘에 대한 요구가 심화되고 있습니다. FP8 은 실현 가능한 저정밀도 대안으로 부상했지만, 4 비트 형식 (예: MXFP4) 을 사용한 학습은 여전히 중요한 과제로 남아 있습니다. 주요 어려움은 활성화 및 기울기 이상치 (outliers) 가 양자화 오차를 증폭시켜 최적화를 불안정하게 만들 수 있다는 점입니다. 초기 연구들은 활성화 이상치를 거칠게 양자화하면 양자화 범위를 지배하여 부동소수점 품질을 심각하게 저하시킬 수 있음을 시사합니다. 혼합 정밀도나 단계별 스케줄을 통해 노이즈를 완화하려는 노력에도 불구하고, 근본적인 질문은 여전히 남아 있습니다: 전방 활성화와 활성화 기울기가 안정적으로 보임에도 불구하고 전체 파이프라인 FP4 학습이 왜 종종 발산하는가? 또한, 학습 파이프라인의 어떤 특정 구성 요소 (전파, 활성화 기울기, 또는 가중치 기울기) 가 수렴 저하에 가장 큰 책임이 있으며, 어떤 안정화 메커니즘이 효과적인지는 명확하지 않습니다.

방법론

저자들은 MLPerf LLM 사전 학습 설정을 준수하며 C4 데이터셋을 사용하여 Llama 3.1–8B의 사전 학습에 대한 통제된 연구를 수행했습니다. 이 연구는 AMD Instinct MI355X GPU에서 네이티브 FP4 텐서 코어 지원을 활용하여 실행되었으며, 이는 소프트웨어 에뮬레이션에 의존했던 이전 연구들과 구별됩니다.

핵심 실험 설계는 MXFP4 활성화에 대한 점진적인 "사다리" 접근법을 포함합니다:

1. 베이스라인: 전체 FP8 학습.
2. 단계별 양자화: 다른 구성 요소들을 FP8 에 고정하면서 세 가지 구성 요소에 걸쳐 MXFP4 를 점진적으로 활성화합니다:
   • 전파 (Fprop) 만.
   • Fprop + 활성화 기울기 (Dgrad).
   • Fprop + Dgrad + 가중치 기울기 (Wgrad).
3. 안정화 평가: 전체 MXFP4 파이프라인 (Fprop + Dgrad + Wgrad) 하에서 안정성을 회복하기 위한 다양한 개입을 테스트했습니다:
   • 확률적 반올림 (Stochastic Rounding): 무작위성을 도입하기 위함.
   • 무작위화된 해다마드 회전 (Randomized Hadamard Rotations): 이상치를 비상관화하기 위함.
   • 결정론적 해다마드 회전 (Deterministic Hadamard Rotations): 구체적으로 $16 \times 16$ (H16) 및 $32 \times 32$ (H32) 행렬을 사용.

지표: 수렴은 3.3의 검증 퍼플렉시티 (perplexity) 에 도달하는 것으로 정의되었습니다. 효율성의 주요 지표는 토큰 오버헤드(FP8 베이스라인 대비 목표 도달에 필요한 토큰의 상대적 증가분) 였습니다. 보조 지표로는 학습 단계 처리량 및 종단간 속도 향상이 포함되었습니다.

주요 기여 및 발견

1. 불안정성 원인 규명

이 연구는 가중치 기울기 (Wgrad) 양자화를 수렴 저하의 주된 원인으로 규명했습니다.

• Fprop 만 MXFP4 로 활성화한 경우 **8–9%**의 토큰 오버헤드가 발생했습니다.
• Fprop + Dgrad를 활성화한 경우 유사한 **10–11%**의 오버헤드가 발생했습니다.
• **전체 파이프라인 (Fprop + Dgrad + Wgrad)**을 안정화 없이 활성화한 경우, 토큰 오버헤드가 **26–27%**로 급증하거나 완전히 발산 ("수렴 안 함") 했습니다.
• 결론: 불안정성은 Fprop 또는 Dgrad 양자화에 내재된 것이 아니라 Wgrad 양자화에 의해 지배됩니다.

2. 안정화 전략 평가

저자들은 확률성 추가 또는 구조적 회전이 전체 파이프라인을 안정화할 수 있는지 테스트했습니다:

• 확률적 반올림: 전체 파이프라인을 안정화하지 못했으며, 실행이 발산했습니다.
• 무작위화된 해다마드 회전: 전체 파이프라인에 적용했을 때 마찬가지로 수렴에 실패했습니다.
• 결정론적 해다마드 회전: 성공적으로 안정적인 최적화를 회복시켰습니다. 전체 MXFP4 파이프라인에 결정론적 해다마드 변환 (구체적으로 H16) 을 적용한 결과, 토큰 오버헤드가 **8–9%**로 감소하여 FP8 베이스라인과 밀접하게 추세를 보였습니다.
• 통찰: 확률적 개입은 이상치 효과를 완화하기보다는 유효 양자화 오차를 증폭시키는 것으로 보입니다. 불안정성은 민감한 기울기 경로를 따라 발생하는 구조적 마이크로 스케일링 오차에 의해 주도되며, 이는 결정론적 회전을 통해 교정될 수 있습니다.

3. 효율성 및 처리량

결정론적 해다마드 (H16) 를 통해 안정성 문제가 해결됨에 따라, 저자들은 FP4 의 실질적인 효율성 향상을 입증했습니다:

• 처리량: MI355X GPU 에서 FP8 베이스라인 대비 MXFP4 + H16 구성은 학습 단계 처리량이 20% 더 높았습니다.
• 종단간 속도 향상: 토큰 오버헤드가 최소화 (8–9%) 되었기 때문에, 처리량 우위는 9–10% 의 종단간 속도 향상으로 이어졌습니다.
• 핵심 통찰: FP4 속도 향상은 "안정성 게이트"를 거칩니다. Wgrad 불안정성을 제어하지 않으면 수렴에 필요한 토큰 수 증가가 저정밀도의 원시 연산 이점을 상쇄합니다.

중요성 및 주장

이 논문은 소프트웨어 에뮬레이션 없이 MI355X 급 하드웨어에서 전체 MXFP4 사전 학습을 최초로 시연했다고 주장합니다. 그 중요성은 FP4 학습 역동에 대한 이해를 전환하는 데 있습니다:

• 실패 메커니즘: FP4 사전 학습의 불안정성은 불충분한 확률성 때문이 아니라, 특히 Wgrad 에서 민감한 기울기 경로를 따라 발생하는 구조적 마이크로 스케일링 오차 때문입니다.
• 실용적 레시피: 결정론적 해다마드 회전은 전체 파이프라인 MXFP4 학습에 필수적이고 효과적인 안정화제이며, 하드웨어 효율성 향상을 실현할 수 있게 합니다.
• 문맥적 한계: 저자들은 겸손하게 이 안정화 레시피가 설정 의존적임을 지적합니다. C4 데이터셋에서 Llama 3.1–8B 에 효과적이지만, 다른 모델, 데이터셋, 또는 파인튜닝 방법에서는 행동이 다를 수 있으며, 추가 증거 없이 보편적이라고 가정해서는 안 됩니다.

요약하자면, 이 연구는 Wgrad 양자화 오차를 확률적 노이즈 주입이 아닌 결정론적 구조적 개입을 통해 제어할 경우, 대규모 사전 학습에 FP8 을 MXFP4 로 대체하는 것이 가능함을 입증했습니다.