[transformers] Hugging Face Transformers: MoE 및 FP8 커널 최적화를 통한 성능 향상

PR 링크: huggingface/transformers#45621 상태: Merged | 변경: +None / -None

들어가며

최근 Hugging Face의 transformers 라이브러리에서 제출된 PR(#45621)은 Mixture-of-Experts (MoE) 모델과 FP8(8비트 부동소수점) 연산의 성능 및 안정성을 크게 향상시키는 중요한 변경사항을 포함하고 있습니다. 특히, MoE 모델에서 사용되는 Expert ID를 처리하는 방식과 FP8 연산을 위한 커널 로딩 및 사용 방식을 개선하여, 기존에 발생할 수 있었던 성능 저하 및 잠재적인 NaN(Not a Number) 오류를 해결하는 데 중점을 두었습니다. 본 글에서는 이 PR의 핵심 변경 사항을 분석하고, 왜 이러한 개선이 중요한지, 그리고 실제 코드 diff를 통해 구체적으로 어떻게 구현되었는지 살펴보겠습니다.

이 PR은 특히 MoE 모델에서 '센티넬(sentinel)' 토큰을 처리하는 로직을 수정하여, 불필요한 연산을 건너뛰도록 함으로써 성능을 높이는 것을 목표로 합니다. 기존에는 Expert ID를 일정 범위로 제한(clamping)하여 센티넬 토큰이 마지막 Expert로 라우팅되는 것처럼 처리했지만, 이 PR에서는 센티넬 토큰을 명시적으로 구분하고 연산에서 제외함으로써 실제 연산량을 줄입니다. 또한, FP8 연산을 위한 Triton 및 DeepGEMM 커널 로딩 방식을 개선하고, 커널 사용 시 발생할 수 있는 잠재적인 메모리 접근 오류 및 NaN 문제를 해결하기 위한 방안을 제시합니다.

코드 분석

1. 커널 로딩 방식 개선 (finegrained_fp8.py, hub_kernels.py)

이 PR은 FP8 연산을 위한 Triton 및 DeepGEMM 커널을 로드하는 방식을 보다 효율적이고 안정적으로 변경했습니다. 기존에는 전역 변수를 사용하여 커널 로딩 상태를 관리했지만, 이제는 functools.cache와 dataclass를 활용하여 커널 로딩을 단 한 번만 수행하고, 로딩된 커널의 엔트리 포인트를 구조화된 방식으로 관리합니다.

Before:

finegrained_fp8.py 파일에서 전역 변수를 사용하여 Triton 커널 로딩 상태를 관리했습니다.

# Lazily-loaded finegrained-fp8 Triton kernel functions (populated by _load_triton_kernel)
triton_fp8_matmul = None
triton_fp8_act_quant = None
triton_batched_fp8_matmul = None
triton_grouped_fp8_matmul = None
# _triton_available: None = not yet attempted, True = loaded, False = failed (won't retry)
_triton_available = None

# Lazily-loaded DeepGEMM kernel functions (populated by _load_deepgemm_kernel)
deepgemm_fp8_matmul = None
deepgemm_grouped_fp8_matmul = None
deepgemm_per_token_cast_to_fp8 = None
# _deepgemm_available: None = not yet attempted, True = loaded, False = failed (won't retry)
_deepgemm_available = None

def _load_triton_kernel():
    # ... (기존 로딩 로직)

def _load_deepgemm_kernel():
    # ... (기존 로딩 로직)

After:

@functools.cache 데코레이터와 @dataclass를 사용하여 커널 로딩 함수를 재정의하고, 로딩된 커널의 함수들을 FineGrainedFP8 및 DeepGEMM 데이터 클래스에 묶어 반환합니다. 이는 커널 로딩 시도 횟수를 제한하고, 상태 관리를 단순화하며, 코드의 가독성을 높입니다.

@dataclass(frozen=True)
class FineGrainedFP8:
    fp8_matmul: Callable
    fp8_act_quant: Callable
    batched_fp8_matmul: Callable
    grouped_fp8_matmul: Callable

@functools.cache
def _load_finegrained_fp8_kernel() -> FineGrainedFP8:
    # ... (새로운 로딩 로직)
    return FineGrainedFP8(...)

@dataclass(frozen=True)
class DeepGEMM:
    fp8_matmul: Callable
    grouped_fp8_matmul: Callable
    per_token_cast_to_fp8: Callable

@functools.cache
def _load_deepgemm_kernel() -> DeepGEMM:
    # ... (새로운 로딩 로직)
    return DeepGEMM(...)

is_kernels_available() 함수를 사용하여 kernels 패키지 설치 여부를 먼저 확인하고, lazy_load_kernel 함수를 통해 커널을 로드합니다. 필요한 함수들이 누락된 경우 ImportError를 발생시켜 문제점을 명확히 합니다.

2. MoE 센티넬 토큰 처리 개선 (moe.py)

MoE 모델에서 센티넬 토큰은 실제 연산에 참여하지 않아야 하지만, 기존 구현에서는 Expert ID를 일정 범위로 제한(clamping)하는 과정에서 센티넬 토큰이 마지막 Expert로 라우팅되는 것처럼 처리되었습니다. 이 PR은 센티넬 토큰을 명시적으로 식별하고 연산에서 제외함으로써 불필요한 계산을 줄입니다.

PR 설명에 따르면, 센티넬 토큰을 처리하는 방식 변경으로 인해 다음과 같은 성능 향상이 관찰되었습니다:

• offsets[-1] = 16384 (100%): 22.7 ms/iter (1.00x)
• offsets[-1] = 8192 (50%): 13.0 ms/iter (0.57x)
• offsets[-1] = 2048 (12.5%): 4.5 ms/iter (0.20x)

이는 센티넬 토큰이 연산에서 제외될 때마다 상당한 속도 향상을 가져옴을 보여줍니다. 센티넬 토큰은 실제로는 라우팅되지 않아야 하므로, 이들을 마지막 Expert로 보내는 것은 잘못된 동작이며 연산 낭비입니다. 이 PR에서는 센티넬 토큰을 히스토그램 계산 시 제외하거나, 정렬(sorting) 과정에서 가장 큰 값으로 처리하여 자연스럽게 제외되도록 합니다.

3. FP8 GEMM 연산 시 NaN 오류 수정 (finegrained_fp8.py)

리뷰 과정에서 AmineDiro는 MoE 모델의 grouped_mm_experts 연산에서 NaN 오류가 발생하는 문제를 지적했습니다. 특히, 커널이 센티넬 토큰에 해당하는 출력(gradient)을 초기화하지 않고 남겨두어, 역전파 과정에서 NaN 값이 전파되는 현상이 발견되었습니다.

문제 상황 재현 코드 (AmineDiro 제공):

# ... (코드 생략)

run_mode(

## 참고 자료
- https://pytorch.org/docs/stable/generated/torch.compile.html
- https://github.com/huggingface/transformers/blob/main/src/transformers/integrations/finegrained_fp8.py
- https://github.com/huggingface/transformers/blob/main/src/transformers/integrations/moe.py
- https://github.com/huggingface/kernels/blob/main/kernels_community/finegrained_fp8/__init__.py
- https://github.com/huggingface/kernels/blob/main/kernels_community/deep-gemm/__init__.py

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