TransformerEngine
v1.13
QuickStart | 설치 | 사용자 안내서 | 예 | FP8 수렴 | 통합 | 릴리스 노트
Transformer Engine (TE)은 Hopper GPU의 8 비트 플로팅 포인트 (FP8) 정밀도를 사용하여 NVIDIA GPU의 변압기 모델을 가속화하기위한 라이브러리로, 훈련 및 추론 모두에서 더 낮은 메모리 활용도로 더 나은 성능을 제공합니다. TE는 인기있는 변압기 아키텍처를위한 고도로 최적화 된 빌딩 블록 모음과 프레임 워크 별 코드와 함께 원활하게 사용할 수있는 자동 혼합 정밀도와 같은 API를 제공합니다. TE에는 또한 다른 딥 러닝 라이브러리와 통합되어 변압기에 대한 FP8 지원을 가능하게하는 프레임 워크 아그네틱 C ++ API가 포함되어 있습니다.
변압기 모델의 매개 변수 수가 계속 증가함에 따라 Bert, GPT 및 T5와 같은 아키텍처의 훈련 및 추론은 매우 기억력이 높아지고 계산 집약적입니다. 대부분의 딥 러닝 프레임 워크는 기본적으로 FP32로 훈련합니다. 그러나 많은 딥 러닝 모델에 대해 완전한 정확도를 달성하는 데 필수적인 것은 아닙니다. 단일 정제 (FP32)와 모델을 훈련시 더 낮은 정밀 (예 : FP16) 형식을 결합한 혼합-정밀 훈련을 사용하면 FP32 훈련에 비해 정확도가 최소화되는 상당한 속도가 높아집니다. 호퍼 GPU 아키텍처 FP8 정밀도가 도입되었으며, 이는 정확도가 떨어지지 않고 FP16에 비해 성능이 향상됩니다. 모든 주요 딥 러닝 프레임 워크는 FP16을 지원하지만 FP8 지원은 오늘날 프레임 워크에서 기본적으로 사용할 수 없습니다.
TE는 인기있는 대형 언어 모델 (LLM) 라이브러리와 통합되는 API를 제공하여 FP8 지원 문제를 해결합니다. FP8 지원에 필요한 스트러크 및 커널을 포함하여 C ++의 프레임 워크 공유 라이브러리뿐만 아니라 변압기 층을 쉽게 구축 할 수있는 모듈로 구성된 Python API를 제공합니다. TE에서 제공하는 모듈은 내부적으로 스케일링 요소 및 FP8 교육에 필요한 기타 값을 유지하므로 사용자를위한 혼합 정밀 교육을 크게 단순화합니다.
import torch
import transformer_engine . pytorch as te
from transformer_engine . common import recipe
# Set dimensions.
in_features = 768
out_features = 3072
hidden_size = 2048
# Initialize model and inputs.
model = te . Linear ( in_features , out_features , bias = True )
inp = torch . randn ( hidden_size , in_features , device = "cuda" )
# Create an FP8 recipe. Note: All input args are optional.
fp8_recipe = recipe . DelayedScaling ( margin = 0 , fp8_format = recipe . Format . E4M3 )
# Enable autocasting for the forward pass
with te . fp8_autocast ( enabled = True , fp8_recipe = fp8_recipe ):
out = model ( inp )
loss = out . sum ()
loss . backward () import flax
import jax
import jax . numpy as jnp
import transformer_engine . jax as te
import transformer_engine . jax . flax as te_flax
from transformer_engine . common import recipe
BATCH = 32
SEQLEN = 128
HIDDEN = 1024
# Initialize RNG and inputs.
rng = jax . random . PRNGKey ( 0 )
init_rng , data_rng = jax . random . split ( rng )
inp = jax . random . normal ( data_rng , [ BATCH , SEQLEN , HIDDEN ], jnp . float32 )
# Create an FP8 recipe. Note: All input args are optional.
fp8_recipe = recipe . DelayedScaling ( margin = 0 , fp8_format = recipe . Format . HYBRID )
# Enable autocasting for the forward pass
with te . fp8_autocast ( enabled = True , fp8_recipe = fp8_recipe ):
model = te_flax . DenseGeneral ( features = HIDDEN )
def loss_fn ( params , other_vars , inp ):
out = model . apply ({ 'params' : params , ** other_vars }, inp )
return jnp . mean ( out )
# Initialize models.
variables = model . init ( init_rng , inp )
other_variables , params = flax . core . pop ( variables , 'params' )
# Construct the forward and backward function
fwd_bwd_fn = jax . value_and_grad ( loss_fn , argnums = ( 0 , 1 ))
for _ in range ( 10 ):
loss , ( param_grads , other_grads ) = fwd_bwd_fn ( params , other_variables , inp )Transformer Engine을 시작하는 가장 빠른 방법은 NVIDIA GPU Cloud (NGC) 카탈로그에서 Docker Images를 사용하는 것입니다. 예를 들어 NGC Pytorch 컨테이너를 대화식으로 사용하려면
docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:23.10-py3여기서 23.10은 컨테이너 버전입니다. 예를 들어, 2023 년 10 월 릴리스의 경우 23.10.
Transformer 엔진의 최신 안정 버전을 설치하려면
pip install git+https://github.com/NVIDIA/TransformerEngine.git@stable지원되는 딥 러닝 프레임 워크가 설치되어 있는지 자동으로 감지하고 변압기 엔진 지원을 구축합니다. 프레임 워크를 명시 적으로 지정하려면 환경 변수 NVTE_Framework를 쉼표로 구분 된 목록으로 설정하십시오 (예 : NVTE_FRAMEWORK = JAX, PYTORCH, PADDLE).
또는 패키지는 변압기 엔진의 PYPI에서 직접 설치할 수 있습니다.
pip install transformer_engine[pytorch]변압기 엔진에 필요한 파이썬 바인딩을 얻으려면 필요한 프레임 워크는 쉼표로 구분 된 목록 (예 : [Jax, Pytorch, 패들])에서 추가 종속성으로 명시 적으로 지정되어야합니다. 트랜스포머 엔진은 핵심 라이브러리를위한 바퀴와 패들 라디스 확장을 선적합니다. 소스 분포는 JAX 및 Pytorch 확장에 배송됩니다.
설치 안내서를 참조하십시오.
Transformer Engine Release V0.11.0은 성능 향상을 위해 Pytorch의 FlashAttention-2에 대한 지원을 추가합니다.
FlashAttention-2 컴파일은 리소스 집약적이며 많은 양의 RAM (버그 참조)이 필요하며 변압기 엔진을 설치하는 동안 메모리 오류가 발생할 수 있습니다. 환경에서 Max_Jobs = 1을 설정하여 문제를 우회하십시오.
NGC Pytorch 23.08+ 컨테이너에는 FlashAttention-2가 포함됩니다.
Transformer 엔진의 세 가지 프레임 워크에서주의 마스크의 정의와 사용을 통일하기 위해 패딩 마스크는 Pytorch 구현에서 해당 위치를 배제하는 데주의를 기울여 해당 위치의 진정한 의미로 변경되었습니다. v1.7 이후, 모든주의 마스크 유형은 동일한 정의를 따릅니다. 여기서 실제는 해당 위치를 포함하여 해당 위치를 포함하여 해당 위치를 포함하여 해당 위치를 가려줍니다.
이 변화의 예는
# for a batch of 3 sequences where `a`s, `b`s and `c`s are the useful tokens
# and `0`s are the padding tokens,
[a, a, a, 0, 0,
b, b, 0, 0, 0,
c, c, c, c, 0]
# the padding mask for this batch before v1.7 is,
[ True, True, True, False, False,
True, True, False, False, False,
True, True, True, True, False]
# and for v1.7 onwards it should be,
[False, False, False, True, True,
False, False, True, True, True,
False, False, False, False, True]FP8은 다른 모델 아키텍처 및 구성에서 광범위하게 테스트되었으며 FP8과 BF16 훈련 손실 곡선 사이에 유의 한 차이가 없음을 발견했습니다. FP8은 또한 다운 스트림 LLM 작업 (예 : Lambada 및 Wikitext)의 정확성에 대해 검증되었습니다. 다음은 다른 프레임 워크에서 수렴에 대해 테스트 된 모델의 예입니다.
| 모델 | 뼈대 | 원천 |
|---|---|---|
| T5-770m | jax/t5x | https://github.com/nvidia/jax-toolbox/tree/main/rosetta/rosetta/projects/t5x#convergence-and-performance |
| MPT-1.3B | 모자이크 작곡가 | https://www.mosaicml.com/blog/coreweave-nvidia-h100-part-1 |
| GPT-5B | JAX/PAXML | https://github.com/nvidia/jax-toolbox/tree/main/rosetta/rosetta/projects/pax#h100-20 |
| GPT-5B | NEMO 프레임 워크 | 요청시 가능합니다 |
| llama2-7b | 알리바바 파이 | https://mp.weixin.qq.com/s/nqt0ukxlbxyh55031zbdebq |
| T5-11B | jax/t5x | 요청시 가능합니다 |
| MPT-13B | 모자이크 작곡가 | https://www.databricks.com/blog/turbocharged-training-optimizing-databricks-mosaic-astack-fp8 |
| GPT-22B | NEMO 프레임 워크 | 요청시 가능합니다 |
| llama2-70b | 알리바바 파이 | https://mp.weixin.qq.com/s/nqt0ukxlbxyh55031zbdebq |
| GPT-175B | JAX/PAXML | https://github.com/nvidia/jax-toolbox/tree/main/rosetta/rosetta/projects/pax#h100-20 |
Transformer Engine은 다음과 같은 인기있는 LLM 프레임 워크와 통합되었습니다.
우리는 Transformer 엔진에 대한 기여를 환영합니다! Transformer 엔진에 기여하고 풀 요청을 수행하려면 Contributing.rst Guide에 요약 된 지침을 따르십시오.
