transformerengine

Transformer Engine(TE)은 NVIDIA GPU에서 Transformer 모델을 가속하기 위한 NVIDIA 라이브러리입니다. 핵심 기능은 저정밀도 연산으로, Hopper, Ada 및 Blackwell GPU에서 8비트 부동소수점(FP8)을, Blackwell에서 MXFP8 및 NVFP4 형식을 지원하며, 학습과 추론 모두에서 성능을 개선하고 메모리 사용을 줄이는 것을 목표로 합니다. 또한 Ampere 이후 아키텍처에서 FP16 및 BF16 전반의 최적화도 지원합니다. 이 라이브러리는 일반적인 Transformer 아키텍처를 위한 고도로 최적화된 빌딩 블록과 프레임워크별 코드와 함께 사용할 수 있는 자동 혼합 정밀도 유사 API를 제공합니다. 프레임워크에 구애받지 않는 C++ API도 포함되어 있어 다른 딥러닝 라이브러리가 Transformer용 FP8 지원을 추가할 수 있습니다. TE 모듈은 내부적으로 FP8 학습에 필요한 스케일링 인자와 관련 값을 유지하므로 사용자의 혼합 정밀도 워크플로가 단순해집니다. 프로젝트에서 나열한 주요 특징으로는 FP8 지원으로 Transformer 레이어를 구축하기 위한 사용하기 쉬운 모듈, 융합 커널 최적화, Hopper, Ada 및 Blackwell에서의 FP8 지원, Blackwell에서의 MXFP8 및 NVFP4 지원, Ampere 이후에서의 FP16/BF16 전반 최적화가 있습니다. PyTorch 및 JAX/Flax에 대한 사용 예시가 제공됩니다. PyTorch에서는 transformer_engine.pytorch를 가져오고, 선택한 FP8 형식으로 DelayedScaling 같은 레시피를 생성한 뒤, 순전파를 te.autocast로 감쌉니다. JAX/Flax에서는 te_flax 모듈과 레시피를 사용한 유사한 autocast 사용법이 제시됩니다. 더 자세한 튜토리얼을 위해 시작하기 가이드가 링크되어 있습니다. 설치 옵션에는 NGC Docker 컨테이너(권장), PyTorch 및/또는 JAX용 extras가 포함된 pip 패키지, PyTorch 통합을 위한 conda-forge 패키지, 소스 빌드가 있습니다. 시스템 요구 사항으로는 Blackwell, Hopper, Grace Hopper/Blackwell, Ada 및 Ampere 하드웨어, 공식 OS인 Linux와 제한적인 WSL2 지원, CUDA 12.1+(Blackwell의 경우 12.8+), cuDNN 9.12+, C++17을 지원하는 GCC 9+ 또는 Clang 10+, 권장 Python 3.12가 언급됩니다. FP8 기능은 컴퓨트 능력 8.9 이상이 필요합니다. CUDA_PATH, CUDNN_PATH, CXX, NVTE_FRAMEWORK, MAX_JOBS, NVTE_CUDA_ARCHS 같은 빌드 관련 환경 변수가 문서화되어 있습니다. README는 PyTorch에서의 FlashAttention-2 및 FlashAttention-3 지원을 다루며, 둘 다 있을 경우 FlashAttention-3가 우선되고, FlashAttention-2 컴파일은 메모리를 많이 사용할 수 있다고 설명합니다. 문제 해결 섹션에서는 ABI 호환성 가져오기 오류, 누락된 헤더 또는 라이브러리, 빌드 리소스 문제, 상세 빌드 로깅, cuDNN 하위 라이브러리 로딩 실패를 포함한 UV/가상 환경 문제, JAX FFI 등록 오류를 다룹니다. v1.7의 호환성 깨지는 변경 사항이 문서화되어 있습니다. PyTorch 패딩 마스크 정의가 변경되어 True가 이제 위치를 포함하는 것이 아니라 마스킹함을 의미하게 되어, 프레임워크 전반에서 마스크 의미가 통일되었습니다. 수렴 관련 노트에서는 테스트된 구성에서 FP8 및 MXFP8이 BF16 학습 손실 곡선과 유의미한 차이를 보이지 않았고, 다운스트림 LLM 작업에서 검증되었으며, Mosaic Composer, Alibaba Pai, Megatron Core를 포함한 프레임워크 전반에서 MPT-1.3B, Llama2-7B, LLM-8B, MPT-13B, MoE-16B, Llama2-70B 같은 모델을 나열합니다. 나열된 통합에는 DeepSpeed, Hugging Face Accelerate, Lightning, MosaicML Composer, NVIDIA JAX Toolbox, NVIDIA Megatron-LM, NVIDIA NeMo Megatron Bridge, Amazon SageMaker Model Parallel Library, Levanter, GPT-NeoX, Hugging Face Nanotron이 있습니다. 이 프로젝트는 Apache-2.0 라이선스이며 CONTRIBUTING 가이드를 통해 기여를 환영합니다.