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.

NVIDIA Transformer Engine

개요 – 고도로 최적화된 커널과 자동 혼합 정밀도 API를 제공하여 NVIDIA GPU에서 Transformer 스타일의 신경망을 가속화하는 라이브러리입니다. FP8, MXFP8, NVFP4와 같은 저정밀도 형식을 사용하여 대규모 언어 모델(LLM), Mixture-of-Experts(MoE) 모델 및 멀티모달 Transformer를 학습하고 실행할 수 있으며, 이를 통해 정확도를 FP16/BF16과 유사하게 유지하면서 메모리 사용량을 줄이고 처리량을 높일 수 있습니다.

주요 기능

FP8 우선 지원: Hopper, Ada 및 Blackwell GPU에서의 FP8 지원과 Blackwell에서의 새로운 MXFP8/NVFP4 형식 지원.
프레임워크에 독립적인 C++ 코어: PyTorch 및 JAX/Flax를 위한 가벼운 Python 바인딩 제공.
Fused kernels: 여러 연산을 단일 GPU 실행으로 결합하여 속도를 높임 (예: FlashAttention-2/-3).
자동 스케일링 계수 처리: 사용자가 te.autocast를 활성화하면 라이브러리가 FP8 학습에 필요한 계산을 관리함.
주요 LLM 스택 통합 훅: DeepSpeed, Hugging Face Accelerate, PyTorch Lightning, MosaicML Composer 등.
병렬화 패턴 지원: Tensor, Sequence, Context 병렬성 및 MoE 워크로드 지원.

전형적인 워크플로우 (PyTorch 예시)

import torch, transformer_engine.pytorch as te
from transformer_engine.common import recipe

model = te.Linear(768, 3072, bias=True)
inp   = torch.randn(2048, 768, device='cuda')
fp8_recipe = recipe.DelayedScaling(margin=0, fp8_format=recipe.Format.E4M3)

with te.autocast(enabled=True, recipe=fp8_recipe):
    out = model(inp)
loss = out.sum(); loss.backward()

JAX/Flax에서도 유사한 패턴을 사용할 수 있으며, te.autocast가 포워드 패스를 래핑합니다.

설치

Docker (권장): NVIDIA NGC 컨테이너 (nvcr.io/nvidia/pytorch:26.01-py3 또는 nvcr.io/nvidia/jax:26.01-py3)를 풀(pull)하세요. 엔진은 /opt/transformerengine 내에 사전 설치되어 있습니다.
pip: pip install --no-build-isolation transformer_engine[pytorch] (또는 [jax] / 둘 다). 일반적인 CUDA, cuDNN, C++17 툴체인을 사용하여 소스 설치가 가능합니다.
conda: conda install -c conda-forge transformer-engine-torch (JAX 지원 예정).

사용 시점

GPU 메모리 또는 연산 대역폭 제한으로 인해 학습이 어려운 LLM 또는 MoE 모델을 학습할 때.
지연 시간(latency)과 메모리 사용량이 중요한 추론 파이프라인을 배포할 때 (특히 FP8 하드웨어를 지원하는 Hopper/Blackwell GPU 사용 시).
모델 코드를 재작성하지 않고 저정밀도 업그레이드를 즉시 적용하고자 하는 PyTorch 또는 JAX 사용 프로젝트.

제한 사항 / 주의 사항

FP8 기능은 Compute Capability 8.9 이상의 GPU (Ada/Hopper/Blackwell)가 필요합니다.
소스 빌드는 메모리 집약적일 수 있습니다 (FlashAttention-2 컴파일 등) – OOM 발생 시 MAX_JOBS=1로 설정하세요.
PyTorch와 엔진 간의 ABI 불일치는 임포트 오류를 일으킬 수 있습니다. 두 라이브러리가 동일한 C++ ABI로 빌드되었는지 확인하십시오.

리소스

전체 사용자 가이드: https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/index.html
예제 노트북: https://github.com/NVIDIA/TransformerEngine/tree/main/examples
최신 뉴스 (2026년 6월): MoE 처리량, NVFP4 및 Nemotron-3 Ultra 모델 관련.

Transformer Engine은 최첨단 저정밀도 하드웨어를 사용하여 현대적인 Transformer 워크로드를 가속화하는 데 중점을 둔 NVIDIA 유지 관리 오픈 소스 프로젝트입니다.

NVIDIA/TransformerEngine

NVIDIA Transformer Engine

관련