Transformer Engine
  • Home

Getting Started

  • Installation
    • Prerequisites
    • Transformer Engine in NGC Containers
    • pip - from PyPI
    • pip - from GitHub
      • Additional Prerequisites
      • Installation (stable release)
      • Installation (development build)
      • Installation (from source)
  • Getting Started
    • Overview
    • Let’s build a Transformer layer!
    • Meet Transformer Engine
    • Fused TE Modules
    • Enabling FP8
  • Frequently Asked Questions (FAQ)
    • FP8 checkpoint compatibility

Python API documentation

  • Common API
    • Format
    • DelayedScaling
  • Framework-specific API
    • pyTorch
      • Linear
        • Linear.forward()
        • Linear.set_tensor_parallel_group()
      • GroupedLinear
        • GroupedLinear.forward()
        • GroupedLinear.set_tensor_parallel_group()
      • LayerNorm
      • RMSNorm
      • LayerNormLinear
        • LayerNormLinear.forward()
        • LayerNormLinear.set_tensor_parallel_group()
      • LayerNormMLP
        • LayerNormMLP.forward()
        • LayerNormMLP.set_tensor_parallel_group()
      • DotProductAttention
        • DotProductAttention.forward()
        • DotProductAttention.set_context_parallel_group()
      • MultiheadAttention
        • MultiheadAttention.forward()
        • MultiheadAttention.set_context_parallel_group()
        • MultiheadAttention.set_tensor_parallel_group()
      • TransformerLayer
        • TransformerLayer.forward()
        • TransformerLayer.set_context_parallel_group()
        • TransformerLayer.set_tensor_parallel_group()
      • InferenceParams
      • CudaRNGStatesTracker
        • CudaRNGStatesTracker.add()
        • CudaRNGStatesTracker.fork()
        • CudaRNGStatesTracker.get_states()
        • CudaRNGStatesTracker.reset()
        • CudaRNGStatesTracker.set_states()
      • fp8_autocast()
      • fp8_model_init()
      • checkpoint()
      • onnx_export()
      • make_graphed_callables()
      • get_cpu_offload_context()
      • moe_permute()
      • moe_unpermute()
    • Jax
      • Pre-defined Variable of Logical Axes
      • Modules
        • TransformerLayerType
        • MeshResource
        • fp8_autocast()
        • update_collections()
        • LayerNorm
        • DenseGeneral
        • LayerNormDenseGeneral
        • LayerNormMLP
        • RelativePositionBiases
        • DotProductAttention
        • MultiHeadAttention
        • TransformerLayer
        • extend_logical_axis_rules()
    • paddle
      • Linear
        • Linear.forward()
      • LayerNorm
      • LayerNormLinear
        • LayerNormLinear.forward()
      • LayerNormMLP
        • LayerNormMLP.forward()
      • FusedScaleMaskSoftmax
        • FusedScaleMaskSoftmax.forward()
      • DotProductAttention
        • DotProductAttention.forward()
      • MultiHeadAttention
        • MultiHeadAttention.forward()
      • TransformerLayer
        • TransformerLayer.forward()
      • fp8_autocast()
      • recompute()

Examples and Tutorials

  • Using FP8 with Transformer Engine
    • Introduction to FP8
      • Structure
      • Mixed precision training - a quick introduction
      • Mixed precision training with FP8
    • Using FP8 with Transformer Engine
      • FP8 recipe
      • FP8 autocasting
      • Handling backward pass
      • Precision
  • Performance Optimizations
    • Multi-GPU training
    • Gradient accumulation fusion
    • FP8 weight caching
  • Accelerating a Hugging Face Llama 2 and Llama 3 models with Transformer Engine
    • Dependencies for this tutorial
    • Table of contents
    • From “Transformer” to “Llama”
    • Hugging Face’s LlamaModel
      • Hugging Face’s LlamaDecoderLayer
        • Self_Attn Layer
        • MLP Layer
    • [Baseline] Running HF LlamaModel (Precision: BF16)
    • [Improvement 1] Replace HF’s LlamaDecoderLayer with TE’s TransformerLayer (Precision: BF16)
      • Transformer Engine’s TransformerLayer
      • TransformerLayer options explained
      • Mapping weights from HF’s LlamaDecoderLayer to TE’s TransformerLayer
    • [Improvement 2] Replace HF’s LlamaDecoderLayer with TE’s TransformerLayer (Precision: FP8)
      • How to run the model in FP8 precision
      • Llama 3 performance results
    • Conclusion

Advanced

  • C/C++ API
    • activation.h
      • NVTE_Activation_Type
        • NVTE_Activation_Type::GELU
        • NVTE_Activation_Type::GEGLU
        • NVTE_Activation_Type::SILU
        • NVTE_Activation_Type::SWIGLU
        • NVTE_Activation_Type::RELU
        • NVTE_Activation_Type::REGLU
        • NVTE_Activation_Type::QGELU
        • NVTE_Activation_Type::QGEGLU
        • NVTE_Activation_Type::SRELU
        • NVTE_Activation_Type::SREGLU
      • nvte_gelu()
      • nvte_silu()
      • nvte_relu()
      • nvte_qgelu()
      • nvte_srelu()
      • nvte_dgelu()
      • nvte_dsilu()
      • nvte_drelu()
      • nvte_dqgelu()
      • nvte_dsrelu()
      • nvte_geglu()
      • nvte_swiglu()
      • nvte_reglu()
      • nvte_qgeglu()
      • nvte_sreglu()
      • nvte_dgeglu()
      • nvte_dswiglu()
      • nvte_dreglu()
      • nvte_dqgeglu()
      • nvte_dsreglu()
    • cast.h
      • nvte_fp8_quantize()
      • nvte_fp8_dequantize()
    • gemm.h
      • nvte_cublas_gemm()
      • nvte_cublas_atomic_gemm()
      • nvte_multi_stream_cublas_gemm()
      • transformer_engine
        • transformer_engine::num_streams
    • fused_attn.h
      • NVTE_QKV_Layout
        • NVTE_QKV_Layout::NVTE_SB3HD
        • NVTE_QKV_Layout::NVTE_SBH3D
        • NVTE_QKV_Layout::NVTE_SBHD_SB2HD
        • NVTE_QKV_Layout::NVTE_SBHD_SBH2D
        • NVTE_QKV_Layout::NVTE_SBHD_SBHD_SBHD
        • NVTE_QKV_Layout::NVTE_BS3HD
        • NVTE_QKV_Layout::NVTE_BSH3D
        • NVTE_QKV_Layout::NVTE_BSHD_BS2HD
        • NVTE_QKV_Layout::NVTE_BSHD_BSH2D
        • NVTE_QKV_Layout::NVTE_BSHD_BSHD_BSHD
        • NVTE_QKV_Layout::NVTE_T3HD
        • NVTE_QKV_Layout::NVTE_TH3D
        • NVTE_QKV_Layout::NVTE_THD_T2HD
        • NVTE_QKV_Layout::NVTE_THD_TH2D
        • NVTE_QKV_Layout::NVTE_THD_THD_THD
      • NVTE_QKV_Layout_Group
        • NVTE_QKV_Layout_Group::NVTE_3HD
        • NVTE_QKV_Layout_Group::NVTE_H3D
        • NVTE_QKV_Layout_Group::NVTE_HD_2HD
        • NVTE_QKV_Layout_Group::NVTE_HD_H2D
        • NVTE_QKV_Layout_Group::NVTE_HD_HD_HD
      • NVTE_QKV_Format
        • NVTE_QKV_Format::NVTE_SBHD
        • NVTE_QKV_Format::NVTE_BSHD
        • NVTE_QKV_Format::NVTE_THD
      • NVTE_Bias_Type
        • NVTE_Bias_Type::NVTE_NO_BIAS
        • NVTE_Bias_Type::NVTE_PRE_SCALE_BIAS
        • NVTE_Bias_Type::NVTE_POST_SCALE_BIAS
        • NVTE_Bias_Type::NVTE_ALIBI
      • NVTE_Mask_Type
        • NVTE_Mask_Type::NVTE_NO_MASK
        • NVTE_Mask_Type::NVTE_PADDING_MASK
        • NVTE_Mask_Type::NVTE_CAUSAL_MASK
        • NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK
        • NVTE_Mask_Type::NVTE_CAUSAL_BOTTOM_RIGHT_MASK
        • NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK
      • NVTE_Fused_Attn_Backend
        • NVTE_Fused_Attn_Backend::NVTE_No_Backend
        • NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen
        • NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen
        • NVTE_Fused_Attn_Backend::NVTE_FP8
      • nvte_get_qkv_layout_group()
      • nvte_get_qkv_format()
      • nvte_get_fused_attn_backend()
      • nvte_fused_attn_fwd_qkvpacked()
      • nvte_fused_attn_bwd_qkvpacked()
      • nvte_fused_attn_fwd_kvpacked()
      • nvte_fused_attn_bwd_kvpacked()
      • nvte_fused_attn_fwd()
      • nvte_fused_attn_bwd()
    • layer_norm.h
      • nvte_layernorm_fwd()
      • nvte_layernorm1p_fwd()
      • nvte_layernorm_bwd()
      • nvte_layernorm1p_bwd()
    • rmsnorm.h
      • nvte_rmsnorm_fwd()
      • nvte_rmsnorm1p_fwd()
      • nvte_rmsnorm_bwd()
      • nvte_rmsnorm1p_bwd()
    • softmax.h
      • nvte_scaled_softmax_forward()
      • nvte_scaled_softmax_backward()
      • nvte_scaled_masked_softmax_forward()
      • nvte_scaled_masked_softmax_backward()
      • nvte_scaled_upper_triang_masked_softmax_forward()
      • nvte_scaled_upper_triang_masked_softmax_backward()
      • nvte_scaled_aligned_causal_masked_softmax_forward()
      • nvte_scaled_aligned_causal_masked_softmax_backward()
    • transformer_engine.h
      • NVTETensor
      • NVTEDType
        • NVTEDType::kNVTEByte
        • NVTEDType::kNVTEInt32
        • NVTEDType::kNVTEInt64
        • NVTEDType::kNVTEFloat32
        • NVTEDType::kNVTEFloat16
        • NVTEDType::kNVTEBFloat16
        • NVTEDType::kNVTEFloat8E4M3
        • NVTEDType::kNVTEFloat8E5M2
        • NVTEDType::kNVTENumTypes
      • nvte_create_tensor()
      • nvte_destroy_tensor()
      • nvte_tensor_type()
      • nvte_tensor_shape()
      • nvte_tensor_data()
      • nvte_tensor_amax()
      • nvte_tensor_scale()
      • nvte_tensor_scale_inv()
      • nvte_tensor_pack_create()
      • nvte_tensor_pack_destroy()
      • NVTEShape
        • NVTEShape::data
        • NVTEShape::ndim
      • NVTETensorPack
        • NVTETensorPack::tensors
        • NVTETensorPack::size
        • NVTETensorPack::MAX_SIZE
      • transformer_engine
        • transformer_engine::DType
        • transformer_engine::TensorWrapper
    • transpose.h
      • nvte_cast_transpose()
      • nvte_transpose()
      • nvte_cast_transpose_dbias()
      • nvte_fp8_transpose_dbias()
      • nvte_multi_cast_transpose()
      • nvte_cast_transpose_dbias_dgelu()
      • nvte_cast_transpose_dbias_dsilu()
      • nvte_cast_transpose_dbias_drelu()
      • nvte_cast_transpose_dbias_dqgelu()
      • nvte_cast_transpose_dbias_dsrelu()
      • nvte_dgeglu_cast_transpose()
      • nvte_dswiglu_cast_transpose()
      • nvte_dreglu_cast_transpose()
      • nvte_dqgeglu_cast_transpose()
      • nvte_dsreglu_cast_transpose()
  • Attention Is All You Need!
    • 1. Attention Backends
      • 1.1 Flash vs. Non-Flash
      • 1.2 flash-attention
      • 1.3 cuDNN Attention
    • 2. Backend Selection
      • 2.1 Debug Information
      • 2.2 User Control
      • 2.3 Example Tests
    • 3. Backend Support
      • 3.1 QKV Layout
      • 3.2 Attention Mask
      • 3.3 Attention Bias
      • 3.4 FP8 Attention
Transformer Engine
  • Search

Privacy Policy | Manage My Privacy | Do Not Sell or Share My Data | Terms of Service | Accessibility | Corporate Policies | Product Security | Contact

© Copyright 2022-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved..