NVIDIA/TransformerEngine

A library for accelerating Transformer models on NVIDIA GPUs, including using 8-bit and 4-bit floating point (FP8 and FP4) precision on Hopper, Ada and Blackwell GPUs, to provide better performance with lower memory utilization in both training and inference.

What it solves

Transformer Engine (TE)은 Transformer 모델(LLM, MoE, 멀티모달 모델 등)을 수천억 개의 파라미터로 확장할 때 발생하는 높은 메모리 및 연산 요구 사항을 해결합니다. 저정밀도 수치 형식을 활용하여 훈련 및 추론 과정 모두에서 메모리 사용량을 줄이고 처리량을 높입니다.

How it works

TE는 Transformer 아키텍처를 위한 고도로 최적화된 빌딩 블록과 융합 커널(fused kernels)을 제공합니다. 사용자가 기존 프레임워크 코드에 저정밀도 형식을 원활하게 통합할 수 있도록 자동 혼합 정밀도 API를 구현합니다.

주요 기술 역량은 다음과 같습니다:

  • Low-Precision Support: Hopper, Ada, Blackwell GPU에서의 8비트 부동 소수점(FP8) 네이티브 지원, 그리고 Blackwell GPU에서의 MXFP8 및 NVFP4를 통한 추가적인 효율성 향상.
  • Framework Integration: PyTorch와 JAX (Flax)를 위한 Python API, 그리고 다른 딥러닝 라이브러리와의 통합을 위한 프레임워크에 구애받지 않는 C++ API를 제공합니다.
  • Automated Scaling: FP8 훈련에 필요한 스케일링 인자(scaling factors)를 내부적으로 관리하여 사용자의 혼합 정밀도 프로세스를 단순화합니다.
  • Optimized Kernels: Mixture-of-Experts (MoE) 및 다양한 병렬화 전략(tensor, sequence, context)을 위한 융합 연산 및 최적화를 포함합니다.

Who it’s for

NVIDIA GPU(Ampere 아키텍처 또는 이후 버전)를 사용하여 대규모 Transformer 모델을 구축하고 훈련하는 연구자 및 엔지니어를 위해 설계되었습니다. 하드웨어 효율성과 훈련 속도를 극대화하고자 하는 사용자에게 적합합니다.

Highlights

  • Broad Precision Support: 최신 NVIDIA GPU에서 FP8, MXFP8, NVFP4를 지원하며, Ampere 및 이후 버전에서 FP16/BF16을 지원합니다.

  • Easy Integration: FP8 지원을 포함한 Transformer 레이어를 구축하기 위한 즉시 사용 가능한 모듈을 제공합니다.

  • C++ API: 딥러닝 라이브러리 개발자를 위한 프레임워크에 구애받지 않는 C++ 라이브러리를 포함합니다.

  • Performance Optimizations: 융합 커널과 PyTorch에서의 FlashAttention-2 및 FlashAttention-3 지원을 특징으로 합니다.