NVIDIA/TransformerEngine

NVIDIA Transformer Engine – 대규모 Transformer를 위한 빠른 저정밀도 학습

무엇인가요 – NVIDIA GPU에서 Transformer 스타일 모델 학습을 위한 고도로 최적화된 C++/CUDA 라이브러리(파이썬 바인딩 포함). PyTorch 및 JAX용으로 FP8, MXFP8, NVFP4, BF16, FP16 학습을 위한 융합 커널과 자동 혼합 정밀도 API를 제공합니다. 스케일링 계수 관리, 양자화 메타데이터, 분산 통신을 처리하여 코드를 변경하지 않고 저정밀도 학습을 채택할 수 있게 합니다.

왜 중요한가요 – 현대의 LLM 및 기타 대규모 모델은 학습에 필요한 컴퓨팅 및 메모리 양에 제한을 받고 있습니다. Hopper/Blackwell/Rubin GPU에서 지원하는 새로운 저정밀도 형식(특히 FP8)으로 대부분의 연산을 이동하고, 여러 연산(어텐션, GEMM, 정규화, MoE 라우팅 등)을 융합함으로써 처리량을 크게 증가시키면서도 모델 품질을 유지할 수 있습니다. 또한 저정밀도 세부 사항을 추상화하여 기존 학습 코드를 그대로 유지할 수 있습니다.

주요 기능

• NVIDIA Hopper, Ada, Blackwell 및 Rubin GPU에서 FP8/MXFP8/NVFP4 학습(이 형식을 지원하는 최신 하드웨어).
• 어텐션, 선형 계층, 레이어 정규화, MoE 라우팅 및 통신(텐서, 시퀀스, 컨텍스트, 전문가 병렬)을 위한 최적화된 융합 커널.
• 자동 스케일링 계수 관리 – amax 기록 및 양자화 메타데이터를 추적하며, PyTorch의 네이티브 AMP와 유사한 autocast 스타일 컨텍스트를 제공.
• 프레임워크 지원 – PyTorch(torch.nn.Module 래퍼) 및 JAX(플랙스 호환 레이어)용 얇고 자연스러운 API.
• 분산 학습 대응 – NCCL과 통합, 통신 오버랩, 제로-코피 전문가 병렬, CUDA 그래프 가속 지원.
• 광범위한 생태계 통합 – Hugging Face Accelerate, PyTorch Lightning, NVIDIA Megatron‑LM, NVIDIA NeMo, Amazon SageMaker Model Parallel Library 등과 즉시 호환.

일반적인 워크플로우

1. 프레임워크용 패키지 설치:

   pip install "transformer_engine[pytorch]"   # 또는 [jax]
   

2. 모델의 레이어를 TE의 대응 버전으로 감싸기 (예: te.Linear, te.LayerNorm, te.MultiHeadAttention).
3. 전방 전파를 autocast 컨텍스트로 감싸서 FP8/FP16/BF16를 자동으로 활성화:

   with transformer_engine.autocast(dtype=torch.float8):
       output = model(input)
   

4. 일반적인 학습 루프를 실행; TE가 스케일링, 양자화, 통신을 내부에서 처리.

시작 및 리소스

• 빠른 시작 문서 – https://docs.nvidia.com/deeplearning/transformer-engine/getting_started/index.html
• API 참조 – PyTorch: https://docs.nvidia.com/deeplearning/transformer-engine/api/pytorch.html ; JAX: https://docs.nvidia.com/deeplearning/transformer-engine/api/jax.html
• 예제 – 리포지토리의 examples/ 폴더에 있는 실행 가능한 노트북 및 스크립트 모음.
• 심층 가이드 – 저정밀도 학습 가이드, FP8 소개, 원본 FP8 논문(arXiv:2209.05433).
• 릴리스 노트 – 최신 v2.19는 Rubin GPU 지원, 하이브리드 양자화, 확장된 FP8 어텐션 추가.

누가 사용해야 하나요

• 최신 NVIDIA GPU를 보유하고 대규모 언어 모델, 비전 Transformer, MoE 모델을 학습하는 연구자 및 엔지니어.
• PyTorch/JAX 코드베이스를 재작성하지 않고도 학습 시간과 GPU 메모리 사용량을 줄이고자 하는 팀.
• Megatron‑LM, Hugging Face, NeMo 등의 인기 LLM 스택 위에서 구축하고 FP8 가속을 활용하고자 하는 사람.

라이선스 – Apache 2.0 (오픈소스).