From ed4d06908dbe3a584aca794d3995075a53f9151f Mon Sep 17 00:00:00 2001 From: Przemek Tredak Date: Mon, 16 Mar 2026 11:37:38 -0700 Subject: [PATCH 01/20] First pass Signed-off-by: Przemek Tredak --- docs/developer/architecture_overview.rst | 161 +++++++++++ .../developer/attention/backend_selection.rst | 135 +++++++++ docs/developer/attention/backends.rst | 130 +++++++++ docs/developer/attention/context_parallel.rst | 87 ++++++ .../attention/fused_attn_kernels.rst | 94 +++++++ .../attention/img/attention_backends.svg | 124 +++++++++ .../attention/img/context_parallel_ring.svg | 102 +++++++ docs/developer/attention/index.rst | 21 ++ docs/developer/attention/jax_attention.rst | 71 +++++ .../developer/attention/pytorch_attention.rst | 105 +++++++ docs/developer/cpp_core/build_system.rst | 109 ++++++++ .../cpp_core/img/scaling_modes_comparison.svg | 89 ++++++ docs/developer/cpp_core/img/tensor_struct.svg | 114 ++++++++ docs/developer/cpp_core/index.rst | 19 ++ docs/developer/cpp_core/kernel_areas.rst | 136 +++++++++ docs/developer/cpp_core/scaling_modes.rst | 185 +++++++++++++ docs/developer/cpp_core/type_system.rst | 236 ++++++++++++++++ .../distributed/comm_gemm_overlap.rst | 129 +++++++++ .../distributed/fsdp_integration.rst | 84 ++++++ .../distributed/img/column_parallel.svg | 110 ++++++++ .../distributed/img/comm_gemm_overlap.svg | 103 +++++++ .../distributed/img/row_parallel.svg | 121 +++++++++ docs/developer/distributed/index.rst | 19 ++ .../developer/distributed/jax_distributed.rst | 104 +++++++ .../distributed/sequence_parallel.rst | 94 +++++++ .../developer/distributed/tensor_parallel.rst | 118 ++++++++ docs/developer/img/architecture_layers.svg | 85 ++++++ docs/developer/img/data_flow_forward.svg | 72 +++++ docs/developer/img/linear_e2e_flow.svg | 125 +++++++++ docs/developer/index.rst | 43 +++ .../img/jax_vs_pytorch_binding.svg | 71 +++++ docs/developer/jax_frontend/index.rst | 23 ++ docs/developer/jax_frontend/module_system.rst | 113 ++++++++ docs/developer/jax_frontend/quantization.rst | 112 ++++++++ docs/developer/jax_frontend/sharding.rst | 100 +++++++ .../jax_frontend/xla_ffi_primitives.rst | 134 +++++++++ docs/developer/linear_walkthrough.rst | 257 ++++++++++++++++++ .../pytorch_frontend/autograd_integration.rst | 113 ++++++++ .../pytorch_frontend/cpp_extensions.rst | 125 +++++++++ .../pytorch_frontend/img/ops_fusion.svg | 67 +++++ .../img/pytorch_module_hierarchy.svg | 110 ++++++++ .../img/transformer_layer.svg | 98 +++++++ docs/developer/pytorch_frontend/index.rst | 20 ++ .../pytorch_frontend/module_hierarchy.rst | 162 +++++++++++ .../pytorch_frontend/ops_framework.rst | 102 +++++++ .../pytorch_frontend/transformer_layer.rst | 84 ++++++ .../quantization/adding_new_type.rst | 128 +++++++++ .../quantization/class_hierarchy.rst | 195 +++++++++++++ .../quantization/img/quantizer_hierarchy.svg | 182 +++++++++++++ .../quantization/img/rowwise_columnwise.svg | 141 ++++++++++ docs/developer/quantization/index.rst | 25 ++ .../quantization/rowwise_columnwise.rst | 116 ++++++++ .../quantization/scaling_recipes.rst | 136 +++++++++ 53 files changed, 5739 insertions(+) create mode 100644 docs/developer/architecture_overview.rst create mode 100644 docs/developer/attention/backend_selection.rst create mode 100644 docs/developer/attention/backends.rst create mode 100644 docs/developer/attention/context_parallel.rst create mode 100644 docs/developer/attention/fused_attn_kernels.rst create mode 100644 docs/developer/attention/img/attention_backends.svg create mode 100644 docs/developer/attention/img/context_parallel_ring.svg create mode 100644 docs/developer/attention/index.rst create mode 100644 docs/developer/attention/jax_attention.rst create mode 100644 docs/developer/attention/pytorch_attention.rst create mode 100644 docs/developer/cpp_core/build_system.rst create mode 100644 docs/developer/cpp_core/img/scaling_modes_comparison.svg create mode 100644 docs/developer/cpp_core/img/tensor_struct.svg create mode 100644 docs/developer/cpp_core/index.rst create mode 100644 docs/developer/cpp_core/kernel_areas.rst create mode 100644 docs/developer/cpp_core/scaling_modes.rst create mode 100644 docs/developer/cpp_core/type_system.rst create mode 100644 docs/developer/distributed/comm_gemm_overlap.rst create mode 100644 docs/developer/distributed/fsdp_integration.rst create mode 100644 docs/developer/distributed/img/column_parallel.svg create mode 100644 docs/developer/distributed/img/comm_gemm_overlap.svg create mode 100644 docs/developer/distributed/img/row_parallel.svg create mode 100644 docs/developer/distributed/index.rst create mode 100644 docs/developer/distributed/jax_distributed.rst create mode 100644 docs/developer/distributed/sequence_parallel.rst create mode 100644 docs/developer/distributed/tensor_parallel.rst create mode 100644 docs/developer/img/architecture_layers.svg create mode 100644 docs/developer/img/data_flow_forward.svg create mode 100644 docs/developer/img/linear_e2e_flow.svg create mode 100644 docs/developer/index.rst create mode 100644 docs/developer/jax_frontend/img/jax_vs_pytorch_binding.svg create mode 100644 docs/developer/jax_frontend/index.rst create mode 100644 docs/developer/jax_frontend/module_system.rst create mode 100644 docs/developer/jax_frontend/quantization.rst create mode 100644 docs/developer/jax_frontend/sharding.rst create mode 100644 docs/developer/jax_frontend/xla_ffi_primitives.rst create mode 100644 docs/developer/linear_walkthrough.rst create mode 100644 docs/developer/pytorch_frontend/autograd_integration.rst create mode 100644 docs/developer/pytorch_frontend/cpp_extensions.rst create mode 100644 docs/developer/pytorch_frontend/img/ops_fusion.svg create mode 100644 docs/developer/pytorch_frontend/img/pytorch_module_hierarchy.svg create mode 100644 docs/developer/pytorch_frontend/img/transformer_layer.svg create mode 100644 docs/developer/pytorch_frontend/index.rst create mode 100644 docs/developer/pytorch_frontend/module_hierarchy.rst create mode 100644 docs/developer/pytorch_frontend/ops_framework.rst create mode 100644 docs/developer/pytorch_frontend/transformer_layer.rst create mode 100644 docs/developer/quantization/adding_new_type.rst create mode 100644 docs/developer/quantization/class_hierarchy.rst create mode 100644 docs/developer/quantization/img/quantizer_hierarchy.svg create mode 100644 docs/developer/quantization/img/rowwise_columnwise.svg create mode 100644 docs/developer/quantization/index.rst create mode 100644 docs/developer/quantization/rowwise_columnwise.rst create mode 100644 docs/developer/quantization/scaling_recipes.rst diff --git a/docs/developer/architecture_overview.rst b/docs/developer/architecture_overview.rst new file mode 100644 index 0000000000..85c044408f --- /dev/null +++ b/docs/developer/architecture_overview.rst @@ -0,0 +1,161 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +.. _architecture-overview: + +Architecture Overview +===================== + +Transformer Engine is structured as a layered system: framework-specific Python frontends +sit on top of a shared C++ core that dispatches to optimized CUDA kernels. + +.. figure:: ./img/architecture_layers.svg + :align: center + :width: 80% + + The four layers of Transformer Engine. + +.. + Diagram description for ``architecture_layers.svg``: + Four horizontal bands stacked vertically. + Top band: "Python API" with two sub-boxes "PyTorch Frontend" and "JAX Frontend". + Second band: "Binding Layer" with sub-boxes "pybind11 (PyTorch)" and "XLA FFI (JAX)". + Third band: "C++ Core Library" (single box spanning full width). + Bottom band: "CUDA Kernels" with sub-boxes for major areas: GEMM, Normalization, + Activation, Attention, Cast/Transpose, Fused Rope, Fused Softmax. + Arrows flow downward between bands. A side annotation shows "transformer_engine/common/" + pointing at the bottom two bands. + +Layer Diagram +------------- + +**Python API** (``transformer_engine/pytorch/``, ``transformer_engine/jax/``) + Framework-specific modules, autograd functions, quantized tensor types, and distributed + utilities. This is what users import and interact with. + +**Binding Layer** (``transformer_engine/pytorch/csrc/``, ``transformer_engine/jax/csrc/``) + pybind11 extensions (PyTorch) and XLA custom-call registrations (JAX) that bridge + Python to C++. These translate framework tensor types into the C API's ``NVTETensor``. + +**C++ Core** (``transformer_engine/common/``) + Framework-agnostic library implementing the actual computation logic. Exposes a C API + (``transformer_engine.h``) for stability across framework bindings. + +**CUDA Kernels** (within ``transformer_engine/common/``) + Hand-optimized CUDA kernels and cuDNN/cuBLASLt integrations organized by functional + area. + +Directory Map +------------- + +.. code-block:: text + + transformer_engine/ + ├── common/ # C++ core (framework-agnostic) + │ ├── include/transformer_engine/ # Public C API headers + │ │ ├── transformer_engine.h # Base types: NVTEDType, NVTETensor, NVTEScalingMode + │ │ ├── fused_attn.h # Attention C API + │ │ ├── gemm.h # GEMM C API + │ │ └── ... # Other kernel APIs + │ ├── common.h # C++ types: DType, SimpleTensor, Tensor + │ ├── gemm/ # cuBLASLt GEMM kernels + │ ├── fused_attn/ # Fused attention (cuDNN + custom) + │ ├── normalization/ # LayerNorm, RMSNorm kernels + │ ├── activation/ # GeLU, SiLU, etc. + │ ├── cast/ # Quantization cast kernels + │ ├── transpose/ # Transpose + cast fusion + │ ├── fused_rope/ # Rotary positional embeddings + │ ├── fused_softmax/ # Scaled softmax variants + │ ├── fused_router/ # MoE router kernels + │ └── comm_gemm/ # Communication-GEMM overlap + │ + ├── pytorch/ # PyTorch frontend + │ ├── module/ # nn.Module subclasses (Linear, LN, Attention, etc.) + │ ├── tensor/ # Quantized tensor implementations + │ ├── attention/ # Attention backends and dispatch + │ ├── ops/ # Op fusion framework + │ ├── cpp_extensions/ # Python wrappers for C++ extensions + │ ├── csrc/ # pybind11 C++ source + │ ├── distributed.py # TP/SP/FSDP utilities + │ ├── quantized_tensor.py # Base quantization classes + │ └── transformer.py # TransformerLayer composition + │ + └── jax/ # JAX frontend + ├── flax/ # Flax nn.Module wrappers + ├── cpp_extensions/ # XLA FFI primitive registration + ├── quantize/ # JAX quantizer and ScaledTensor + ├── attention.py # JAX attention entry points + ├── sharding.py # Mesh and sharding utilities + └── distributed.py # JAX distributed collectives + +Data Flow: Forward Pass +----------------------- + +A typical forward pass through a quantized linear layer follows this path: + +.. figure:: ./img/data_flow_forward.svg + :align: center + :width: 90% + + Data flow through a quantized linear forward pass. + +.. + Diagram description for ``data_flow_forward.svg``: + Left-to-right flow diagram. + 1. Box "Input (BF16/FP16)" with arrow to + 2. Box "Quantize (cast kernel)" which outputs "FP8 data + scale" arrow to + 3. Box "GEMM (cuBLASLt)" which also receives "Weight (FP8 + scale)" from above, arrow to + 4. Box "Output (BF16/FP16)". + A dashed feedback arrow from GEMM back to Quantize labeled "amax history (delayed scaling)". + Below the main flow, a note: "Scale computation depends on NVTEScalingMode". + +1. **Input arrives** in high precision (BF16 or FP16). +2. **Quantization** casts the input to FP8 (or other low-precision format), producing + quantized data and associated scaling factors. The specific cast kernel and scale + granularity depend on the :ref:`scaling mode `. +3. **GEMM** executes in low precision via cuBLASLt, consuming quantized activations and + weights along with their scale inverses. +4. **Output** is produced in high precision. + +For delayed tensor scaling, an amax (absolute maximum) feedback loop maintains scaling +factor history across iterations. + +The backward pass follows the same pattern twice: once for the activation gradient (dgrad) +and once for the weight gradient (wgrad). The wgrad path uses *columnwise* quantized data +to satisfy cuBLASLt's transposed layout requirements — see +:doc:`quantization/rowwise_columnwise` for details. + +Key Design Decisions +-------------------- + +**C API boundary** + The C++ core exposes a C API (not C++) for ABI stability. This means framework bindings + interact with opaque ``NVTETensor`` handles rather than C++ objects directly. The C++ + ``Tensor`` and ``SimpleTensor`` structs in ``common.h`` are internal implementation + details. + +**Dual type systems** + The C API uses ``NVTEDType`` (a plain C enum), while the C++ core uses + ``DType`` (a C++ ``enum class`` in the ``transformer_engine`` namespace). These are + related but not implicitly convertible — see :doc:`cpp_core/type_system` for conversion + patterns. + +**Framework-agnostic kernels** + All CUDA kernels live in ``common/`` and know nothing about PyTorch or JAX. Framework + frontends are responsible for converting their tensor types to ``NVTETensor`` before + calling into the C API. + +**Quantizer pattern** + Quantization is not a single function but a stateful builder (``Quantizer``) that + produces ``QuantizedTensorStorage`` (lightweight, no autograd) or ``QuantizedTensor`` + (full torch.Tensor subclass). This separation keeps internal kernel code free from + autograd overhead — see :doc:`quantization/class_hierarchy`. + +Next Steps +---------- + +- :doc:`linear_walkthrough` — End-to-end trace through a Linear forward and backward pass +- :doc:`cpp_core/index` — C++ core type system, kernels, and build system +- :doc:`quantization/index` — Quantization architecture and adding new types diff --git a/docs/developer/attention/backend_selection.rst b/docs/developer/attention/backend_selection.rst new file mode 100644 index 0000000000..006d3c48fd --- /dev/null +++ b/docs/developer/attention/backend_selection.rst @@ -0,0 +1,135 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +.. _backend-selection: + +Backend Selection +================= + +Transformer Engine automatically selects the best attention backend based on the +hardware, input configuration, and user preferences. This page documents the selection +logic and how to override it. + +Selection Flow +-------------- + +The backend is selected at the start of each ``DotProductAttention.forward()`` call. +The logic (in ``transformer_engine/pytorch/attention/dot_product_attention/``) follows +a priority order: + +.. code-block:: text + + 1. Check user override (env vars, constructor args) + 2. Check FP8 — if FP8 enabled, try cuDNN FP8 fused + 3. Check cuDNN F16 fused — if supported by config + 4. Check FlashAttention — if installed and supported + 5. Fall back to unfused attention + +At each step, the logic checks whether the backend supports the current configuration: +head dimension, sequence length, mask type, dropout, GQA groups, etc. + +Environment Variables +--------------------- + +.. list-table:: + :header-rows: 1 + :widths: 35 65 + + * - Variable + - Description + * - ``NVTE_FUSED_ATTN`` + - ``0`` to disable cuDNN fused attention, ``1`` to enable (default: ``1``) + * - ``NVTE_FLASH_ATTN`` + - ``0`` to disable FlashAttention, ``1`` to enable (default: ``1``) + * - ``NVTE_FUSED_ATTN_BACKEND`` + - Force a specific cuDNN fused attention sub-backend (integer ID) + * - ``NVTE_FUSED_ATTN_FORCE_WORKSPACE_OPT`` + - Force workspace optimization for fused attention + +Constructor Override +-------------------- + +``DotProductAttention`` accepts a ``backend`` parameter: + +.. code-block:: python + + attn = te.DotProductAttention( + num_attention_heads=32, + kv_channels=128, + attention_type="self", + backend="flash_attention", # Force FlashAttention + ) + +Feature Compatibility Matrix +----------------------------- + +Not all backends support all features. Key restrictions: + +.. list-table:: + :header-rows: 1 + :widths: 30 15 15 15 15 + + * - Feature + - Unfused + - Flash + - cuDNN F16 + - cuDNN FP8 + * - Causal mask + - Yes + - Yes + - Yes + - Yes + * - Padding mask + - Yes + - Yes + - Yes + - Yes + * - Sliding window + - No + - Yes + - Yes + - Yes + * - GQA/MQA + - Yes + - Yes + - Yes + - Yes + * - Dropout + - Yes + - Yes + - Yes + - Yes + * - FP8 Q/K/V + - No + - No + - No + - Yes + * - Context parallel + - No + - Limited + - Yes + - Yes + * - Head dim > 256 + - Yes + - No + - Yes + - Limited + +Debugging Selection +------------------- + +To see which backend was selected, set: + +.. code-block:: bash + + NVTE_DEBUG=1 + +This logs the backend selection decision and the reason for any fallbacks. + +See Also +-------- + +- :doc:`backends` — Detailed description of each backend +- :doc:`context_parallel` — How CP affects backend selection diff --git a/docs/developer/attention/backends.rst b/docs/developer/attention/backends.rst new file mode 100644 index 0000000000..0a5260851c --- /dev/null +++ b/docs/developer/attention/backends.rst @@ -0,0 +1,130 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +.. _attention-backends: + +Attention Backends +================== + +Transformer Engine supports multiple attention backends, each optimized for different +hardware, precision, and feature combinations. + +.. figure:: ./img/attention_backends.svg + :align: center + :width: 80% + + Taxonomy of attention backends with architecture annotations. + +.. + Diagram description for ``attention_backends.svg``: + Tree/taxonomy diagram: + "Attention Backends" + ├── "Unfused" (FlashAttention via torch) + │ └── Any GPU, no FP8, fallback + ├── "Flash Attention" (Tri Dao's FlashAttention) + │ └── Ampere+, BF16/FP16, fastest for non-FP8 + └── "Fused Attention" (cuDNN-based) + ├── "F16 Fused" — cuDNN with BF16/FP16 + │ └── Ampere+, supports arbitrary seq len + ├── "FP8 Fused" — cuDNN with FP8 + │ └── Hopper+, best perf with FP8 training + └── "Custom" — TE's own CUDA kernels + └── Legacy, limited feature support + +Overview +-------- + +.. list-table:: + :header-rows: 1 + :widths: 20 15 15 15 35 + + * - Backend + - Precision + - Min Arch + - cuDNN Required + - Notes + * - Unfused + - BF16/FP16 + - Any + - No + - Fallback; uses PyTorch's native attention or manual QK^TV + * - FlashAttention + - BF16/FP16 + - Ampere + - No + - Tri Dao's FlashAttention; fast, memory-efficient + * - Fused (F16) + - BF16/FP16 + - Ampere + - Yes + - cuDNN graph-based; supports all mask types + * - Fused (FP8) + - FP8 + - Hopper + - Yes + - FP8 Q/K/V with cuDNN; best throughput for FP8 training + * - Custom/THD + - BF16/FP16 + - Hopper + - No + - TE custom kernels; used for specific context-parallel patterns + +Backend Details +--------------- + +Unfused Attention +^^^^^^^^^^^^^^^^^ + +The simplest backend. Computes attention as separate matrix multiplications: + +.. code-block:: text + + scores = Q @ K^T / sqrt(d) + scores = mask(scores) + weights = softmax(scores) + output = weights @ V + +Used as a fallback when no fused backend supports the requested configuration. Also +useful for debugging (produces identical numerics to textbook attention). + +FlashAttention +^^^^^^^^^^^^^^ + +Uses Tri Dao's FlashAttention algorithm, which tiles the computation to avoid +materializing the full attention matrix. Key properties: + +- O(N) memory instead of O(N²) for the attention matrix. +- IO-aware: optimized for GPU memory hierarchy. +- Supports causal masks, variable-length sequences. +- No FP8 support (BF16/FP16 only). + +cuDNN Fused Attention (F16 and FP8) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Uses cuDNN's graph API to fuse the entire attention computation (QK^T, scaling, masking, +softmax, dropout, V multiply) into a single cuDNN operation. Key properties: + +- **F16 variant**: BF16/FP16 inputs, Ampere and later. +- **FP8 variant**: FP8 Q/K/V inputs with per-tensor or block scales, Hopper and later. +- Supports: arbitrary masks (causal, padding, sliding window), GQA/MQA, dropout. +- Requires cuDNN 9.3+ for FP8. + +The cuDNN backend is the recommended choice for production FP8 training. + +Custom/THD Kernels +^^^^^^^^^^^^^^^^^^ + +TE's own CUDA attention kernels, used for specific context-parallel patterns +(THD = Token-Head-Dimension layout). These handle cases where cuDNN doesn't support the +required distributed communication pattern. + +See :doc:`fused_attn_kernels` for the C++ kernel organization. + +See Also +-------- + +- :doc:`backend_selection` — How the backend is chosen at runtime +- :doc:`pytorch_attention` — PyTorch DotProductAttention and MultiheadAttention modules +- :doc:`fused_attn_kernels` — C++ kernel organization and cuDNN integration diff --git a/docs/developer/attention/context_parallel.rst b/docs/developer/attention/context_parallel.rst new file mode 100644 index 0000000000..c435a2fda6 --- /dev/null +++ b/docs/developer/attention/context_parallel.rst @@ -0,0 +1,87 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +.. _context-parallel: + +Context Parallelism +=================== + +Context Parallelism (CP) distributes the sequence dimension across multiple GPUs, +enabling training with very long sequences that don't fit in a single GPU's memory. + +.. figure:: ./img/context_parallel_ring.svg + :align: center + :width: 70% + + Ring of 4 GPUs passing KV chunks in context parallelism. + +.. + Diagram description for ``context_parallel_ring.svg``: + Four GPU boxes arranged in a ring (square layout). + GPU 0 has "seq[0:L/4]", GPU 1 has "seq[L/4:L/2]", GPU 2 has "seq[L/2:3L/4]", + GPU 3 has "seq[3L/4:L]". + Arrows form a ring: GPU0 → GPU1 → GPU2 → GPU3 → GPU0. + Each arrow is labeled "KV chunk". + Center text: "Each GPU computes local Q × all KV via ring passes" + +Overview +-------- + +In CP, each GPU holds a shard of the sequence for Q, K, and V. To compute full attention, +each GPU needs access to the K and V from *all* sequence shards. CP achieves this by +passing KV chunks between GPUs in a ring or all-gather pattern. + +Strategies +---------- + +**Ring Attention** + +KV chunks are passed around a ring of GPUs. Each GPU: + +1. Computes attention for its local Q against the currently-held KV chunk. +2. Sends its KV chunk to the next GPU in the ring. +3. Receives a new KV chunk from the previous GPU. +4. Repeats until all KV chunks have been seen. + +This overlaps communication with computation: while computing attention on the current +KV chunk, the next chunk is being transferred. + +**All-Gather KV** + +All GPUs perform an all-gather to collect the full K and V tensors, then compute +attention locally. Simpler than ring attention but requires more memory (full K and V +on each GPU). + +**THD (Token-Head-Dimension) Layout** + +A specialized layout where tokens from different sequence positions are interleaved +across GPUs in a way that enables efficient attention computation with minimal +communication. Used with TE's custom CUDA attention kernels. + +Integration with Attention Backends +------------------------------------ + +CP support varies by attention backend: + +- **cuDNN Fused**: Full CP support via both ring and all-gather strategies. +- **FlashAttention**: Limited CP support. +- **Unfused**: No CP support. + +The CP strategy is configured via ``DotProductAttention``: + +.. code-block:: python + + attn = te.DotProductAttention( + num_attention_heads=32, + kv_channels=128, + cp_group=cp_process_group, + cp_stream=cuda_stream, + ) + +See Also +-------- + +- :doc:`backends` — Backend support for context parallelism +- :doc:`/developer/distributed/sequence_parallel` — Sequence parallelism (different from CP) diff --git a/docs/developer/attention/fused_attn_kernels.rst b/docs/developer/attention/fused_attn_kernels.rst new file mode 100644 index 0000000000..9c69f535e2 --- /dev/null +++ b/docs/developer/attention/fused_attn_kernels.rst @@ -0,0 +1,94 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +.. _fused-attn-kernels: + +Fused Attention Kernels +======================= + +The C++ fused attention implementation lives in ``transformer_engine/common/fused_attn/`` +and represents the most complex kernel area in Transformer Engine. + +C API +----- + +**Header**: ``transformer_engine/common/include/transformer_engine/fused_attn.h`` + +Key functions: + +- ``nvte_fused_attn_fwd()`` — Forward pass +- ``nvte_fused_attn_bwd()`` — Backward pass +- ``nvte_fused_attn_fwd_kvpacked()`` / ``nvte_fused_attn_bwd_kvpacked()`` — KV-packed + variants (K and V in a single tensor) +- ``nvte_fused_attn_fwd_qkvpacked()`` / ``nvte_fused_attn_bwd_qkvpacked()`` — QKV-packed + variants + +Each function accepts: + +- Q, K, V tensors (as ``NVTETensor``) +- Bias tensor (optional) +- Softmax auxiliary data (for backward — stores softmax statistics) +- Attention configuration (heads, scaling factor, dropout, mask type, etc.) +- Workspace tensor (pre-allocated scratch memory) + +Directory Structure +------------------- + +.. code-block:: text + + fused_attn/ + ├── fused_attn.h # Internal C++ API + ├── fused_attn.cpp # C API implementation, backend dispatch + ├── fused_attn_f16_max512_seqlen.h/.cu # cuDNN F16 (short sequences) + ├── fused_attn_f16_arbitrary_seqlen.h/.cu # cuDNN F16 (any sequence length) + ├── fused_attn_fp8.h/.cu # cuDNN FP8 + ├── thd_utils.h # THD layout utilities for CP + └── utils.h # Shared utilities + +cuDNN Integration +----------------- + +The cuDNN-based backends use cuDNN's **graph API** (via the cuDNN Frontend library in +``3rdparty/cudnn-frontend``). The approach: + +1. **Build a cuDNN graph** describing the attention computation (matmuls, softmax, + dropout, scaling). +2. **Compile the graph** for the target GPU architecture. +3. **Execute the graph** with input tensors. + +Graph compilation is expensive, so compiled graphs are cached for reuse across iterations. +The cache key includes: sequence lengths, head dimensions, number of heads, data types, +mask type, and dropout configuration. + +FP8 Fused Attention +-------------------- + +The FP8 variant (``fused_attn_fp8.cu``) extends the cuDNN graph with quantization nodes: + +- Q, K, V arrive in FP8 with associated scale inverses. +- Internal computation happens in higher precision (cuDNN manages this). +- Output can be produced in FP8 (with a new scale) or BF16. +- Softmax statistics are saved in FP32 for the backward pass. + +Workspace Management +-------------------- + +Fused attention kernels require significant scratch memory. The workspace size is queried +before execution: + +.. code-block:: cpp + + // Query workspace size + nvte_fused_attn_fwd(q, k, v, ..., workspace, stream); + // First call with workspace.data == nullptr returns required size + // Second call with allocated workspace executes the kernel + +This two-pass pattern (query size, then execute) is common across TE's C API. + +See Also +-------- + +- :doc:`backends` — How fused attention fits in the backend taxonomy +- :doc:`/developer/cpp_core/kernel_areas` — Other kernel areas diff --git a/docs/developer/attention/img/attention_backends.svg b/docs/developer/attention/img/attention_backends.svg new file mode 100644 index 0000000000..de3013407d --- /dev/null +++ b/docs/developer/attention/img/attention_backends.svg @@ -0,0 +1,124 @@ + + + + + + + + + + + + + + + + + + + + + + Attention Backends + + + + + + + + + + + + + + + Unfused + + + + + + + PyTorch Native + Manual QKᵀV matmuls + + Any GPU + BF16 / FP16, fallback path + + + + SM70+ (any arch) + + + + FlashAttention + + + + + + + Tri Dao's FlashAttention + Tiled, IO-aware algorithm + + Ampere+, BF16 / FP16 + O(N) memory, no materialization + + + + SM80+ (Ampere) + + + + Fused (cuDNN) + + + + + + + + + + + F16 Fused + + + + Ampere+ | SM80+ + cuDNN graph API + + + + + FP8 Fused + + + + Hopper+ | SM90+ + Best FP8 performance + + + + + + Custom / THD + + + + Hopper+ | SM90+ + CP-specific kernels + + + + Legend: + + Unfused (fallback) + + FlashAttention + + cuDNN Fused + SM80 = Ampere (A100) | SM90 = Hopper (H100) | SM70 = Volta (V100) + \ No newline at end of file diff --git a/docs/developer/attention/img/context_parallel_ring.svg b/docs/developer/attention/img/context_parallel_ring.svg new file mode 100644 index 0000000000..4783af6b1e --- /dev/null +++ b/docs/developer/attention/img/context_parallel_ring.svg @@ -0,0 +1,102 @@ + + + + + + + + + + + + + + + Context Parallelism: Ring Attention + + + + + + + GPU 0 + seq[0 : L/4] + Q + + Q + + K + + V + + + + + + GPU 1 + seq[L/4 : L/2] + + Q + + K + + V + + + + + + GPU 2 + seq[L/2 : 3L/4] + + Q + + K + + V + + + + + + GPU 3 + seq[3L/4 : L] + + Q + + K + + V + + + + + + + KV chunk + + + + + KV + chunk + + + + + KV chunk + + + + + KV + chunk + + + + Ring Attention + Each GPU computes local Q x all KV + + + KV chunks rotate through the ring until every GPU has seen all KV pairs. + \ No newline at end of file diff --git a/docs/developer/attention/index.rst b/docs/developer/attention/index.rst new file mode 100644 index 0000000000..5bf432ee19 --- /dev/null +++ b/docs/developer/attention/index.rst @@ -0,0 +1,21 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +Attention +========= + +Attention is the most complex subsystem in Transformer Engine, spanning multiple backends, +hardware-specific optimizations, and distributed strategies. This section covers the +backend taxonomy, selection logic, context parallelism, and kernel organization. + +.. toctree:: + :maxdepth: 1 + + backends + backend_selection + context_parallel + pytorch_attention + jax_attention + fused_attn_kernels diff --git a/docs/developer/attention/jax_attention.rst b/docs/developer/attention/jax_attention.rst new file mode 100644 index 0000000000..9309320440 --- /dev/null +++ b/docs/developer/attention/jax_attention.rst @@ -0,0 +1,71 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +JAX Attention +============= + +The JAX frontend provides attention through functional APIs that wrap the same C++ core +as the PyTorch frontend, but use JAX's XLA FFI mechanism for dispatch. + +Entry Points +------------ + +**Location**: ``transformer_engine/jax/attention.py`` + +The primary functions are: + +- ``fused_attn()`` — Fused attention using cuDNN (equivalent to PyTorch's cuDNN fused + backend). +- ``fused_attn_fwd()`` / ``fused_attn_bwd()`` — Forward and backward primitives. + +These are JAX primitives registered via the XLA FFI mechanism +(see :doc:`/developer/jax_frontend/xla_ffi_primitives`). + +Differences from PyTorch +------------------------- + +.. list-table:: + :header-rows: 1 + :widths: 30 35 35 + + * - Aspect + - PyTorch + - JAX + * - API style + - ``nn.Module`` (``DotProductAttention``) + - Functional (``fused_attn()``) + * - Backend selection + - Automatic with fallback chain + - Primarily cuDNN fused + * - Autograd + - ``torch.autograd.Function`` + - JAX custom VJP via primitives + * - Context parallel + - Ring attention with NCCL + - Integrated with XLA sharding + * - KV cache + - ``InferenceParams`` class + - Managed externally by user + +Usage Example +------------- + +.. code-block:: python + + from transformer_engine.jax.attention import fused_attn + + output = fused_attn( + query, key, value, + bias=attn_bias, + mask=attn_mask, + scaling_factor=1.0 / math.sqrt(head_dim), + is_training=True, + ) + +See Also +-------- + +- :doc:`backends` — Backend descriptions (shared with PyTorch) +- :doc:`/developer/jax_frontend/xla_ffi_primitives` — How JAX primitives are registered diff --git a/docs/developer/attention/pytorch_attention.rst b/docs/developer/attention/pytorch_attention.rst new file mode 100644 index 0000000000..be2b5da1f4 --- /dev/null +++ b/docs/developer/attention/pytorch_attention.rst @@ -0,0 +1,105 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +PyTorch Attention Modules +========================= + +The PyTorch frontend provides two attention modules at different abstraction levels. + +DotProductAttention +------------------- + +**Location**: ``transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py`` + +The core attention computation module. Given pre-projected Q, K, V tensors, it computes +scaled dot-product attention with automatic backend selection. + +**Key parameters:** + +- ``num_attention_heads`` / ``num_gqa_groups`` — Head configuration +- ``kv_channels`` — Per-head dimension +- ``attention_type`` — ``"self"`` or ``"cross"`` +- ``attn_mask_type`` — ``"causal"``, ``"padding"``, ``"no_mask"``, ``"arbitrary"`` +- ``window_size`` — For sliding window attention +- ``cp_group`` / ``cp_stream`` — Context parallelism configuration + +**Forward signature** (simplified): + +.. code-block:: python + + def forward(self, query, key, value, + attn_mask=None, + attention_bias=None, + inference_params=None): + +The module handles: + +1. Backend selection (see :doc:`backend_selection`). +2. KV caching for autoregressive inference (via ``InferenceParams``). +3. Context parallel communication if ``cp_group`` is set. +4. FP8 quantization of Q/K/V if FP8 is enabled. + +MultiheadAttention +------------------ + +**Location**: ``transformer_engine/pytorch/attention/multi_head_attention.py`` + +A higher-level module that wraps QKV projection + ``DotProductAttention`` + output +projection: + +.. code-block:: text + + Input + │ + ├── QKV Projection (LayerNormLinear, column-parallel) + │ └── Produces Q, K, V tensors + │ + ├── DotProductAttention + │ └── Computes attention output + │ + └── Output Projection (Linear, row-parallel) + └── Projects back to hidden_size + +**Key parameters** (in addition to DotProductAttention params): + +- ``hidden_size`` — Input/output dimension +- ``input_layernorm`` — Whether to include LayerNorm before QKV projection +- ``return_layernorm_output`` — Return LN output for residual connection + +InferenceParams +--------------- + +**Location**: ``transformer_engine/pytorch/attention/`` + +Manages KV cache for autoregressive generation: + +.. code-block:: python + + class InferenceParams: + max_sequence_length: int + max_batch_size: int + # Pre-allocated KV cache tensors + key_value_memory_dict: dict # layer_idx → (K_cache, V_cache) + +During inference, ``DotProductAttention`` appends new KV to the cache and attends over +the full cached sequence. + +Softmax Variants +---------------- + +TE provides fused softmax implementations for attention scores: + +- ``ScaledSoftmax`` — Standard scaled softmax +- ``ScaledMaskedSoftmax`` — Softmax with arbitrary mask +- ``ScaledUpperTriangMaskedSoftmax`` — Causal mask (upper triangular) + +These are used by the unfused attention backend. The fused backends (cuDNN, Flash) +handle softmax internally. + +See Also +-------- + +- :doc:`backends` — Which backend executes the attention +- :doc:`/developer/pytorch_frontend/module_hierarchy` — Where attention fits in the module tree diff --git a/docs/developer/cpp_core/build_system.rst b/docs/developer/cpp_core/build_system.rst new file mode 100644 index 0000000000..33f60429b2 --- /dev/null +++ b/docs/developer/cpp_core/build_system.rst @@ -0,0 +1,109 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +.. _build-system: + +Build System +============ + +Transformer Engine uses a hybrid build system: a Python ``setup.py`` orchestrates the +top-level install, while CMake handles the C++/CUDA compilation of the core library. + +Build Pipeline +-------------- + +.. code-block:: text + + pip install -e . -v --no-build-isolation + │ + ▼ + setup.py + ├── Checks/fetches git submodules (CUTLASS, cuDNN-frontend) + ├── Detects CUDA toolkit and GPU architectures + ├── Invokes CMake to build transformer_engine/common/ + └── Builds framework-specific extensions: + ├── PyTorch: pybind11 extensions via torch.utils.cpp_extension + └── JAX: XLA FFI extensions + +Key Environment Variables +------------------------- + +.. list-table:: + :header-rows: 1 + :widths: 30 70 + + * - Variable + - Description + * - ``NVTE_FRAMEWORK`` + - Select framework: ``pytorch``, ``jax``, or auto-detect (default) + * - ``NVTE_CUDA_ARCHS`` + - Semicolon-separated GPU architectures (e.g., ``"80;90;100"``) + * - ``MAX_JOBS`` + - Maximum parallel compilation jobs + * - ``CUDA_PATH`` + - Custom CUDA toolkit path + * - ``CUDNN_PATH`` + - Custom cuDNN path + * - ``NVTE_RELEASE_BUILD`` + - Enable release-mode optimizations + +Git Submodules +-------------- + +The build requires two submodules in ``3rdparty/``: + +- **CUTLASS** (``3rdparty/cutlass``) — NVIDIA's CUDA Templates for Linear Algebra + Subroutines. Used for custom GEMM kernels and attention kernels. +- **cuDNN Frontend** (``3rdparty/cudnn-frontend``) — C++ frontend for cuDNN graph API. + Used for fused attention and normalization. + +``setup.py`` automatically initializes these if they are missing, but manual setup is +sometimes needed: + +.. code-block:: bash + + git submodule update --init --recursive + +CMake Structure +--------------- + +The CMake build lives under ``transformer_engine/common/CMakeLists.txt`` and produces +``libtransformer_engine.so``. Key configuration: + +- **Architecture flags**: ``NVTE_CUDA_ARCHS`` maps to CMake's + ``CMAKE_CUDA_ARCHITECTURES``. Only specified architectures are compiled, which + significantly affects build time. +- **Conditional compilation**: Some features are gated on CUDA version (e.g., FP4 support + requires CUDA 12.8+). +- **Header paths**: cuDNN headers are found via ``CUDNN_PATH`` or system paths. + +Developer Build Tips +-------------------- + +**Fast incremental rebuild** (C++ only, skip Python): + +.. code-block:: bash + + cmake --build build/cmake --parallel 4 + +**Minimal architecture for fast iteration**: + +.. code-block:: bash + + NVTE_CUDA_ARCHS="90" pip install -e . -v --no-build-isolation + +**Verbose CMake output** (for debugging build issues): + +.. code-block:: bash + + cmake --build build/cmake --verbose --parallel 4 + +**Common build issues:** + +- ``fatal error: cudnn.h: No such file or directory`` → Set ``CUDNN_PATH`` +- ``nvcc fatal: Unsupported gpu architecture`` → Check ``NVTE_CUDA_ARCHS`` matches + your GPU +- Submodule errors → Run ``git submodule update --init --recursive`` +- Out of memory during compilation → Reduce ``MAX_JOBS`` (try ``MAX_JOBS=1``) diff --git a/docs/developer/cpp_core/img/scaling_modes_comparison.svg b/docs/developer/cpp_core/img/scaling_modes_comparison.svg new file mode 100644 index 0000000000..b3073647b5 --- /dev/null +++ b/docs/developer/cpp_core/img/scaling_modes_comparison.svg @@ -0,0 +1,89 @@ + + + + + + + + + + Scaling Modes Comparison + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + Mode Name + Scale Granularity + Block Size + Data Type + + + Delayed Tensor Scaling + 1 per tensor + N/A + FP8 E4M3 / E5M2 + + + MXFP8 1D Scaling + 1 per 32 elements + 32 + FP8 + E8M0 scales + + + Block Scaling 1D + 1 per 1xN block + configurable + FP8 E4M3 / E5M2 + + + Block Scaling 2D + 1 per NxN block + configurable + FP8 E4M3 / E5M2 + + + NVFP4 1D Scaling + 1 per 16 elements + 16 + FP4 E2M1 + E8M0 + \ No newline at end of file diff --git a/docs/developer/cpp_core/img/tensor_struct.svg b/docs/developer/cpp_core/img/tensor_struct.svg new file mode 100644 index 0000000000..d4adb519d1 --- /dev/null +++ b/docs/developer/cpp_core/img/tensor_struct.svg @@ -0,0 +1,114 @@ + + + + + + + + + + + + + + + + + + + + + + Tensor + + + + data + : SimpleTensor + + + columnwise_data + : SimpleTensor + + + amax + : SimpleTensor + + + columnwise_amax + : SimpleTensor + + + scale + : SimpleTensor + + + scale_inv + : SimpleTensor + + + columnwise_scale_inv + : SimpleTensor + + + + + + scaling_mode + : NVTEScalingMode + + + nvte_tensor + : NVTETensor + + + with_gemm_swizzled_scales + : bool + + + + + + + + + + + + + + + + SimpleTensor + + + + dptr + : void* + + + shape + : vector<size_t> + + + dtype + : DType + + + + + + + + + + + + + + + + + + + + 7 + + \ No newline at end of file diff --git a/docs/developer/cpp_core/index.rst b/docs/developer/cpp_core/index.rst new file mode 100644 index 0000000000..66c01c90b3 --- /dev/null +++ b/docs/developer/cpp_core/index.rst @@ -0,0 +1,19 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +C++ Core +======== + +The C++ core (``transformer_engine/common/``) is the framework-agnostic foundation of +Transformer Engine. It provides the C API, type system, CUDA kernels, and build +infrastructure shared by all frontends. + +.. toctree:: + :maxdepth: 1 + + type_system + kernel_areas + scaling_modes + build_system diff --git a/docs/developer/cpp_core/kernel_areas.rst b/docs/developer/cpp_core/kernel_areas.rst new file mode 100644 index 0000000000..c2be48859a --- /dev/null +++ b/docs/developer/cpp_core/kernel_areas.rst @@ -0,0 +1,136 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +.. _kernel-areas: + +Kernel Areas +============ + +The C++ core organizes CUDA kernels into functional areas, each in its own subdirectory +under ``transformer_engine/common/``. Every area exposes a C API header in +``include/transformer_engine/`` and implements one or more CUDA kernels. + +GEMM (``gemm/``) +----------------- + +Matrix multiplication via cuBLASLt with FP8/MXFP8/NVFP4 support. + +- **Header**: ``include/transformer_engine/gemm.h`` +- **Key file**: ``gemm/cublaslt_gemm.cu`` +- **Entry point**: ``nvte_general_gemm()`` +- Handles: scale/scale-inverse application, bias addition, pre-GeLU fusion, grouped GEMM +- Dispatches to cuBLASLt with appropriate compute types based on input precision and + scaling mode. + +Normalization (``normalization/``) +---------------------------------- + +LayerNorm and RMSNorm with optional FP8 output quantization. + +- **Header**: ``include/transformer_engine/normalization.h`` +- **Key files**: ``normalization/layernorm/``, ``normalization/rmsnorm/`` +- Variants: forward, backward, fused with quantization (cast output to FP8 in the same + kernel to save memory bandwidth) +- Architecture dispatch: separate kernel implementations for different GPU architectures + (e.g., Hopper uses warp-specialized kernels). + +Activation (``activation/``) +----------------------------- + +Element-wise activation functions with optional FP8 output. + +- **Header**: ``include/transformer_engine/activation.h`` +- Supports: GeLU, SiLU/Swish, ReLU, QuickGeLU, SReLU, and their gated variants +- Fused quantization: activation + cast to FP8 in a single kernel pass. + +Cast (``cast/``) +----------------- + +Type-casting and quantization kernels. + +- **Header**: ``include/transformer_engine/cast.h`` +- ``nvte_quantize()`` — cast high-precision data to quantized format with scale computation +- Handles all scaling modes (tensor, block, MXFP8, NVFP4) with mode-specific kernels. + +Transpose (``transpose/``) +--------------------------- + +Transpose operations, often fused with casting. + +- **Header**: ``include/transformer_engine/transpose.h`` +- ``nvte_transpose()`` — standalone transpose +- Fused variants: cast + transpose in a single kernel (critical for producing columnwise + data efficiently during forward pass). + +Fused Attention (``fused_attn/``) +--------------------------------- + +The largest and most complex kernel area. See :doc:`/developer/attention/fused_attn_kernels` +for detailed coverage. + +- **Header**: ``include/transformer_engine/fused_attn.h`` +- Multiple backends: cuDNN-based (F16 and FP8), custom CUDA kernels +- Supports: MHA, GQA, MQA, arbitrary head dims, causal/padding masks, dropout, + sliding window, FP8 quantization. + +Fused RoPE (``fused_rope/``) +----------------------------- + +Rotary positional embedding applied as a fused CUDA kernel. + +- **Header**: ``include/transformer_engine/fused_rope.h`` +- Applies rotary embedding in-place during forward and backward passes. + +Fused Softmax (``fused_softmax/``) +----------------------------------- + +Scaled softmax variants optimized for attention score computation. + +- **Header**: ``include/transformer_engine/softmax.h`` +- Variants: scaled softmax, scaled masked softmax, scaled upper-triangular masked softmax. + +Fused Router (``fused_router/``) +--------------------------------- + +Mixture-of-Experts (MoE) routing kernels. + +- Permutation and unpermutation of tokens across experts. +- Includes its own type-switch macros (note: these need DType casting, see + :doc:`type_system`). + +Communication-GEMM Overlap (``comm_gemm/``) +-------------------------------------------- + +Overlapping NCCL collectives with GEMM computation using CUDA multi-stream. + +- See :doc:`/developer/distributed/comm_gemm_overlap` for the design. +- Uses CUDA events and user-allocated buffers (``UserBuffers``) to pipeline communication + and computation. + +Multi-Tensor Operations (``multi_tensor/``) +------------------------------------------- + +Fused operations over lists of tensors (e.g., multi-tensor scale for optimizer steps). + +- **Header**: ``include/transformer_engine/multi_tensor.h`` + +Architecture Dispatch +--------------------- + +Many kernel areas provide architecture-specific implementations. The typical pattern: + +.. code-block:: text + + normalization/ + ├── common.h # Shared types and helpers + ├── layernorm/ + │ ├── ln_api.cpp # C API entry point, dispatches by arch + │ ├── ln_sm80.cu # Ampere (SM80) kernel + │ ├── ln_sm90.cu # Hopper (SM90) kernel + │ └── ln_sm100.cu # Blackwell (SM100) kernel + +The C API function queries the GPU architecture at runtime and dispatches to the +appropriate kernel. See :doc:`build_system` for how ``NVTE_CUDA_ARCHS`` controls which +architectures are compiled. diff --git a/docs/developer/cpp_core/scaling_modes.rst b/docs/developer/cpp_core/scaling_modes.rst new file mode 100644 index 0000000000..c3a68ab6c0 --- /dev/null +++ b/docs/developer/cpp_core/scaling_modes.rst @@ -0,0 +1,185 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +.. _scaling-modes: + +Scaling Modes +============= + +Transformer Engine supports multiple quantization scaling strategies, each offering +different trade-offs between precision, performance, and hardware requirements. The +scaling mode is selected via the ``NVTEScalingMode`` enum and propagated through the +``Tensor.scaling_mode`` field. + +.. figure:: ./img/scaling_modes_comparison.svg + :align: center + :width: 80% + + Comparison of scale granularity across scaling modes. + +.. + Diagram description for ``scaling_modes_comparison.svg``: + A table/matrix with 5 rows (one per scaling mode) and columns: + Mode Name | Enum Value | Scale Granularity | Block Size | Data Type | Min Arch. + DELAYED_TENSOR_SCALING | 0 | 1 scale per tensor | N/A | FP8 E4M3/E5M2 | Hopper + MXFP8_1D_SCALING | 1 | 1 scale per 32 elements | 32 | FP8 + E8M0 scales | Blackwell + BLOCK_SCALING_1D | 2 | 1 scale per 1×N block | configurable | FP8 E4M3/E5M2 | Blackwell + BLOCK_SCALING_2D | 3 | 1 scale per N×N block | configurable | FP8 E4M3/E5M2 | Blackwell + NVFP4_1D_SCALING | 4 | 1 scale per 16 elements | 16 | FP4 E2M1 + E8M0 | Blackwell + +Overview +-------- + +.. list-table:: + :header-rows: 1 + :widths: 25 10 25 15 25 + + * - Mode + - Enum + - Scale Granularity + - Block Size + - Data Type + * - Delayed Tensor Scaling + - ``0`` + - 1 scale per tensor + - N/A + - FP8 E4M3 / E5M2 + * - MXFP8 1D Scaling + - ``1`` + - 1 scale per 32 elements + - 32 + - FP8 + E8M0 scales + * - Block Scaling 1D + - ``2`` + - 1 scale per 1×N block + - Configurable + - FP8 E4M3 / E5M2 + * - Block Scaling 2D + - ``3`` + - 1 scale per N×N block + - Configurable + - FP8 E4M3 / E5M2 + * - NVFP4 1D Scaling + - ``4`` + - 1 scale per 16 elements + - 16 + - FP4 E2M1 + E8M0 + +Delayed Tensor Scaling (``NVTE_DELAYED_TENSOR_SCALING``) +-------------------------------------------------------- + +The original FP8 scaling mode. A single scale factor applies to the entire tensor, +computed from the *previous iteration's* amax (absolute maximum) value. + +**How it works:** + +1. After each forward/backward pass, the kernel records the amax of the output. +2. The ``FP8GlobalStateManager`` maintains an amax history window (typically 1024 values). +3. Before the next iteration, the scale is computed as: + ``scale = fp8_max / amax_history.max()`` +4. This scale is applied uniformly to all elements during quantization. + +**Trade-offs:** + +- Simple and fast (no per-element scale overhead in GEMM). +- Scale is one iteration stale — can cause overflow/underflow on loss spikes. +- Requires the amax feedback loop between iterations. + +**C++ helpers:** + +.. code-block:: cpp + + bool is_tensor_scaling(const NVTEScalingMode &mode); + bool is_delayed_tensor_scaling(const NVTEScalingMode &mode); + +MXFP8 1D Scaling (``NVTE_MXFP8_1D_SCALING``) +---------------------------------------------- + +Microscaling FP8 format per the OCP MX specification. Each block of 32 contiguous +elements shares a single E8M0 (8-bit exponent-only) scale factor. + +**Key details:** + +- Block size is fixed at 32 elements. +- Scales are E8M0 format (power-of-two only, no mantissa bits). +- Scales may need "GEMM swizzling" — a specific memory layout that cuBLASLt expects. + The ``with_gemm_swizzled_scales`` flag on ``Tensor`` tracks this. +- Current scaling (no amax history needed). + +**C++ helpers:** + +.. code-block:: cpp + + bool is_mxfp8_scaling(const NVTEScalingMode &mode); + bool is_mxfp_scaling(const NVTEScalingMode &mode); // alias + +Block Scaling (``NVTE_BLOCK_SCALING_1D`` / ``NVTE_BLOCK_SCALING_2D``) +---------------------------------------------------------------------- + +Per-block FP8 scaling with configurable block dimensions. ``1D`` uses 1×N blocks +(one scale per row-slice of N elements), while ``2D`` uses N×N blocks. + +**Key details:** + +- Block size is configurable via the quantizer's ``block_scaling_dim`` property. +- Scales are FP32, not E8M0 (more precise than MXFP8). +- Current scaling (computed just-in-time, no history). +- Available on Blackwell and later architectures. + +**C++ helpers:** + +.. code-block:: cpp + + bool is_block_scaling(const NVTEScalingMode &mode); + // Returns true for all non-tensor-scaling modes + +NVFP4 1D Scaling (``NVTE_NVFP4_1D_SCALING``) +---------------------------------------------- + +4-bit floating point with E8M0 block scales, targeting maximum compression. + +**Key details:** + +- Data is FP4 E2M1 format (2 exponent bits, 1 mantissa bit). +- Each block of 16 elements shares one E8M0 scale. +- Like MXFP8, may require GEMM-swizzled scales. +- Blackwell and later only. + +**C++ helpers:** + +.. code-block:: cpp + + bool is_nvfp4_scaling(const NVTEScalingMode &mode); + bool is_nvfp_scaling(const NVTEScalingMode &mode); // alias + +Scale Storage and Layout +------------------------ + +For all block-scaling modes, scale tensors are stored alongside the data tensor. +The ``Tensor`` struct provides separate scale inverse fields for rowwise and columnwise +layouts: + +- ``scale_inv`` — scale inverses for decoding rowwise data (used in forward GEMM). +- ``columnwise_scale_inv`` — scale inverses for decoding columnwise data (used in wgrad + GEMM). + +See :doc:`/developer/quantization/rowwise_columnwise` for how these layouts interact +with GEMM execution. + +Checking Scaling Mode in Code +----------------------------- + +The ``common.h`` header provides inline helpers that should be preferred over direct +enum comparison: + +.. code-block:: cpp + + // Prefer this: + if (is_mxfp8_scaling(tensor.scaling_mode)) { ... } + + // Over this: + if (tensor.scaling_mode == NVTE_MXFP8_1D_SCALING) { ... } + +This makes code resilient to future additions to the ``NVTEScalingMode`` enum. diff --git a/docs/developer/cpp_core/type_system.rst b/docs/developer/cpp_core/type_system.rst new file mode 100644 index 0000000000..9030adc08a --- /dev/null +++ b/docs/developer/cpp_core/type_system.rst @@ -0,0 +1,236 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +.. _type-system: + +Type System +=========== + +Transformer Engine maintains two parallel type systems for historical and ABI-stability +reasons: a C type system for the public API and a C++ type system for internal use. + +C Types (Public API) +-------------------- + +Defined in ``transformer_engine/common/include/transformer_engine/transformer_engine.h``. + +NVTEDType +^^^^^^^^^ + +A plain C ``enum`` representing data types at the API boundary: + +.. code-block:: c + + enum NVTEDType { + kNVTEByte = 0, + kNVTEInt16 = 1, + kNVTEInt32 = 2, + kNVTEInt64 = 3, + kNVTEFloat32 = 4, + kNVTEFloat16 = 5, + kNVTEBFloat16 = 6, + kNVTEFloat8E4M3 = 7, + kNVTEFloat8E5M2 = 8, + kNVTEFloat8E8M0 = 9, + kNVTEFloat4E2M1 = 10, + kNVTENumTypes + }; + +NVTEBasicTensor +^^^^^^^^^^^^^^^ + +A lightweight, non-owning tensor descriptor used to populate ``NVTETensor`` parameters: + +.. code-block:: c + + struct NVTEBasicTensor { + void *data_ptr; // Raw pointer to data buffer + NVTEDType dtype; // Data type + NVTEShape shape; // Shape (up to 15 dimensions) + }; + +NVTETensor +^^^^^^^^^^ + +An opaque handle (``typedef void *NVTETensor``) to the internal ``Tensor`` object. +Created with ``nvte_create_tensor()`` and populated via ``nvte_tensor_set()`` with +``NVTETensorParam`` keys: + +.. code-block:: c + + enum NVTETensorParam { + kNVTERowwiseData = 0, + kNVTEColumnwiseData = 1, + kNVTEScale = 2, + kNVTEAmax = 3, + kNVTERowwiseScaleInv = 4, + kNVTEColumnwiseScaleInv = 5, + kNVTEColumnwiseAmax = 6, + kNVTEWithGEMMSwizzledScales = 7, + }; + +NVTEScalingMode +^^^^^^^^^^^^^^^ + +Controls the quantization granularity — see :doc:`scaling_modes` for details. + +C++ Types (Internal) +-------------------- + +Defined in ``transformer_engine/common/common.h``. + +DType +^^^^^ + +A C++ ``enum class`` in the ``transformer_engine`` namespace, mirroring ``NVTEDType``: + +.. code-block:: cpp + + namespace transformer_engine { + enum class DType { + kByte, kInt16, kInt32, kInt64, + kFloat32, kFloat16, kBFloat16, + kFloat8E4M3, kFloat8E5M2, kFloat8E8M0, + kFloat4E2M1, kNumTypes + }; + } + +.. warning:: + + ``DType`` and ``NVTEDType`` have the same numeric values but are **not** implicitly + convertible. Use ``static_cast`` for conversions. + +SimpleTensor +^^^^^^^^^^^^ + +A C++ wrapper around ``NVTEBasicTensor`` that provides constructors, implicit conversions, +and utility methods: + +.. code-block:: cpp + + struct SimpleTensor { + void *dptr; + std::vector shape; + DType dtype; + + // Implicit conversion from NVTEBasicTensor + SimpleTensor(const NVTEBasicTensor &tensor); + // Implicit conversion to NVTEBasicTensor + operator NVTEBasicTensor() const; + + size_t numel() const; + bool has_data() const; + size_t buffer_size_bytes() const; + }; + +``SimpleTensor`` converts freely to/from ``NVTEBasicTensor``, handling the +``NVTEDType`` ↔ ``DType`` cast internally. + +Tensor +^^^^^^ + +The main internal tensor type, composed of multiple ``SimpleTensor`` members: + +.. figure:: ./img/tensor_struct.svg + :align: center + :width: 60% + + UML-style diagram of the Tensor struct and its SimpleTensor members. + +.. + Diagram description for ``tensor_struct.svg``: + UML struct box labeled "Tensor": + Fields: + + data : SimpleTensor [rowwise quantized data] + + columnwise_data : SimpleTensor [columnwise quantized data] + + amax : SimpleTensor [absolute maximum] + + columnwise_amax : SimpleTensor [columnwise amax] + + scale : SimpleTensor [quantization scale] + + scale_inv : SimpleTensor [rowwise scale inverse] + + columnwise_scale_inv : SimpleTensor [columnwise scale inverse] + + scaling_mode : NVTEScalingMode + + nvte_tensor : NVTETensor (opaque handle) + + with_gemm_swizzled_scales : bool + Below, a smaller UML box for "SimpleTensor": + + dptr : void* + + shape : vector + + dtype : DType + Arrow from each SimpleTensor field in Tensor to the SimpleTensor box (composition). + +.. code-block:: cpp + + struct Tensor { + SimpleTensor data; // Rowwise data + SimpleTensor columnwise_data; // Columnwise data (for wgrad) + SimpleTensor amax; // Absolute maximum + SimpleTensor columnwise_amax; // Columnwise amax + SimpleTensor scale; // Quantization scale + SimpleTensor scale_inv; // Rowwise scale inverse + SimpleTensor columnwise_scale_inv; // Columnwise scale inverse + + NVTEScalingMode scaling_mode; + NVTETensor nvte_tensor; // Opaque handle for C API + bool with_gemm_swizzled_scales; // MXFP8/NVFP4 scale format flag + }; + +This structure holds all the metadata needed for a quantized tensor: the data itself, +scaling factors for both rowwise and columnwise layouts, and amax values for delayed +scaling. + +DType / NVTEDType Conversion Patterns +-------------------------------------- + +When working across the C/C++ boundary, conversions are frequently needed. Here are the +common patterns: + +**NVTEDType → DType** (e.g., reading from ``NVTEBasicTensor.dtype``): + +.. code-block:: cpp + + NVTEDType nvte_type = tensor.dtype; + DType dtype = static_cast(nvte_type); + +**DType → NVTEDType** (e.g., writing to ``NVTEBasicTensor.dtype``): + +.. code-block:: cpp + + DType dtype = DType::kFloat8E4M3; + basic_tensor.dtype = static_cast(dtype); + +**Comparison**: Overloaded ``==`` and ``!=`` operators exist, so direct comparison works: + +.. code-block:: cpp + + if (nvte_dtype == DType::kFloat8E4M3) { ... } // OK + +**Function overloads**: Many utility functions accept both types: + +.. code-block:: cpp + + // Both of these work: + bool a = is_fp8_dtype(DType::kFloat8E4M3); + bool b = is_fp8_dtype(kNVTEFloat8E4M3); + +Overloaded functions include: ``is_fp8_dtype``, ``typeToSize``, ``typeToNumBits``, +``get_buffer_size_bytes``, ``get_cuda_dtype``, ``get_cudnn_dtype``, ``get_cudnn_fe_dtype``, +and ``to_string``. + +Common Pitfalls +^^^^^^^^^^^^^^^ + +1. **auto deduction**: ``auto dtype = tensor.dtype;`` deduces ``NVTEDType`` for + ``NVTEBasicTensor`` fields. If passed to a function expecting ``DType``, add an + explicit cast. + +2. **Ternary expressions**: ``condition ? nvte_val : dtype_val`` won't compile — wrap + one side in ``static_cast``. + +3. **ADL (Argument-Dependent Lookup)**: Functions in the ``transformer_engine`` namespace + won't be found by ADL when called with ``NVTEDType`` arguments (which are in global + scope). Use the namespace qualifier: ``transformer_engine::to_string(nvte_type)``. + +4. **Custom switch macros**: Macros like ``TRANSFORMER_ENGINE_TYPE_SWITCH_ALL`` cast to + ``DType`` internally, but custom switch macros (e.g., in ``fused_router/utils.h``) + may need the same treatment. diff --git a/docs/developer/distributed/comm_gemm_overlap.rst b/docs/developer/distributed/comm_gemm_overlap.rst new file mode 100644 index 0000000000..9a607142e6 --- /dev/null +++ b/docs/developer/distributed/comm_gemm_overlap.rst @@ -0,0 +1,129 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +.. _comm-gemm-overlap: + +Communication-GEMM Overlap +=========================== + +Transformer Engine can overlap NCCL collective communication with GEMM computation to +hide communication latency. This is critical for scaling TP across nodes where +inter-node bandwidth is limited. + +.. figure:: ./img/comm_gemm_overlap.svg + :align: center + :width: 80% + + Timeline showing GEMM compute overlapping with all-gather/reduce-scatter. + +.. + Diagram description for ``comm_gemm_overlap.svg``: + Two horizontal timelines: + Top "Without overlap": + [All-gather (full)] → [GEMM (full)] — sequential, total time = comm + compute + Bottom "With overlap": + Chunk 0: [AG chunk 0] [GEMM chunk 0] + Chunk 1: [AG chunk 1] [GEMM chunk 1] + Chunk 2: [AG chunk 2] [GEMM chunk 2] + Total time ≈ max(comm, compute) — pipelined + +Concept +------- + +In tensor parallelism, a column-parallel linear requires an all-gather before GEMM, and +a row-parallel linear requires a reduce-scatter after GEMM. Normally these are +sequential: + +.. code-block:: text + + all-gather → GEMM (column-parallel forward) + GEMM → reduce-scatter (row-parallel forward) + +With overlap, the all-gather/reduce-scatter is split into chunks that execute on a +communication stream while GEMM chunks execute on the compute stream: + +.. code-block:: text + + comm stream: [AG chunk 0][AG chunk 1][AG chunk 2]... + compute stream: [GEMM 0] [GEMM 1] [GEMM 2]... + +UserBuffers +----------- + +**Location**: ``transformer_engine/common/comm_gemm/`` + +The overlap implementation uses pre-allocated ``UserBuffers`` — persistent GPU memory +buffers that are registered with NCCL for direct access. This avoids per-iteration +buffer allocation overhead and enables fine-grained pipelining. + +Key components: + +- ``UserBuffersManager``: Manages lifetime and allocation of persistent buffers. +- Multi-stream coordination via CUDA events for synchronization between communication + and computation streams. + +Overlap Modes +------------- + +The ``CommOverlapAlgo`` enum defines the available overlap strategies: + +.. list-table:: + :header-rows: 1 + :widths: 35 65 + + * - Algorithm + - Description + * - ``BULK_OVERLAP_AG`` + - All-gather runs fully in parallel with GEMM (simplest) + * - ``BULK_OVERLAP_RS`` + - Reduce-scatter runs fully in parallel with GEMM + * - ``SPLIT_PIPELINED_AG_P2P`` + - All-gather split into fine-grained chunks with P2P, pipelined with GEMM chunks + * - ``SPLIT_PIPELINED_RS`` + - Reduce-scatter split into chunks, pipelined with GEMM chunks + * - ``SPLIT_PIPELINED_RS_P2P`` + - Reduce-scatter with P2P, pipelined with GEMM + * - ``ATOMIC_GEMM_RS`` + - Uses CUDA atomics with reduce-scatter for producer-consumer overlap + * - ``ATOMIC_GEMM_AG_P2P`` + - Uses CUDA atomics with all-gather P2P + +Enabling Overlap +---------------- + +.. code-block:: python + + model = te.TransformerLayer( + hidden_size=4096, + ffn_hidden_size=16384, + num_attention_heads=32, + tp_group=tp_group, + ub_overlap_rs=True, # Overlap reduce-scatter + ub_overlap_ag=True, # Overlap all-gather + ) + +Environment variables: + +- ``NVTE_UB_OVERLAP``: Enable/disable UserBuffers overlap +- ``NVTE_UB_SPLIT_RS``: Enable split-pipelined reduce-scatter +- ``NVTE_UB_SPLIT_AG``: Enable split-pipelined all-gather + +Performance Considerations +-------------------------- + +Overlap is most beneficial when: + +- TP spans across nodes (high communication latency). +- GEMM is large enough to hide communication time. +- Multiple GEMM chunks can execute efficiently (not too small). + +Overlap adds complexity (buffer management, stream synchronization) and may not help when +communication is already fast (e.g., within NVLink-connected nodes). + +See Also +-------- + +- :doc:`tensor_parallel` — The TP collectives that are being overlapped +- :doc:`/developer/cpp_core/kernel_areas` — The comm_gemm kernel area diff --git a/docs/developer/distributed/fsdp_integration.rst b/docs/developer/distributed/fsdp_integration.rst new file mode 100644 index 0000000000..efa6cda678 --- /dev/null +++ b/docs/developer/distributed/fsdp_integration.rst @@ -0,0 +1,84 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +.. _fsdp-integration: + +FSDP Integration +================ + +Transformer Engine integrates with PyTorch's Fully Sharded Data Parallelism (FSDP) to +enable combined data and tensor parallelism with FP8 quantization. + +Challenges +---------- + +FSDP shards model parameters across data-parallel ranks and gathers them on-demand +during forward/backward. This creates challenges for FP8: + +1. **FP8 weight storage**: Weights may be stored in FP8 for memory savings, but FSDP's + gather/scatter operates on raw tensors — it doesn't know about FP8 scales. +2. **Amax synchronization**: Delayed scaling requires amax values to be consistent across + FSDP ranks (all ranks must agree on the scale). +3. **Quantized gradients**: FSDP's gradient reduction must handle FP8 gradients correctly. + +Solution: FP8-Aware Hooks +-------------------------- + +TE registers FSDP hooks that handle FP8 tensor conversions: + +**Pre-gather hook**: Before FSDP gathers a parameter shard, convert FP8 storage to a +format that can be gathered (e.g., pack data + scale into a single buffer). + +**Post-gather hook**: After FSDP gathers the full parameter, reconstruct the FP8 tensor +with proper scales. + +**Pre-scatter hook**: Before FSDP scatters gradients, ensure FP8 gradient data and scales +are properly packed. + +Usage +----- + +.. code-block:: python + + import torch.distributed.fsdp as fsdp + import transformer_engine.pytorch as te + + model = te.TransformerLayer(...) + + # Wrap with FSDP — TE hooks are registered automatically + model = fsdp.FullyShardedDataParallel( + model, + auto_wrap_policy=..., + ) + + # FP8 training works normally + with te.fp8_autocast(enabled=True): + output = model(input) + +FP8 Weight Storage with FSDP +----------------------------- + +When FP8 weights are enabled, the parameter is stored in FP8 format even when sharded +by FSDP. The lifecycle: + +1. **Initialization**: Weight is created in high precision, then optionally cast to FP8. +2. **FSDP shard**: The FP8 data (and associated scale) is sharded across ranks. +3. **FSDP gather**: Full FP8 weight is gathered; TE reconstructs the quantized tensor. +4. **Forward/backward**: GEMM operates on the gathered FP8 weight. +5. **FSDP scatter**: Updated FP8 weight shard is stored locally. + +Limitations +----------- + +- FSDP2 (``torch.distributed._composable.fsdp``) has better support than FSDP1. +- FP8 weight caching interacts with FSDP gather/scatter — weights may be re-quantized + on each gather. +- Mixed FSDP + TP configurations require careful process group setup. + +See Also +-------- + +- :doc:`tensor_parallel` — TP can be combined with FSDP +- :doc:`/developer/quantization/class_hierarchy` — FP8 tensor types that FSDP must handle diff --git a/docs/developer/distributed/img/column_parallel.svg b/docs/developer/distributed/img/column_parallel.svg new file mode 100644 index 0000000000..37ccdefacf --- /dev/null +++ b/docs/developer/distributed/img/column_parallel.svg @@ -0,0 +1,110 @@ + + + + + + + + + + + + + + Column-Parallel Linear + + + + Input X + (same on all GPUs) + + + + @ + + + Weight W (split along output dim) + + + + W₀ (GPU 0) + + + + W₁ (GPU 1) + + + + W₂ (GPU 2) + + + + W₃ (GPU 3) + + + + + + Per-GPU compute + + + + GPU 0 + Y₀ = X @ W₀ᵀ + + + + GPU 1 + Y₁ = X @ W₁ᵀ + + + + GPU 2 + Y₂ = X @ W₂ᵀ + + + + GPU 3 + Y₃ = X @ W₃ᵀ + + + + + + + + + Partial outputs + + + Y₀ + + + Y₁ + + + Y₂ + + + Y₃ + + + (split along hidden dim) + + + + All-gather input if sequence parallel; output stays partitioned across GPUs + diff --git a/docs/developer/distributed/img/comm_gemm_overlap.svg b/docs/developer/distributed/img/comm_gemm_overlap.svg new file mode 100644 index 0000000000..8a8bc3865b --- /dev/null +++ b/docs/developer/distributed/img/comm_gemm_overlap.svg @@ -0,0 +1,103 @@ + + + + + + + + + + + + + + + + + Communication-GEMM Overlap (Split Pipelined) + + + Without Overlap (Sequential) + + + + GPU + + + + All-gather + + + + GEMM + + + + + + Total time = comm + compute + + + With Overlap (Split Pipelined) + + + Comm + stream + + + Compute + stream + + + + AG chunk 0 + + + AG chunk 1 + + + AG chunk 2 + + + + GEMM 0 + + + GEMM 1 + + + GEMM 2 + + + + + + + + + + + + + + Total time ≈ max(comm, compute) + + + + Communication (all-gather) + + Compute (GEMM) + diff --git a/docs/developer/distributed/img/row_parallel.svg b/docs/developer/distributed/img/row_parallel.svg new file mode 100644 index 0000000000..429e0f95e5 --- /dev/null +++ b/docs/developer/distributed/img/row_parallel.svg @@ -0,0 +1,121 @@ + + + + + + + + + + + + + + Row-Parallel Linear + + + Input X + (partitioned) + + + X₀ (GPU 0) + + + X₁ (GPU 1) + + + X₂ (GPU 2) + + + X₃ (GPU 3) + + + + + + + + + Weight W + (split along input dim) + + + W₀ (GPU 0) + + + W₁ (GPU 1) + + + W₂ (GPU 2) + + + W₃ (GPU 3) + + + @ + @ + @ + @ + + + + + + + + + Partial results + + + GPU 0 + Y₀ = X₀ @ W₀ᵀ + + + GPU 1 + Y₁ = X₁ @ W₁ᵀ + + + GPU 2 + Y₂ = X₂ @ W₂ᵀ + + + GPU 3 + Y₃ = X₃ @ W₃ᵀ + + + + + + + + + + All-reduce / + Reduce-scatter + Y = ΣYᵢ + + + + + + + Output Y + (complete) + + + + Input partitioned across GPUs; output requires reduction (all-reduce or reduce-scatter) + diff --git a/docs/developer/distributed/index.rst b/docs/developer/distributed/index.rst new file mode 100644 index 0000000000..884cb4276d --- /dev/null +++ b/docs/developer/distributed/index.rst @@ -0,0 +1,19 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +Distributed Training +==================== + +Transformer Engine provides built-in support for several distributed training strategies, +with optimizations specific to low-precision computation. + +.. toctree:: + :maxdepth: 1 + + tensor_parallel + sequence_parallel + fsdp_integration + comm_gemm_overlap + jax_distributed diff --git a/docs/developer/distributed/jax_distributed.rst b/docs/developer/distributed/jax_distributed.rst new file mode 100644 index 0000000000..b63323a55a --- /dev/null +++ b/docs/developer/distributed/jax_distributed.rst @@ -0,0 +1,104 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +JAX Distributed +=============== + +JAX uses a fundamentally different approach to distributed training than PyTorch. +Instead of explicit collective calls, JAX relies on XLA's SPMD (Single Program Multiple +Data) partitioner to automatically insert communication based on tensor sharding +annotations. + +Sharding Model +-------------- + +In JAX + Transformer Engine: + +- A ``Mesh`` defines the physical device topology (e.g., 2×4 grid of GPUs). +- ``PartitionSpec`` annotations on tensors specify how each dimension maps to mesh axes. +- XLA's compiler automatically inserts all-gather, reduce-scatter, etc. + +.. code-block:: python + + from jax.sharding import Mesh, PartitionSpec as P + + # 2D mesh: 2 data-parallel × 4 tensor-parallel + mesh = Mesh(devices, axis_names=("dp", "tp")) + + # Weight sharded along TP axis + weight_spec = P(None, "tp") # [hidden, hidden/tp] + +MeshResource +------------ + +**Location**: ``transformer_engine/jax/sharding.py`` + +TE's ``MeshResource`` wraps JAX's mesh configuration for use with TE modules: + +.. code-block:: python + + from transformer_engine.jax.sharding import MeshResource + + mesh_resource = MeshResource( + dp_resource="dp", + tp_resource="tp", + pp_resource="pp", # optional pipeline parallel + ) + +This tells TE modules which mesh axes correspond to data, tensor, and pipeline +parallelism. + +Comparison with PyTorch TP +-------------------------- + +.. list-table:: + :header-rows: 1 + :widths: 25 37 38 + + * - Aspect + - PyTorch + - JAX + * - Communication + - Explicit (``allreduce()``, etc.) + - Implicit (XLA compiler inserts) + * - Weight sharding + - Manual (``parallel_mode``) + - Via ``PartitionSpec`` + * - Process groups + - ``torch.distributed`` groups + - ``jax.sharding.Mesh`` axes + * - FSDP + - ``torch.distributed.fsdp`` + - XLA SPMD (automatic) + * - Configuration + - Module constructor params + - ``MeshResource`` + partition specs + +TP in JAX TE Modules +--------------------- + +JAX TE Flax modules accept sharding annotations: + +.. code-block:: python + + from transformer_engine.jax.flax import TransformerLayer + + layer = TransformerLayer( + hidden_size=4096, + mlp_hidden_size=16384, + num_attention_heads=32, + mesh_resource=mesh_resource, + ) + +The module uses the mesh resource to: + +1. Shard weight parameters according to column/row parallel patterns. +2. Annotate intermediate tensors for XLA to optimize communication. + +See Also +-------- + +- :doc:`tensor_parallel` — PyTorch TP concepts (same mathematical operations) +- :doc:`/developer/jax_frontend/sharding` — Detailed JAX sharding utilities diff --git a/docs/developer/distributed/sequence_parallel.rst b/docs/developer/distributed/sequence_parallel.rst new file mode 100644 index 0000000000..6ff01adc32 --- /dev/null +++ b/docs/developer/distributed/sequence_parallel.rst @@ -0,0 +1,94 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +.. _sequence-parallel: + +Sequence Parallelism +==================== + +Sequence Parallelism (SP) extends Tensor Parallelism by partitioning the sequence +dimension for operations that don't require the full hidden dimension (LayerNorm, +dropout, residual connections). This reduces memory by distributing activations that +TP alone would replicate. + +How It Works +------------ + +In standard TP, operations like LayerNorm process the full sequence on every GPU +(replicated). With SP: + +- **LayerNorm** and **residual additions** operate on sequence-partitioned data + (each GPU handles ``seq_len / tp_size`` tokens). +- **Column-parallel linear** begins with an **all-gather** along the sequence dimension + to reconstruct the full sequence before the GEMM. +- **Row-parallel linear** ends with a **reduce-scatter** along the sequence dimension + instead of all-reduce, producing sequence-partitioned output. + +.. code-block:: text + + Without SP: With SP: + [Full seq, Full hidden] [Seq/TP, Full hidden] + │ │ + LayerNorm LayerNorm (local) + │ │ + Column-parallel All-gather → Column-parallel + │ │ + Row-parallel Row-parallel → Reduce-scatter + │ │ + [Full seq, Full hidden] [Seq/TP, Full hidden] + +Memory Savings +-------------- + +SP reduces activation memory by a factor of ``tp_size`` for: + +- LayerNorm activations +- Dropout masks +- Residual connection buffers + +These savings are significant for long-sequence training where activations dominate +memory. + +Enabling SP +----------- + +SP is enabled alongside TP: + +.. code-block:: python + + layer = te.TransformerLayer( + hidden_size=4096, + ffn_hidden_size=16384, + num_attention_heads=32, + tp_group=tp_group, + sequence_parallel=True, + ) + +**Requirements:** + +- TP must be enabled (``tp_group`` must be set). +- Sequence length must be divisible by ``tp_size``. +- All ranks in the TP group process the same batch. + +Implementation +-------------- + +**Location**: ``transformer_engine/pytorch/distributed.py`` + +Key functions: + +- ``gather_along_first_dim(input, tp_group)``: All-gather along sequence dim + (before column-parallel GEMM). +- ``reduce_scatter_along_first_dim(input, tp_group)``: Reduce-scatter along sequence dim + (after row-parallel GEMM). + +These are inserted automatically by the TE modules when ``sequence_parallel=True``. + +See Also +-------- + +- :doc:`tensor_parallel` — TP fundamentals that SP builds on +- :doc:`/developer/attention/context_parallel` — Context parallelism (different strategy + for long sequences) diff --git a/docs/developer/distributed/tensor_parallel.rst b/docs/developer/distributed/tensor_parallel.rst new file mode 100644 index 0000000000..50c1fe9bef --- /dev/null +++ b/docs/developer/distributed/tensor_parallel.rst @@ -0,0 +1,118 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +.. _tensor-parallel: + +Tensor Parallelism +================== + +Tensor Parallelism (TP) splits individual weight matrices across multiple GPUs, +enabling models too large for a single GPU's memory. Transformer Engine implements +Megatron-style column-parallel and row-parallel linear layers. + +.. figure:: ./img/column_parallel.svg + :align: center + :width: 70% + + Column-parallel linear: input is broadcast, weight columns are split across GPUs. + +.. + Diagram description for ``column_parallel.svg``: + Left: "Input X" (full tensor, same on all GPUs). + Center: Weight matrix W split vertically into W_0, W_1, W_2, W_3 (4 GPUs). + Each GPU computes Y_i = X @ W_i^T (partial output). + Right: Partial outputs Y_0..Y_3 are concatenated (or kept split for next layer). + Label: "All-gather input if needed; output is split along hidden dim" + +.. figure:: ./img/row_parallel.svg + :align: center + :width: 70% + + Row-parallel linear: input is split, weight rows are split, output is reduced. + +.. + Diagram description for ``row_parallel.svg``: + Left: "Input X" split into X_0, X_1, X_2, X_3 (one per GPU). + Center: Weight matrix W split horizontally into W_0, W_1, W_2, W_3. + Each GPU computes Y_i = X_i @ W_i^T (partial result). + Right: "All-reduce (or reduce-scatter)" to produce final output Y = sum(Y_i). + Label: "Input is partitioned; output requires reduction" + +Column-Parallel Linear +---------------------- + +Used for the **first** linear layer in a pair (e.g., QKV projection, MLP FC1): + +.. code-block:: python + + linear = te.Linear( + in_features=hidden_size, + out_features=4 * hidden_size, + parallel_mode="column", + tp_group=tp_process_group, + ) + +**Forward**: Each GPU holds columns ``W[:, start:end]`` and computes partial output. +If the input is not already partitioned, an all-gather collects the full input first. + +**Backward**: Gradient of the output is split; weight gradient is computed locally. +Input gradient requires an all-reduce (or reduce-scatter with sequence parallelism). + +Row-Parallel Linear +------------------- + +Used for the **second** linear layer in a pair (e.g., attention output projection, +MLP FC2): + +.. code-block:: python + + linear = te.Linear( + in_features=4 * hidden_size, + out_features=hidden_size, + parallel_mode="row", + tp_group=tp_process_group, + ) + +**Forward**: Each GPU holds rows ``W[start:end, :]`` and computes a partial result. +The partial results are summed via all-reduce (or reduce-scatter). + +**Backward**: Gradient is broadcast (or all-gathered); weight gradient is computed locally. + +Column + Row Pairing +-------------------- + +In a Transformer layer, column-parallel and row-parallel layers always appear in pairs: + +.. code-block:: text + + MLP: FC1 (column-parallel) → Activation → FC2 (row-parallel) + Attn: QKV (column-parallel) → Attention → Out (row-parallel) + +This pairing ensures that only one all-reduce per pair is needed (at the row-parallel +layer's output), minimizing communication. + +Implementation Details +---------------------- + +**Location**: ``transformer_engine/pytorch/distributed.py`` and +``transformer_engine/pytorch/module/linear.py`` + +Key functions: + +- ``allreduce()``: Synchronous all-reduce across TP group. +- ``reduce_scatter_along_first_dim()``: For sequence parallelism integration. +- ``gather_along_first_dim()``: For sequence parallelism integration. + +The ``parallel_mode`` parameter on ``Linear`` controls: + +1. How the weight is sharded at initialization. +2. Which collective (all-gather, all-reduce, reduce-scatter) is inserted in forward/backward. +3. Whether sequence parallelism communication is used. + +See Also +-------- + +- :doc:`sequence_parallel` — SP extends TP to the sequence dimension +- :doc:`comm_gemm_overlap` — Overlapping GEMM with TP communication diff --git a/docs/developer/img/architecture_layers.svg b/docs/developer/img/architecture_layers.svg new file mode 100644 index 0000000000..9de855ec77 --- /dev/null +++ b/docs/developer/img/architecture_layers.svg @@ -0,0 +1,85 @@ + + + + + + + + + + + + + + + + TransformerEngine Architecture Layers + + + + Python API + + + PyTorch Frontend + + + JAX Frontend + + + + + + + Binding Layer + + + pybind11 (PyTorch) + + + XLA FFI (JAX) + + + + + + + C++ Core Library + + + + + + + CUDA Kernels + + + GEMM + + + Normali- + zation + + + Activation + + + Attention + + + Cast / + Transpose + + + + + + transformer_engine/ + common/ + \ No newline at end of file diff --git a/docs/developer/img/data_flow_forward.svg b/docs/developer/img/data_flow_forward.svg new file mode 100644 index 0000000000..ff539cd198 --- /dev/null +++ b/docs/developer/img/data_flow_forward.svg @@ -0,0 +1,72 @@ + + + + + + + + + + + + + + + + FP8 Forward Data Flow + + + + Input + (BF16 / FP16) + + + + + + + Quantize + (cast kernel) + + + + FP8 data + + scale + + + + GEMM + (cuBLASLt) + + + + + + + Output + (BF16 / FP16) + + + + Weight + (FP8 + scale) + + + + + + + amax history (delayed scaling) + + + + Scale computation depends on NVTEScalingMode + \ No newline at end of file diff --git a/docs/developer/img/linear_e2e_flow.svg b/docs/developer/img/linear_e2e_flow.svg new file mode 100644 index 0000000000..b5056c0b4d --- /dev/null +++ b/docs/developer/img/linear_e2e_flow.svg @@ -0,0 +1,125 @@ + + + + + + + + + + + + + + + + + + + Linear Module End-to-End Flow + + + + Forward + + + + Input + (BF16) + + + + + + + All-gather + (col-parallel + + SP) + + + + + + + Quantize + (FP8 + row + col) + + + + + + + GEMM + + + + + + + Reduce-scatter / + All-reduce + (row-parallel) + + + + Backward + + + + grad_output + + + + + + + Quantize + + + + + + + Dgrad + GEMM + + + + + + + Wgrad + GEMM + + + + + + + All-reduce + grad_weight + + + + + + saved columnwise + input + + + + + saved weight + \ No newline at end of file diff --git a/docs/developer/index.rst b/docs/developer/index.rst new file mode 100644 index 0000000000..cf74227043 --- /dev/null +++ b/docs/developer/index.rst @@ -0,0 +1,43 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +Developer Guide +=============== + +This section documents the internal architecture, design decisions, and implementation +details of Transformer Engine. It is intended for contributors, maintainers, and AI +agents working on the codebase. + +.. note:: + + This guide describes the *internal* structure of Transformer Engine, not the + user-facing API. For the public API, see :doc:`/api/pytorch` and :doc:`/api/jax`. + +How to Use This Guide +--------------------- + +- **New contributors**: Start with :doc:`architecture_overview` for the big picture, + then read :doc:`linear_walkthrough` for an end-to-end trace through a concrete module. +- **Working on quantization**: See :doc:`quantization/index` for the Quantizer/Storage/Tensor + design and :doc:`cpp_core/type_system` for the underlying C/C++ types. +- **Working on attention**: See :doc:`attention/index` for backend selection and kernel + organization. +- **Working on distributed**: See :doc:`distributed/index` for tensor/sequence parallelism + and communication overlap. +- **Framework-specific work**: See :doc:`pytorch_frontend/index` or + :doc:`jax_frontend/index`. + +.. toctree:: + :caption: Developer Guide + :maxdepth: 2 + + architecture_overview + linear_walkthrough + cpp_core/index + quantization/index + pytorch_frontend/index + jax_frontend/index + attention/index + distributed/index diff --git a/docs/developer/jax_frontend/img/jax_vs_pytorch_binding.svg b/docs/developer/jax_frontend/img/jax_vs_pytorch_binding.svg new file mode 100644 index 0000000000..0e16e35809 --- /dev/null +++ b/docs/developer/jax_frontend/img/jax_vs_pytorch_binding.svg @@ -0,0 +1,71 @@ + + + + + + + + + + + + + + + + + + + + + + + PyTorch nn.Module + + + + torch.Tensor + + + + pybind11 Extensions + + + + + + + + + JAX Flax nn.Module + + + + jax.Array + + + + XLA FFI Custom Calls + + + + + + + + + C++ Core Library + transformer_engine/common/ + + + + + + + CUDA Kernels + diff --git a/docs/developer/jax_frontend/index.rst b/docs/developer/jax_frontend/index.rst new file mode 100644 index 0000000000..78c4c3574c --- /dev/null +++ b/docs/developer/jax_frontend/index.rst @@ -0,0 +1,23 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +JAX Frontend +============ + +.. note:: + + The JAX frontend documentation is a **draft**. JAX support in Transformer Engine is + actively evolving and some details may change. + +The JAX frontend (``transformer_engine/jax/``) provides Flax ``nn.Module`` wrappers, +XLA FFI custom-call primitives, and sharding utilities for distributed training. + +.. toctree:: + :maxdepth: 1 + + module_system + xla_ffi_primitives + quantization + sharding diff --git a/docs/developer/jax_frontend/module_system.rst b/docs/developer/jax_frontend/module_system.rst new file mode 100644 index 0000000000..820b4ac691 --- /dev/null +++ b/docs/developer/jax_frontend/module_system.rst @@ -0,0 +1,113 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +JAX Module System +================= + +The JAX frontend provides Flax ``nn.Module`` wrappers that mirror the PyTorch frontend's +module hierarchy but use JAX idioms (functional transforms, XLA compilation, sharding +annotations). + +.. figure:: ./img/jax_vs_pytorch_binding.svg + :align: center + :width: 70% + + PyTorch and JAX frontends share the same C++ core via different binding mechanisms. + +.. + Diagram description for ``jax_vs_pytorch_binding.svg``: + Two parallel vertical paths: + Left path: "PyTorch nn.Module" → "pybind11" → "C++ Core" + Right path: "JAX Flax nn.Module" → "XLA FFI" → "C++ Core" + Both paths converge at the "C++ Core" box at the bottom. + Labels: "torch.Tensor" on left, "jax.Array" on right. + +Module Location +--------------- + +**Location**: ``transformer_engine/jax/flax/module.py`` + +Flax Modules +------------ + +The JAX frontend provides these Flax modules. Core modules are in +``transformer_engine/jax/flax/module.py``, while attention and transformer composition +are in ``transformer_engine/jax/flax/transformer.py``. + +.. list-table:: + :header-rows: 1 + :widths: 30 30 40 + + * - JAX Module + - PyTorch Equivalent + - Description + * - ``DenseGeneral`` + - ``Linear`` + - General linear layer with FP8 and flexible contracting dimensions + * - ``LayerNorm`` + - ``LayerNorm`` / ``RMSNorm`` + - Normalization (supports both ``layernorm`` and ``rmsnorm`` types) + * - ``LayerNormDenseGeneral`` + - ``LayerNormLinear`` + - Fused LN + Linear + * - ``LayerNormMLP`` + - ``LayerNormMLP`` + - Fused LN + MLP + * - ``DotProductAttention`` + - ``DotProductAttention`` + - Attention (fused/unfused auto-selection) + * - ``MultiHeadAttention`` + - ``MultiheadAttention`` + - Full MHA block with projections + * - ``TransformerLayer`` + - ``TransformerLayer`` + - Complete Transformer block + +All FP8-capable modules inherit from ``TransformerEngineBase``, which provides +``generate_quantizer_set()`` for recipe-based quantizer creation. + +Key Differences from PyTorch +----------------------------- + +**Functional style**: JAX modules are stateless. Parameters are passed explicitly, not +stored as attributes: + +.. code-block:: python + + # PyTorch: stateful + layer = te.Linear(768, 3072) + output = layer(input) + + # JAX: functional + layer = te_jax.DenseGeneral(features=3072) + params = layer.init(rng, input) + output = layer.apply(params, input) + +**Sharding via annotations**: Instead of ``parallel_mode`` constructor args, JAX modules +use ``MeshResource`` and rely on XLA SPMD for communication: + +.. code-block:: python + + layer = te_jax.DenseGeneral( + features=3072, + mesh_resource=MeshResource(tp_resource="tp"), + ) + +**Quantization context**: Instead of ``fp8_autocast()``, JAX uses explicit quantizer +arguments: + +.. code-block:: python + + output = layer.apply( + params, input, + quantizer=fp8_quantizer, + ) + +See Also +-------- + +- :doc:`/developer/pytorch_frontend/module_hierarchy` — PyTorch module hierarchy for comparison +- :doc:`xla_ffi_primitives` — How JAX modules call into C++ kernels +- :doc:`sharding` — Detailed sharding configuration diff --git a/docs/developer/jax_frontend/quantization.rst b/docs/developer/jax_frontend/quantization.rst new file mode 100644 index 0000000000..e68af1bdda --- /dev/null +++ b/docs/developer/jax_frontend/quantization.rst @@ -0,0 +1,112 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +JAX Quantization +================ + +The JAX frontend has its own quantization system that parallels the PyTorch design but +uses JAX idioms (immutable arrays, functional transforms). + +Quantizer Hierarchy +------------------- + +**Location**: ``transformer_engine/jax/quantize/quantizer.py`` + +The JAX quantizer hierarchy mirrors PyTorch's: + +.. list-table:: + :header-rows: 1 + :widths: 35 35 30 + + * - JAX Quantizer + - PyTorch Equivalent + - Scaling Mode + * - ``DelayedScaleQuantizer`` + - ``Float8Quantizer`` + - Delayed tensor + * - ``CurrentScaleQuantizer`` + - ``Float8CurrentScalingQuantizer`` + - Current tensor + * - ``BlockScaleQuantizer`` + - ``MXFP8Quantizer`` + - MXFP8 block + * - ``NVFP4Quantizer`` + - ``NVFP4Quantizer`` + - NVFP4 + * - ``GroupedQuantizer`` + - (no direct equivalent) + - Group-wise FP8 + +ScaledTensor Types +------------------ + +**Location**: ``transformer_engine/jax/quantize/tensor.py`` + +Instead of PyTorch's ``QuantizedTensor`` (which is a ``torch.Tensor`` subclass), JAX uses +``ScaledTensor`` types registered as JAX pytrees. There are two core types: + +- ``ScaledTensor1x`` — Single-layout quantized tensor (either rowwise or columnwise). +- ``ScaledTensor2x`` — Dual-layout tensor wrapping both a rowwise and columnwise + ``ScaledTensor1x``. + +A ``NoScaleTensor`` type wraps unquantized data with optional amax tracking. + +Since JAX arrays are immutable, these are simple data classes (not JAX array subclasses), +registered as pytree nodes so JAX can trace through them: + +.. code-block:: python + + @register_pytree_node_class + @dataclass + class ScaledTensor1x: + data: jax.Array # Quantized data + scale_inv: jax.Array # Scale inverse for dequantization + scaling_mode: ScalingMode + dq_dtype: jnp.dtype # Target dtype for dequantization + is_colwise: bool # Whether this is a columnwise layout + + @register_pytree_node_class + @dataclass + class ScaledTensor2x: + rowwise_tensor: ScaledTensor1x + colwise_tensor: ScaledTensor1x + +Key Differences from PyTorch +----------------------------- + +1. **No ``__torch_dispatch__``**: JAX doesn't have dispatch machinery. Operations must + explicitly handle ``ScaledTensor`` inputs. + +2. **No autograd integration**: JAX uses ``custom_vjp`` on the primitive level, not on + the tensor type. + +3. **Immutable**: Quantized tensors cannot be modified in-place. Amax history for delayed + scaling is managed through JAX's stateful mechanisms (Flax variable collections). + +4. **XLA-compatible**: ``ScaledTensor`` fields are regular ``jax.Array`` values that XLA + can trace through and compile. + +Usage with Modules +------------------ + +Quantizers are typically created via the ``QuantizerFactory`` from a recipe: + +.. code-block:: python + + from transformer_engine.jax.quantize.helper import get_quantize_config_with_recipe + from transformer_engine.common.recipe import DelayedScaling + + recipe = DelayedScaling() + config = get_quantize_config_with_recipe(recipe) + + # Inside a Flax module, quantizer sets are auto-created + # via TransformerEngineBase.generate_quantizer_set() + # Metadata (scales, amax_history) stored in Flax variable collections + +See Also +-------- + +- :doc:`/developer/quantization/class_hierarchy` — PyTorch quantization design (same concepts) +- :doc:`/developer/quantization/scaling_recipes` — Recipe system shared between frontends diff --git a/docs/developer/jax_frontend/sharding.rst b/docs/developer/jax_frontend/sharding.rst new file mode 100644 index 0000000000..e47a266f4c --- /dev/null +++ b/docs/developer/jax_frontend/sharding.rst @@ -0,0 +1,100 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +Sharding +======== + +The JAX frontend uses XLA's SPMD partitioner for distributed training. Transformer +Engine provides utilities to configure sharding for TE modules. + +MeshResource +------------ + +**Location**: ``transformer_engine/jax/sharding.py`` + +``MeshResource`` maps logical parallelism dimensions to physical mesh axes: + +.. code-block:: python + + from transformer_engine.jax.sharding import MeshResource + + mesh_resource = MeshResource( + dp_resource="data", # Data parallel axis name + tp_resource="model", # Tensor parallel axis name + tpsp_resource="model", # Tensor + sequence parallel axis name (optional) + fsdp_resource="fsdp", # FSDP axis name (optional) + pp_resource="pipe", # Pipeline parallel axis name (optional) + cp_resource="context", # Context parallel axis name (optional) + ) + +This is passed to TE modules, which use it to determine how to partition weights and +activations. The ``global_shard_guard(resource)`` context manager sets the active +``MeshResource`` globally. + +Logical Axis Annotations +------------------------- + +TE modules annotate their tensors with logical axes (e.g., "hidden", "heads", +"sequence"). These are mapped to physical mesh axes via ``MeshResource``: + +.. code-block:: text + + Logical axis "hidden" → physical axis "model" (TP) + Logical axis "batch" → physical axis "data" (DP) + Logical axis "sequence" → may be split across CP or SP + +The mapping is defined by TE's sharding rules registered with each XLA primitive +(see :doc:`xla_ffi_primitives`). + +Weight Sharding +--------------- + +For tensor parallelism, weights are sharded along the appropriate dimension: + +**Column-parallel** (QKV projection, MLP FC1): + +.. code-block:: python + + # Weight shape: [hidden, heads * head_dim] + # Sharded: [hidden, heads * head_dim / tp] + # PartitionSpec: P(None, "tp") + +**Row-parallel** (output projection, MLP FC2): + +.. code-block:: python + + # Weight shape: [heads * head_dim, hidden] + # Sharded: [heads * head_dim / tp, hidden] + # PartitionSpec: P("tp", None) + +XLA automatically inserts all-gather and reduce-scatter operations based on these +partition specs. + +Activation Sharding +------------------- + +Activations are annotated to minimize communication: + +- **Before column-parallel**: Full tensor, replicated across TP. +- **After column-parallel**: Sharded along hidden dim. +- **After row-parallel**: Full tensor (XLA inserts all-reduce). + +With sequence parallelism, activations outside the parallel region are sharded along the +sequence dimension. + +Custom Sharding Rules +--------------------- + +Each TE XLA primitive registers sharding rules via ``infer_sharding_from_operands()`` +and ``partition()`` methods on the ``BasePrimitive`` class. These tell XLA: + +1. Given input sharding, what is the output sharding? +2. How to partition the operation across devices. + +See Also +-------- + +- :doc:`/developer/distributed/jax_distributed` — Overview of JAX distributed training +- :doc:`xla_ffi_primitives` — Where sharding rules are registered diff --git a/docs/developer/jax_frontend/xla_ffi_primitives.rst b/docs/developer/jax_frontend/xla_ffi_primitives.rst new file mode 100644 index 0000000000..f04500839e --- /dev/null +++ b/docs/developer/jax_frontend/xla_ffi_primitives.rst @@ -0,0 +1,134 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +XLA FFI Primitives +================== + +The JAX frontend communicates with the C++ core through XLA's Foreign Function Interface +(FFI). Each TE operation is registered as a custom call that XLA can invoke during +compiled execution. + +BasePrimitive +------------- + +**Location**: ``transformer_engine/jax/cpp_extensions/base.py`` + +``BasePrimitive`` is the base class for all JAX TE primitives. It provides the machinery +to: + +1. **Register** a custom call with XLA (both forward and backward). +2. **Define abstract evaluation** (shape/dtype inference without running the kernel). +3. **Define the custom VJP** (JAX's equivalent of PyTorch's autograd backward). +4. **Handle sharding** (propagate partition specs through the operation). + +.. code-block:: python + + class BasePrimitive(metaclass=ABCMeta): + + @classmethod + def enabled(cls) -> bool: ... # Enable/disable via NVTE_JAX_CUSTOM_CALLS env var + + @staticmethod + @abstractmethod + def abstract(*args, **kwargs): ... # Shape/dtype inference + + @staticmethod + @abstractmethod + def lowering(ctx, *args, **kwargs): ... # MLIR lowering to XLA custom call + + @staticmethod + @abstractmethod + def impl(*args, **kwargs): ... # Eager execution implementation + + @staticmethod + @abstractmethod + def batcher(*args): ... # vmap batch rules + + @staticmethod + @abstractmethod + def partition(*args): ... # Shardy multi-device partitioning + + @staticmethod + @abstractmethod + def shardy_sharding_rule(*args): ... # Shardy sharding spec string + +Registration Flow +----------------- + +The ``register_primitive(cls)`` function registers each primitive in two forms: + +- **Inner primitive**: Single-device execution, no sharding awareness. +- **Outer primitive**: Multi-device aware, uses Shardy custom partitioning. + +When the JAX TE module is imported, each primitive: + +1. Creates a JAX ``core.Primitive`` object (both inner and outer). +2. Registers ``abstract`` as the abstract evaluation rule. +3. Registers ``lowering`` as the MLIR lowering rule (maps to XLA custom call via FFI). +4. Registers ``batcher`` for ``jax.vmap`` support. +5. Registers ``partition`` and ``shardy_sharding_rule`` for XLA SPMD. +6. Backward is registered via ``jax.custom_vjp`` on the outer primitive. + +Primitives can be selectively enabled/disabled via the ``NVTE_JAX_CUSTOM_CALLS`` +environment variable. + +Custom Call Execution +--------------------- + +When XLA encounters a TE custom call during execution: + +.. code-block:: text + + JAX trace → MLIR lowering → XLA custom call → C++ handler → CUDA kernel + +The C++ handler: + +1. Receives raw GPU buffer pointers from XLA. +2. Wraps them in ``NVTETensor`` handles. +3. Calls the same C API functions as the PyTorch frontend. +4. Returns results via pre-allocated output buffers. + +Comparison with PyTorch Bridge +------------------------------- + +.. list-table:: + :header-rows: 1 + :widths: 25 37 38 + + * - Aspect + - PyTorch (pybind11) + - JAX (XLA FFI) + * - Binding mechanism + - Python C extension + - XLA custom call + * - When called + - Eagerly per operation + - During XLA compilation/execution + * - Autodiff + - ``torch.autograd.Function`` + - ``jax.custom_vjp`` + * - Shape inference + - Implicit (PyTorch is eager) + - Explicit ``abstract()`` method + * - Parallelism + - Manual collectives + - XLA SPMD via sharding rules + +Key Primitives +-------------- + +Major primitives registered in ``transformer_engine/jax/cpp_extensions/``: + +- ``gemm`` — Matrix multiplication with FP8 +- ``layernorm_fwd`` / ``layernorm_bwd`` — Normalization +- ``fused_attn_fwd`` / ``fused_attn_bwd`` — Fused attention +- ``cast_fp8`` / ``cast_fp8_transpose`` — FP8 quantization +- ``activation_fwd`` / ``activation_bwd`` — Activation functions + +See Also +-------- + +- :doc:`/developer/pytorch_frontend/cpp_extensions` — PyTorch equivalent +- :doc:`/developer/cpp_core/kernel_areas` — The C++ kernels being called diff --git a/docs/developer/linear_walkthrough.rst b/docs/developer/linear_walkthrough.rst new file mode 100644 index 0000000000..c85334267a --- /dev/null +++ b/docs/developer/linear_walkthrough.rst @@ -0,0 +1,257 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +.. _linear-walkthrough: + +Linear Module: End-to-End Walkthrough +====================================== + +This page traces a complete forward and backward pass through ``te.Linear``, the most +fundamental Transformer Engine module. It connects concepts from across the codebase: +quantization, GEMM, distributed communication, and autograd. + +.. figure:: ./img/linear_e2e_flow.svg + :align: center + :width: 95% + + End-to-end flow through Linear forward and backward with FP8 and tensor parallelism. + +.. + Diagram description for ``linear_e2e_flow.svg``: + Two horizontal swim lanes labeled "Forward" and "Backward". + Forward lane (left to right): + 1. "Input (BF16)" → + 2. "All-gather (if column-parallel + SP)" → + 3. "input_quantizer(input)" producing "FP8 rowwise + columnwise" → + 4. "general_gemm(weight, input)" producing "Output (BF16)" → + 5. "Reduce-scatter (if row-parallel + SP) or All-reduce (if row-parallel)" + Backward lane (right to left): + 1. "grad_output (BF16)" → + 2. "grad_output_quantizer(grad)" → + 3. "Dgrad GEMM: general_gemm(grad, weight)" producing "grad_input" → + 4. "Wgrad GEMM: general_gemm(input_columnwise, grad)" producing "grad_weight" → + 5. "All-reduce grad_weight (if column-parallel)" + Dotted arrows from forward boxes 3→backward box 4 labeled "saved columnwise input" + and forward weight→backward box 3 labeled "saved weight". + +Setup +----- + +.. code-block:: python + + import transformer_engine.pytorch as te + from transformer_engine.common.recipe import DelayedScaling + + # Create a Linear module with FP8 + linear = te.Linear(4096, 16384, bias=True) + + # Enable FP8 + recipe = DelayedScaling(fp8_format=te.recipe.Format.HYBRID) + with te.fp8_autocast(enabled=True, fp8_recipe=recipe): + output = linear(input) # Triggers the flow below + +Phase 1: Module Forward Entry +------------------------------ + +**File**: ``transformer_engine/pytorch/module/linear.py``, ``Linear.forward()`` (line ~1343) + +1. ``pre_forward()`` is called (inherited from ``TransformerEngineBaseModule``): + + - Checks if FP8 is enabled via ``FP8GlobalStateManager``. + - Creates or refreshes quantizer instances from the active recipe. + - For delayed scaling: computes new scales from amax history. + +2. Quantizers are prepared: + + - ``input_quantizer`` — for the activation tensor + - ``weight_quantizer`` — for the weight parameter + - ``grad_output_quantizer`` — for the backward gradient + +3. ``_Linear.apply()`` is called, entering the autograd function. + +Phase 2: _Linear.forward() — Input Preparation +------------------------------------------------- + +**File**: ``transformer_engine/pytorch/module/linear.py``, ``_Linear.forward()`` (line ~87) + +The input tensor is prepared based on the parallel mode and FP8 state: + +**Column-parallel with sequence parallelism** (most common for QKV/FC1): + +.. code-block:: python + + # 1. Set quantizer usage: need rowwise for forward GEMM, + # columnwise for backward wgrad GEMM + input_quantizer.set_usage(rowwise=True, columnwise=backward_needs_input) + + # 2. Quantize local input shard + inputmat = input_quantizer(inputmat) # Returns QuantizedTensorStorage + + # 3. All-gather along sequence dimension + inputmat_total, _ = gather_along_first_dim(inputmat, tp_group) + +**No parallelism** (standalone Linear): + +.. code-block:: python + + input_quantizer.set_usage(rowwise=True, columnwise=backward_needs_input) + inputmat = input_quantizer(inputmat) + inputmat_total = inputmat + +The ``input_quantizer(inputmat)`` call produces a ``QuantizedTensorStorage`` containing: + +- Rowwise FP8 data + ``scale_inv`` (for forward GEMM) +- Columnwise FP8 data + ``columnwise_scale_inv`` (saved for wgrad GEMM) + +See :doc:`quantization/rowwise_columnwise` for why both layouts are needed. + +Phase 3: _Linear.forward() — Weight Preparation +-------------------------------------------------- + +.. code-block:: python + + # Configure weight quantizer + weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) + + # Get quantized weight (may use cached version) + weightmat = module.get_weight_workspace( + tensor=weight, + quantizer=weight_quantizer, + cache_name="weight", # Cache across microbatches + update_workspace=is_first_microbatch, + ) + +Weight caching is important: the same FP8 weight can be reused across gradient +accumulation microbatches, avoiding redundant quantization. + +Phase 4: _Linear.forward() — GEMM +----------------------------------- + +.. code-block:: python + + # Forward GEMM: output = input @ weight^T + gemm_out, *_ = general_gemm( + weightmat, # A operand (FP8) + inputmat_total, # B operand (FP8) + out_dtype=activation_dtype, + bias=bias, + use_split_accumulator=use_split_accumulator, + ub=ub_obj, # Userbuffers overlap (optional) + ub_type=ub_type, + ) + +This call chain: + +1. **Python**: ``general_gemm()`` in ``transformer_engine/pytorch/cpp_extensions/gemm.py`` + extracts raw tensors and scales from the quantized inputs. +2. **pybind11**: ``te_general_gemm()`` in ``transformer_engine/pytorch/csrc/extensions/gemm.cpp`` + constructs ``NVTETensor`` handles and calls the C API. +3. **C API**: ``nvte_general_gemm()`` in ``transformer_engine/common/gemm/`` + dispatches to cuBLASLt with FP8 compute types. +4. **cuBLASLt**: Executes the FP8 matrix multiplication on the GPU. + +Phase 5: _Linear.forward() — Output Communication +---------------------------------------------------- + +For **row-parallel** (output projection, FC2): + +.. code-block:: python + + if sequence_parallel: + out, _ = reduce_scatter_along_first_dim(out, tp_group) + elif tensor_parallel: + out, _ = allreduce(out, tp_group) + +For **column-parallel**: no communication needed (output is naturally partitioned). + +Phase 6: _Linear.forward() — Save for Backward +-------------------------------------------------- + +.. code-block:: python + + # Save quantized input for wgrad (columnwise data) + # Save quantized weight for dgrad + ctx.save_for_backward(inputmat, weightmat, weight, ...) + ctx.weight_quantizer = weight_quantizer + +Key optimization: only the *columnwise* portion of the input is kept (the rowwise +data is discarded after the forward GEMM). This halves the activation memory for FP8. + +Phase 7: _Linear.backward() — Dgrad +-------------------------------------- + +.. code-block:: python + + # Quantize grad_output + grad_output_quantizer.set_usage(rowwise=True, columnwise=True) + qgrad_output = grad_output_quantizer(grad_output) + + # Dgrad GEMM: grad_input = grad_output @ weight + dgrad, *_ = general_gemm( + weightmat, # Columnwise weight from forward + qgrad_output, # Rowwise grad_output + out_dtype=activation_dtype, + ) + +Phase 8: _Linear.backward() — Wgrad +-------------------------------------- + +.. code-block:: python + + # Wgrad GEMM: grad_weight = input^T @ grad_output + # Uses columnwise input saved from forward + wgrad, *_ = general_gemm( + inputmat, # Columnwise input from forward + qgrad_output, # Columnwise grad_output + out_dtype=activation_dtype, + grad=True, # Indicates wgrad GEMM + ) + +This is why the forward pass computed both rowwise and columnwise input: the columnwise +data is consumed here by the wgrad GEMM. + +Phase 9: Post-Backward +------------------------ + +1. ``post_forward()`` records the output amax for delayed scaling. +2. For distributed training, amax values are all-reduced across TP ranks. +3. The amax history is updated for the next iteration's scale computation. + +Summary of Files Touched +------------------------- + +.. list-table:: + :header-rows: 1 + :widths: 50 50 + + * - File + - Role in Linear + * - ``pytorch/module/linear.py`` + - ``Linear.forward()``, ``_Linear`` autograd + * - ``pytorch/module/base.py`` + - ``pre_forward()``, FP8 state management + * - ``pytorch/quantized_tensor.py`` + - ``Quantizer``, ``QuantizedTensorStorage`` base classes + * - ``pytorch/tensor/float8_tensor.py`` + - ``Float8Quantizer``, FP8 cast implementation + * - ``pytorch/cpp_extensions/gemm.py`` + - ``general_gemm()`` Python wrapper + * - ``pytorch/csrc/extensions/gemm.cpp`` + - pybind11 GEMM binding + * - ``common/gemm/cublaslt_gemm.cu`` + - cuBLASLt GEMM kernel dispatch + * - ``pytorch/distributed.py`` + - ``allreduce()``, ``gather_along_first_dim()``, etc. + * - ``pytorch/quantization.py`` + - ``FP8GlobalStateManager``, ``fp8_autocast`` + +See Also +-------- + +- :doc:`architecture_overview` — High-level system architecture +- :doc:`quantization/class_hierarchy` — Quantizer/Storage/Tensor design +- :doc:`quantization/rowwise_columnwise` — Why both layouts exist +- :doc:`pytorch_frontend/autograd_integration` — Autograd patterns +- :doc:`distributed/tensor_parallel` — TP communication in Linear diff --git a/docs/developer/pytorch_frontend/autograd_integration.rst b/docs/developer/pytorch_frontend/autograd_integration.rst new file mode 100644 index 0000000000..21d89f07e0 --- /dev/null +++ b/docs/developer/pytorch_frontend/autograd_integration.rst @@ -0,0 +1,113 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +.. _autograd-integration: + +Autograd Integration +==================== + +Transformer Engine modules use custom ``torch.autograd.Function`` subclasses to control +what is saved for backward, how FP8 tensors flow through the autograd graph, and when +activation recomputation occurs. + +The _Linear Pattern +------------------- + +The canonical example is ``_Linear`` in ``transformer_engine/pytorch/module/linear.py``. +This autograd function wraps the forward GEMM and defines the corresponding backward +(dgrad + wgrad). + +**Forward** (pseudocode): + +.. code-block:: python + + class _Linear(torch.autograd.Function): + @staticmethod + def forward(ctx, input, weight, bias, ...): + # 1. Quantize input (produces rowwise + columnwise) + qinput = input_quantizer(input) + + # 2. Get quantized weight + qweight = weight_quantizer(weight) + + # 3. Forward GEMM: output = qinput @ qweight^T + output = general_gemm(qinput, qweight, bias=bias) + + # 4. Save for backward + ctx.save_for_backward(qinput, qweight, ...) + + return output + +**Backward** (pseudocode): + +.. code-block:: python + + @staticmethod + def backward(ctx, grad_output): + qinput, qweight, ... = ctx.saved_tensors + + # 1. Quantize grad_output + qgrad = grad_quantizer(grad_output) + + # 2. Dgrad GEMM: grad_input = qgrad @ qweight + grad_input = general_gemm(qgrad, qweight) + + # 3. Wgrad GEMM: grad_weight = qinput^T @ qgrad + # Uses columnwise qinput saved from forward + grad_weight = general_gemm(qinput.columnwise, qgrad) + + return grad_input, grad_weight, grad_bias, ... + +Saved Tensor Strategy +--------------------- + +What gets saved for backward depends on the configuration: + +**FP8 disabled**: Standard PyTorch behavior — save full-precision input and weight. + +**FP8 enabled (no recompute)**: Save the ``QuantizedTensor`` objects from forward. These +contain both rowwise data (used by dgrad GEMM) and columnwise data (used by wgrad GEMM). +Memory cost: ~2× the FP8 data size (rowwise + columnwise). + +**FP8 enabled + activation recompute**: Only save the high-precision input (or a stashed +copy). During backward, re-run the quantization to produce fresh FP8 data. Saves memory +at the cost of recomputation. + +Activation Recomputation +------------------------ + +TE supports activation recomputation (gradient checkpointing) at multiple granularities: + +- **Full recompute**: Re-run the entire forward pass during backward. Standard PyTorch + ``checkpoint()`` works with TE modules. +- **Selective recompute**: Only recompute the quantization (cast) operations, keeping + the GEMM results. This is cheaper than full recompute because casting is + memory-bandwidth-bound while GEMM is compute-bound. + +The ``TransformerLayer`` module provides a ``activation_checkpointing`` parameter that +controls this behavior. + +FP8 Tensors in Autograd +------------------------ + +``QuantizedTensor`` (a ``torch.Tensor`` subclass) can be saved via +``ctx.save_for_backward()`` like any other tensor. Key behaviors: + +- **No implicit dequantization** during save — the FP8 data is stored as-is. +- During backward, the saved ``QuantizedTensor`` is retrieved and its ``.get_storage()`` + method provides the raw FP8 data + scales for GEMM. +- If a non-TE operation receives a ``QuantizedTensor``, ``__torch_dispatch__`` + automatically dequantizes it. + +Interaction with torch.compile +------------------------------- + +TE modules are compatible with ``torch.compile`` but require care: + +- ``NVTE_TORCH_COMPILE=1`` enables compile-friendly code paths. +- Some autograd functions use ``torch.compiler.is_compiling()`` to select between + eager and compile-compatible implementations. +- FP8 state management (amax updates, scale refreshes) must happen outside compiled + regions because they involve in-place mutation of module state. diff --git a/docs/developer/pytorch_frontend/cpp_extensions.rst b/docs/developer/pytorch_frontend/cpp_extensions.rst new file mode 100644 index 0000000000..71cfd76796 --- /dev/null +++ b/docs/developer/pytorch_frontend/cpp_extensions.rst @@ -0,0 +1,125 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +.. _cpp-extensions: + +C++ Extensions Bridge +===================== + +The PyTorch frontend communicates with the C++ core through a pybind11 extension module. +This bridge converts PyTorch tensors to ``NVTETensor`` handles and calls the C API. + +Architecture +------------ + +.. code-block:: text + + Python (transformer_engine/pytorch/cpp_extensions/) + │ torch.Tensor → NVTETensor conversion + ▼ + pybind11 (transformer_engine/pytorch/csrc/extensions/) + │ Call C API functions + ▼ + C API (transformer_engine/common/include/transformer_engine/) + │ Unpack NVTETensor → internal Tensor struct + ▼ + CUDA kernels (transformer_engine/common/) + +Python Side +----------- + +**Location**: ``transformer_engine/pytorch/cpp_extensions/`` + +Each kernel area has a corresponding Python wrapper module: + +- ``gemm.py`` — ``general_gemm()``, ``grouped_gemm()`` +- ``normalization.py`` — ``layernorm_fwd()``, ``rmsnorm_fwd()``, etc. +- ``activation.py`` — ``gelu()``, ``silu()``, etc. +- ``cast.py`` — ``quantize()``, ``dequantize()`` +- ``transpose.py`` — ``transpose()``, ``cast_transpose()`` +- ``attention.py`` — ``fused_attn_fwd()``, ``fused_attn_bwd()`` + +These wrappers: + +1. Extract raw data tensors and scales from ``QuantizedTensor`` / ``QuantizedTensorStorage`` + objects. +2. Construct ``NVTETensor`` handles via helper functions. +3. Call the pybind11-exposed C++ function. +4. Wrap outputs back into Python tensor types. + +**Example** — ``general_gemm()`` (simplified): + +.. code-block:: python + + def general_gemm(A, B, bias=None, ...): + # Convert QuantizedTensor → NVTETensor components + A_data, A_scale_inv, A_dtype = _extract_tensors(A) + B_data, B_scale_inv, B_dtype = _extract_tensors(B) + + # Call pybind11 extension + output = tex.general_gemm( + A_data, A_scale_inv, A_dtype, + B_data, B_scale_inv, B_dtype, + bias, output_dtype, ... + ) + return output + +C++ Side (pybind11) +------------------- + +**Location**: ``transformer_engine/pytorch/csrc/extensions/`` + +The pybind11 module (``transformer_engine/pytorch/csrc/ts_fp8_op.cpp`` or similar) +registers Python-callable functions that: + +1. Accept ``torch::Tensor`` and scalar arguments from Python. +2. Construct ``NVTETensor`` handles using ``makeTransformerEngineTensor()``. +3. Call the C API function (e.g., ``nvte_general_gemm()``). +4. Return ``torch::Tensor`` outputs. + +**Example** — GEMM binding (simplified): + +.. code-block:: cpp + + void te_general_gemm( + at::Tensor A, at::Tensor A_scale_inv, int A_dtype, + at::Tensor B, at::Tensor B_scale_inv, int B_dtype, + at::Tensor D, at::Tensor bias, ...) { + + // Create NVTETensor handles + auto te_A = makeTransformerEngineTensor( + A.data_ptr(), A.sizes(), A_dtype, A_scale_inv, ...); + auto te_B = makeTransformerEngineTensor( + B.data_ptr(), B.sizes(), B_dtype, B_scale_inv, ...); + + // Call C API + nvte_general_gemm(te_A.data(), te_B.data(), te_D.data(), + te_bias.data(), stream); + } + +Tensor Conversion Helpers +------------------------- + +The ``makeTransformerEngineTensor()`` family of functions (in +``transformer_engine/pytorch/csrc/common.h``) handles the conversion from PyTorch +tensors to ``NVTETensor``: + +- Sets up the ``NVTEBasicTensor`` parameters (data pointer, shape, dtype). +- Attaches scaling metadata (scale, scale_inv, amax) via ``nvte_tensor_set()``. +- Sets the ``NVTEScalingMode`` on the tensor. + +The reverse conversion (``NVTETensor`` → PyTorch tensor) is typically not needed because +C API functions write into pre-allocated output buffers. + +Adding a New Extension +---------------------- + +To expose a new C API function to Python: + +1. Add the C API function to the appropriate header in ``include/transformer_engine/``. +2. Add a pybind11 wrapper in ``csrc/extensions/``. +3. Register it in the pybind11 module definition. +4. Add a Python wrapper in ``cpp_extensions/``. +5. Update ``__init__.py`` exports as needed. diff --git a/docs/developer/pytorch_frontend/img/ops_fusion.svg b/docs/developer/pytorch_frontend/img/ops_fusion.svg new file mode 100644 index 0000000000..23b7811275 --- /dev/null +++ b/docs/developer/pytorch_frontend/img/ops_fusion.svg @@ -0,0 +1,67 @@ + + + + + + + + + + + + + + + + Before Fusion + + + + LayerNorm + kernel launch + + + + + + + Quantize (FP8) + kernel launch + + + + + + + Linear (GEMM) + kernel launch + + + 3 kernel launches + + + + Op Fuser + + + After Fusion + + + + LayerNormLinear + (fused kernel) + single kernel launch + + + 1 kernel launch + + + + Fusion reduces overhead from kernel launch latency and memory traffic + diff --git a/docs/developer/pytorch_frontend/img/pytorch_module_hierarchy.svg b/docs/developer/pytorch_frontend/img/pytorch_module_hierarchy.svg new file mode 100644 index 0000000000..05242e70db --- /dev/null +++ b/docs/developer/pytorch_frontend/img/pytorch_module_hierarchy.svg @@ -0,0 +1,110 @@ + + + + + + + + + + + + + + torch.nn.Module + + + + + + + + + + + + + + + + TransformerEngineBaseModule + + + + TransformerLayer + + + (composes TE modules) + + + + + + + + + + + + + + + + + + + + + + + + + + Linear + + + + LayerNorm + + + + RMSNorm + + + + LayerNormLinear + + + + LayerNormMLP + + + + GroupedLinear + + + + DotProductAttention + + + + MultiheadAttention + + + + PyTorch base class + + + TE abstract base module + + + TE concrete modules + + + Composite module (composes TE modules) + diff --git a/docs/developer/pytorch_frontend/img/transformer_layer.svg b/docs/developer/pytorch_frontend/img/transformer_layer.svg new file mode 100644 index 0000000000..8ac0fae501 --- /dev/null +++ b/docs/developer/pytorch_frontend/img/transformer_layer.svg @@ -0,0 +1,98 @@ + + + + + + + + + + + + + + + + + Input + + + + + + + LayerNorm + + + + + + + + Self-Attention + (MultiheadAttention) + + + + + + + + + + + + + residual + + + + + + + LayerNorm + + + + + + + + MLP + (FC1 -> Act -> FC2) + + + + + + + + + + + + residual + + + + + + + Output + + + + Compute block + + + + + Residual add + + + Skip connection + diff --git a/docs/developer/pytorch_frontend/index.rst b/docs/developer/pytorch_frontend/index.rst new file mode 100644 index 0000000000..8cfc721340 --- /dev/null +++ b/docs/developer/pytorch_frontend/index.rst @@ -0,0 +1,20 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +PyTorch Frontend +================ + +The PyTorch frontend (``transformer_engine/pytorch/``) provides ``nn.Module`` subclasses, +autograd integration, quantized tensor types, and distributed utilities. It is the most +widely used frontend and the primary development target. + +.. toctree:: + :maxdepth: 1 + + module_hierarchy + autograd_integration + cpp_extensions + ops_framework + transformer_layer diff --git a/docs/developer/pytorch_frontend/module_hierarchy.rst b/docs/developer/pytorch_frontend/module_hierarchy.rst new file mode 100644 index 0000000000..84150817bf --- /dev/null +++ b/docs/developer/pytorch_frontend/module_hierarchy.rst @@ -0,0 +1,162 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +.. _module-hierarchy: + +Module Hierarchy +================ + +All PyTorch modules in Transformer Engine extend a common base class that provides +FP8 state management, quantizer lifecycle, and distributed training hooks. + +.. figure:: ./img/pytorch_module_hierarchy.svg + :align: center + :width: 80% + + Inheritance tree of TransformerEngine PyTorch modules. + +.. + Diagram description for ``pytorch_module_hierarchy.svg``: + Tree diagram with torch.nn.Module at the top. + Below it: TransformerEngineBaseModule. + Below that, branching to: + ├── Linear + ├── LayerNorm + ├── RMSNorm + ├── LayerNormLinear + ├── LayerNormMLP + ├── GroupedLinear + ├── DotProductAttention + └── MultiheadAttention + A separate branch from torch.nn.Module: TransformerLayer (composes the above modules + rather than inheriting from BaseModule). + +TransformerEngineBaseModule +--------------------------- + +**Location**: ``transformer_engine/pytorch/module/base.py`` + +``TransformerEngineBaseModule`` extends ``torch.nn.Module`` and is the foundation for all +TE modules that support FP8 quantization. Key responsibilities: + +**FP8 State Management** + +- ``init_fp8_metadata()``: Creates quantizer instances and amax history buffers based on + the active recipe. +- ``pre_forward()``: Called at the start of each forward pass to update FP8 state (refresh + scales from amax history, check if FP8 is enabled). +- ``post_forward()``: Called after forward to record amax values. +- FP8 metadata is registered as module buffers for serialization. + +**Quantizer Access** + +- Maintains quantizers for each quantized tensor (input, weight, gradient). +- Quantizer instances are recreated when the recipe changes. + +**Distributed Hooks** + +- ``set_tensor_parallel_group()``: Configure TP process group. +- ``set_sequence_parallel()``: Enable sequence parallelism. +- Manages all-reduce of amax values across distributed ranks. + +**Parameter Management** + +- ``weight`` parameter with optional FP8 storage. +- ``bias`` parameter (optional). +- Weight caching for FP8 weights that persist across iterations. + +Module Subclasses +----------------- + +Linear +^^^^^^ + +**Location**: ``transformer_engine/pytorch/module/linear.py`` + +The workhorse module. Performs ``output = input × weight^T + bias`` with optional FP8 +quantization of both input and weight. + +- Defines a ``_Linear`` autograd function (see :doc:`autograd_integration`). +- Supports tensor parallelism (column-parallel and row-parallel modes). +- Can fuse with preceding LayerNorm (via ``LayerNormLinear``). + +LayerNorm / RMSNorm +^^^^^^^^^^^^^^^^^^^^ + +**Location**: ``transformer_engine/pytorch/module/layernorm.py``, +``transformer_engine/pytorch/module/rmsnorm.py`` + +Normalization modules that can fuse their output quantization with the normalization +kernel (single kernel pass for norm + cast to FP8). + +LayerNormLinear +^^^^^^^^^^^^^^^ + +**Location**: ``transformer_engine/pytorch/module/layernorm_linear.py`` + +Fuses LayerNorm (or RMSNorm) with a subsequent Linear. The normalization output is +produced directly in FP8, avoiding a round-trip through high precision. + +LayerNormMLP +^^^^^^^^^^^^ + +**Location**: ``transformer_engine/pytorch/module/layernorm_mlp.py`` + +Fuses LayerNorm + Linear + Activation + Linear (the full MLP block). This is the most +aggressively fused module, combining up to 4 operations. + +GroupedLinear +^^^^^^^^^^^^^ + +**Location**: ``transformer_engine/pytorch/module/grouped_linear.py`` + +Batched linear operations for Mixture-of-Experts (MoE) where multiple smaller linear +layers execute as a single grouped GEMM. + +DotProductAttention +^^^^^^^^^^^^^^^^^^^ + +**Location**: ``transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py`` + +Computes scaled dot-product attention with automatic backend selection +(see :doc:`/developer/attention/backends`). Supports FP8 attention on Hopper+. + +MultiheadAttention +^^^^^^^^^^^^^^^^^^ + +**Location**: ``transformer_engine/pytorch/attention/multi_head_attention.py`` + +Composes QKV projection (Linear), DotProductAttention, and output projection (Linear) +into a complete multi-head attention block. Handles the split into heads and optional +key-value caching for inference. + +TransformerLayer +^^^^^^^^^^^^^^^^ + +**Location**: ``transformer_engine/pytorch/transformer.py`` + +Composes a full Transformer block (see :doc:`transformer_layer`). Note that +``TransformerLayer`` does **not** inherit from ``TransformerEngineBaseModule`` — it +composes TE modules rather than being one. + +Module Lifecycle +---------------- + +A typical forward pass through a TE module: + +.. code-block:: text + + 1. fp8_autocast() sets global FP8 state + 2. module.forward() called + a. pre_forward() — refresh FP8 scales, create quantizers + b. Quantize input via input_quantizer(input) + c. Get quantized weight (cached or quantize via weight_quantizer) + d. Call _Linear.forward() autograd function + i. general_gemm(quantized_input, quantized_weight) + ii. Save tensors for backward + e. post_forward() — record amax values + 3. Return output + +See :doc:`autograd_integration` for details on step (d). diff --git a/docs/developer/pytorch_frontend/ops_framework.rst b/docs/developer/pytorch_frontend/ops_framework.rst new file mode 100644 index 0000000000..68f8f0cfeb --- /dev/null +++ b/docs/developer/pytorch_frontend/ops_framework.rst @@ -0,0 +1,102 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +.. _ops-framework: + +Op Fusion Framework +=================== + +Transformer Engine includes an operation fusion framework (``transformer_engine/pytorch/ops/``) +that enables composing and fusing operations for better performance. + +Overview +-------- + +.. figure:: ./img/ops_fusion.svg + :align: center + :width: 70% + + Before and after op fusion: separate ops become a single fused op. + +.. + Diagram description for ``ops_fusion.svg``: + Left side "Before Fusion": + Three sequential boxes: "LayerNorm" → "Quantize" → "Linear" + Each box has a separate kernel launch arrow. + Right side "After Fusion": + Single box: "LayerNormLinear (fused)" + Single kernel launch arrow. + Arrow between left and right labeled "Op Fuser". + +The framework defines three levels of abstraction: + +- **FusibleOperation** (``ops/op.py``): Abstract base for any operation that can + participate in fusion. +- **BasicOperation** (``ops/op.py``): Single operations (GEMM, bias, normalization, etc.) + that serve as building blocks. Holds parameters and quantization state. +- **FusedOperation** (``ops/op.py``): Compound operations composed of multiple basic ops. + Delegates parameter management to its constituent basic ops. +- **OperationFuser** (``ops/fuser.py``): Manages the fusion pipeline, using a custom + autograd function (``_OperationFuserAutogradFunction``) to coordinate forward/backward. +- **Sequential** (``ops/sequential.py``): A container that groups consecutive + ``FusibleOperation`` instances for automatic fusion. + +Basic Ops +--------- + +Basic ops extend ``BasicOperation`` and implement ``op_forward()`` / ``op_backward()``: + +.. code-block:: text + + ops/basic/ + ├── linear.py # BasicLinear — matrix multiplication + ├── bias.py # Bias — bias addition + ├── normalization.py # LayerNorm / RMSNorm + ├── activation.py # GELU, GEGLU, SwiGLU, GLU + ├── quantization.py # Quantize — cast to/from FP8 + ├── all_gather.py # AllGather (distributed) + ├── all_reduce.py # AllReduce (distributed) + ├── reduce_scatter.py # ReduceScatter (distributed) + ├── dropout.py # Dropout + ├── reshape.py # Reshape + └── identity.py # Identity + +Fused Ops +--------- + +Fused ops extend ``FusedOperation`` and compose basic ops into optimized sequences. +The most important fused op is ``Linear`` (``ops/linear.py``), which composes +``BasicLinear`` + ``Bias`` + optional ``AllReduce``/``ReduceScatter`` for tensor +parallelism: + +.. code-block:: python + + # ops/linear.py — row-parallel example + Linear = FusedOperation([BasicLinear, Bias, ReduceScatter]) + +Op Fuser +-------- + +The ``OperationFuser`` (``ops/fuser.py``) manages a pipeline of fusible operations. It +wraps them in ``_OperationFuserAutogradFunction``, which: + +- **Forward**: Applies ops sequentially with potential fusion, passing quantizers between + ops. +- **Backward**: Reverses op order, applies fused/unfused backward, using per-op saved + contexts. + +``Sequential`` (``ops/sequential.py``) is the user-facing container. It groups +consecutive ``FusibleOperation`` instances via ``_make_module_groups()`` and applies +the fuser automatically. + +Relationship to Modules +----------------------- + +The ops framework is an alternative to the monolithic module approach +(``LayerNormLinear``, ``LayerNormMLP``). Modules like ``te.Linear`` (in ``module/``) use +the ``_Linear`` autograd function directly, while the ops framework allows more flexible +composition and automatic fusion discovery. + +Both approaches ultimately call the same C++ extensions and CUDA kernels. diff --git a/docs/developer/pytorch_frontend/transformer_layer.rst b/docs/developer/pytorch_frontend/transformer_layer.rst new file mode 100644 index 0000000000..3dbc580415 --- /dev/null +++ b/docs/developer/pytorch_frontend/transformer_layer.rst @@ -0,0 +1,84 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +.. _transformer-layer: + +TransformerLayer +================ + +``TransformerLayer`` (``transformer_engine/pytorch/transformer.py``) composes TE modules +into a complete Transformer block. Unlike other TE modules, it does **not** inherit from +``TransformerEngineBaseModule`` — it is a pure composition of TE sub-modules. + +.. figure:: ./img/transformer_layer.svg + :align: center + :width: 40% + + Data flow through a TransformerLayer. + +.. + Diagram description for ``transformer_layer.svg``: + Vertical flow diagram: + "Input" → "LayerNorm" → "Self-Attention (MHA)" → "+" (residual add) → + "LayerNorm" → "MLP (LayerNormMLP)" → "+" (residual add) → "Output" + Residual connections shown as arrows bypassing the attention and MLP blocks. + +Architecture +------------ + +A standard ``TransformerLayer`` contains: + +.. code-block:: text + + TransformerLayer + ├── self_attention: MultiheadAttention + │ ├── qkv_projection: LayerNormLinear # Fused LN + QKV projection + │ ├── core_attention: DotProductAttention # Attention computation + │ └── output_projection: Linear # Output projection + ├── layernorm_mlp: LayerNormMLP # Fused LN + FC1 + Act + FC2 + └── (optional) cross_attention: MultiheadAttention + +Configuration +------------- + +Key constructor parameters: + +.. list-table:: + :header-rows: 1 + :widths: 30 70 + + * - Parameter + - Description + * - ``hidden_size`` + - Model hidden dimension + * - ``ffn_hidden_size`` + - Feed-forward network intermediate size + * - ``num_attention_heads`` + - Number of attention heads + * - ``num_gqa_groups`` + - Number of GQA groups (for grouped-query attention) + * - ``layer_type`` + - ``"encoder"`` or ``"decoder"`` (controls cross-attention) + * - ``self_attn_mask_type`` + - ``"causal"``, ``"padding"``, ``"no_mask"``, etc. + * - ``normalization`` + - ``"LayerNorm"`` or ``"RMSNorm"`` + * - ``activation`` + - ``"gelu"``, ``"swiglu"``, ``"geglu"``, etc. + +Distributed Training +-------------------- + +When ``set_tensor_parallel_group()`` is called, the sub-modules automatically configure: + +- QKV projection as **column-parallel** (split heads across ranks). +- Attention output projection as **row-parallel** (each rank has partial output, reduced + via all-reduce or reduce-scatter). +- MLP FC1 as **column-parallel**, FC2 as **row-parallel**. + +With sequence parallelism enabled, the LayerNorm and residual additions operate on +sequence-partitioned data. + +See :doc:`/developer/distributed/tensor_parallel` for details on the parallelism patterns. diff --git a/docs/developer/quantization/adding_new_type.rst b/docs/developer/quantization/adding_new_type.rst new file mode 100644 index 0000000000..57e54402e4 --- /dev/null +++ b/docs/developer/quantization/adding_new_type.rst @@ -0,0 +1,128 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +Adding a New Quantization Type +============================== + +This guide walks through the steps to add a new low-precision format to Transformer +Engine, using the existing types (FP8, MXFP8, NVFP4) as templates. + +Prerequisites +------------- + +Before starting, you need: + +- A new data type (or an existing one with a new scaling strategy). +- CUDA kernel(s) for quantization (cast) and dequantization. +- cuBLASLt support for the type in GEMM (or a custom GEMM kernel). + +Step 1: Add the Data Type to the C API +--------------------------------------- + +If your format uses a new data type (not already in ``NVTEDType``): + +1. Add an entry to the ``NVTEDType`` enum in + ``transformer_engine/common/include/transformer_engine/transformer_engine.h``: + + .. code-block:: c + + enum NVTEDType { + // ... existing types ... + kNVTENewType = 11, + kNVTENumTypes + }; + +2. Add the corresponding entry to ``DType`` in ``transformer_engine/common/common.h``. + +3. Add overloads for: ``typeToSize``, ``typeToNumBits``, ``is_fp8_dtype`` (if applicable), + ``get_cuda_dtype``, ``get_cudnn_dtype``, and ``to_string``. + +Step 2: Add a Scaling Mode (if needed) +--------------------------------------- + +If your format uses a new scaling strategy: + +1. Add an entry to ``NVTEScalingMode`` in ``transformer_engine.h``. +2. Add helper functions in ``common.h`` (e.g., ``is_new_scaling()``). + +Step 3: Implement Cast Kernels +------------------------------- + +Create a new file or extend existing files in ``transformer_engine/common/cast/``: + +1. Implement the quantization kernel (high precision → new format). +2. Implement the dequantization kernel (new format → high precision). +3. Register C API entry points in the cast header. + +Step 4: Extend GEMM Support +---------------------------- + +In ``transformer_engine/common/gemm/``: + +1. Add handling for the new type in cuBLASLt GEMM dispatch. +2. Handle scale/scale_inv for the new scaling mode. +3. Update the compute type selection logic. + +Step 5: Create Python Quantizer and Tensor Classes +--------------------------------------------------- + +Create a new file ``transformer_engine/pytorch/tensor/new_type_tensor.py`` with: + +1. **``NewTypeQuantizer(Quantizer)``**: Implements ``__call__()`` and ``quantize()`` + to produce quantized tensors in the new format. + +2. **``NewTypeTensorStorage(QuantizedTensorStorage)``**: Holds the raw quantized data + and scale inverses. + +3. **``NewTypeTensor(QuantizedTensor)``**: Full ``torch.Tensor`` subclass with + ``__torch_dispatch__`` support. + +Use ``transformer_engine/pytorch/tensor/float8_tensor.py`` as a reference implementation. + +Step 6: Register in the Recipe System +-------------------------------------- + +1. Add a recipe option that selects your new quantizer. +2. Update the recipe → quantizer mapping in the module initialization code. + +Step 7: Update C++ Extensions +------------------------------ + +In ``transformer_engine/pytorch/cpp_extensions/``: + +1. Update tensor conversion utilities to handle the new quantized tensor type. +2. Ensure ``NVTETensor`` construction properly sets scaling mode and parameters. + +Step 8: Test +------------ + +Minimum test coverage: + +- [ ] Cast round-trip: quantize → dequantize produces acceptable error. +- [ ] GEMM: FP8 GEMM with new type matches BF16 GEMM within tolerance. +- [ ] Linear forward/backward: ``te.Linear`` produces correct gradients. +- [ ] Distributed: quantized tensors survive all-gather/reduce-scatter. +- [ ] Serialization: ``torch.save`` / ``torch.load`` round-trip. + +Checklist +--------- + +.. code-block:: text + + [ ] NVTEDType entry (if new type) + [ ] DType entry (if new type) + [ ] Type utility overloads (typeToSize, etc.) + [ ] NVTEScalingMode entry (if new scaling) + [ ] Scaling mode helpers (is_new_scaling, etc.) + [ ] Cast kernel (quantize) + [ ] Cast kernel (dequantize) + [ ] GEMM support + [ ] NewTypeQuantizer class + [ ] NewTypeTensorStorage class + [ ] NewTypeTensor class + [ ] Recipe integration + [ ] C++ extension updates + [ ] Unit tests + [ ] Lint passes (Black, clang-format, cpplint) diff --git a/docs/developer/quantization/class_hierarchy.rst b/docs/developer/quantization/class_hierarchy.rst new file mode 100644 index 0000000000..6d9af28ea5 --- /dev/null +++ b/docs/developer/quantization/class_hierarchy.rst @@ -0,0 +1,195 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +.. _quantization-class-hierarchy: + +Class Hierarchy +=============== + +The quantization system is built around three parallel class hierarchies that work +together: **Quantizers** (builders), **Storage** (lightweight internal), and **Tensors** +(full PyTorch integration). + +.. figure:: ./img/quantizer_hierarchy.svg + :align: center + :width: 90% + + The three-column quantization class hierarchy. + +.. + Diagram description for ``quantizer_hierarchy.svg``: + Three columns side by side: + Column 1 "Quantizer (Builder)": + Quantizer (abstract base) + ├── Float8Quantizer (delayed tensor scaling) + ├── Float8CurrentScalingQuantizer (current tensor scaling) + ├── MXFP8Quantizer (MXFP8 block scaling) + ├── Float8BlockQuantizer (1D/2D block scaling) + └── NVFP4Quantizer (NVFP4 scaling) + Column 2 "Storage (Internal)": + QuantizedTensorStorage (abstract base) + ├── Float8TensorStorage + ├── MXFP8TensorStorage + ├── Float8BlockwiseQTensorStorage + └── NVFP4TensorStorage + Column 3 "Tensor (PyTorch)": + QuantizedTensor (torch.Tensor subclass) + ├── Float8Tensor + ├── MXFP8Tensor + ├── Float8BlockwiseQTensor + └── NVFP4Tensor + Arrows from each Quantizer to its corresponding Storage and Tensor. + +Design Rationale +---------------- + +The three-class design exists because quantized data has two distinct consumers: + +1. **Internal kernel code** needs lightweight access to raw data + scales, without + PyTorch autograd overhead. This is served by ``QuantizedTensorStorage``. +2. **User-facing code** needs a full ``torch.Tensor`` subclass that participates in + autograd, supports ``__torch_dispatch__``, and can be passed to standard PyTorch APIs. + This is served by ``QuantizedTensor``. +3. **Configuration and creation** needs a stateful builder that knows how to compute + scales, manage amax history, and produce both storage and tensor forms. This is the + ``Quantizer``. + +Quantizer (Base: ``transformer_engine/pytorch/quantized_tensor.py``) +--------------------------------------------------------------------- + +The ``Quantizer`` abstract class is the entry point for all quantization. Key +responsibilities: + +- **Configure** what kind of quantization to apply (data type, scaling mode, block size). +- **Track usage** mode (rowwise, columnwise, or both) via ``set_usage()``. +- **Produce** quantized tensors via ``__call__()`` (or ``quantize()``). +- **Manage state** for delayed scaling (amax history, scale update). + +.. code-block:: python + + class Quantizer: + # Core interface + def __call__(self, tensor, *, noop=False) -> QuantizedTensor: ... + def quantize(self, tensor, *, noop=False) -> QuantizedTensorStorage: ... + + # Usage control + def set_usage(self, *, rowwise=False, columnwise=False): ... + + # Properties + @property + def dtype(self) -> torch.dtype: ... + @property + def scaling_mode(self) -> str: ... + +Concrete Quantizers +^^^^^^^^^^^^^^^^^^^ + +.. list-table:: + :header-rows: 1 + :widths: 30 20 50 + + * - Class + - Location + - Description + * - ``Float8Quantizer`` + - ``tensor/float8_tensor.py`` + - Delayed tensor scaling with amax history + * - ``Float8CurrentScalingQuantizer`` + - ``tensor/float8_tensor.py`` + - Current (just-in-time) tensor scaling + * - ``MXFP8Quantizer`` + - ``tensor/mxfp8_tensor.py`` + - MXFP8 with 32-element E8M0 block scales + * - ``Float8BlockQuantizer`` + - ``tensor/float8_blockwise_tensor.py`` + - 1D or 2D block scaling with configurable block size + * - ``NVFP4Quantizer`` + - ``tensor/nvfp4_tensor.py`` + - NVFP4 with 16-element E8M0 block scales + +QuantizedTensorStorage (Base: ``transformer_engine/pytorch/quantized_tensor.py``) +---------------------------------------------------------------------------------- + +``QuantizedTensorStorage`` is a lightweight wrapper around quantized data and its +associated scales. It is **not** a ``torch.Tensor`` subclass — it has no autograd +support and minimal overhead. + +.. code-block:: python + + class QuantizedTensorStorage: + @property + def data(self) -> torch.Tensor: ... # Raw quantized data + @property + def scale_inv(self) -> torch.Tensor: ... # Scale inverse for dequantization + @property + def dtype(self) -> torch.dtype: ... + @property + def scaling_mode(self) -> str: ... + + def dequantize(self) -> torch.Tensor: ... # Convert back to high precision + +Storage objects are used internally by: + +- C++ extension calls (converted to ``NVTETensor`` for kernel dispatch) +- The GEMM pipeline (directly passes data + scale_inv) +- Op fusion internals (avoids autograd overhead in fused kernels) + +QuantizedTensor (Base: ``transformer_engine/pytorch/quantized_tensor.py``) +--------------------------------------------------------------------------- + +``QuantizedTensor`` is a ``torch.Tensor`` subclass that wraps a +``QuantizedTensorStorage`` and adds PyTorch integration: + +.. code-block:: python + + class QuantizedTensor(torch.Tensor): + _storage: QuantizedTensorStorage + + # torch.Tensor subclass machinery + def __torch_dispatch__(cls, func, types, args, kwargs): ... + + # Access underlying storage + def get_storage(self) -> QuantizedTensorStorage: ... + + # Dequantize + def dequantize(self) -> torch.Tensor: ... + def float(self) -> torch.Tensor: ... + +Key behaviors: + +- **Dequantize on access**: Most ``torch`` operations trigger automatic dequantization + via ``__torch_dispatch__``, so quantized tensors "just work" in standard PyTorch code + (at the cost of a dequantize). +- **GEMM fast path**: The GEMM implementation checks for ``QuantizedTensor`` inputs and + extracts the storage directly, avoiding dequantization. +- **Autograd compatible**: Can be saved in autograd's ``ctx.save_for_backward()``. + +Lifecycle +--------- + +A typical quantization lifecycle: + +.. code-block:: python + + # 1. Create quantizer (usually done once per module) + quantizer = Float8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) + quantizer.set_usage(rowwise=True, columnwise=True) + + # 2. Quantize a tensor (returns QuantizedTensor for autograd) + qinput = quantizer(input_tensor) + + # 3. Inside GEMM, extract storage for kernel call + storage = qinput.get_storage() + # Pass storage.data and storage.scale_inv to C++ kernel + + # 4. Dequantize if needed for non-optimized ops + fp32_data = qinput.dequantize() + +See Also +-------- + +- :doc:`scaling_recipes` — How scaling parameters (amax, scale) are managed globally +- :doc:`rowwise_columnwise` — Why both layouts exist and when each is used +- :doc:`adding_new_type` — Step-by-step guide to implementing a new quantization format diff --git a/docs/developer/quantization/img/quantizer_hierarchy.svg b/docs/developer/quantization/img/quantizer_hierarchy.svg new file mode 100644 index 0000000000..9ea425c615 --- /dev/null +++ b/docs/developer/quantization/img/quantizer_hierarchy.svg @@ -0,0 +1,182 @@ + + + + + + + + + + + + + Quantization Class Hierarchy + + + + + + + + Quantizer (Builder) + + + + Quantizer + abstract base + + + + + + + + + + Float8Quantizer + + + + + Float8CurrentScaling + Quantizer + + + + + MXFP8Quantizer + + + + + Float8BlockQuantizer + + + + + NVFP4Quantizer + + + + + + + + Storage (Internal) + + + + QuantizedTensorStorage + abstract base + + + + + + + + + + Float8TensorStorage + + + + + MXFP8TensorStorage + + + + + Float8BlockwiseQ + TensorStorage + + + + + NVFP4TensorStorage + + + + + + + + Tensor (PyTorch) + + + + QuantizedTensor + torch.Tensor subclass + + + + + + + + + + Float8Tensor + + + + + MXFP8Tensor + + + + + Float8BlockwiseQTensor + + + + + NVFP4Tensor + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + produces / wraps + + + shares Float8 storage + \ No newline at end of file diff --git a/docs/developer/quantization/img/rowwise_columnwise.svg b/docs/developer/quantization/img/rowwise_columnwise.svg new file mode 100644 index 0000000000..81fcdc08b6 --- /dev/null +++ b/docs/developer/quantization/img/rowwise_columnwise.svg @@ -0,0 +1,141 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + Forward GEMM + + + + + + + + + A + Activation + (rowwise, FP8) + scale_inv per row + + + × + + + + + + + + + B + Weight + (columnwise, FP8) + columnwise_scale_inv + per column + + + = + + + + Output + Output + (BF16) + + + Y = A × B + Row-scaled activations × column-scaled weights + + + + + + + + Backward Wgrad GEMM + + + + + + + + + Aᵀ + Activationᵀ + (columnwise from fwd) + + + × + + + + + + + + + dY + grad_output + (rowwise, FP8) + + + = + + + + Weight + grad + + + dW = Aᵀ × dY + Transposed activations reused from forward pass + + + + + + + + + + same data, different layout + + + + + + Key Insight + Forward pass quantizes activations rowwise (for fwd GEMM) and columnwise (saved for bwd wgrad GEMM). + \ No newline at end of file diff --git a/docs/developer/quantization/index.rst b/docs/developer/quantization/index.rst new file mode 100644 index 0000000000..d72dd90126 --- /dev/null +++ b/docs/developer/quantization/index.rst @@ -0,0 +1,25 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +Quantization +============ + +Transformer Engine supports multiple low-precision formats (FP8 E4M3/E5M2, MXFP8, +block-scaled FP8, NVFP4) through a unified quantization architecture. This section +describes the internal design of that system. + +The quantization system has three main axes: + +- **What** to quantize (the Quantizer/Storage/Tensor class hierarchy) +- **How** to scale (the scaling recipe and global FP8 state) +- **Where** to store (rowwise vs. columnwise layouts for GEMM) + +.. toctree:: + :maxdepth: 1 + + class_hierarchy + scaling_recipes + rowwise_columnwise + adding_new_type diff --git a/docs/developer/quantization/rowwise_columnwise.rst b/docs/developer/quantization/rowwise_columnwise.rst new file mode 100644 index 0000000000..40345d9a48 --- /dev/null +++ b/docs/developer/quantization/rowwise_columnwise.rst @@ -0,0 +1,116 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +.. _rowwise-columnwise: + +Rowwise vs. Columnwise Usage +============================= + +Quantized tensors in Transformer Engine can be stored in two layouts: **rowwise** and +**columnwise**. This dual-layout design is driven by cuBLASLt's requirements for FP8 +GEMM operands. + +.. figure:: ./img/rowwise_columnwise.svg + :align: center + :width: 80% + + Forward GEMM uses rowwise data; backward wgrad GEMM uses columnwise data. + +.. + Diagram description for ``rowwise_columnwise.svg``: + Two side-by-side GEMM diagrams: + Left "Forward GEMM (dgrad)": + Matrix A (activation, rowwise) × Matrix B (weight, columnwise) = Output + A is labeled "row-major, scale_inv per row" + B is labeled "col-major, columnwise_scale_inv per column" + Right "Backward GEMM (wgrad)": + Matrix A^T (activation transposed, columnwise) × Matrix B (grad_output, rowwise) = Weight grad + A^T is labeled "using columnwise_data from forward" + Arrow from left A to right A^T labeled "same data, different layout view" + +Why Two Layouts? +---------------- + +cuBLASLt's FP8 GEMM computes ``C = A × B`` where: + +- ``A`` is read in **row-major** order (consecutive elements along the last dimension). +- ``B`` is read in **column-major** order (consecutive elements along the first dimension). + +In the **forward pass**: + +- ``output = activation × weight^T`` +- Activation is the ``A`` operand → needs **rowwise** quantization. +- Weight is the ``B`` operand → needs **columnwise** quantization. + +In the **backward wgrad pass**: + +- ``weight_grad = activation^T × grad_output`` +- Activation transposed is the ``A`` operand → needs **columnwise** data from forward. +- Grad output is the ``B`` operand → needs **columnwise** quantization. + +This is why the forward pass must produce *both* rowwise and columnwise quantized data +for activations: the rowwise data is consumed immediately by the forward GEMM, while the +columnwise data is saved for the backward wgrad GEMM. + +Setting Usage on Quantizers +---------------------------- + +The ``Quantizer.set_usage()`` method controls which layouts are produced: + +.. code-block:: python + + # Forward activation quantizer: needs both layouts + fwd_quantizer.set_usage(rowwise=True, columnwise=True) + + # Weight quantizer: only columnwise for forward GEMM + weight_quantizer.set_usage(rowwise=False, columnwise=True) + + # Backward grad_output quantizer: needs both for dgrad and wgrad + bwd_quantizer.set_usage(rowwise=True, columnwise=True) + +When ``columnwise=True``, the quantization kernel produces an additional transposed +copy of the data (or, for block-scaling modes, computes column-oriented block scales). + +Impact on the Tensor Struct +---------------------------- + +In the C++ ``Tensor`` struct, the layout is reflected in separate fields: + +.. code-block:: cpp + + struct Tensor { + SimpleTensor data; // Rowwise quantized data + SimpleTensor columnwise_data; // Columnwise quantized data + SimpleTensor scale_inv; // Rowwise scale inverses + SimpleTensor columnwise_scale_inv; // Columnwise scale inverses + // ... + }; + +For **tensor scaling** (delayed/current), ``data`` and ``columnwise_data`` point to +different memory buffers (the data is physically transposed). + +For **block scaling** modes, the data may be the same buffer but with different scale +tensors (row-oriented vs. column-oriented block scales). + +Performance Implications +------------------------ + +Producing both layouts doubles the quantization work and memory for activations. This is +a deliberate trade-off: + +- **Without dual layout**: The wgrad GEMM would need to transpose and requantize at + backward time, adding latency to the critical path. +- **With dual layout**: Extra memory and compute during forward, but the backward wgrad + can proceed immediately with pre-computed columnwise data. + +For memory-constrained scenarios, activation recomputation can be used — the forward +pass discards the saved columnwise data and recomputes it during backward. See +:doc:`/developer/pytorch_frontend/autograd_integration` for details. + +See Also +-------- + +- :doc:`class_hierarchy` — How Storage objects hold both layouts +- :doc:`/developer/cpp_core/scaling_modes` — Scale granularity per mode diff --git a/docs/developer/quantization/scaling_recipes.rst b/docs/developer/quantization/scaling_recipes.rst new file mode 100644 index 0000000000..f338b04103 --- /dev/null +++ b/docs/developer/quantization/scaling_recipes.rst @@ -0,0 +1,136 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +.. _scaling-recipes: + +Scaling Recipes +=============== + +Scaling recipes control *how* quantization parameters (scales, amax values) are computed +and updated across training iterations. They sit between the user-facing ``recipe`` API +and the per-tensor ``Quantizer`` instances. + +Recipe Configuration +-------------------- + +Users configure quantization via recipe classes from ``transformer_engine.common.recipe``. +The recipe hierarchy: + +- ``Recipe`` — abstract base class +- ``DelayedScaling`` — delayed tensor scaling with amax history +- ``Float8CurrentScaling`` — just-in-time per-tensor scaling +- ``MXFP8BlockScaling`` — MXFP8 with 32-element blocks +- ``Float8BlockScaling`` — configurable 1D/2D block scaling +- ``NVFP4BlockScaling`` — NVFP4 with 16-element blocks +- ``CustomRecipe`` — user-provided quantizer factory + +Recipe classes specify: + +- **FP8 format**: Which FP8 types to use for forward (typically E4M3) and backward + (typically E5M2). +- **Amax history length**: How many past amax values to track (default: 1024). +- **Amax compute algorithm**: How to derive the scale from history (typically ``max``). +- **Scaling mode**: Delayed tensor, current tensor, MXFP8, block, or NVFP4. + +.. code-block:: python + + import transformer_engine.pytorch as te + from transformer_engine.common.recipe import DelayedScaling + + recipe = DelayedScaling( + fp8_format=te.recipe.Format.HYBRID, # E4M3 fwd, E5M2 bwd + amax_history_len=1024, + amax_compute_algo="max", + ) + + with te.fp8_autocast(enabled=True, fp8_recipe=recipe): + output = model(input) + +FP8GlobalStateManager +--------------------- + +The ``FP8GlobalStateManager`` (in ``transformer_engine/pytorch/quantization.py``) is a singleton +that coordinates FP8 state across all modules in a model. Key responsibilities: + +- **Track global FP8 enabled/disabled state** for the current forward pass. +- **Distribute recipe configuration** to all ``TransformerEngineBaseModule`` instances. +- **Coordinate amax reduction** across distributed ranks (for delayed scaling, all ranks + must agree on the global amax). + +The manager is activated by the ``fp8_autocast`` context manager: + +.. code-block:: python + + # Pseudocode for fp8_autocast + @contextmanager + def fp8_autocast(enabled, fp8_recipe): + FP8GlobalStateManager.set_enabled(enabled) + FP8GlobalStateManager.set_recipe(fp8_recipe) + try: + yield + finally: + FP8GlobalStateManager.reset() + +Per-Module State +---------------- + +Each ``TransformerEngineBaseModule`` maintains its own quantization state: + +- **Quantizer instances**: One per quantized tensor (input, weight, output gradient, weight + gradient). Created based on the active recipe. +- **Amax history buffers**: Registered as module buffers for delayed scaling. Updated + after each forward/backward pass. +- **Scale tensors**: Computed from amax history before each forward pass. + +The module's ``init_fp8_metadata()`` method creates this state, and +``pre_forward()`` / ``post_forward()`` manage the per-iteration lifecycle. + +Amax Update Flow (Delayed Scaling) +----------------------------------- + +For delayed tensor scaling, the amax update follows this sequence each iteration: + +1. **Pre-forward**: Compute new scales from existing amax history. +2. **Forward**: GEMM kernels record the amax of their FP8 outputs into a buffer. +3. **Backward**: Same as forward for gradient GEMMs. +4. **Post-backward**: Roll the amax history window (shift old values, insert new amax). +5. **All-reduce**: If using distributed training, reduce amax across ranks. + +For current-scaling and block-scaling modes, scales are computed just-in-time during +the cast kernel, so no amax history is maintained. + +Recipe → Quantizer Mapping +-------------------------- + +The recipe's scaling mode determines which ``Quantizer`` subclass is instantiated: + +.. list-table:: + :header-rows: 1 + :widths: 40 30 30 + + * - Recipe / Mode + - Forward Quantizer + - Backward Quantizer + * - ``DelayedScaling`` + - ``Float8Quantizer`` + - ``Float8Quantizer`` + * - Current scaling + - ``Float8CurrentScalingQuantizer`` + - ``Float8CurrentScalingQuantizer`` + * - MXFP8 + - ``MXFP8Quantizer`` + - ``MXFP8Quantizer`` + * - Block scaling + - ``Float8BlockQuantizer`` + - ``Float8BlockQuantizer`` + * - NVFP4 + - ``NVFP4Quantizer`` + - ``NVFP4Quantizer`` + +See Also +-------- + +- :doc:`class_hierarchy` — The Quantizer/Storage/Tensor class design +- :doc:`/developer/pytorch_frontend/module_hierarchy` — How modules manage FP8 state From a0c674c964af07d9b9c6ffd4f484989982788383 Mon Sep 17 00:00:00 2001 From: Przemek Tredak Date: Mon, 16 Mar 2026 11:38:28 -0700 Subject: [PATCH 02/20] Some fixes Signed-off-by: Przemek Tredak --- docs/developer/architecture_overview.rst | 66 +-- docs/developer/cpp_core/type_system.rst | 14 +- docs/developer/index.rst | 6 +- docs/developer/linear_walkthrough.rst | 334 +++++++++++---- .../pytorch_frontend/ops_framework.rst | 404 +++++++++++++++--- docs/index.rst | 8 +- 6 files changed, 666 insertions(+), 166 deletions(-) diff --git a/docs/developer/architecture_overview.rst b/docs/developer/architecture_overview.rst index 85c044408f..7e928c6288 100644 --- a/docs/developer/architecture_overview.rst +++ b/docs/developer/architecture_overview.rst @@ -127,31 +127,47 @@ and once for the weight gradient (wgrad). The wgrad path uses *columnwise* quant to satisfy cuBLASLt's transposed layout requirements — see :doc:`quantization/rowwise_columnwise` for details. -Key Design Decisions --------------------- - -**C API boundary** - The C++ core exposes a C API (not C++) for ABI stability. This means framework bindings - interact with opaque ``NVTETensor`` handles rather than C++ objects directly. The C++ - ``Tensor`` and ``SimpleTensor`` structs in ``common.h`` are internal implementation - details. - -**Dual type systems** - The C API uses ``NVTEDType`` (a plain C enum), while the C++ core uses - ``DType`` (a C++ ``enum class`` in the ``transformer_engine`` namespace). These are - related but not implicitly convertible — see :doc:`cpp_core/type_system` for conversion - patterns. - -**Framework-agnostic kernels** - All CUDA kernels live in ``common/`` and know nothing about PyTorch or JAX. Framework - frontends are responsible for converting their tensor types to ``NVTETensor`` before - calling into the C API. - -**Quantizer pattern** - Quantization is not a single function but a stateful builder (``Quantizer``) that - produces ``QuantizedTensorStorage`` (lightweight, no autograd) or ``QuantizedTensor`` - (full torch.Tensor subclass). This separation keeps internal kernel code free from - autograd overhead — see :doc:`quantization/class_hierarchy`. +Guiding Principles +------------------- + +These principles apply across the entire codebase and inform design decisions at every +layer. + +**Three API surfaces, one core** + Transformer Engine exposes three distinct API surfaces — a C API + (``transformer_engine.h``), PyTorch ``nn.Module`` subclasses, and JAX XLA primitives — + all backed by the same framework-agnostic C++ core. CUDA kernels live in ``common/`` + and know nothing about PyTorch or JAX; framework frontends are responsible for + converting their tensor types before calling into the C API. + +**Feel native to each framework** + Each frontend should feel like a natural extension of its framework, not a foreign + library. PyTorch modules follow ``nn.Module`` conventions (autograd, ``state_dict``, + ``torch.compile``). JAX primitives register with XLA's custom-call and sharding + machinery. Users should be able to mix and match TE components with native framework + code freely — a ``te.Linear`` should be a drop-in replacement for ``torch.nn.Linear``. + +**Leverage the ecosystem, don't reinvent it** + TE provides custom CUDA kernels where they deliver clear performance wins (fused + attention, quantized GEMM, fused normalization + cast). For everything else, we prefer + the framework's native capabilities — ``torch.compile`` for fusion, XLA for graph + optimization, cuBLASLt for matrix multiplication. A custom kernel must justify its + existence against the ecosystem alternative. + +**No memory allocation in the C++ layer** + The C++ core never allocates GPU memory directly. All memory is allocated by the + framework frontend and passed down to the C API as pointers within ``NVTETensor`` + handles. This keeps memory management in the framework's control (important for + memory pools, garbage collection, and distributed training) and avoids hidden + allocation surprises. + +**Hide complexity, but stay extensible** + Low-precision training involves many moving parts (scaling modes, amax history, block + sizes, format selection). TE hides this complexity behind simple APIs — a user can + pass a recipe to ``fp8_autocast`` and get FP8 training without understanding the + internals. But we do not want to be a black box: the quantization system is designed + to be extensible (see :doc:`quantization/adding_new_type`), scaling recipes are + configurable, and advanced users can provide custom quantizers. Next Steps ---------- diff --git a/docs/developer/cpp_core/type_system.rst b/docs/developer/cpp_core/type_system.rst index 9030adc08a..afa0373424 100644 --- a/docs/developer/cpp_core/type_system.rst +++ b/docs/developer/cpp_core/type_system.rst @@ -8,8 +8,18 @@ Type System =========== -Transformer Engine maintains two parallel type systems for historical and ABI-stability -reasons: a C type system for the public API and a C++ type system for internal use. +Transformer Engine maintains two parallel type systems: a C type system for the public +API and a C++ type system for internal use. + +The C++ core exposes a **C API** (not C++) for ABI stability. Framework bindings interact +with opaque ``NVTETensor`` handles rather than C++ objects directly. The C++ ``Tensor`` +and ``SimpleTensor`` structs in ``common.h`` are internal implementation details — +they are never exposed across the API boundary. + +This design means there are **two dtype enums** (``NVTEDType`` in C, ``DType`` in C++) +and **two tensor abstractions** (``NVTEBasicTensor`` in C, ``SimpleTensor`` in C++). They +mirror each other but are not implicitly convertible. The rest of this page documents both +systems and the conversion patterns between them. C Types (Public API) -------------------- diff --git a/docs/developer/index.rst b/docs/developer/index.rst index cf74227043..c2485a8719 100644 --- a/docs/developer/index.rst +++ b/docs/developer/index.rst @@ -19,9 +19,9 @@ How to Use This Guide --------------------- - **New contributors**: Start with :doc:`architecture_overview` for the big picture, - then read :doc:`linear_walkthrough` for an end-to-end trace through a concrete module. -- **Working on quantization**: See :doc:`quantization/index` for the Quantizer/Storage/Tensor - design and :doc:`cpp_core/type_system` for the underlying C/C++ types. + then read :doc:`linear_walkthrough` for an end-to-end trace through a concrete PyTorch module. +- **Working on new quantization recipe**: See :doc:`quantization/index` for the + Quantizer/Storage/Tensor design and :doc:`cpp_core/type_system` for the underlying C/C++ types. - **Working on attention**: See :doc:`attention/index` for backend selection and kernel organization. - **Working on distributed**: See :doc:`distributed/index` for tensor/sequence parallelism diff --git a/docs/developer/linear_walkthrough.rst b/docs/developer/linear_walkthrough.rst index c85334267a..87dc8959bc 100644 --- a/docs/developer/linear_walkthrough.rst +++ b/docs/developer/linear_walkthrough.rst @@ -12,6 +12,11 @@ This page traces a complete forward and backward pass through ``te.Linear``, the fundamental Transformer Engine module. It connects concepts from across the codebase: quantization, GEMM, distributed communication, and autograd. +The goal is not just to show *what* happens at each step, but *why* the code is structured +this way — what constraints drive the ordering, what trade-offs are being made, and how +low-precision quantization fundamentally changes the data flow compared to a standard +``torch.nn.Linear``. + .. figure:: ./img/linear_e2e_flow.svg :align: center :width: 95% @@ -36,19 +41,45 @@ quantization, GEMM, distributed communication, and autograd. Dotted arrows from forward boxes 3→backward box 4 labeled "saved columnwise input" and forward weight→backward box 3 labeled "saved weight". +Why This Walkthrough Matters +----------------------------- + +A standard ``torch.nn.Linear`` is straightforward: ``output = input @ weight.T + bias``. +TE's Linear does the same math, but introduces two major dimensions of complexity: + +1. **Low-precision quantization** — Inputs and weights must be cast to a low-precision + format (FP8, MXFP8, NVFP4) *before* GEMM, and the backward pass needs the data in a + *different layout* (columnwise) than the forward pass (rowwise). This means the forward + pass must proactively prepare data for backward. + +2. **Tensor parallelism** — The weight matrix is sharded across GPUs, so collective + communication (all-gather, reduce-scatter) must be interleaved with computation. + The placement of communication relative to GEMM depends on whether this is a + column-parallel or row-parallel linear. + +These two concerns interact with each other at every step — for instance, quantizing +*before* an all-gather halves the communication volume. The walkthrough below makes these +interactions explicit. + +The flow described here is **recipe-agnostic**: the same phases execute regardless of +whether the active recipe is MXFP8, current scaling, block scaling, or delayed scaling. +The quantizer abstraction (see :doc:`quantization/class_hierarchy`) hides recipe-specific +details — ``_Linear`` simply calls ``quantizer(tensor)`` and receives quantized data with +the appropriate scales, no matter how those scales were computed. + Setup ----- .. code-block:: python import transformer_engine.pytorch as te - from transformer_engine.common.recipe import DelayedScaling + from transformer_engine.common.recipe import MXFP8BlockScaling - # Create a Linear module with FP8 + # Create a Linear module linear = te.Linear(4096, 16384, bias=True) - # Enable FP8 - recipe = DelayedScaling(fp8_format=te.recipe.Format.HYBRID) + # Enable FP8 with MXFP8 recipe + recipe = MXFP8BlockScaling() with te.fp8_autocast(enabled=True, fp8_recipe=recipe): output = linear(input) # Triggers the flow below @@ -57,167 +88,316 @@ Phase 1: Module Forward Entry **File**: ``transformer_engine/pytorch/module/linear.py``, ``Linear.forward()`` (line ~1343) -1. ``pre_forward()`` is called (inherited from ``TransformerEngineBaseModule``): +**Why this phase exists**: Before any math can happen, TE needs to set up the quantization +infrastructure. Unlike BF16 where you just call GEMM directly, low-precision formats +require quantizer objects that know how to cast data and manage scales. This setup phase +bridges the gap between the user-facing recipe configuration and the quantizer objects +that actually perform casts. + +1. ``prepare_forward()`` is called (inherited from ``TransformerEngineBaseModule``): + + - Checks if FP8 is enabled via ``FP8GlobalStateManager`` — a global singleton that + tracks whether we're inside an ``fp8_autocast`` context. This global state exists + because FP8 behavior must be coordinated across all TE modules in the model. + - Creates or refreshes **quantizer** instances from the active recipe. Each recipe + type produces a different quantizer class (``MXFP8Quantizer``, + ``Float8CurrentScalingQuantizer``, ``Float8BlockQuantizer``, etc.). + - Recipe-specific customization is applied — for example, current scaling configures + ``amax_epsilon`` and ``power_2_scale`` on quantizers, while MXFP8 needs no such + configuration. - - Checks if FP8 is enabled via ``FP8GlobalStateManager``. - - Creates or refreshes quantizer instances from the active recipe. - - For delayed scaling: computes new scales from amax history. +2. Three quantizers are prepared, one for each tensor that will be cast to FP8: -2. Quantizers are prepared: + - ``input_quantizer`` — for the activation tensor (forward GEMM input) + - ``weight_quantizer`` — for the weight parameter (forward GEMM input) + - ``grad_output_quantizer`` — for the backward gradient (backward GEMM input) - - ``input_quantizer`` — for the activation tensor - - ``weight_quantizer`` — for the weight parameter - - ``grad_output_quantizer`` — for the backward gradient + **Why three separate quantizers?** Each tensor has independent scaling state because + their value distributions differ significantly — activations, weights, and gradients + have different magnitudes and ranges. A single shared scale would waste dynamic range. -3. ``_Linear.apply()`` is called, entering the autograd function. +3. ``_Linear.apply()`` is called, entering the custom ``torch.autograd.Function``. This + boundary separates the ``nn.Module`` API from the autograd machinery — everything + beyond this point runs inside PyTorch's autograd graph and must carefully manage what + is saved for backward. -Phase 2: _Linear.forward() — Input Preparation -------------------------------------------------- +Phase 2: _Linear.forward() — Input Quantization and Communication +------------------------------------------------------------------- -**File**: ``transformer_engine/pytorch/module/linear.py``, ``_Linear.forward()`` (line ~87) +**File**: ``transformer_engine/pytorch/module/linear.py``, ``_Linear.forward()`` (line ~167) -The input tensor is prepared based on the parallel mode and FP8 state: +**Why this is the most complex phase**: Input preparation must solve a chicken-and-egg +problem. The GEMM needs the *full* (ungathered) input, but communication is expensive. +With FP8, we also need both a rowwise *and* columnwise quantized version. The code must +decide: quantize before or after communication? The answer depends on the parallel mode. -**Column-parallel with sequence parallelism** (most common for QKV/FC1): +**Column-parallel with sequence parallelism** (QKV projection, FC1): + +Each GPU holds a shard of the input along the sequence dimension. The GEMM needs the +full sequence, so an all-gather is required. The key decision: **quantize locally first, +then all-gather the quantized data** — this communicates 1 byte per element instead of 2 +(BF16), cutting communication volume in half. .. code-block:: python - # 1. Set quantizer usage: need rowwise for forward GEMM, - # columnwise for backward wgrad GEMM + # 1. Configure what quantization layouts we need. + # Rowwise: for the forward GEMM (this phase). + # Columnwise: for the backward wgrad GEMM (Phase 8, later). + # We quantize both now because the original BF16 input will be discarded. input_quantizer.set_usage(rowwise=True, columnwise=backward_needs_input) - # 2. Quantize local input shard - inputmat = input_quantizer(inputmat) # Returns QuantizedTensorStorage + # 2. Quantize the local shard — produces a QuantizedTensorStorage containing + # quantized rowwise data, columnwise data, and their scale inverses. + inputmat = input_quantizer(inputmat) - # 3. All-gather along sequence dimension + # 3. All-gather the quantized data across TP ranks. + # Each rank sends its local shard; all ranks receive the full tensor. inputmat_total, _ = gather_along_first_dim(inputmat, tp_group) +.. note:: + + The columnwise data is NOT all-gathered along with the rowwise data in all cases. + Some quantizer types (e.g. delayed scaling ``Float8Quantizer``) don't support + columnwise all-gather, so columnwise is disabled before the gather and the input is + re-quantized columnwise in the backward pass. This is a trade-off: less communication + complexity at the cost of redundant quantization later. + **No parallelism** (standalone Linear): +No communication needed — the full input is already local. Quantization still produces +both layouts because the backward pass will need columnwise data regardless. + .. code-block:: python input_quantizer.set_usage(rowwise=True, columnwise=backward_needs_input) inputmat = input_quantizer(inputmat) - inputmat_total = inputmat + inputmat_total = inputmat # No gather needed -The ``input_quantizer(inputmat)`` call produces a ``QuantizedTensorStorage`` containing: +**What ``backward_needs_input`` means**: This flag is ``True`` when ``weight.requires_grad`` +and gradients are enabled — i.e., during training. During inference or when the weight is +frozen, the backward wgrad GEMM won't run, so there's no point producing columnwise data. +This avoids a significant portion of the quantization cost during inference. -- Rowwise FP8 data + ``scale_inv`` (for forward GEMM) -- Columnwise FP8 data + ``columnwise_scale_inv`` (saved for wgrad GEMM) +Phase 3: _Linear.forward() — Weight Quantization +-------------------------------------------------- -See :doc:`quantization/rowwise_columnwise` for why both layouts are needed. +**File**: ``transformer_engine/pytorch/module/linear.py``, ``_Linear.forward()`` (line ~250) -Phase 3: _Linear.forward() — Weight Preparation --------------------------------------------------- +**Why weight handling differs from input**: Weights are persistent parameters that don't +change between microbatches during gradient accumulation. Inputs change every microbatch. +This asymmetry enables an important optimization: **quantized weight caching**. .. code-block:: python - # Configure weight quantizer + # Configure weight quantizer — rowwise for forward GEMM, + # columnwise for backward dgrad GEMM (where the weight is the "other" operand). weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) - # Get quantized weight (may use cached version) + # Get quantized weight, potentially from cache. + # On the first microbatch: quantize and cache. + # On subsequent microbatches: return cached quantized weight. weightmat = module.get_weight_workspace( tensor=weight, quantizer=weight_quantizer, - cache_name="weight", # Cache across microbatches - update_workspace=is_first_microbatch, + cache_name=(None if is_first_microbatch is None else "weight"), + update_workspace=is_first_microbatch is None or is_first_microbatch, ) -Weight caching is important: the same FP8 weight can be reused across gradient -accumulation microbatches, avoiding redundant quantization. +**Why cache quantized weights?** In gradient accumulation with *N* microbatches, the weight +doesn't change until the optimizer step. Without caching, we'd quantize the same weight +*N* times. With caching, we quantize once and reuse the quantized version for all *N* +forward passes. The ``is_first_microbatch`` parameter controls this: ``True`` triggers a +fresh quantization, ``False`` returns the cached version, ``None`` disables caching +entirely (quantize every time). + +**Columnwise weight**: The weight also needs a columnwise layout for the backward dgrad +GEMM (Phase 7). The quantizer produces both during the same quantization call, just like +the input quantizer. The columnwise weight data is saved for backward. Phase 4: _Linear.forward() — GEMM ----------------------------------- +**File**: ``transformer_engine/pytorch/cpp_extensions/gemm.py`` + +**Why this isn't just a matmul**: Low-precision GEMM requires additional metadata (scale +inverses) that standard matrix multiplication APIs don't handle. TE wraps cuBLASLt with +scale-aware logic and optional communication overlap. + .. code-block:: python # Forward GEMM: output = input @ weight^T - gemm_out, *_ = general_gemm( - weightmat, # A operand (FP8) - inputmat_total, # B operand (FP8) + # The result is produced in high precision (BF16/FP16) — cuBLASLt accumulates + # in FP32 internally and casts the output down. This is why low-precision GEMM + # doesn't lose as much accuracy as you might expect from 8-bit math. + gemm_out, *_, reduce_scatter_out = general_gemm( + weightmat, # A operand (quantized + scale_inv) + inputmat_total, # B operand (quantized + scale_inv) out_dtype=activation_dtype, - bias=bias, + bias=bias, # Fused bias add (avoids separate kernel launch) use_split_accumulator=use_split_accumulator, ub=ub_obj, # Userbuffers overlap (optional) ub_type=ub_type, ) -This call chain: +The call traverses four layers, matching the :doc:`architecture_overview`: -1. **Python**: ``general_gemm()`` in ``transformer_engine/pytorch/cpp_extensions/gemm.py`` - extracts raw tensors and scales from the quantized inputs. -2. **pybind11**: ``te_general_gemm()`` in ``transformer_engine/pytorch/csrc/extensions/gemm.cpp`` - constructs ``NVTETensor`` handles and calls the C API. -3. **C API**: ``nvte_general_gemm()`` in ``transformer_engine/common/gemm/`` - dispatches to cuBLASLt with FP8 compute types. -4. **cuBLASLt**: Executes the FP8 matrix multiplication on the GPU. +1. **Python** (``cpp_extensions/gemm.py``): Extracts raw data pointers and scale tensors + from the ``QuantizedTensorStorage``. The C++ layer can't understand Python quantized + tensor types, so this translation is mandatory. +2. **pybind11** (``csrc/extensions/gemm.cpp``): Constructs opaque ``NVTETensor`` handles — + the C API's tensor abstraction. This is where TE crosses the Python/C++ boundary. +3. **C API** (``common/gemm/``): Selects the cuBLASLt algorithm and configures compute + types based on the input precision and scaling mode. +4. **cuBLASLt**: Executes the actual matrix multiply on the GPU. + +**``use_split_accumulator``**: When ``True``, cuBLASLt uses higher-precision intermediate +accumulators, trading some performance for numerical accuracy. This is controlled by the +recipe and defaults to ``False`` for the forward pass (where some accumulation error is +acceptable) and ``True`` for backward (where gradient accuracy matters more). Phase 5: _Linear.forward() — Output Communication ---------------------------------------------------- +**Why output communication depends on parallel mode**: Tensor parallelism shards the +weight matrix, and the parallel mode determines which dimension is sharded. This directly +determines what communication is needed to produce the correct output. + For **row-parallel** (output projection, FC2): +Each GPU computes a partial sum (because the input is split across the inner dimension). +These partial results must be summed across GPUs to produce the correct output. + .. code-block:: python if sequence_parallel: + # Reduce-scatter: sum partial results AND scatter the output along the + # sequence dimension, so each GPU only stores its local sequence shard. + # This is more memory-efficient than all-reduce because the output is + # smaller on each GPU. out, _ = reduce_scatter_along_first_dim(out, tp_group) elif tensor_parallel: + # All-reduce: sum partial results and replicate the full output. + # Used when sequence parallelism is disabled. out, _ = allreduce(out, tp_group) -For **column-parallel**: no communication needed (output is naturally partitioned). +For **column-parallel** (QKV, FC1): No communication needed. Each GPU holds a different +slice of the output features, and these slices are independent — they don't need to be +summed or gathered until a downstream operation requires the full output. Phase 6: _Linear.forward() — Save for Backward -------------------------------------------------- +**Why this phase is critical for memory efficiency**: PyTorch autograd requires saving +tensors from the forward pass to compute gradients. For a standard ``nn.Linear``, this +means saving the full input and weight in BF16. With FP8, TE saves *quantized* tensors — +half the memory — but must carefully track which quantization layouts to keep. + .. code-block:: python - # Save quantized input for wgrad (columnwise data) - # Save quantized weight for dgrad - ctx.save_for_backward(inputmat, weightmat, weight, ...) - ctx.weight_quantizer = weight_quantizer + # Discard rowwise data — it was only needed for the forward GEMM (Phase 4) + # and is no longer useful. Keep only columnwise data for the wgrad GEMM. + if backward_needs_input and own_quantized_input: + inputmat.update_usage(rowwise_usage=False, columnwise_usage=True) + + ctx.save_for_backward(inputmat, weightmat, weight, bias) + +**The key insight**: The forward GEMM consumes *rowwise* input, but the backward wgrad +GEMM needs *columnwise* input (see :doc:`quantization/rowwise_columnwise`). Since TE +quantized both layouts in Phase 2, it can now discard the rowwise data, keeping only +the columnwise data in the autograd cache. This means activation memory for FP8 is +roughly **half** of what it would be in BF16 — one byte per element (FP8 columnwise) +instead of two (BF16). -Key optimization: only the *columnwise* portion of the input is kept (the rowwise -data is discarded after the forward GEMM). This halves the activation memory for FP8. +For column-parallel with sequence parallelism, only the *local shard* of the input is +saved, not the full all-gathered tensor. The all-gather will be repeated in the backward +pass. This trades communication for memory — a favorable trade-off in large model training +where activation memory is the binding constraint. -Phase 7: _Linear.backward() — Dgrad --------------------------------------- +Phase 7: _Linear.backward() — Dgrad (Activation Gradient) +----------------------------------------------------------- + +**File**: ``transformer_engine/pytorch/module/linear.py``, ``_Linear.backward()`` (line ~495) + +**Why dgrad runs before wgrad**: The dgrad GEMM (``grad_input = grad_output @ weight``) +is on the critical path of backpropagation — downstream layers are waiting for +``grad_input`` to continue their own backward passes. The wgrad GEMM (``grad_weight``) +only updates this layer's parameters and doesn't block anything. Running dgrad first +minimizes the time other layers spend waiting. .. code-block:: python - # Quantize grad_output + # Quantize grad_output for the backward GEMMs. + # Rowwise: needed for dgrad GEMM (this phase). + # Columnwise: needed for wgrad GEMM (Phase 8). + # Both are computed now because the BF16 grad_output will be discarded. grad_output_quantizer.set_usage(rowwise=True, columnwise=True) - qgrad_output = grad_output_quantizer(grad_output) + grad_output = grad_output_preprocess(ctx, grad_output, ...) # Dgrad GEMM: grad_input = grad_output @ weight + # Uses rowwise grad_output and columnwise weight (saved from forward). dgrad, *_ = general_gemm( - weightmat, # Columnwise weight from forward - qgrad_output, # Rowwise grad_output + weight_fp8, # Columnwise weight saved in Phase 6 + grad_output, # Rowwise grad_output + layout="NN", + grad=True, # Signals this is a backward GEMM (affects accumulator) out_dtype=activation_dtype, ) -Phase 8: _Linear.backward() — Wgrad --------------------------------------- +**Communication overlap opportunity**: For column-parallel, the dgrad output must be +reduce-scattered (or all-reduced) back to each GPU's local shard. This communication is +launched *asynchronously* — it runs on a separate CUDA stream while the wgrad GEMM +(Phase 8) executes on the compute stream. Similarly, if the wgrad needs a re-gathered +input (column-parallel + SP), that all-gather is launched asynchronously during dgrad +and synchronized before wgrad begins. + +Phase 8: _Linear.backward() — Wgrad (Weight Gradient) +------------------------------------------------------- + +**Why wgrad uses columnwise data**: The wgrad GEMM computes ``grad_weight = input^T @ +grad_output``. The transpose of rowwise data is columnwise data. Rather than transposing +at this point (which would require a separate kernel launch and temporary memory), the +forward pass pre-computed the columnwise layout in Phase 2. This is the payoff for the +"dual layout" strategy. .. code-block:: python + # Synchronize any async communication from Phase 7 + if inputmat_total_work is not None: + inputmat_total_work.wait() + + # Ensure columnwise data is available + inputmat_total.update_usage(columnwise_usage=True) + # Wgrad GEMM: grad_weight = input^T @ grad_output - # Uses columnwise input saved from forward wgrad, *_ = general_gemm( - inputmat, # Columnwise input from forward - qgrad_output, # Columnwise grad_output + inputmat_total, # Columnwise input saved from forward + grad_output, # Columnwise grad_output from Phase 7 out_dtype=activation_dtype, - grad=True, # Indicates wgrad GEMM + grad=True, + ub=ub_obj_wgrad, + ub_type=ub_type_wgrad, ) -This is why the forward pass computed both rowwise and columnwise input: the columnwise -data is consumed here by the wgrad GEMM. +**``fuse_wgrad_accumulation``**: When enabled (common in Megatron-LM), the wgrad GEMM +accumulates directly into ``weight.main_grad`` instead of allocating a separate gradient +tensor. This avoids an extra memory allocation and addition kernel, which matters when +gradients are large (e.g., 4096 × 16384 = 67M parameters per Linear). + +Phase 9: Post-Backward Cleanup +-------------------------------- + +After both GEMMs complete, the backward pass handles recipe-specific bookkeeping. -Phase 9: Post-Backward ------------------------- +For most recipes (MXFP8, current scaling, block scaling), there is no post-backward state +to update — scales are computed inline during quantization and don't carry state across +iterations. This is one of the advantages of these recipes: they are stateless. -1. ``post_forward()`` records the output amax for delayed scaling. -2. For distributed training, amax values are all-reduced across TP ranks. -3. The amax history is updated for the next iteration's scale computation. +**Delayed scaling only**: The delayed scaling recipe is an exception. It maintains an +**amax history** — a rolling window of observed maximum values — that feeds forward into +the next iteration's scale computation. After the backward pass completes, the first FP8 +module in the model triggers ``FP8GlobalStateManager.reduce_and_update_fp8_tensors()``, +which all-reduces amax values across TP ranks and updates the history buffer. This ensures +all ranks agree on scales. If you're not using delayed scaling, this code path is skipped. Summary of Files Touched ------------------------- @@ -231,11 +411,13 @@ Summary of Files Touched * - ``pytorch/module/linear.py`` - ``Linear.forward()``, ``_Linear`` autograd * - ``pytorch/module/base.py`` - - ``pre_forward()``, FP8 state management + - ``prepare_forward()``, quantizer setup * - ``pytorch/quantized_tensor.py`` - ``Quantizer``, ``QuantizedTensorStorage`` base classes + * - ``pytorch/tensor/mxfp8_tensor.py`` + - ``MXFP8Quantizer``, block-scaled FP8 cast * - ``pytorch/tensor/float8_tensor.py`` - - ``Float8Quantizer``, FP8 cast implementation + - ``Float8Quantizer``, ``Float8CurrentScalingQuantizer`` * - ``pytorch/cpp_extensions/gemm.py`` - ``general_gemm()`` Python wrapper * - ``pytorch/csrc/extensions/gemm.cpp`` diff --git a/docs/developer/pytorch_frontend/ops_framework.rst b/docs/developer/pytorch_frontend/ops_framework.rst index 68f8f0cfeb..a7419a9a4e 100644 --- a/docs/developer/pytorch_frontend/ops_framework.rst +++ b/docs/developer/pytorch_frontend/ops_framework.rst @@ -5,14 +5,14 @@ .. _ops-framework: -Op Fusion Framework -=================== +Op Fusion Framework (``te.ops``) +================================ Transformer Engine includes an operation fusion framework (``transformer_engine/pytorch/ops/``) -that enables composing and fusing operations for better performance. +that enables composing and automatically fusing operations for better performance. -Overview --------- +For **user-facing** documentation (basic usage, examples), see :doc:`/examples/op_fuser/op_fuser`. +This page covers the **internal architecture** for developers working on the framework itself. .. figure:: ./img/ops_fusion.svg :align: center @@ -30,73 +30,359 @@ Overview Single kernel launch arrow. Arrow between left and right labeled "Op Fuser". -The framework defines three levels of abstraction: +Motivation +---------- + +TE's monolithic modules (``LayerNormLinear``, ``LayerNormMLP``) achieve high performance +through hand-tuned fusion, but they are rigid — adding a new model architecture or +experimenting with different fusion patterns requires modifying large, complex modules. + +The ops framework takes the opposite approach: users compose small, atomic operations +and the framework **automatically discovers and applies fusions** via a sliding-window +pattern matcher. The same fused kernels are used either way. + +Architecture +------------ + +The framework has five key components: + +.. code-block:: text + + Sequential User-facing container (like torch.nn.Sequential) + └── OperationFuser Manages fusion pipeline for a group of ops + └── _OperationFuserAutogradFunction Custom autograd for fused forward/backward + └── FusibleOperation Abstract base for all ops + ├── BasicOperation Atomic op (holds params, implements fwd/bwd) + └── FusedOperation Compound op (wraps multiple BasicOperations) + +**Sequential** (``ops/sequential.py``) + Drop-in replacement for ``torch.nn.Sequential``. On first forward call, lazily groups + consecutive ``FusibleOperation`` instances into ``OperationFuser`` objects. Non-fusible + ``torch.nn.Module`` instances remain standalone between groups. + +**OperationFuser** (``ops/fuser.py``) + Manages one group of fusible operations. On first call, applies registered fusion + functions to discover fusible patterns. Executes the fused pipeline via a custom + ``torch.autograd.Function``. + +**BasicOperation** (``ops/op.py``) + An atomic operation that holds parameters and implements ``op_forward()`` / + ``op_backward()`` with an ``OperationContext`` (similar to ``torch.autograd.Function``'s + ``ctx``). All quantization state (``Quantizer`` instances, FP8 metadata) lives here. + +**FusedOperation** (``ops/op.py``) + A compound operation that replaces one or more ``BasicOperation`` instances with a + fused implementation. It is **stateless** — it accesses parameters and quantizers + through its constituent basic ops. Forward and backward may use different fusions + (e.g., ``ForwardLinearBiasActivation`` for forward, ``BackwardActivationBias`` for + backward). + +**OperationContext** (``ops/op.py``) + Lightweight dataclass for caching state between forward and backward, per basic op: + + .. code-block:: python + + @dataclasses.dataclass + class OperationContext: + saved_tensors: Optional[tuple[Optional[torch.Tensor], ...]] + to_save: Optional[tuple[Optional[torch.Tensor], ...]] + requires_grad: bool = True -- **FusibleOperation** (``ops/op.py``): Abstract base for any operation that can - participate in fusion. -- **BasicOperation** (``ops/op.py``): Single operations (GEMM, bias, normalization, etc.) - that serve as building blocks. Holds parameters and quantization state. -- **FusedOperation** (``ops/op.py``): Compound operations composed of multiple basic ops. - Delegates parameter management to its constituent basic ops. -- **OperationFuser** (``ops/fuser.py``): Manages the fusion pipeline, using a custom - autograd function (``_OperationFuserAutogradFunction``) to coordinate forward/backward. -- **Sequential** (``ops/sequential.py``): A container that groups consecutive - ``FusibleOperation`` instances for automatic fusion. + def save_for_backward(self, *tensors): ... -Basic Ops ---------- +Execution Flow +-------------- -Basic ops extend ``BasicOperation`` and implement ``op_forward()`` / ``op_backward()``: +**First call** through ``Sequential.forward()``: .. code-block:: text - ops/basic/ - ├── linear.py # BasicLinear — matrix multiplication - ├── bias.py # Bias — bias addition - ├── normalization.py # LayerNorm / RMSNorm - ├── activation.py # GELU, GEGLU, SwiGLU, GLU - ├── quantization.py # Quantize — cast to/from FP8 - ├── all_gather.py # AllGather (distributed) - ├── all_reduce.py # AllReduce (distributed) - ├── reduce_scatter.py # ReduceScatter (distributed) - ├── dropout.py # Dropout - ├── reshape.py # Reshape - └── identity.py # Identity - -Fused Ops ---------- - -Fused ops extend ``FusedOperation`` and compose basic ops into optimized sequences. -The most important fused op is ``Linear`` (``ops/linear.py``), which composes -``BasicLinear`` + ``Bias`` + optional ``AllReduce``/``ReduceScatter`` for tensor -parallelism: + Sequential.forward(input) + │ + ├─ _make_module_groups(modules) + │ Group consecutive FusibleOps → [OperationFuser, Module, OperationFuser, ...] + │ + └─ for group in module_groups: + if OperationFuser: + group(input) + │ + ├─ maybe_fuse_ops(...) + │ ├─ Reset recipe state for all basic ops + │ ├─ Apply forward_fusion_functions (sliding window pattern match) + │ └─ Apply backward_fusion_functions + │ + ├─ op.pre_fuser_forward() for each op + │ + └─ _OperationFuserAutogradFunction.apply(input, fuser, ...) + │ + ├─ Forward: for (op, basic_op_idxs) in forward_ops: + │ op.fuser_forward(ctxs, input, quantizers_from_neighbors, ...) + │ + └─ Backward: for (op, basic_op_idxs) in reversed(backward_ops): + op.fuser_backward(ctxs, grad_output, ...) + +**Subsequent calls** reuse the cached fusion state unless invalidated (recipe change, +gradient requirements change, etc.). + +Fusion Pipeline +--------------- + +Fusion is driven by **registered fusion functions** — callables that scan a list of ops +with a sliding window and replace matching patterns: .. code-block:: python - # ops/linear.py — row-parallel example - Linear = FusedOperation([BasicLinear, Bias, ReduceScatter]) + # Type signature for fusion functions + OperationFusionFunction = Callable[ + [list[FusibleOperation], ...], # ops to scan + list[FusibleOperation], # ops after fusion (same or fewer) + ] -Op Fuser --------- +Fusion functions are registered globally and applied in order: + +.. code-block:: python + + from transformer_engine.pytorch.ops import register_forward_fusion + + register_forward_fusion(my_fusion_function) + register_forward_fusion(another_fusion, prepend=True) # Higher priority + +**Forward and backward have separate fusion registries.** This is important because the +optimal fusion pattern can differ between passes. For example, a forward pass might fuse +``[BasicLinear, Bias, GELU]`` into ``ForwardLinearBiasActivation``, while the backward +fuses ``[GELU_bwd, Bias_bwd]`` into ``BackwardActivationBias``. + +**Sliding window pattern matching** is the standard approach. Each fusion function scans +with a window of 2-3 ops, checks if the pattern matches (including constraints like +dtype, parallel mode, etc.), and replaces with a ``FusedOperation`` if so. + +Quantizer Passing Between Ops +------------------------------ + +A key mechanism for enabling fused quantization: the fuser passes quantizers between +adjacent ops during execution. + +.. code-block:: python + + # In _OperationFuserAutogradFunction.forward(): + for i, (op, idxs) in enumerate(forward_ops): + # Get quantizer from previous op's backward + prev_quantizer = forward_ops[i-1].get_grad_output_quantizer() + # Get quantizer from next op's forward + next_quantizer = forward_ops[i+1].get_input_quantizer() -The ``OperationFuser`` (``ops/fuser.py``) manages a pipeline of fusible operations. It -wraps them in ``_OperationFuserAutogradFunction``, which: + op.fuser_forward( + ..., + prev_op_grad_output_quantizer=prev_quantizer, + next_op_input_quantizer=next_quantizer, + ) -- **Forward**: Applies ops sequentially with potential fusion, passing quantizers between - ops. -- **Backward**: Reverses op order, applies fused/unfused backward, using per-op saved - contexts. +This allows an op to **quantize its output eagerly** using the next op's quantizer, +enabling fused cast kernels. For example, a LayerNorm op receiving the next Linear's +input quantizer can fuse the normalization and FP8 cast into a single kernel. -``Sequential`` (``ops/sequential.py``) is the user-facing container. It groups -consecutive ``FusibleOperation`` instances via ``_make_module_groups()`` and applies -the fuser automatically. +Catalog of Basic Operations +---------------------------- -Relationship to Modules ------------------------ +.. list-table:: + :header-rows: 1 + :widths: 25 40 35 -The ops framework is an alternative to the monolithic module approach -(``LayerNormLinear``, ``LayerNormMLP``). Modules like ``te.Linear`` (in ``module/``) use -the ``_Linear`` autograd function directly, while the ops framework allows more flexible -composition and automatic fusion discovery. + * - Operation + - File + - Description + * - ``BasicLinear`` + - ``basic/basic_linear.py`` + - Core GEMM (no bias). Supports FP8, TP, SP. + * - ``Bias`` + - ``basic/bias.py`` + - Additive bias parameter. + * - ``LayerNorm`` + - ``basic/layer_norm.py`` + - Layer normalization with learnable scale/bias. + * - ``RMSNorm`` + - ``basic/rmsnorm.py`` + - Root mean square normalization. + * - ``GELU``, ``SiLU``, ``ReLU``, ``QGELU`` + - ``basic/activation.py`` + - Element-wise activation functions. + * - ``GEGLU``, ``SwiGLU``, ``ReGLU``, ``SReGLU``, ``QGEGLU``, ``GLU`` + - ``basic/activation.py``, ``basic/swiglu.py`` + - Gated activation functions (2× input split). + * - ``Quantize`` + - ``basic/quantize.py`` + - Explicit FP8 cast (identity if FP8 disabled). + * - ``AllGather`` + - ``basic/all_gather.py`` + - Distributed all-gather along first dim. + * - ``AllReduce`` + - ``basic/all_reduce.py`` + - Distributed all-reduce. + * - ``ReduceScatter`` + - ``basic/reduce_scatter.py`` + - Distributed reduce-scatter along first dim. + * - ``AddExtraInput`` + - ``basic/add_extra_input.py`` + - Add an extra tensor input (for residual connections). + * - ``MakeExtraOutput`` + - ``basic/make_extra_output.py`` + - Export intermediate tensor as extra output. + * - ``Dropout`` + - ``basic/dropout.py`` + - Random zeroing. + * - ``Reshape`` + - ``basic/reshape.py`` + - Tensor reshape. + * - ``Identity`` + - ``basic/identity.py`` + - Pass-through (useful as fusion anchor). + * - ``ConstantScale`` + - ``basic/constant_scale.py`` + - Multiply by constant factor. + * - ``L2Normalization`` + - ``basic/l2normalization.py`` + - L2 normalization. + * - ``GroupedLinear`` + - ``basic/grouped_linear.py`` + - Multiple parallel GEMMs (MoE). + +Catalog of Fused Operations +---------------------------- + +**Forward fusions** (registered in ``ops/fused/__init__.py``): + +.. list-table:: + :header-rows: 1 + :widths: 30 35 35 + + * - Fused Operation + - Pattern + - Description + * - ``ForwardLinearBiasActivation`` + - ``BasicLinear`` + ``Bias`` + activation + - Fused GEMM + bias + activation + * - ``ForwardLinearBiasAdd`` + - ``BasicLinear`` + ``Bias`` + ``AddExtraInput`` + - Fused GEMM + bias + residual add + * - ``ForwardLinearScaleAdd`` + - ``BasicLinear`` + ``ConstantScale`` + ``AddExtraInput`` + - Fused GEMM + scale + add + * - ``UserbuffersForwardLinear`` + - ``BasicLinear`` + ``Bias`` + ``ReduceScatter`` + - GEMM overlapped with TP communication + +**Backward fusions**: + +.. list-table:: + :header-rows: 1 + :widths: 30 35 35 + + * - Fused Operation + - Pattern + - Description + * - ``BackwardActivationBias`` + - Backward of activation + bias + - Fused activation and bias gradients + * - ``BackwardAddRMSNorm`` + - Backward of add + RMSNorm + - Fused residual add and norm gradients + * - ``BackwardLinearAdd`` + - ``MakeExtraOutput`` + ``BasicLinear`` + - Fused dgrad GEMM + residual add + * - ``BackwardLinearScale`` + - Backward GEMM + scale + - Fused backward GEMM and scaling + * - ``UserbuffersBackwardLinear`` + - Backward with Userbuffers + - Backward GEMM overlapped with TP communication + +Extra Inputs and Outputs +------------------------- + +Most ops are purely sequential (one input, one output). Some ops need branching: + +- ``AddExtraInput`` (``num_extra_inputs=1``): Accepts an additional tensor and adds it + to the intermediate value. Used for residual connections. +- ``MakeExtraOutput`` (``num_extra_outputs=1``): Copies the intermediate value to an + extra output. Used to extract residual before a transformation. + +Extra inputs are passed as additional arguments to ``Sequential.forward()``, and extra +outputs are returned as additional values: + +.. code-block:: python + + # MLP with residual + mlp = te.ops.Sequential( + te.ops.LayerNorm(4096), + te.ops.MakeExtraOutput(), # Branch: emit residual + te.ops.Linear(4096, 16384), + te.ops.SwiGLU(), + te.ops.Linear(8192, 4096), + te.ops.AddExtraInput(), # Merge: add residual + ) + + y, residual = mlp(x) # Extra output from MakeExtraOutput + # Or split across two Sequentials: + # y = mlp2(mlp1_out, residual) # Extra input to AddExtraInput + +Ops with extra inputs/outputs must override ``fuser_forward`` / ``fuser_backward`` +(the default ``BasicOperation`` wrappers raise an error for these). + +The ``te.ops.Linear`` Composition +---------------------------------- + +``te.ops.Linear`` (``ops/linear.py``) is a ``FusedOperation`` that serves as a drop-in +replacement for ``torch.nn.Linear``. It composes basic ops based on the parallelism +configuration: + +.. code-block:: text + + No TP: BasicLinear → Bias + Column TP: BasicLinear → Bias + Row TP: BasicLinear → Bias → ReduceScatter (if SP) or AllReduce + +The weight and bias parameters are registered on the ``Linear`` and forwarded to the +underlying ``BasicLinear`` and ``Bias`` ops. State dict methods are overridden for +backward compatibility with ``torch.nn.Linear`` checkpoints. + +Relationship to Monolithic Modules +------------------------------------ + +The ops framework and the monolithic modules (``module/linear.py``, +``module/layernorm_mlp.py``, etc.) are **two parallel implementations** of the same +functionality: + +.. list-table:: + :header-rows: 1 + :widths: 25 37 38 + + * - Aspect + - Monolithic Modules (``module/``) + - Ops Framework (``ops/``) + * - Fusion + - Hand-coded in ``_Linear``, ``_LayerNormMLP`` + - Automatic via pattern-matching fuser + * - Flexibility + - Fixed compositions + - Arbitrary op sequences + * - Performance + - Maximum (hand-tuned) + - Comparable (same kernels) + * - Complexity + - High (400+ line autograd functions) + - Distributed across small ops + * - Maturity + - Production + - Experimental + +Both approaches call the same C++ extensions and CUDA kernels. The ops framework is +intended to eventually replace the monolithic modules as the primary implementation, +once it reaches feature parity. + +See Also +-------- -Both approaches ultimately call the same C++ extensions and CUDA kernels. +- :doc:`/examples/op_fuser/op_fuser` — User-facing documentation with usage examples + and how to implement new operations +- :doc:`autograd_integration` — How ``_Linear`` autograd works (monolithic approach) +- :doc:`/developer/quantization/class_hierarchy` — Quantizer classes used by the ops diff --git a/docs/index.rst b/docs/index.rst index 7389553679..b1dd93c0ed 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -13,7 +13,7 @@ Transformer Engine documentation To see the documentation for the latest stable release, refer to: * `Release Notes `_ - * `Developer Guide `_ (stable version of this page) + * `User Guide `_ (stable version of this page) .. include:: ../README.rst :start-after: overview-begin-marker-do-not-remove @@ -69,3 +69,9 @@ Transformer Engine documentation envvars examples/attention/attention.ipynb examples/attention/cp_ag_thd_dpa_jax_deep_dive.ipynb + +.. toctree:: + :hidden: + :caption: Developer Guide + + developer/index From c941076677941fff3654df6a2bc6d757913e476c Mon Sep 17 00:00:00 2001 From: Przemek Tredak Date: Mon, 16 Mar 2026 13:50:35 -0700 Subject: [PATCH 03/20] Next round of fixes Signed-off-by: Przemek Tredak --- docs/developer/attention/backends.rst | 11 +- .../developer/attention/pytorch_attention.rst | 1 + docs/developer/cpp_core/kernel_areas.rst | 72 +++++++- docs/developer/cpp_core/scaling_modes.rst | 7 + .../distributed/comm_gemm_overlap.rst | 20 +- .../distributed/fsdp_integration.rst | 173 ++++++++++++++---- .../jax_frontend/xla_ffi_primitives.rst | 52 +++++- docs/developer/linear_walkthrough.rst | 2 +- .../pytorch_frontend/cpp_extensions.rst | 18 +- .../pytorch_frontend/module_hierarchy.rst | 36 +++- .../pytorch_frontend/ops_framework.rst | 3 +- .../quantization/class_hierarchy.rst | 51 ++++-- 12 files changed, 347 insertions(+), 99 deletions(-) diff --git a/docs/developer/attention/backends.rst b/docs/developer/attention/backends.rst index 0a5260851c..ffa411b9ac 100644 --- a/docs/developer/attention/backends.rst +++ b/docs/developer/attention/backends.rst @@ -55,16 +55,21 @@ Overview - Ampere - No - Tri Dao's FlashAttention; fast, memory-efficient - * - Fused (F16) + * - Fused F16 (max512) - BF16/FP16 - Ampere - Yes - - cuDNN graph-based; supports all mask types + - cuDNN graph-based; sequence length ≤ 512 + * - Fused F16 (arbitrary) + - BF16/FP16 + - Ampere + - Yes + - cuDNN graph-based; arbitrary sequence length, supports all mask types * - Fused (FP8) - FP8 - Hopper - Yes - - FP8 Q/K/V with cuDNN; best throughput for FP8 training + - FP8 Q/K/V with cuDNN; sequence length ≤ 512; best throughput for FP8 training * - Custom/THD - BF16/FP16 - Hopper diff --git a/docs/developer/attention/pytorch_attention.rst b/docs/developer/attention/pytorch_attention.rst index be2b5da1f4..09edfcbb29 100644 --- a/docs/developer/attention/pytorch_attention.rst +++ b/docs/developer/attention/pytorch_attention.rst @@ -94,6 +94,7 @@ TE provides fused softmax implementations for attention scores: - ``ScaledSoftmax`` — Standard scaled softmax - ``ScaledMaskedSoftmax`` — Softmax with arbitrary mask - ``ScaledUpperTriangMaskedSoftmax`` — Causal mask (upper triangular) +- ``ScaledAlignedCausalMaskedSoftmax`` — Aligned causal mask variant These are used by the unfused attention backend. The fused backends (cuDNN, Flash) handle softmax internally. diff --git a/docs/developer/cpp_core/kernel_areas.rst b/docs/developer/cpp_core/kernel_areas.rst index c2be48859a..e8267294a1 100644 --- a/docs/developer/cpp_core/kernel_areas.rst +++ b/docs/developer/cpp_core/kernel_areas.rst @@ -116,21 +116,73 @@ Fused operations over lists of tensors (e.g., multi-tensor scale for optimizer s - **Header**: ``include/transformer_engine/multi_tensor.h`` +Hadamard Transform (``hadamard_transform/``) +-------------------------------------------- + +Hadamard and group Hadamard transforms, with optional fused quantization. + +- **Header**: ``include/transformer_engine/hadamard_transform.h`` +- Variants: standard, group, fused with cast, graph-safe implementations. + +Recipe (``recipe/``) +-------------------- + +Scaling recipe kernels that compute quantization scales for each recipe type. + +- **Header**: ``include/transformer_engine/recipe.h`` +- Per-recipe files: ``current_scaling.cu``, ``delayed_scaling.cu``, + ``fp8_block_scaling.cu``, ``mxfp8_scaling.cu``, ``nvfp4.cu``. + +Dropout (``dropout/``) +----------------------- + +Dropout kernel. + +- **Header**: ``include/transformer_engine/dropout.h`` + +Swizzle (``swizzle/``) +----------------------- + +Scale swizzling for GEMM-compatible layouts. + +- **Header**: ``include/transformer_engine/swizzle.h`` +- Includes block scaling variant (``swizzle_block_scaling.cu``). + +Permutation (``permutation/``) +------------------------------- + +Token permutation/unpermutation for MoE expert routing. + +- **Header**: ``include/transformer_engine/permutation.h`` + +Padding (``util/padding.cu``) +------------------------------ + +Utility kernel for padding tensors to alignment requirements. + +- **Header**: ``include/transformer_engine/padding.h`` + Architecture Dispatch --------------------- -Many kernel areas provide architecture-specific implementations. The typical pattern: +Many kernel areas provide architecture-specific implementations. Unlike some CUDA +libraries that use separate source files per architecture (e.g., ``kernel_sm80.cu``, +``kernel_sm90.cu``), TE uses **template specialization and compile-time dispatch** within +shared source files: .. code-block:: text normalization/ - ├── common.h # Shared types and helpers + ├── common.h # Shared types and helpers + ├── common.cpp # KernelRegistry, runtime dispatch ├── layernorm/ - │ ├── ln_api.cpp # C API entry point, dispatches by arch - │ ├── ln_sm80.cu # Ampere (SM80) kernel - │ ├── ln_sm90.cu # Hopper (SM90) kernel - │ └── ln_sm100.cu # Blackwell (SM100) kernel - -The C API function queries the GPU architecture at runtime and dispatches to the -appropriate kernel. See :doc:`build_system` for how ``NVTE_CUDA_ARCHS`` controls which -architectures are compiled. + │ ├── ln_fwd_cuda_kernel.cu # Forward kernels (all archs) + │ ├── ln_bwd_semi_cuda_kernel.cu # Backward kernels + │ ├── ln_fwd_kernels.cuh # Templated kernel implementations + │ └── ln_bwd_kernels.cuh + +Kernels are heavily templated (data types, hidden sizes, warp configurations) and the +appropriate specialization is selected at runtime via a ``KernelRegistry``. CMake compile +flags (``NVTE_CUDA_ARCHS``) control which GPU architectures are compiled — this affects +which template instantiations are generated, not which source files are included. See +:doc:`build_system` for details. diff --git a/docs/developer/cpp_core/scaling_modes.rst b/docs/developer/cpp_core/scaling_modes.rst index c3a68ab6c0..c521cdd337 100644 --- a/docs/developer/cpp_core/scaling_modes.rst +++ b/docs/developer/cpp_core/scaling_modes.rst @@ -183,3 +183,10 @@ enum comparison: if (tensor.scaling_mode == NVTE_MXFP8_1D_SCALING) { ... } This makes code resilient to future additions to the ``NVTEScalingMode`` enum. + +Invalid Scaling Sentinel +^^^^^^^^^^^^^^^^^^^^^^^^ + +The enum includes a sentinel value ``NVTE_INVALID_SCALING = 100`` used to detect +uninitialized or misconfigured scaling modes. Code that switches on scaling mode should +handle this case explicitly (e.g., with an error or assertion). diff --git a/docs/developer/distributed/comm_gemm_overlap.rst b/docs/developer/distributed/comm_gemm_overlap.rst index 9a607142e6..1906edea36 100644 --- a/docs/developer/distributed/comm_gemm_overlap.rst +++ b/docs/developer/distributed/comm_gemm_overlap.rst @@ -89,10 +89,17 @@ The ``CommOverlapAlgo`` enum defines the available overlap strategies: - Uses CUDA atomics with reduce-scatter for producer-consumer overlap * - ``ATOMIC_GEMM_AG_P2P`` - Uses CUDA atomics with all-gather P2P + * - ``ATOMIC_GEMM_RS_P2P`` + - Uses CUDA atomics with reduce-scatter P2P + * - ``EXTERNAL_BULK_OVERLAP_AG`` + - All-gather with externally managed buffer overlap Enabling Overlap ---------------- +Overlap is enabled per-layer via constructor parameters on ``te.Linear`` or +``te.TransformerLayer``: + .. code-block:: python model = te.TransformerLayer( @@ -100,15 +107,16 @@ Enabling Overlap ffn_hidden_size=16384, num_attention_heads=32, tp_group=tp_group, - ub_overlap_rs=True, # Overlap reduce-scatter - ub_overlap_ag=True, # Overlap all-gather + ub_overlap_rs=True, # Overlap reduce-scatter in forward (row-parallel) + ub_overlap_ag=True, # Overlap all-gather in forward (column-parallel) ) -Environment variables: +At the ``_Linear`` autograd level, finer-grained control is available via parameters that +specify overlap for each pass direction: -- ``NVTE_UB_OVERLAP``: Enable/disable UserBuffers overlap -- ``NVTE_UB_SPLIT_RS``: Enable split-pipelined reduce-scatter -- ``NVTE_UB_SPLIT_AG``: Enable split-pipelined all-gather +- ``ub_overlap_rs_fprop`` / ``ub_overlap_rs_dgrad`` — reduce-scatter in forward / dgrad +- ``ub_overlap_ag_fprop`` / ``ub_overlap_ag_dgrad`` — all-gather in forward / dgrad +- ``ub_bulk_dgrad`` / ``ub_bulk_wgrad`` — bulk overlap during backward Performance Considerations -------------------------- diff --git a/docs/developer/distributed/fsdp_integration.rst b/docs/developer/distributed/fsdp_integration.rst index efa6cda678..5ffdf44d8e 100644 --- a/docs/developer/distributed/fsdp_integration.rst +++ b/docs/developer/distributed/fsdp_integration.rst @@ -9,73 +9,174 @@ FSDP Integration ================ Transformer Engine integrates with PyTorch's Fully Sharded Data Parallelism (FSDP) to -enable combined data and tensor parallelism with FP8 quantization. +enable combined data and tensor parallelism with FP8 quantization. TE supports both +**FSDP1** (the wrapper-based ``FullyShardedDataParallel``) and **FSDP2** (the composable +``fully_shard`` API), using different mechanisms for each. Challenges ---------- FSDP shards model parameters across data-parallel ranks and gathers them on-demand -during forward/backward. This creates challenges for FP8: +during forward/backward. This creates two challenges for quantized training: -1. **FP8 weight storage**: Weights may be stored in FP8 for memory savings, but FSDP's - gather/scatter operates on raw tensors — it doesn't know about FP8 scales. -2. **Amax synchronization**: Delayed scaling requires amax values to be consistent across - FSDP ranks (all ranks must agree on the scale). -3. **Quantized gradients**: FSDP's gradient reduction must handle FP8 gradients correctly. +1. **Quantized weight all-gather**: FSDP gathers weight shards into full parameters + before each forward/backward pass. When weights are stored in FP8, FSDP must gather + the quantized data *and* its associated scales, then reconstruct a valid quantized + tensor on the other side. -Solution: FP8-Aware Hooks --------------------------- +2. **Activation sharding**: Activations saved for backward consume memory proportional to + model size. Sharding them across FSDP ranks reduces per-rank memory, but requires + handling quantized tensor types during scatter/gather. -TE registers FSDP hooks that handle FP8 tensor conversions: +These two challenges are handled by separate mechanisms, described below. -**Pre-gather hook**: Before FSDP gathers a parameter shard, convert FP8 storage to a -format that can be gathered (e.g., pack data + scale into a single buffer). +FSDP2: Quantized Weight All-Gather +------------------------------------ -**Post-gather hook**: After FSDP gathers the full parameter, reconstruct the FP8 tensor -with proper scales. +**FSDP2** (``torch.distributed._composable.fsdp.fully_shard``) is the preferred approach. +It uses PyTorch's **tensor subclass protocol** — FSDP2 automatically discovers and calls +``fsdp_pre_all_gather()`` and ``fsdp_post_all_gather()`` methods on tensor subclasses +during weight all-gather. No explicit setup or registration is needed. -**Pre-scatter hook**: Before FSDP scatters gradients, ensure FP8 gradient data and scales -are properly packed. +**Location**: Methods on ``QuantizedTensor`` subclasses in ``transformer_engine/pytorch/tensor/``. -Usage ------ +``fsdp_pre_all_gather(self, mesh, orig_size, contiguous_orig_stride, module, mp_policy)`` + Called **before** FSDP all-gathers a sharded weight parameter. Extracts the raw + quantized data (and optionally scale_inv tensors) into plain ``torch.Tensor`` objects + that FSDP can all-gather normally. Returns a tuple of tensors to gather and a metadata + tuple (dtype, usage flags, scales) that will be passed to ``fsdp_post_all_gather``. + + The method determines which layouts to include based on training state: during the + forward pass it sends rowwise data; during backward it sends columnwise data (or both, + depending on the ``reshard_after_forward`` policy). + +``fsdp_post_all_gather(self, all_gather_outputs, metadata, param_dtype, *, out=None)`` + Called **after** FSDP has gathered the full tensors. Reconstructs the ``QuantizedTensor`` + subclass from the gathered data and the metadata. On the first call, creates a new + tensor; on subsequent calls, reuses the ``out`` tensor (passed by FSDP2) and updates + it in place. + +**Supported tensor types**: + +.. list-table:: + :header-rows: 1 + :widths: 30 70 + + * - Tensor Class + - Notes + * - ``Float8Tensor`` + - Gathers uint8 data + scale_inv. Handles amax reduction for current scaling. + * - ``MXFP8Tensor`` + - Gathers data + scale_inv with padding/unpadding for block-aligned scales. + +.. note:: + + FSDP2 also requires that certain ``torch.ops.aten`` operations preserve the tensor + subclass (e.g., ``slice``, ``view``, ``copy_``). ``Float8Tensor`` defines + ``_ops_to_preserve_subclass_in_fsdp2`` to list these operations. + +**Usage**: .. code-block:: python - import torch.distributed.fsdp as fsdp + from torch.distributed._composable.fsdp import fully_shard import transformer_engine.pytorch as te model = te.TransformerLayer(...) - # Wrap with FSDP — TE hooks are registered automatically - model = fsdp.FullyShardedDataParallel( - model, - auto_wrap_policy=..., - ) + # FSDP2 automatically handles quantized weight all-gather + fully_shard(model) - # FP8 training works normally with te.fp8_autocast(enabled=True): - output = model(input) + output = model(input) # Weights gathered/scattered automatically + +FSDP1: Activation Scatter/Gather +----------------------------------- + +**FSDP1** uses a different mechanism for a different purpose: sharding *activations* (not +weights) saved for backward. This is implemented via direct function calls within the +``_Linear`` autograd function. + +**Location**: ``transformer_engine/pytorch/distributed.py`` + +``_fsdp_scatter_tensors(fsdp_group, *tensors)`` + Called in ``_Linear.forward()`` after the forward GEMM. Shards saved activations + (quantized input and optionally quantized weight) across FSDP ranks using + ``split_tensor_into_1d_equal_chunks()``. This reduces per-rank activation memory + proportionally to the FSDP world size. Returns the original shapes for reconstruction. -FP8 Weight Storage with FSDP ------------------------------ +``_fsdp_gather_tensors(fsdp_group, shapes, *tensors)`` + Called in ``_Linear.backward()`` before the dgrad/wgrad GEMMs. Reconstructs full + activations from shards using ``gather_split_1d_tensor()`` and reshapes to original + dimensions. -When FP8 weights are enabled, the parameter is stored in FP8 format even when sharded -by FSDP. The lifecycle: +``prepare_te_modules_for_fsdp(fsdp_root)`` + Must be called **after** wrapping with FSDP1. Injects the FSDP process group reference + into each TE module so that ``_fsdp_scatter_tensors`` / ``_fsdp_gather_tensors`` know + which group to communicate with. + +TE also provides a convenience wrapper: + +.. code-block:: python + + # TE's FSDP1 wrapper calls prepare_te_modules_for_fsdp automatically + from transformer_engine.pytorch.distributed import FullyShardedDataParallel + + model = FullyShardedDataParallel(te_model, ...) + +**Usage (explicit setup)**: + +.. code-block:: python + + from torch.distributed.fsdp import FullyShardedDataParallel as FSDP + from transformer_engine.pytorch.distributed import prepare_te_modules_for_fsdp + + model = FSDP(te_model, ...) + prepare_te_modules_for_fsdp(model) + + with te.fp8_autocast(enabled=True): + output = model(input) -1. **Initialization**: Weight is created in high precision, then optionally cast to FP8. -2. **FSDP shard**: The FP8 data (and associated scale) is sharded across ranks. -3. **FSDP gather**: Full FP8 weight is gathered; TE reconstructs the quantized tensor. -4. **Forward/backward**: GEMM operates on the gathered FP8 weight. -5. **FSDP scatter**: Updated FP8 weight shard is stored locally. +Summary: Which Mechanism Does What +------------------------------------ + +.. list-table:: + :header-rows: 1 + :widths: 20 40 40 + + * - Aspect + - FSDP2 (composable) + - FSDP1 (wrapper) + * - What it handles + - Quantized **weight** all-gather + - **Activation** scatter/gather + * - Mechanism + - ``fsdp_pre_all_gather`` / ``fsdp_post_all_gather`` on tensor subclasses + - ``_fsdp_scatter_tensors`` / ``_fsdp_gather_tensors`` called from ``_Linear`` + * - Setup required + - None (automatic discovery) + - ``prepare_te_modules_for_fsdp()`` + * - Called by + - PyTorch's FSDP2 internals + - ``_Linear.forward()`` / ``_Linear.backward()`` + +.. note:: + + The activation scatter/gather from FSDP1 (``_fsdp_scatter_tensors`` / + ``_fsdp_gather_tensors``) is also used in FSDP2 setups — it is called from + ``_Linear`` regardless of which FSDP version manages the weights. The two mechanisms + are complementary, not mutually exclusive. Limitations ----------- -- FSDP2 (``torch.distributed._composable.fsdp``) has better support than FSDP1. +- FSDP sharding is not valid for models initialized with ``primary_weights_in_fp8=True``. - FP8 weight caching interacts with FSDP gather/scatter — weights may be re-quantized on each gather. - Mixed FSDP + TP configurations require careful process group setup. +- ``fsdp_pre_all_gather`` / ``fsdp_post_all_gather`` are currently implemented only for + ``Float8Tensor`` and ``MXFP8Tensor``, not for ``Float8BlockwiseQTensor`` or + ``NVFP4Tensor``. See Also -------- diff --git a/docs/developer/jax_frontend/xla_ffi_primitives.rst b/docs/developer/jax_frontend/xla_ffi_primitives.rst index f04500839e..edbe0cc5bf 100644 --- a/docs/developer/jax_frontend/xla_ffi_primitives.rst +++ b/docs/developer/jax_frontend/xla_ffi_primitives.rst @@ -119,13 +119,53 @@ Comparison with PyTorch Bridge Key Primitives -------------- -Major primitives registered in ``transformer_engine/jax/cpp_extensions/``: +Major primitives registered in ``transformer_engine/jax/cpp_extensions/``, organized by +functional area: -- ``gemm`` — Matrix multiplication with FP8 -- ``layernorm_fwd`` / ``layernorm_bwd`` — Normalization -- ``fused_attn_fwd`` / ``fused_attn_bwd`` — Fused attention -- ``cast_fp8`` / ``cast_fp8_transpose`` — FP8 quantization -- ``activation_fwd`` / ``activation_bwd`` — Activation functions +**GEMM** (``gemm.py``): + +- ``te_gemm_v2_ffi`` — Matrix multiplication with FP8 +- ``te_grouped_gemm_ffi`` — Grouped GEMM for MoE +- ``te_grouped_gemm_d2h_group_sizes_ffi`` — Device-to-host group size transfer + +**Quantization** (``quantization.py``): + +- ``te_dbias_quantize_ffi`` — Quantize with optional bias gradient +- ``te_grouped_quantize_ffi`` — Grouped quantization + +Exported as: ``quantize``, ``quantize_dbias``, ``grouped_quantize``, ``grouped_dbias``. + +**Activation** (``activation.py``): + +- ``te_act_lu_ffi`` — Forward activation (GeLU, SiLU, etc.) +- ``te_dact_dbias_quantize_ffi`` — Fused backward activation + bias gradient + quantize + +Exported as: ``act_lu``, ``dact_lu``, ``quantize_dact_dbias``. + +**Normalization** (``normalization.py``): + +- ``te_norm_forward_ffi`` — LayerNorm / RMSNorm forward +- ``te_norm_backward_ffi`` — LayerNorm / RMSNorm backward + +**Attention** (``attention.py``): + +- ``te_fused_attn_forward_ffi`` — Fused attention forward +- ``te_fused_attn_backward_ffi`` — Fused attention backward + +**Softmax** (``softmax.py``): + +- ``te_scaled_softmax_forward_ffi`` / ``te_scaled_softmax_backward_ffi`` +- ``te_scaled_masked_softmax_forward_ffi`` / ``te_scaled_masked_softmax_backward_ffi`` +- ``te_scaled_upper_triang_masked_softmax_forward_ffi`` / ``te_scaled_upper_triang_masked_softmax_backward_ffi`` + +**Router / MoE** (``router.py``): + +- ``te_fused_topk_with_score_function_forward_ffi`` / ``te_fused_topk_with_score_function_backward_ffi`` +- ``te_fused_moe_aux_loss_forward_ffi`` / ``te_fused_moe_aux_loss_backward_ffi`` + +**Amax** (``amax.py``): + +- ``te_rht_amax_ffi`` — Randomized Hadamard transform + amax calculation See Also -------- diff --git a/docs/developer/linear_walkthrough.rst b/docs/developer/linear_walkthrough.rst index 87dc8959bc..90051c0e2c 100644 --- a/docs/developer/linear_walkthrough.rst +++ b/docs/developer/linear_walkthrough.rst @@ -124,7 +124,7 @@ that actually perform casts. Phase 2: _Linear.forward() — Input Quantization and Communication ------------------------------------------------------------------- -**File**: ``transformer_engine/pytorch/module/linear.py``, ``_Linear.forward()`` (line ~167) +**File**: ``transformer_engine/pytorch/module/linear.py``, ``_Linear.forward()`` (line ~87) **Why this is the most complex phase**: Input preparation must solve a chicken-and-egg problem. The GEMM needs the *full* (ungathered) input, but communication is expensive. diff --git a/docs/developer/pytorch_frontend/cpp_extensions.rst b/docs/developer/pytorch_frontend/cpp_extensions.rst index 71cfd76796..341ebfe9ef 100644 --- a/docs/developer/pytorch_frontend/cpp_extensions.rst +++ b/docs/developer/pytorch_frontend/cpp_extensions.rst @@ -32,16 +32,18 @@ Python Side **Location**: ``transformer_engine/pytorch/cpp_extensions/`` -Each kernel area has a corresponding Python wrapper module: +Only a subset of C++ extensions have dedicated Python wrapper modules: - ``gemm.py`` — ``general_gemm()``, ``grouped_gemm()`` -- ``normalization.py`` — ``layernorm_fwd()``, ``rmsnorm_fwd()``, etc. -- ``activation.py`` — ``gelu()``, ``silu()``, etc. -- ``cast.py`` — ``quantize()``, ``dequantize()`` -- ``transpose.py`` — ``transpose()``, ``cast_transpose()`` -- ``attention.py`` — ``fused_attn_fwd()``, ``fused_attn_bwd()`` +- ``fused_attn.py`` — ``fused_attn_fwd()``, ``fused_attn_bwd()`` -These wrappers: +Most other C++ extensions (normalization, activation, cast, transpose, etc.) are exposed +**directly** through the compiled ``transformer_engine_torch`` pybind11 module (imported +as ``tex``) and called without a Python wrapper layer. For example, normalization is +called as ``tex.layernorm_fwd(...)`` rather than through a ``normalization.py`` wrapper. + +The Python wrappers that do exist serve as the translation layer between PyTorch's +quantized tensor types and the C++ API: 1. Extract raw data tensors and scales from ``QuantizedTensor`` / ``QuantizedTensorStorage`` objects. @@ -71,7 +73,7 @@ C++ Side (pybind11) **Location**: ``transformer_engine/pytorch/csrc/extensions/`` -The pybind11 module (``transformer_engine/pytorch/csrc/ts_fp8_op.cpp`` or similar) +The pybind11 module (``transformer_engine/pytorch/csrc/extensions/pybind.cpp``) registers Python-callable functions that: 1. Accept ``torch::Tensor`` and scalar arguments from Python. diff --git a/docs/developer/pytorch_frontend/module_hierarchy.rst b/docs/developer/pytorch_frontend/module_hierarchy.rst index 84150817bf..843678de6d 100644 --- a/docs/developer/pytorch_frontend/module_hierarchy.rst +++ b/docs/developer/pytorch_frontend/module_hierarchy.rst @@ -20,18 +20,20 @@ FP8 state management, quantizer lifecycle, and distributed training hooks. .. Diagram description for ``pytorch_module_hierarchy.svg``: Tree diagram with torch.nn.Module at the top. - Below it: TransformerEngineBaseModule. - Below that, branching to: + Below it, two branches: + Branch 1: TransformerEngineBaseModule, with children: ├── Linear ├── LayerNorm ├── RMSNorm ├── LayerNormLinear ├── LayerNormMLP ├── GroupedLinear - ├── DotProductAttention - └── MultiheadAttention - A separate branch from torch.nn.Module: TransformerLayer (composes the above modules - rather than inheriting from BaseModule). + └── DotProductAttention + Branch 2: Directly from torch.nn.Module: + ├── MultiheadAttention (composes DotProductAttention + Linear projections) + ├── TransformerLayer (composes the above modules) + ├── Fp8Padding + └── Fp8Unpadding TransformerEngineBaseModule --------------------------- @@ -132,14 +134,30 @@ Composes QKV projection (Linear), DotProductAttention, and output projection (Li into a complete multi-head attention block. Handles the split into heads and optional key-value caching for inference. +.. note:: + + ``MultiheadAttention`` extends ``torch.nn.Module`` directly, **not** + ``TransformerEngineBaseModule``. It delegates FP8 behavior to its child modules + (``Linear``, ``DotProductAttention``) rather than managing FP8 state itself. + +Fp8Padding / Fp8Unpadding +^^^^^^^^^^^^^^^^^^^^^^^^^^ + +**Location**: ``transformer_engine/pytorch/module/fp8_padding.py``, +``transformer_engine/pytorch/module/fp8_unpadding.py`` + +Modules for padding/unpadding sequences to token-count multiples required by FP8 kernels. +These extend ``torch.nn.Module`` directly (not ``TransformerEngineBaseModule``) and use +custom autograd functions internally. + TransformerLayer ^^^^^^^^^^^^^^^^ **Location**: ``transformer_engine/pytorch/transformer.py`` -Composes a full Transformer block (see :doc:`transformer_layer`). Note that -``TransformerLayer`` does **not** inherit from ``TransformerEngineBaseModule`` — it -composes TE modules rather than being one. +Composes a full Transformer block (see :doc:`transformer_layer`). Like +``MultiheadAttention``, it extends ``torch.nn.Module`` directly and delegates FP8 +behavior to its child TE modules. Module Lifecycle ---------------- diff --git a/docs/developer/pytorch_frontend/ops_framework.rst b/docs/developer/pytorch_frontend/ops_framework.rst index a7419a9a4e..19ee5e2de3 100644 --- a/docs/developer/pytorch_frontend/ops_framework.rst +++ b/docs/developer/pytorch_frontend/ops_framework.rst @@ -206,7 +206,8 @@ Catalog of Basic Operations * - ``GELU``, ``SiLU``, ``ReLU``, ``QGELU`` - ``basic/activation.py`` - Element-wise activation functions. - * - ``GEGLU``, ``SwiGLU``, ``ReGLU``, ``SReGLU``, ``QGEGLU``, ``GLU`` + * - ``GEGLU``, ``SwiGLU``, ``ReGLU``, ``SReGLU``, ``QGEGLU``, ``GLU``, + ``ClampedSwiGLU``, ``ScaledSwiGLU`` - ``basic/activation.py``, ``basic/swiglu.py`` - Gated activation functions (2× input split). * - ``Quantize`` diff --git a/docs/developer/quantization/class_hierarchy.rst b/docs/developer/quantization/class_hierarchy.rst index 6d9af28ea5..71695cf053 100644 --- a/docs/developer/quantization/class_hierarchy.rst +++ b/docs/developer/quantization/class_hierarchy.rst @@ -119,23 +119,31 @@ support and minimal overhead. .. code-block:: python class QuantizedTensorStorage: - @property - def data(self) -> torch.Tensor: ... # Raw quantized data - @property - def scale_inv(self) -> torch.Tensor: ... # Scale inverse for dequantization - @property - def dtype(self) -> torch.dtype: ... - @property - def scaling_mode(self) -> str: ... - - def dequantize(self) -> torch.Tensor: ... # Convert back to high precision + # Key internal attributes (accessed directly, not via properties) + _data: torch.Tensor # Raw quantized data + _scale_inv: torch.Tensor # Scale inverse for dequantization + + # Public interface + def update_usage(self, *, rowwise_usage=None, columnwise_usage=None): ... + def get_usages(self) -> tuple[bool, bool]: ... + def prepare_for_saving(self) -> tuple: ... # For autograd save_for_backward + def restore_from_saved(cls, saved) -> "QuantizedTensorStorage": ... + def quantize_(self, tensor, quantizer): ... # In-place quantization + def copy_from_storage(self, other): ... Storage objects are used internally by: - C++ extension calls (converted to ``NVTETensor`` for kernel dispatch) -- The GEMM pipeline (directly passes data + scale_inv) +- The GEMM pipeline (directly passes internal data + scale tensors) - Op fusion internals (avoids autograd overhead in fused kernels) +.. note:: + + ``_data`` and ``_scale_inv`` are internal attributes, not public properties. Subclasses + (e.g., ``Float8TensorStorage``) access them directly. External code should use the + ``get_data_tensors()`` method on ``QuantizedTensor`` or pass storage objects to GEMM + APIs that know how to extract the internal data. + QuantizedTensor (Base: ``transformer_engine/pytorch/quantized_tensor.py``) --------------------------------------------------------------------------- @@ -145,17 +153,22 @@ QuantizedTensor (Base: ``transformer_engine/pytorch/quantized_tensor.py``) .. code-block:: python class QuantizedTensor(torch.Tensor): - _storage: QuantizedTensorStorage - # torch.Tensor subclass machinery def __torch_dispatch__(cls, func, types, args, kwargs): ... - # Access underlying storage - def get_storage(self) -> QuantizedTensorStorage: ... + # Data access + def get_data_tensors(self) -> dict: ... # Extract raw data for kernel calls + def get_metadata(self) -> dict: ... # Scaling metadata # Dequantize def dequantize(self) -> torch.Tensor: ... def float(self) -> torch.Tensor: ... + def bfloat16(self) -> torch.Tensor: ... + def half(self) -> torch.Tensor: ... + + # Lifecycle + def quantize_(self, tensor, quantizer): ... + def clear(self): ... Key behaviors: @@ -163,7 +176,7 @@ Key behaviors: via ``__torch_dispatch__``, so quantized tensors "just work" in standard PyTorch code (at the cost of a dequantize). - **GEMM fast path**: The GEMM implementation checks for ``QuantizedTensor`` inputs and - extracts the storage directly, avoiding dequantization. + uses ``get_data_tensors()`` to extract raw data directly, avoiding dequantization. - **Autograd compatible**: Can be saved in autograd's ``ctx.save_for_backward()``. Lifecycle @@ -180,9 +193,9 @@ A typical quantization lifecycle: # 2. Quantize a tensor (returns QuantizedTensor for autograd) qinput = quantizer(input_tensor) - # 3. Inside GEMM, extract storage for kernel call - storage = qinput.get_storage() - # Pass storage.data and storage.scale_inv to C++ kernel + # 3. Inside GEMM, extract data tensors for kernel call + data_tensors = qinput.get_data_tensors() + # Pass raw data and scale_inv to C++ kernel # 4. Dequantize if needed for non-optimized ops fp32_data = qinput.dequantize() From 294ef606e0c302430b688d82c0b9197cf4b83165 Mon Sep 17 00:00:00 2001 From: Przemek Tredak Date: Thu, 19 Mar 2026 15:20:51 -0700 Subject: [PATCH 04/20] First round of reviews Signed-off-by: Przemek Tredak --- docs/developer/linear_walkthrough.rst | 291 +++++++++--------- .../pytorch_frontend/autograd_integration.rst | 31 ++ .../quantization/class_hierarchy.rst | 15 + .../quantization/rowwise_columnwise.rst | 8 +- .../quantization/scaling_recipes.rst | 8 + 5 files changed, 208 insertions(+), 145 deletions(-) diff --git a/docs/developer/linear_walkthrough.rst b/docs/developer/linear_walkthrough.rst index 90051c0e2c..b7a48f5abc 100644 --- a/docs/developer/linear_walkthrough.rst +++ b/docs/developer/linear_walkthrough.rst @@ -12,11 +12,6 @@ This page traces a complete forward and backward pass through ``te.Linear``, the fundamental Transformer Engine module. It connects concepts from across the codebase: quantization, GEMM, distributed communication, and autograd. -The goal is not just to show *what* happens at each step, but *why* the code is structured -this way — what constraints drive the ordering, what trade-offs are being made, and how -low-precision quantization fundamentally changes the data flow compared to a standard -``torch.nn.Linear``. - .. figure:: ./img/linear_e2e_flow.svg :align: center :width: 95% @@ -28,16 +23,16 @@ low-precision quantization fundamentally changes the data flow compared to a sta Two horizontal swim lanes labeled "Forward" and "Backward". Forward lane (left to right): 1. "Input (BF16)" → - 2. "All-gather (if column-parallel + SP)" → - 3. "input_quantizer(input)" producing "FP8 rowwise + columnwise" → + 2. "input_quantizer(input)" producing "FP8 rowwise + columnwise" → + 3. "All-gather (if column-parallel + SP)" → 4. "general_gemm(weight, input)" producing "Output (BF16)" → 5. "Reduce-scatter (if row-parallel + SP) or All-reduce (if row-parallel)" Backward lane (right to left): 1. "grad_output (BF16)" → 2. "grad_output_quantizer(grad)" → 3. "Dgrad GEMM: general_gemm(grad, weight)" producing "grad_input" → - 4. "Wgrad GEMM: general_gemm(input_columnwise, grad)" producing "grad_weight" → - 5. "All-reduce grad_weight (if column-parallel)" + 4. "Wgrad GEMM: general_gemm(input_columnwise, grad_columnwise)" producing "grad_weight" → + 5. "Reduce-scatter (if row-parallel + SP) or All-reduce (if row-parallel)grad_input" Dotted arrows from forward boxes 3→backward box 4 labeled "saved columnwise input" and forward weight→backward box 3 labeled "saved weight". @@ -53,15 +48,15 @@ TE's Linear does the same math, but introduces two major dimensions of complexit pass must proactively prepare data for backward. 2. **Tensor parallelism** — The weight matrix is sharded across GPUs, so collective - communication (all-gather, reduce-scatter) must be interleaved with computation. - The placement of communication relative to GEMM depends on whether this is a + communication (all-gather, reduce-scatter or all-reduce) must be interleaved with computation. + The type and placement of communication relative to GEMM depends on whether this is a column-parallel or row-parallel linear. These two concerns interact with each other at every step — for instance, quantizing -*before* an all-gather halves the communication volume. The walkthrough below makes these +before an all-gather halves the communication volume. The walkthrough below makes these interactions explicit. -The flow described here is **recipe-agnostic**: the same phases execute regardless of +The flow described here is recipe-agnostic: the same phases execute regardless of whether the active recipe is MXFP8, current scaling, block scaling, or delayed scaling. The quantizer abstraction (see :doc:`quantization/class_hierarchy`) hides recipe-specific details — ``_Linear`` simply calls ``quantizer(tensor)`` and receives quantized data with @@ -70,6 +65,8 @@ the appropriate scales, no matter how those scales were computed. Setup ----- +Here is the basic example of how one could create and call the Linear layer: + .. code-block:: python import transformer_engine.pytorch as te @@ -86,9 +83,9 @@ Setup Phase 1: Module Forward Entry ------------------------------ -**File**: ``transformer_engine/pytorch/module/linear.py``, ``Linear.forward()`` (line ~1343) +**File**: ``transformer_engine/pytorch/module/linear.py``, ``Linear.forward()`` -**Why this phase exists**: Before any math can happen, TE needs to set up the quantization +Before any math can happen, TE needs to set up the quantization infrastructure. Unlike BF16 where you just call GEMM directly, low-precision formats require quantizer objects that know how to cast data and manage scales. This setup phase bridges the gap between the user-facing recipe configuration and the quantizer objects @@ -97,70 +94,98 @@ that actually perform casts. 1. ``prepare_forward()`` is called (inherited from ``TransformerEngineBaseModule``): - Checks if FP8 is enabled via ``FP8GlobalStateManager`` — a global singleton that - tracks whether we're inside an ``fp8_autocast`` context. This global state exists - because FP8 behavior must be coordinated across all TE modules in the model. - - Creates or refreshes **quantizer** instances from the active recipe. Each recipe - type produces a different quantizer class (``MXFP8Quantizer``, + tracks whether we're inside an ``autocast`` context. This global state exists + because FP8 behavior must be coordinated across all TE modules in the model and + we want to be able to pass the information about the recipe the TE modules which + could be multiple levels deep in the user's model hierarchy. + - Creates or refreshes ``Quantizer`` instances from the active recipe. Each recipe + type produces a different ``Quantizer`` subclass (``MXFP8Quantizer``, ``Float8CurrentScalingQuantizer``, ``Float8BlockQuantizer``, etc.). - Recipe-specific customization is applied — for example, current scaling configures ``amax_epsilon`` and ``power_2_scale`` on quantizers, while MXFP8 needs no such configuration. -2. Three quantizers are prepared, one for each tensor that will be cast to FP8: +2. Six quantizers are prepared — three for the forward pass and three for the backward + pass: - ``input_quantizer`` — for the activation tensor (forward GEMM input) - ``weight_quantizer`` — for the weight parameter (forward GEMM input) - - ``grad_output_quantizer`` — for the backward gradient (backward GEMM input) + - ``output_quantizer`` — for the forward output (optional, when ``fp8_output=True``) + - ``grad_output_quantizer`` — for the backward gradient + - ``grad_input_quantizer`` — for the backward gradient input (optional, when ``fp8_grad=True``) + - ``grad_weight_quantizer`` — for the backward weight gradient (currently unused) + + All six are created here because the backward pass executes outside the + ``fp8_autocast`` context and therefore no longer has access to the recipe + configuration. By creating the backward quantizers during forward setup, we + capture the recipe information while it is still available. - **Why three separate quantizers?** Each tensor has independent scaling state because - their value distributions differ significantly — activations, weights, and gradients - have different magnitudes and ranges. A single shared scale would waste dynamic range. + Each tensor gets its own quantizer because activations, weights, and gradients have + different value distributions — sharing a single scale would waste dynamic range. -3. ``_Linear.apply()`` is called, entering the custom ``torch.autograd.Function``. This - boundary separates the ``nn.Module`` API from the autograd machinery — everything - beyond this point runs inside PyTorch's autograd graph and must carefully manage what - is saved for backward. +3. ``_Linear.apply()`` is called, entering the custom ``torch.autograd.Function``. Using + custom autograd functions enables us to hide the details of lower precision from + PyTorch's autograd and implement a custom backward pass. + + All non-Tensor attributes (quantizers, flags, configuration) are packed into a single + ``non_tensor_args`` tuple and passed as one parameter. This is a deliberate + optimization: each additional parameter to an autograd function adds CPU overhead + because PyTorch must inspect it. Packing everything into a single object avoids that + cost. Phase 2: _Linear.forward() — Input Quantization and Communication ------------------------------------------------------------------- -**File**: ``transformer_engine/pytorch/module/linear.py``, ``_Linear.forward()`` (line ~87) +**File**: ``transformer_engine/pytorch/module/linear.py``, ``_Linear.forward()`` -**Why this is the most complex phase**: Input preparation must solve a chicken-and-egg -problem. The GEMM needs the *full* (ungathered) input, but communication is expensive. -With FP8, we also need both a rowwise *and* columnwise quantized version. The code must -decide: quantize before or after communication? The answer depends on the parallel mode. +This is the most complex phase. The GEMM needs the full (ungathered) input, but +communication is expensive and we want to minimize it. A natural idea is to quantize +locally first and communicate the quantized data (1 byte per element for FP8 instead of 2 +for BF16). However, the backward wgrad GEMM needs columnwise data while the forward GEMM +needs rowwise data — so both copies must exist somewhere. Because both layouts are needed, +the total data volume does not shrink if we communicate both together. The key insight is +that we can split the communication into two phases: all-gather only the rowwise data +during the forward pass, and all-gather only the columnwise data during the backward pass. +This also reduces the amount of data each GPU must save between forward and backward, +since only the local shard (before all-gather) needs to be retained. **Column-parallel with sequence parallelism** (QKV projection, FC1): Each GPU holds a shard of the input along the sequence dimension. The GEMM needs the -full sequence, so an all-gather is required. The key decision: **quantize locally first, -then all-gather the quantized data** — this communicates 1 byte per element instead of 2 -(BF16), cutting communication volume in half. +full sequence, so an all-gather is required. The key decision: quantize locally first, +then all-gather the quantized data — for FP8, this communicates 1 byte per element +instead of 2 (BF16), cutting communication volume in half. .. code-block:: python # 1. Configure what quantization layouts we need. # Rowwise: for the forward GEMM (this phase). # Columnwise: for the backward wgrad GEMM (Phase 8, later). - # We quantize both now because the original BF16 input will be discarded. + # We quantize both now because the original BF16 input will be discarded + # to save GPU memory. input_quantizer.set_usage(rowwise=True, columnwise=backward_needs_input) - # 2. Quantize the local shard — produces a QuantizedTensorStorage containing - # quantized rowwise data, columnwise data, and their scale inverses. + # 2. Quantize the local shard. + # The quantizer's `internal` flag is set to True here, so it returns a lightweight + # QuantizedTensorStorage rather than a full QuantizedTensor (torch.Tensor subclass). + # See :doc:`quantization/class_hierarchy` for the internal flag and the distinction + # between these two types. inputmat = input_quantizer(inputmat) - # 3. All-gather the quantized data across TP ranks. - # Each rank sends its local shard; all ranks receive the full tensor. - inputmat_total, _ = gather_along_first_dim(inputmat, tp_group) + # 3. All-gather only the rowwise quantized data across TP ranks. + # The quantizer is passed to the gather with columnwise disabled so that only + # the rowwise data is communicated. Columnwise data stays local. + input_quantizer.set_usage(rowwise=True, columnwise=False) + inputmat_total, _ = gather_along_first_dim( + inputmat, tp_group, quantizer=input_quantizer, + ) .. note:: - The columnwise data is NOT all-gathered along with the rowwise data in all cases. - Some quantizer types (e.g. delayed scaling ``Float8Quantizer``) don't support - columnwise all-gather, so columnwise is disabled before the gather and the input is - re-quantized columnwise in the backward pass. This is a trade-off: less communication - complexity at the cost of redundant quantization later. + The columnwise data is never all-gathered together with the rowwise data during the + forward pass. Instead, it is all-gathered separately during the backward pass when + it is actually needed (see Phase 8). This split is what allows each GPU to save only + its local shard between forward and backward. **No parallelism** (standalone Linear): @@ -169,23 +194,24 @@ both layouts because the backward pass will need columnwise data regardless. .. code-block:: python + # backward_needs_input is True when weight.requires_grad and gradients are enabled + # (i.e., during training). During inference, no columnwise data is needed. input_quantizer.set_usage(rowwise=True, columnwise=backward_needs_input) inputmat = input_quantizer(inputmat) inputmat_total = inputmat # No gather needed -**What ``backward_needs_input`` means**: This flag is ``True`` when ``weight.requires_grad`` -and gradients are enabled — i.e., during training. During inference or when the weight is -frozen, the backward wgrad GEMM won't run, so there's no point producing columnwise data. -This avoids a significant portion of the quantization cost during inference. - Phase 3: _Linear.forward() — Weight Quantization -------------------------------------------------- -**File**: ``transformer_engine/pytorch/module/linear.py``, ``_Linear.forward()`` (line ~250) +**File**: ``transformer_engine/pytorch/module/linear.py``, ``_Linear.forward()`` -**Why weight handling differs from input**: Weights are persistent parameters that don't +Before quantizing the weight, if CPU offloading is enabled, the code starts offloading +the input activation to CPU memory. This is done as early as possible so that the +host-to-device transfer can overlap with the weight quantization and the upcoming GEMM. + +Weights are persistent parameters that don't change between microbatches during gradient accumulation. Inputs change every microbatch. -This asymmetry enables an important optimization: **quantized weight caching**. +This asymmetry enables an important optimization: quantized weight caching. .. code-block:: python @@ -203,14 +229,14 @@ This asymmetry enables an important optimization: **quantized weight caching**. update_workspace=is_first_microbatch is None or is_first_microbatch, ) -**Why cache quantized weights?** In gradient accumulation with *N* microbatches, the weight +In gradient accumulation with *N* microbatches, the weight doesn't change until the optimizer step. Without caching, we'd quantize the same weight *N* times. With caching, we quantize once and reuse the quantized version for all *N* forward passes. The ``is_first_microbatch`` parameter controls this: ``True`` triggers a fresh quantization, ``False`` returns the cached version, ``None`` disables caching entirely (quantize every time). -**Columnwise weight**: The weight also needs a columnwise layout for the backward dgrad +The weight also needs a columnwise layout for the backward dgrad GEMM (Phase 7). The quantizer produces both during the same quantization call, just like the input quantizer. The columnwise weight data is saved for backward. @@ -219,16 +245,11 @@ Phase 4: _Linear.forward() — GEMM **File**: ``transformer_engine/pytorch/cpp_extensions/gemm.py`` -**Why this isn't just a matmul**: Low-precision GEMM requires additional metadata (scale -inverses) that standard matrix multiplication APIs don't handle. TE wraps cuBLASLt with -scale-aware logic and optional communication overlap. - .. code-block:: python # Forward GEMM: output = input @ weight^T - # The result is produced in high precision (BF16/FP16) — cuBLASLt accumulates - # in FP32 internally and casts the output down. This is why low-precision GEMM - # doesn't lose as much accuracy as you might expect from 8-bit math. + # general_gemm can also accept an output quantizer to produce a quantized output + # directly, avoiding a separate quantization step when the next layer needs FP8 input. gemm_out, *_, reduce_scatter_out = general_gemm( weightmat, # A operand (quantized + scale_inv) inputmat_total, # B operand (quantized + scale_inv) @@ -241,28 +262,26 @@ scale-aware logic and optional communication overlap. The call traverses four layers, matching the :doc:`architecture_overview`: -1. **Python** (``cpp_extensions/gemm.py``): Extracts raw data pointers and scale tensors - from the ``QuantizedTensorStorage``. The C++ layer can't understand Python quantized - tensor types, so this translation is mandatory. -2. **pybind11** (``csrc/extensions/gemm.cpp``): Constructs opaque ``NVTETensor`` handles — - the C API's tensor abstraction. This is where TE crosses the Python/C++ boundary. +1. **Python** (``cpp_extensions/gemm.py``): Passes the ``QuantizedTensorStorage`` objects + through to the C++ layer. No data extraction happens here. +2. **pybind11** (``csrc/extensions/gemm.cpp``): Accepts inputs as ``py::handle`` / + ``py::object`` types and converts them to opaque ``NVTETensor`` handles — the C API's + tensor abstraction. The Python types always own the underlying memory; + ``NVTETensor`` is a non-owning view. 3. **C API** (``common/gemm/``): Selects the cuBLASLt algorithm and configures compute types based on the input precision and scaling mode. 4. **cuBLASLt**: Executes the actual matrix multiply on the GPU. -**``use_split_accumulator``**: When ``True``, cuBLASLt uses higher-precision intermediate +``use_split_accumulator``: When ``True``, cuBLASLt uses higher-precision intermediate accumulators, trading some performance for numerical accuracy. This is controlled by the recipe and defaults to ``False`` for the forward pass (where some accumulation error is -acceptable) and ``True`` for backward (where gradient accuracy matters more). +acceptable) and ``True`` for backward (where gradient accuracy matters more). This setting +is Hopper-specific and is a no-op on Blackwell. Phase 5: _Linear.forward() — Output Communication ---------------------------------------------------- -**Why output communication depends on parallel mode**: Tensor parallelism shards the -weight matrix, and the parallel mode determines which dimension is sharded. This directly -determines what communication is needed to produce the correct output. - -For **row-parallel** (output projection, FC2): +For row-parallel (output projection, FC2): Each GPU computes a partial sum (because the input is split across the inner dimension). These partial results must be summed across GPUs to produce the correct output. @@ -280,17 +299,17 @@ These partial results must be summed across GPUs to produce the correct output. # Used when sequence parallelism is disabled. out, _ = allreduce(out, tp_group) -For **column-parallel** (QKV, FC1): No communication needed. Each GPU holds a different +For column-parallel (QKV, FC1): No communication needed. Each GPU holds a different slice of the output features, and these slices are independent — they don't need to be summed or gathered until a downstream operation requires the full output. Phase 6: _Linear.forward() — Save for Backward -------------------------------------------------- -**Why this phase is critical for memory efficiency**: PyTorch autograd requires saving -tensors from the forward pass to compute gradients. For a standard ``nn.Linear``, this -means saving the full input and weight in BF16. With FP8, TE saves *quantized* tensors — -half the memory — but must carefully track which quantization layouts to keep. +PyTorch autograd requires saving tensors from the forward pass to compute gradients. For a +standard ``nn.Linear``, this means saving the full input and weight in BF16. With FP8, TE +saves quantized tensors — half the memory — but must carefully track which quantization +layouts to keep. .. code-block:: python @@ -299,30 +318,43 @@ half the memory — but must carefully track which quantization layouts to keep. if backward_needs_input and own_quantized_input: inputmat.update_usage(rowwise_usage=False, columnwise_usage=True) - ctx.save_for_backward(inputmat, weightmat, weight, bias) - -**The key insight**: The forward GEMM consumes *rowwise* input, but the backward wgrad -GEMM needs *columnwise* input (see :doc:`quantization/rowwise_columnwise`). Since TE -quantized both layouts in Phase 2, it can now discard the rowwise data, keeping only -the columnwise data in the autograd cache. This means activation memory for FP8 is -roughly **half** of what it would be in BF16 — one byte per element (FP8 columnwise) -instead of two (BF16). - -For column-parallel with sequence parallelism, only the *local shard* of the input is + # Flatten QuantizedTensorStorage objects for save_for_backward. + tensors_to_save, tensor_objects = prepare_for_saving( + inputmat, weightmat, weight, bias, + ) + ctx.save_for_backward(*tensors_to_save) + ctx.tensor_objects = tensor_objects + +The ``prepare_for_saving`` / ``restore_from_saved`` pair handles the fact that +``QuantizedTensorStorage`` is not a ``torch.Tensor`` and therefore cannot be passed +directly to ``ctx.save_for_backward()``. It splits each storage into metadata and raw +tensors so that PyTorch can manage their lifetime correctly. See +:doc:`pytorch_frontend/autograd_integration` for the full explanation of why this is +necessary and how it works. + +The forward GEMM consumes rowwise input, but the backward wgrad GEMM needs columnwise +input (see :doc:`quantization/rowwise_columnwise`). Since TE quantized both layouts in +Phase 2, it can now discard the rowwise data, keeping only the columnwise data. This means +activation memory stored for backward with FP8 is roughly half of what it would be in +BF16 — one byte per element (FP8 columnwise) instead of two (BF16). + +For column-parallel with sequence parallelism, only the local shard of the input is saved, not the full all-gathered tensor. The all-gather will be repeated in the backward -pass. This trades communication for memory — a favorable trade-off in large model training -where activation memory is the binding constraint. +pass. For BF16 training this trades communication for memory — a favorable trade-off in +large model training where activation memory is the binding constraint. For FP8 the total +communication volume is the same as BF16: the forward all-gathers only rowwise data and +the backward all-gathers only columnwise data, so we would need to communicate both pieces +regardless. Phase 7: _Linear.backward() — Dgrad (Activation Gradient) ----------------------------------------------------------- -**File**: ``transformer_engine/pytorch/module/linear.py``, ``_Linear.backward()`` (line ~495) +**File**: ``transformer_engine/pytorch/module/linear.py``, ``_Linear.backward()`` -**Why dgrad runs before wgrad**: The dgrad GEMM (``grad_input = grad_output @ weight``) -is on the critical path of backpropagation — downstream layers are waiting for -``grad_input`` to continue their own backward passes. The wgrad GEMM (``grad_weight``) -only updates this layer's parameters and doesn't block anything. Running dgrad first -minimizes the time other layers spend waiting. +In the default execution mode, both dgrad and wgrad run within the same autograd function, +so PyTorch cannot overlap dgrad consumers with wgrad. However, TE provides a special +option to delay wgrad and run it outside this function entirely, which allows overlapping +wgrad with pipeline-parallel communication bubbles. This is not the default behavior. .. code-block:: python @@ -343,7 +375,7 @@ minimizes the time other layers spend waiting. out_dtype=activation_dtype, ) -**Communication overlap opportunity**: For column-parallel, the dgrad output must be +Communication overlap opportunity: For column-parallel, the dgrad output must be reduce-scattered (or all-reduced) back to each GPU's local shard. This communication is launched *asynchronously* — it runs on a separate CUDA stream while the wgrad GEMM (Phase 8) executes on the compute stream. Similarly, if the wgrad needs a re-gathered @@ -353,11 +385,11 @@ and synchronized before wgrad begins. Phase 8: _Linear.backward() — Wgrad (Weight Gradient) ------------------------------------------------------- -**Why wgrad uses columnwise data**: The wgrad GEMM computes ``grad_weight = input^T @ -grad_output``. The transpose of rowwise data is columnwise data. Rather than transposing -at this point (which would require a separate kernel launch and temporary memory), the -forward pass pre-computed the columnwise layout in Phase 2. This is the payoff for the -"dual layout" strategy. +The wgrad GEMM computes ``grad_weight = input^T @ grad_output``. The transpose of rowwise +data is columnwise data, so the pre-computed columnwise layout from Phase 2 is used here. +Transposing at this point would cause both performance and numerics problems — for +block-wise formats, the transpose requires requantization with different block boundaries, +introducing additional quantization error (see :doc:`quantization/rowwise_columnwise`). .. code-block:: python @@ -378,10 +410,9 @@ forward pass pre-computed the columnwise layout in Phase 2. This is the payoff f ub_type=ub_type_wgrad, ) -**``fuse_wgrad_accumulation``**: When enabled (common in Megatron-LM), the wgrad GEMM +``fuse_wgrad_accumulation``: When enabled (common in Megatron-LM), the wgrad GEMM accumulates directly into ``weight.main_grad`` instead of allocating a separate gradient -tensor. This avoids an extra memory allocation and addition kernel, which matters when -gradients are large (e.g., 4096 × 16384 = 67M parameters per Linear). +tensor. This avoids an extra memory allocation and addition kernel. Phase 9: Post-Backward Cleanup -------------------------------- @@ -392,42 +423,15 @@ For most recipes (MXFP8, current scaling, block scaling), there is no post-backw to update — scales are computed inline during quantization and don't carry state across iterations. This is one of the advantages of these recipes: they are stateless. -**Delayed scaling only**: The delayed scaling recipe is an exception. It maintains an -**amax history** — a rolling window of observed maximum values — that feeds forward into +Delayed scaling only: The delayed scaling recipe is an exception. It maintains an +amax history — a rolling window of observed maximum values — that feeds forward into the next iteration's scale computation. After the backward pass completes, the first FP8 -module in the model triggers ``FP8GlobalStateManager.reduce_and_update_fp8_tensors()``, -which all-reduces amax values across TP ranks and updates the history buffer. This ensures -all ranks agree on scales. If you're not using delayed scaling, this code path is skipped. - -Summary of Files Touched -------------------------- - -.. list-table:: - :header-rows: 1 - :widths: 50 50 - - * - File - - Role in Linear - * - ``pytorch/module/linear.py`` - - ``Linear.forward()``, ``_Linear`` autograd - * - ``pytorch/module/base.py`` - - ``prepare_forward()``, quantizer setup - * - ``pytorch/quantized_tensor.py`` - - ``Quantizer``, ``QuantizedTensorStorage`` base classes - * - ``pytorch/tensor/mxfp8_tensor.py`` - - ``MXFP8Quantizer``, block-scaled FP8 cast - * - ``pytorch/tensor/float8_tensor.py`` - - ``Float8Quantizer``, ``Float8CurrentScalingQuantizer`` - * - ``pytorch/cpp_extensions/gemm.py`` - - ``general_gemm()`` Python wrapper - * - ``pytorch/csrc/extensions/gemm.cpp`` - - pybind11 GEMM binding - * - ``common/gemm/cublaslt_gemm.cu`` - - cuBLASLt GEMM kernel dispatch - * - ``pytorch/distributed.py`` - - ``allreduce()``, ``gather_along_first_dim()``, etc. - * - ``pytorch/quantization.py`` - - ``FP8GlobalStateManager``, ``fp8_autocast`` +module in the autocast region triggers +``FP8GlobalStateManager.reduce_and_update_fp8_tensors()``, which all-reduces amax values +across TP ranks and updates the history buffer. This ensures all ranks agree on scales. If +you're not using delayed scaling, this code path is skipped. See +:doc:`quantization/scaling_recipes` for details on the amax update flow and important +caveats about distributed correctness. See Also -------- @@ -437,3 +441,4 @@ See Also - :doc:`quantization/rowwise_columnwise` — Why both layouts exist - :doc:`pytorch_frontend/autograd_integration` — Autograd patterns - :doc:`distributed/tensor_parallel` — TP communication in Linear +- :doc:`/features/low_precision_training/index` — User-facing recipe documentation and examples diff --git a/docs/developer/pytorch_frontend/autograd_integration.rst b/docs/developer/pytorch_frontend/autograd_integration.rst index 21d89f07e0..810d16b501 100644 --- a/docs/developer/pytorch_frontend/autograd_integration.rst +++ b/docs/developer/pytorch_frontend/autograd_integration.rst @@ -71,6 +71,37 @@ What gets saved for backward depends on the configuration: contain both rowwise data (used by dgrad GEMM) and columnwise data (used by wgrad GEMM). Memory cost: ~2× the FP8 data size (rowwise + columnwise). +Saving QuantizedTensorStorage via prepare_for_saving +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +PyTorch offers two ways to pass data from forward to backward: direct attributes on +``ctx``, and ``ctx.save_for_backward()``. Direct ``ctx`` attributes are not released after +the backward pass — they persist until the *next* forward pass. Storing tensors there +would keep memory alive throughout the entire backward pass and into the next forward, +which is very wasteful. Using ``save_for_backward`` lets PyTorch release the memory +promptly, but it only accepts ``torch.Tensor`` objects. + +Since ``QuantizedTensorStorage`` is not a ``torch.Tensor``, the helper function +``prepare_for_saving()`` (in ``quantized_tensor.py``) splits each storage into: + +- Its **metadata** (with all tensor fields set to ``None``) — stored on ``ctx.tensor_objects``. +- A list of raw **``torch.Tensor`` objects** — passed through ``ctx.save_for_backward()``. + +In the backward pass, ``restore_from_saved()`` reassembles the original +``QuantizedTensorStorage`` objects from the saved tensors and metadata. + +.. code-block:: python + + # Forward: split and save + tensors_to_save, tensor_objects = prepare_for_saving(inputmat, weightmat, weight, bias) + ctx.save_for_backward(*tensors_to_save) + ctx.tensor_objects = tensor_objects + + # Backward: reassemble + inputmat, weightmat, weight, bias = restore_from_saved( + ctx.saved_tensors, ctx.tensor_objects, + ) + **FP8 enabled + activation recompute**: Only save the high-precision input (or a stashed copy). During backward, re-run the quantization to produce fresh FP8 data. Saves memory at the cost of recomputation. diff --git a/docs/developer/quantization/class_hierarchy.rst b/docs/developer/quantization/class_hierarchy.rst index 71695cf053..718e4565f7 100644 --- a/docs/developer/quantization/class_hierarchy.rst +++ b/docs/developer/quantization/class_hierarchy.rst @@ -83,6 +83,21 @@ responsibilities: @property def scaling_mode(self) -> str: ... +The ``internal`` Flag +^^^^^^^^^^^^^^^^^^^^^ + +The ``Quantizer.internal`` flag controls whether ``__call__()`` returns a +``QuantizedTensorStorage`` (lightweight, no autograd overhead) or a full +``QuantizedTensor`` (``torch.Tensor`` subclass with autograd support): + +- ``internal=True`` → returns ``QuantizedTensorStorage``. Used for tensors that stay + within the boundaries of a custom ``torch.autograd.Function`` (e.g., the activation + inside ``_Linear``). Because these tensors never escape to user code or participate in + autograd, the storage-only form avoids the CPU overhead of ``torch.Tensor`` subclass + mechanics. +- ``internal=False`` → returns ``QuantizedTensor``. Used for tensors that are visible + outside the autograd function (e.g., returned to the user as a quantized output). + Concrete Quantizers ^^^^^^^^^^^^^^^^^^^ diff --git a/docs/developer/quantization/rowwise_columnwise.rst b/docs/developer/quantization/rowwise_columnwise.rst index 40345d9a48..5b2e7a781a 100644 --- a/docs/developer/quantization/rowwise_columnwise.rst +++ b/docs/developer/quantization/rowwise_columnwise.rst @@ -101,9 +101,13 @@ Producing both layouts doubles the quantization work and memory for activations. a deliberate trade-off: - **Without dual layout**: The wgrad GEMM would need to transpose and requantize at - backward time, adding latency to the critical path. + backward time. Beyond the performance cost (extra kernel launch and temporary memory), + this is a numerics problem for block-wise formats. For MXFP8, block-wise FP8, and + NVFP4, transposing requires requantization with different block boundaries, which + introduces additional quantization error. - **With dual layout**: Extra memory and compute during forward, but the backward wgrad - can proceed immediately with pre-computed columnwise data. + can proceed immediately with pre-computed columnwise data, avoiding both the latency + and the accuracy loss. For memory-constrained scenarios, activation recomputation can be used — the forward pass discards the saved columnwise data and recomputes it during backward. See diff --git a/docs/developer/quantization/scaling_recipes.rst b/docs/developer/quantization/scaling_recipes.rst index f338b04103..fd858d4c74 100644 --- a/docs/developer/quantization/scaling_recipes.rst +++ b/docs/developer/quantization/scaling_recipes.rst @@ -101,6 +101,14 @@ For delayed tensor scaling, the amax update follows this sequence each iteration For current-scaling and block-scaling modes, scales are computed just-in-time during the cast kernel, so no amax history is maintained. +.. warning:: + + For the amax all-reduce to work correctly, every rank must participate in the + collective. This means each ``fp8_autocast`` region must contain at least one FP8 + module to trigger the reduction. The modules do not need to be the same across + ranks — the reduction operates on amaxes from all layers simultaneously and updates + only those that were touched (identified by checking against the initial zero value). + Recipe → Quantizer Mapping -------------------------- From 31a2719900db38874696aad7033863dafd4a5297 Mon Sep 17 00:00:00 2001 From: Przemek Tredak Date: Thu, 19 Mar 2026 15:42:44 -0700 Subject: [PATCH 05/20] After second round of review Signed-off-by: Przemek Tredak --- docs/developer/architecture_overview.rst | 17 +-- .../distributed/fsdp_integration.rst | 4 +- docs/developer/index.rst | 6 +- docs/developer/jax_frontend/module_system.rst | 2 +- docs/developer/linear_walkthrough.rst | 4 +- .../pytorch_frontend/autograd_integration.rst | 75 ++++++++----- docs/developer/pytorch_frontend/index.rst | 3 +- .../pytorch_frontend/module_hierarchy.rst | 106 ++---------------- .../quantization/scaling_recipes.rst | 12 +- docs/installation.rst | 9 ++ 10 files changed, 91 insertions(+), 147 deletions(-) diff --git a/docs/developer/architecture_overview.rst b/docs/developer/architecture_overview.rst index 7e928c6288..8d4d1c389d 100644 --- a/docs/developer/architecture_overview.rst +++ b/docs/developer/architecture_overview.rst @@ -24,7 +24,7 @@ sit on top of a shared C++ core that dispatches to optimized CUDA kernels. Second band: "Binding Layer" with sub-boxes "pybind11 (PyTorch)" and "XLA FFI (JAX)". Third band: "C++ Core Library" (single box spanning full width). Bottom band: "CUDA Kernels" with sub-boxes for major areas: GEMM, Normalization, - Activation, Attention, Cast/Transpose, Fused Rope, Fused Softmax. + Quantize, Activation, Attention, Other. Arrows flow downward between bands. A side annotation shows "transformer_engine/common/" pointing at the bottom two bands. @@ -41,7 +41,7 @@ Layer Diagram **C++ Core** (``transformer_engine/common/``) Framework-agnostic library implementing the actual computation logic. Exposes a C API - (``transformer_engine.h``) for stability across framework bindings. + (``transformer_engine.h``, etc.) for stability across framework bindings. **CUDA Kernels** (within ``transformer_engine/common/``) Hand-optimized CUDA kernels and cuDNN/cuBLASLt integrations organized by functional @@ -93,7 +93,8 @@ Directory Map Data Flow: Forward Pass ----------------------- -A typical forward pass through a quantized linear layer follows this path: +A typical forward pass through a quantized linear layer follows this path (see +:doc:`linear_walkthrough` for a detailed end-to-end trace): .. figure:: ./img/data_flow_forward.svg :align: center @@ -123,8 +124,9 @@ For delayed tensor scaling, an amax (absolute maximum) feedback loop maintains s factor history across iterations. The backward pass follows the same pattern twice: once for the activation gradient (dgrad) -and once for the weight gradient (wgrad). The wgrad path uses *columnwise* quantized data -to satisfy cuBLASLt's transposed layout requirements — see +and once for the weight gradient (wgrad). Both backward GEMMs use inputs already quantized +during the forward pass (saved via autograd). The wgrad path uses *columnwise* quantized +data to satisfy cuBLASLt's transposed layout requirements — see :doc:`quantization/rowwise_columnwise` for details. Guiding Principles @@ -150,8 +152,7 @@ layer. **Leverage the ecosystem, don't reinvent it** TE provides custom CUDA kernels where they deliver clear performance wins (fused attention, quantized GEMM, fused normalization + cast). For everything else, we prefer - the framework's native capabilities — ``torch.compile`` for fusion, XLA for graph - optimization, cuBLASLt for matrix multiplication. A custom kernel must justify its + the framework's native capabilities. A custom kernel must justify its existence against the ecosystem alternative. **No memory allocation in the C++ layer** @@ -164,7 +165,7 @@ layer. **Hide complexity, but stay extensible** Low-precision training involves many moving parts (scaling modes, amax history, block sizes, format selection). TE hides this complexity behind simple APIs — a user can - pass a recipe to ``fp8_autocast`` and get FP8 training without understanding the + pass a recipe to ``autocast`` and get FP8 training without understanding the internals. But we do not want to be a black box: the quantization system is designed to be extensible (see :doc:`quantization/adding_new_type`), scaling recipes are configurable, and advanced users can provide custom quantizers. diff --git a/docs/developer/distributed/fsdp_integration.rst b/docs/developer/distributed/fsdp_integration.rst index 5ffdf44d8e..c25b5a134b 100644 --- a/docs/developer/distributed/fsdp_integration.rst +++ b/docs/developer/distributed/fsdp_integration.rst @@ -87,7 +87,7 @@ during weight all-gather. No explicit setup or registration is needed. # FSDP2 automatically handles quantized weight all-gather fully_shard(model) - with te.fp8_autocast(enabled=True): + with te.autocast(enabled=True): output = model(input) # Weights gathered/scattered automatically FSDP1: Activation Scatter/Gather @@ -134,7 +134,7 @@ TE also provides a convenience wrapper: model = FSDP(te_model, ...) prepare_te_modules_for_fsdp(model) - with te.fp8_autocast(enabled=True): + with te.autocast(enabled=True): output = model(input) Summary: Which Mechanism Does What diff --git a/docs/developer/index.rst b/docs/developer/index.rst index c2485a8719..ffc1133cf1 100644 --- a/docs/developer/index.rst +++ b/docs/developer/index.rst @@ -19,12 +19,14 @@ How to Use This Guide --------------------- - **New contributors**: Start with :doc:`architecture_overview` for the big picture, - then read :doc:`linear_walkthrough` for an end-to-end trace through a concrete PyTorch module. + then read :doc:`linear_walkthrough` for an end-to-end trace through a concrete PyTorch + module. For building from source, see :doc:`/installation`. For running tests and + contributing guidelines, see ``CONTRIBUTING.rst`` at the repository root. - **Working on new quantization recipe**: See :doc:`quantization/index` for the Quantizer/Storage/Tensor design and :doc:`cpp_core/type_system` for the underlying C/C++ types. - **Working on attention**: See :doc:`attention/index` for backend selection and kernel organization. -- **Working on distributed**: See :doc:`distributed/index` for tensor/sequence parallelism +- **Working on distributed training**: See :doc:`distributed/index` for tensor/sequence parallelism and communication overlap. - **Framework-specific work**: See :doc:`pytorch_frontend/index` or :doc:`jax_frontend/index`. diff --git a/docs/developer/jax_frontend/module_system.rst b/docs/developer/jax_frontend/module_system.rst index 820b4ac691..b25527ed7f 100644 --- a/docs/developer/jax_frontend/module_system.rst +++ b/docs/developer/jax_frontend/module_system.rst @@ -95,7 +95,7 @@ use ``MeshResource`` and rely on XLA SPMD for communication: mesh_resource=MeshResource(tp_resource="tp"), ) -**Quantization context**: Instead of ``fp8_autocast()``, JAX uses explicit quantizer +**Quantization context**: Instead of ``autocast()``, JAX uses explicit quantizer arguments: .. code-block:: python diff --git a/docs/developer/linear_walkthrough.rst b/docs/developer/linear_walkthrough.rst index b7a48f5abc..04dbc5a8f8 100644 --- a/docs/developer/linear_walkthrough.rst +++ b/docs/developer/linear_walkthrough.rst @@ -77,7 +77,7 @@ Here is the basic example of how one could create and call the Linear layer: # Enable FP8 with MXFP8 recipe recipe = MXFP8BlockScaling() - with te.fp8_autocast(enabled=True, fp8_recipe=recipe): + with te.autocast(enabled=True, recipe=recipe): output = linear(input) # Triggers the flow below Phase 1: Module Forward Entry @@ -116,7 +116,7 @@ that actually perform casts. - ``grad_weight_quantizer`` — for the backward weight gradient (currently unused) All six are created here because the backward pass executes outside the - ``fp8_autocast`` context and therefore no longer has access to the recipe + ``autocast`` context and therefore no longer has access to the recipe configuration. By creating the backward quantizers during forward setup, we capture the recipe information while it is still available. diff --git a/docs/developer/pytorch_frontend/autograd_integration.rst b/docs/developer/pytorch_frontend/autograd_integration.rst index 810d16b501..7a6d57c767 100644 --- a/docs/developer/pytorch_frontend/autograd_integration.rst +++ b/docs/developer/pytorch_frontend/autograd_integration.rst @@ -81,14 +81,19 @@ would keep memory alive throughout the entire backward pass and into the next fo which is very wasteful. Using ``save_for_backward`` lets PyTorch release the memory promptly, but it only accepts ``torch.Tensor`` objects. -Since ``QuantizedTensorStorage`` is not a ``torch.Tensor``, the helper function +Since ``QuantizedTensorStorage`` is not a ``torch.Tensor`` (see +:doc:`/developer/quantization/class_hierarchy` for the distinction between +``QuantizedTensorStorage`` and ``QuantizedTensor``), the helper function ``prepare_for_saving()`` (in ``quantized_tensor.py``) splits each storage into: - Its **metadata** (with all tensor fields set to ``None``) — stored on ``ctx.tensor_objects``. - A list of raw **``torch.Tensor`` objects** — passed through ``ctx.save_for_backward()``. In the backward pass, ``restore_from_saved()`` reassembles the original -``QuantizedTensorStorage`` objects from the saved tensors and metadata. +``QuantizedTensorStorage`` objects from the saved tensors and metadata. After +reassembling, ``ctx.tensor_objects`` must be deleted to avoid keeping references to the +reassembled tensors on ``ctx`` (which would defeat the purpose of using +``save_for_backward`` in the first place). .. code-block:: python @@ -97,48 +102,62 @@ In the backward pass, ``restore_from_saved()`` reassembles the original ctx.save_for_backward(*tensors_to_save) ctx.tensor_objects = tensor_objects - # Backward: reassemble + # Backward: reassemble and release ctx references inputmat, weightmat, weight, bias = restore_from_saved( ctx.saved_tensors, ctx.tensor_objects, ) - -**FP8 enabled + activation recompute**: Only save the high-precision input (or a stashed -copy). During backward, re-run the quantization to produce fresh FP8 data. Saves memory -at the cost of recomputation. + del ctx.tensor_objects # Avoid holding tensor references on ctx Activation Recomputation ------------------------ -TE supports activation recomputation (gradient checkpointing) at multiple granularities: +TE provides its own ``checkpoint()`` function (in ``transformer_engine/pytorch/distributed.py``) +rather than relying on the standard PyTorch ``torch.utils.checkpoint.checkpoint()``. The TE +version handles FP8 state correctly across the recompute boundary. -- **Full recompute**: Re-run the entire forward pass during backward. Standard PyTorch - ``checkpoint()`` works with TE modules. -- **Selective recompute**: Only recompute the quantization (cast) operations, keeping - the GEMM results. This is cheaper than full recompute because casting is - memory-bandwidth-bound while GEMM is compute-bound. +TE supports activation recomputation (gradient checkpointing) at multiple granularities: -The ``TransformerLayer`` module provides a ``activation_checkpointing`` parameter that -controls this behavior. +- **Full recompute**: Re-run the entire forward pass during backward. The TE checkpoint + function re-enters the ``autocast`` context during recomputation so that quantizers + and FP8 state are correctly restored. +- **Selective recompute**: Only recompute specific operations (e.g., the core attention + computation) while keeping other activations saved. The ``TransformerLayer`` module + provides an ``activation_checkpointing`` parameter that controls selective + recomputation. FP8 Tensors in Autograd ------------------------ ``QuantizedTensor`` (a ``torch.Tensor`` subclass) can be saved via -``ctx.save_for_backward()`` like any other tensor. Key behaviors: +``ctx.save_for_backward()`` like any other tensor. No implicit dequantization happens +during save — the FP8 data is stored as-is. + +__torch_dispatch__ and Automatic Dequantization +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -- **No implicit dequantization** during save — the FP8 data is stored as-is. -- During backward, the saved ``QuantizedTensor`` is retrieved and its ``.get_storage()`` - method provides the raw FP8 data + scales for GEMM. -- If a non-TE operation receives a ``QuantizedTensor``, ``__torch_dispatch__`` - automatically dequantizes it. +``QuantizedTensor`` implements ``__torch_dispatch__`` to intercept all PyTorch operations. +This is the most complex part of the ``QuantizedTensor`` class. The dispatch logic handles +three cases: + +1. **TE-optimized operations** (e.g., GEMM): These recognize ``QuantizedTensor`` inputs + and extract the raw FP8 data + scales directly via ``get_data_tensors()``, avoiding + any dequantization. + +2. **Non-mutable operations**: For standard PyTorch ops that don't modify their inputs + (e.g., ``torch.add``, ``torch.matmul``), the dispatch automatically dequantizes all + quantized tensor arguments, executes the operation in high precision, and returns the + result. This makes quantized tensors "just work" with arbitrary PyTorch code, at the + cost of a dequantize. + +3. **In-place operations**: These require special care. The dispatch dequantizes the + inputs, executes the in-place op on the dequantized data, then re-quantizes the + result back into the original ``QuantizedTensor``'s storage. This ensures the + quantized representation stays consistent after mutation. Interaction with torch.compile ------------------------------- -TE modules are compatible with ``torch.compile`` but require care: - -- ``NVTE_TORCH_COMPILE=1`` enables compile-friendly code paths. -- Some autograd functions use ``torch.compiler.is_compiling()`` to select between - eager and compile-compatible implementations. -- FP8 state management (amax updates, scale refreshes) must happen outside compiled - regions because they involve in-place mutation of module state. +TE modules work with ``torch.compile`` by disabling compilation for TE-specific code +(which results in graph breaks). The ``NVTE_TORCH_COMPILE=1`` environment variable +enables paths where TE uses ``torch.compile`` internally within some modules. In +CPU-limited training scenarios, the graph breaks may cause slowdowns. diff --git a/docs/developer/pytorch_frontend/index.rst b/docs/developer/pytorch_frontend/index.rst index 8cfc721340..db09b50cee 100644 --- a/docs/developer/pytorch_frontend/index.rst +++ b/docs/developer/pytorch_frontend/index.rst @@ -7,8 +7,7 @@ PyTorch Frontend ================ The PyTorch frontend (``transformer_engine/pytorch/``) provides ``nn.Module`` subclasses, -autograd integration, quantized tensor types, and distributed utilities. It is the most -widely used frontend and the primary development target. +autograd integration, quantized tensor types, and distributed utilities. .. toctree:: :maxdepth: 1 diff --git a/docs/developer/pytorch_frontend/module_hierarchy.rst b/docs/developer/pytorch_frontend/module_hierarchy.rst index 843678de6d..ccdb164a34 100644 --- a/docs/developer/pytorch_frontend/module_hierarchy.rst +++ b/docs/developer/pytorch_frontend/module_hierarchy.rst @@ -54,7 +54,8 @@ TE modules that support FP8 quantization. Key responsibilities: **Quantizer Access** -- Maintains quantizers for each quantized tensor (input, weight, gradient). +- Maintains quantizers for each quantized tensor (input, weight, output, gradient input, + gradient output, gradient weight). - Quantizer instances are recreated when the recipe changes. **Distributed Hooks** @@ -63,101 +64,14 @@ TE modules that support FP8 quantization. Key responsibilities: - ``set_sequence_parallel()``: Enable sequence parallelism. - Manages all-reduce of amax values across distributed ranks. -**Parameter Management** +**Weight Caching** -- ``weight`` parameter with optional FP8 storage. -- ``bias`` parameter (optional). -- Weight caching for FP8 weights that persist across iterations. +- Provides caching infrastructure for quantized weights that persist across microbatches + during gradient accumulation. -Module Subclasses ------------------ - -Linear -^^^^^^ - -**Location**: ``transformer_engine/pytorch/module/linear.py`` - -The workhorse module. Performs ``output = input × weight^T + bias`` with optional FP8 -quantization of both input and weight. - -- Defines a ``_Linear`` autograd function (see :doc:`autograd_integration`). -- Supports tensor parallelism (column-parallel and row-parallel modes). -- Can fuse with preceding LayerNorm (via ``LayerNormLinear``). - -LayerNorm / RMSNorm -^^^^^^^^^^^^^^^^^^^^ - -**Location**: ``transformer_engine/pytorch/module/layernorm.py``, -``transformer_engine/pytorch/module/rmsnorm.py`` - -Normalization modules that can fuse their output quantization with the normalization -kernel (single kernel pass for norm + cast to FP8). - -LayerNormLinear -^^^^^^^^^^^^^^^ - -**Location**: ``transformer_engine/pytorch/module/layernorm_linear.py`` - -Fuses LayerNorm (or RMSNorm) with a subsequent Linear. The normalization output is -produced directly in FP8, avoiding a round-trip through high precision. - -LayerNormMLP -^^^^^^^^^^^^ - -**Location**: ``transformer_engine/pytorch/module/layernorm_mlp.py`` - -Fuses LayerNorm + Linear + Activation + Linear (the full MLP block). This is the most -aggressively fused module, combining up to 4 operations. - -GroupedLinear -^^^^^^^^^^^^^ - -**Location**: ``transformer_engine/pytorch/module/grouped_linear.py`` - -Batched linear operations for Mixture-of-Experts (MoE) where multiple smaller linear -layers execute as a single grouped GEMM. - -DotProductAttention -^^^^^^^^^^^^^^^^^^^ - -**Location**: ``transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py`` - -Computes scaled dot-product attention with automatic backend selection -(see :doc:`/developer/attention/backends`). Supports FP8 attention on Hopper+. - -MultiheadAttention -^^^^^^^^^^^^^^^^^^ - -**Location**: ``transformer_engine/pytorch/attention/multi_head_attention.py`` - -Composes QKV projection (Linear), DotProductAttention, and output projection (Linear) -into a complete multi-head attention block. Handles the split into heads and optional -key-value caching for inference. - -.. note:: - - ``MultiheadAttention`` extends ``torch.nn.Module`` directly, **not** - ``TransformerEngineBaseModule``. It delegates FP8 behavior to its child modules - (``Linear``, ``DotProductAttention``) rather than managing FP8 state itself. - -Fp8Padding / Fp8Unpadding -^^^^^^^^^^^^^^^^^^^^^^^^^^ - -**Location**: ``transformer_engine/pytorch/module/fp8_padding.py``, -``transformer_engine/pytorch/module/fp8_unpadding.py`` - -Modules for padding/unpadding sequences to token-count multiples required by FP8 kernels. -These extend ``torch.nn.Module`` directly (not ``TransformerEngineBaseModule``) and use -custom autograd functions internally. - -TransformerLayer -^^^^^^^^^^^^^^^^ - -**Location**: ``transformer_engine/pytorch/transformer.py`` - -Composes a full Transformer block (see :doc:`transformer_layer`). Like -``MultiheadAttention``, it extends ``torch.nn.Module`` directly and delegates FP8 -behavior to its child TE modules. +Note that individual parameters (``weight``, ``bias``) are defined by each subclass, +not by the base module. See the module hierarchy diagram above for the full list of +subclasses and their docstrings for module-specific details. Module Lifecycle ---------------- @@ -166,12 +80,12 @@ A typical forward pass through a TE module: .. code-block:: text - 1. fp8_autocast() sets global FP8 state + 1. autocast() sets global FP8 state 2. module.forward() called a. pre_forward() — refresh FP8 scales, create quantizers b. Quantize input via input_quantizer(input) c. Get quantized weight (cached or quantize via weight_quantizer) - d. Call _Linear.forward() autograd function + d. Call the module's custom autograd function i. general_gemm(quantized_input, quantized_weight) ii. Save tensors for backward e. post_forward() — record amax values diff --git a/docs/developer/quantization/scaling_recipes.rst b/docs/developer/quantization/scaling_recipes.rst index fd858d4c74..7eacc550b5 100644 --- a/docs/developer/quantization/scaling_recipes.rst +++ b/docs/developer/quantization/scaling_recipes.rst @@ -45,7 +45,7 @@ Recipe classes specify: amax_compute_algo="max", ) - with te.fp8_autocast(enabled=True, fp8_recipe=recipe): + with te.autocast(enabled=True, recipe=recipe): output = model(input) FP8GlobalStateManager @@ -59,15 +59,15 @@ that coordinates FP8 state across all modules in a model. Key responsibilities: - **Coordinate amax reduction** across distributed ranks (for delayed scaling, all ranks must agree on the global amax). -The manager is activated by the ``fp8_autocast`` context manager: +The manager is activated by the ``autocast`` context manager: .. code-block:: python - # Pseudocode for fp8_autocast + # Pseudocode for autocast @contextmanager - def fp8_autocast(enabled, fp8_recipe): + def autocast(enabled, recipe): FP8GlobalStateManager.set_enabled(enabled) - FP8GlobalStateManager.set_recipe(fp8_recipe) + FP8GlobalStateManager.set_recipe(recipe) try: yield finally: @@ -104,7 +104,7 @@ the cast kernel, so no amax history is maintained. .. warning:: For the amax all-reduce to work correctly, every rank must participate in the - collective. This means each ``fp8_autocast`` region must contain at least one FP8 + collective. This means each ``autocast`` region must contain at least one FP8 module to trigger the reduction. The modules do not need to be the same across ranks — the reduction operates on amaxes from all layers simultaneously and updates only those that were touched (identified by checking against the initial zero value). diff --git a/docs/installation.rst b/docs/installation.rst index cc48a0adac..5ac8fdfff3 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -122,6 +122,15 @@ To build the C++ extensions with debug symbols, e.g. with the `-g` flag: NVTE_BUILD_DEBUG=1 pip3 install --no-build-isolation . +To speed up repeated builds (e.g. during development), enable ccache: + +.. code-block:: bash + + NVTE_USE_CCACHE=1 pip3 install --no-build-isolation . + +This uses `ccache` by default. To use a different binary (e.g. `sccache`), set +`NVTE_CCACHE_BIN=sccache`. + .. include:: ../README.rst :start-after: troubleshooting-begin-marker-do-not-remove :end-before: troubleshooting-end-marker-do-not-remove From ff8f1c265e372250ee9749e9aefb4b3e5a21d8a1 Mon Sep 17 00:00:00 2001 From: Przemek Tredak Date: Thu, 19 Mar 2026 16:45:47 -0700 Subject: [PATCH 06/20] And more review changes Signed-off-by: Przemek Tredak --- .../pytorch_frontend/cpp_extensions.rst | 191 +++++++++++++----- docs/developer/quantization/index.rst | 15 +- 2 files changed, 144 insertions(+), 62 deletions(-) diff --git a/docs/developer/pytorch_frontend/cpp_extensions.rst b/docs/developer/pytorch_frontend/cpp_extensions.rst index 341ebfe9ef..1cddfad259 100644 --- a/docs/developer/pytorch_frontend/cpp_extensions.rst +++ b/docs/developer/pytorch_frontend/cpp_extensions.rst @@ -9,7 +9,9 @@ C++ Extensions Bridge ===================== The PyTorch frontend communicates with the C++ core through a pybind11 extension module. -This bridge converts PyTorch tensors to ``NVTETensor`` handles and calls the C API. +This bridge accepts Python objects (``QuantizedTensorStorage``, ``QuantizedTensor``, or +plain ``torch.Tensor``), converts them to ``TensorWrapper`` / ``NVTETensor`` handles, and +calls the C API. Architecture ------------ @@ -17,13 +19,13 @@ Architecture .. code-block:: text Python (transformer_engine/pytorch/cpp_extensions/) - │ torch.Tensor → NVTETensor conversion + │ Optional auxiliary processing (e.g. custom recipe handling) ▼ pybind11 (transformer_engine/pytorch/csrc/extensions/) - │ Call C API functions + │ py::handle → TensorWrapper conversion, GIL release, C API call ▼ C API (transformer_engine/common/include/transformer_engine/) - │ Unpack NVTETensor → internal Tensor struct + │ NVTETensor handle → internal Tensor struct ▼ CUDA kernels (transformer_engine/common/) @@ -38,35 +40,28 @@ Only a subset of C++ extensions have dedicated Python wrapper modules: - ``fused_attn.py`` — ``fused_attn_fwd()``, ``fused_attn_bwd()`` Most other C++ extensions (normalization, activation, cast, transpose, etc.) are exposed -**directly** through the compiled ``transformer_engine_torch`` pybind11 module (imported +directly through the compiled ``transformer_engine_torch`` pybind11 module (imported as ``tex``) and called without a Python wrapper layer. For example, normalization is called as ``tex.layernorm_fwd(...)`` rather than through a ``normalization.py`` wrapper. -The Python wrappers that do exist serve as the translation layer between PyTorch's -quantized tensor types and the C++ API: - -1. Extract raw data tensors and scales from ``QuantizedTensor`` / ``QuantizedTensorStorage`` - objects. -2. Construct ``NVTETensor`` handles via helper functions. -3. Call the pybind11-exposed C++ function. -4. Wrap outputs back into Python tensor types. - -**Example** — ``general_gemm()`` (simplified): +The Python wrappers that do exist provide a hook for auxiliary processing before calling +into C++. For example, ``general_gemm()`` handles +custom recipe dispatch and block-scaling detection, then passes the Python objects +through to C++: .. code-block:: python - def general_gemm(A, B, bias=None, ...): - # Convert QuantizedTensor → NVTETensor components - A_data, A_scale_inv, A_dtype = _extract_tensors(A) - B_data, B_scale_inv, B_dtype = _extract_tensors(B) + def general_gemm(A, B, bias=None, quantization_params=None, ...): + # Auxiliary logic (custom tensor dispatch, debug handling, etc.) + if is_custom(A) or is_custom(B): + return custom_gemm(...) - # Call pybind11 extension - output = tex.general_gemm( - A_data, A_scale_inv, A_dtype, - B_data, B_scale_inv, B_dtype, - bias, output_dtype, ... + # Pass Python objects as-is to the pybind11 layer. + # QuantizedTensorStorage/QuantizedTensor objects are NOT unpacked here. + out, bias_grad, gelu_input, extra_output = tex.generic_gemm( + A, transa, B, transb, out, quantization_params, ... ) - return output + return out, bias_grad, gelu_input, extra_output C++ Side (pybind11) ------------------- @@ -74,46 +69,132 @@ C++ Side (pybind11) **Location**: ``transformer_engine/pytorch/csrc/extensions/`` The pybind11 module (``transformer_engine/pytorch/csrc/extensions/pybind.cpp``) -registers Python-callable functions that: +registers Python-callable functions. These functions: -1. Accept ``torch::Tensor`` and scalar arguments from Python. -2. Construct ``NVTETensor`` handles using ``makeTransformerEngineTensor()``. -3. Call the C API function (e.g., ``nvte_general_gemm()``). -4. Return ``torch::Tensor`` outputs. +1. Accept inputs as ``py::handle`` or ``py::object`` types — not ``at::Tensor``. This + allows them to receive ``QuantizedTensorStorage``, ``QuantizedTensor``, or plain + ``torch.Tensor`` objects without the Python layer needing to unpack them. +2. Convert inputs to ``TensorWrapper`` objects via ``makeTransformerEngineTensor()``. +3. Release the GIL and call the C API function. GIL release is important to prevent + hangs in case of hardware failures on parallel ranks. +4. Return ``py::object`` outputs (which may be ``QuantizedTensorStorage`` or + ``torch.Tensor`` depending on whether an output quantizer was provided). **Example** — GEMM binding (simplified): .. code-block:: cpp - void te_general_gemm( - at::Tensor A, at::Tensor A_scale_inv, int A_dtype, - at::Tensor B, at::Tensor B_scale_inv, int B_dtype, - at::Tensor D, at::Tensor bias, ...) { + std::vector gemm( + py::handle A, bool transa, + py::handle B, bool transb, + py::object D, + py::handle quantizer, ...) { + + // Convert Python objects to TensorWrapper (handles all tensor types). + // The second argument is the quantizer — py::none() here because inputs + // are already quantized and the quantizer is only needed for outputs. + auto none = py::none(); + TensorWrapper A_tensor = makeTransformerEngineTensor(A, none); + TensorWrapper B_tensor = makeTransformerEngineTensor(B, none); + + // Convert the Python quantizer handle to a C++ Quantizer object, + // then use it to create the output buffer. + std::unique_ptr quantizer_cpp = convert_quantizer(quantizer); + auto [out_tensor, out_py] = quantizer_cpp->create_tensor(shape, dtype); + + // Release GIL and call C API + NVTE_SCOPED_GIL_RELEASE({ + nvte_cublas_gemm_v2(transa, transb, &alpha, + A_tensor.data(), B_tensor.data(), + &beta, out_tensor.data(), out_tensor.data(), + te_workspace.data(), config, main_stream); + }); + + return {out_py, bias_grad, gelu_input, extra_output}; + } + +The ``NVTE_SCOPED_GIL_RELEASE`` macro checks ``PyGILState_Check()`` and releases the GIL +only if it is currently held. This ensures C API calls (which may block on CUDA operations) +do not hold the GIL, preventing deadlocks in multi-threaded or multi-rank scenarios. - // Create NVTETensor handles - auto te_A = makeTransformerEngineTensor( - A.data_ptr(), A.sizes(), A_dtype, A_scale_inv, ...); - auto te_B = makeTransformerEngineTensor( - B.data_ptr(), B.sizes(), B_dtype, B_scale_inv, ...); +TensorWrapper and NVTETensor +---------------------------- - // Call C API - nvte_general_gemm(te_A.data(), te_B.data(), te_D.data(), - te_bias.data(), stream); - } +``NVTETensor`` (defined in ``transformer_engine.h``) is an opaque handle (``typedef void +*NVTETensor``). Unlike typical opaque handles in C libraries, ``NVTETensor`` is not a +direct pointer to an internal structure. Instead, it is an integer handle into a +pre-allocated pool of tensor descriptors (see ``transformer_engine.cpp``). + +``TensorWrapper`` (also in ``transformer_engine.h``) is a lightweight C++ RAII wrapper +around an ``NVTETensor`` that provides builder methods (``set_rowwise_data()``, +``set_columnwise_data()``, ``set_scale()``, etc.) and calls ``nvte_destroy_tensor()`` on +destruction. Neither type owns the underlying data memory on the GPU — the Python objects +retain ownership. + +Tensor Conversion: makeTransformerEngineTensor +---------------------------------------------- + +The key overload (in ``transformer_engine/pytorch/csrc/common.h`` and ``common.cpp``) +accepts ``py::handle`` inputs: + +.. code-block:: cpp + + TensorWrapper makeTransformerEngineTensor(py::handle tensor, py::handle quantizer); + +This function performs runtime type detection to handle all Python tensor types: -Tensor Conversion Helpers -------------------------- +1. It iterates over a registry of known quantized types (defined in + ``transformer_engine/pytorch/csrc/extensions/pybind.h`` as + ``custom_types_converters``). Each entry maps a ``PyTypeObject*`` to an extraction + function. This registry is populated during extension initialization + (``init_extension()``), which is why the extensions must be initialized before any + tensor conversion can occur. +2. When a match is found, it calls the corresponding extraction function (e.g., + ``NVTETensorFromFloat8Tensor``) which reads the internal data pointers and scale + tensors from the Python object and populates a ``TensorWrapper``. +3. If no quantized type matches, it falls back to casting the input as a plain + ``at::Tensor`` and constructing a ``TensorWrapper`` from its data pointer and shape. -The ``makeTransformerEngineTensor()`` family of functions (in -``transformer_engine/pytorch/csrc/common.h``) handles the conversion from PyTorch -tensors to ``NVTETensor``: +This design is what allows the pybind11 layer to accept any tensor type as a +``py::handle`` without requiring the Python side to unpack anything. -- Sets up the ``NVTEBasicTensor`` parameters (data pointer, shape, dtype). -- Attaches scaling metadata (scale, scale_inv, amax) via ``nvte_tensor_set()``. -- Sets the ``NVTEScalingMode`` on the tensor. +Output Buffer Creation +---------------------- + +When an output quantizer is provided, the quantizer's ``create_tensor()`` method +(implemented in C++, e.g., +``transformer_engine/pytorch/csrc/quantizer.cpp``) allocates both the data buffer and +the Python wrapper in one step: + +.. code-block:: cpp + + // Simplified from Float8Quantizer::create_tensor() + std::pair Float8Quantizer::create_tensor( + const std::vector& shape, DType dtype) const { + + // Allocate data buffer + at::Tensor data = at::empty(shape, at::TensorOptions().dtype(torch::kUInt8).device(torch::kCUDA)); + + // Allocate scale inverse + at::Tensor scale_inv = at::empty({1}, at::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA)); + + // Create Python wrapper (Float8TensorStorage or Float8Tensor depending on `internal`) + py::object out_py = /* construct Python object from data + scale_inv */; + + // Create C++ TensorWrapper pointing to the same buffers + TensorWrapper out_cpp(this->get_scaling_mode()); + out_cpp.set_rowwise_data(data.data_ptr(), this->dtype, shape); + out_cpp.set_rowwise_scale_inv(scale_inv.data_ptr(), DType::kFloat32, {1}); + + return {std::move(out_cpp), std::move(out_py)}; + } -The reverse conversion (``NVTETensor`` → PyTorch tensor) is typically not needed because -C API functions write into pre-allocated output buffers. +The returned ``TensorWrapper`` and ``py::object`` point to the same underlying GPU memory. +The C API writes into the ``TensorWrapper``'s buffers, and the ``py::object`` is returned +to Python with the results already in place. Because ``TensorWrapper`` is non-owning, the +``py::object`` must be kept alive for as long as the GPU memory is needed — even if only +the ``TensorWrapper`` is passed to ``nvte_`` functions. Some functions also accept +pre-allocated output buffers from the caller instead of creating new ones via a quantizer. Adding a New Extension ---------------------- @@ -123,5 +204,5 @@ To expose a new C API function to Python: 1. Add the C API function to the appropriate header in ``include/transformer_engine/``. 2. Add a pybind11 wrapper in ``csrc/extensions/``. 3. Register it in the pybind11 module definition. -4. Add a Python wrapper in ``cpp_extensions/``. -5. Update ``__init__.py`` exports as needed. +4. (Optional) Add a Python wrapper in ``cpp_extensions/`` if auxiliary processing is + needed before calling the ``tex`` function. diff --git a/docs/developer/quantization/index.rst b/docs/developer/quantization/index.rst index d72dd90126..1480b59a31 100644 --- a/docs/developer/quantization/index.rst +++ b/docs/developer/quantization/index.rst @@ -6,15 +6,16 @@ Quantization ============ -Transformer Engine supports multiple low-precision formats (FP8 E4M3/E5M2, MXFP8, -block-scaled FP8, NVFP4) through a unified quantization architecture. This section -describes the internal design of that system. +This section describes the internal design of the quantization system. For user-facing +documentation on supported formats, recipes, and usage examples, see +:doc:`/features/low_precision_training/index`. -The quantization system has three main axes: +The developer documentation here focuses on the internal implementation: -- **What** to quantize (the Quantizer/Storage/Tensor class hierarchy) -- **How** to scale (the scaling recipe and global FP8 state) -- **Where** to store (rowwise vs. columnwise layouts for GEMM) +- The Quantizer / QuantizedTensorStorage / QuantizedTensor class hierarchy +- How scaling recipes map to quantizer instances +- The rowwise vs. columnwise layout design for GEMM operands +- How to add a new quantization type .. toctree:: :maxdepth: 1 From f03315f0a8976cc1fa5c89f82b80eeed4935750e Mon Sep 17 00:00:00 2001 From: Przemek Tredak Date: Fri, 20 Mar 2026 17:02:47 -0700 Subject: [PATCH 07/20] And another fixes from review Signed-off-by: Przemek Tredak --- docs/developer/cpp_core/build_system.rst | 115 +++++++--- docs/developer/cpp_core/kernel_areas.rst | 216 +++++------------- docs/developer/cpp_core/scaling_modes.rst | 173 +++++++------- docs/developer/cpp_core/type_system.rst | 44 ++-- .../pytorch_frontend/ops_framework.rst | 129 ++--------- .../pytorch_frontend/transformer_layer.rst | 161 +++++++------ .../quantization/rowwise_columnwise.rst | 11 +- transformer_engine/common/common.h | 4 - .../common/gemm/cublaslt_gemm.cu | 8 +- .../common/gemm/cublaslt_grouped_gemm.cu | 8 +- 10 files changed, 378 insertions(+), 491 deletions(-) diff --git a/docs/developer/cpp_core/build_system.rst b/docs/developer/cpp_core/build_system.rst index 33f60429b2..ac99526f83 100644 --- a/docs/developer/cpp_core/build_system.rst +++ b/docs/developer/cpp_core/build_system.rst @@ -8,24 +8,48 @@ Build System ============ -Transformer Engine uses a hybrid build system: a Python ``setup.py`` orchestrates the -top-level install, while CMake handles the C++/CUDA compilation of the core library. +Transformer Engine has two build paths: a monolithic developer build and a split +pip-wheel build for distribution. -Build Pipeline --------------- +Monolithic Developer Build +-------------------------- -.. code-block:: text +This is the build most developers use during day-to-day work. A single ``pip install`` +invocation builds everything — the C++ core library and the framework-specific extensions +— into one package: + +.. code-block:: bash pip install -e . -v --no-build-isolation + +The build pipeline: + +.. code-block:: text + + pip install -e . │ ▼ setup.py ├── Checks/fetches git submodules (CUTLASS, cuDNN-frontend) ├── Detects CUDA toolkit and GPU architectures - ├── Invokes CMake to build transformer_engine/common/ + ├── Invokes CMake to build libtransformer_engine.so (C++ core) └── Builds framework-specific extensions: - ├── PyTorch: pybind11 extensions via torch.utils.cpp_extension - └── JAX: XLA FFI extensions + ├── PyTorch: pybind11 extension via torch.utils.cpp_extension + └── JAX: XLA FFI extension via pybind11 + +The C++ core is built by CMake (``transformer_engine/common/CMakeLists.txt``) and +produces ``libtransformer_engine.so``. Framework-specific extensions are built by each +framework's own extension system (PyTorch uses ``torch.utils.cpp_extension.CppExtension``, +JAX uses pybind11 directly). Both link against the shared core library. + +Pip Wheel Build +--------------- + +The distribution build (used for PyPI releases) creates multiple separate packages +rather than one monolithic install. This is controlled by the ``NVTE_RELEASE_BUILD`` +environment variable. When set, ``setup.py`` produces a metapackage with no compiled +extensions — the core library and framework extensions are packaged and distributed +separately. Key Environment Variables ------------------------- @@ -37,7 +61,7 @@ Key Environment Variables * - Variable - Description * - ``NVTE_FRAMEWORK`` - - Select framework: ``pytorch``, ``jax``, or auto-detect (default) + - Select framework: ``pytorch``, ``jax``, ``all``, ``none``, or auto-detect (default) * - ``NVTE_CUDA_ARCHS`` - Semicolon-separated GPU architectures (e.g., ``"80;90;100"``) * - ``MAX_JOBS`` @@ -47,7 +71,32 @@ Key Environment Variables * - ``CUDNN_PATH`` - Custom cuDNN path * - ``NVTE_RELEASE_BUILD`` - - Enable release-mode optimizations + - Build as split pip-wheel metapackage (for distribution) + * - ``NVTE_USE_CCACHE`` + - Set to ``1`` to enable ccache for faster incremental builds + * - ``NVTE_CCACHE_BIN`` + - Path to ccache binary (default: ``ccache``). E.g. ``sccache`` for CI. + * - ``NVTE_BUILD_DEBUG`` + - Set to ``1`` to build C++ extensions with debug symbols + +NVTE_CUDA_ARCHS +^^^^^^^^^^^^^^^^ + +This variable controls which GPU architectures the CUDA kernels are compiled for. It +directly affects build time (more architectures = longer builds) and which GPUs the +resulting binary supports. + +If not set, the build system auto-detects based on the installed CUDA toolkit version: + +- CUDA 13.0+: ``"75;80;89;90;100;120"`` +- CUDA 12.8–12.9: ``"70;80;89;90;100;120"`` +- CUDA < 12.8: ``"70;80;89;90"`` + +For development, set this to only your target GPU to minimize build time: + +.. code-block:: bash + + NVTE_CUDA_ARCHS="90" pip install -e . -v --no-build-isolation Git Submodules -------------- @@ -66,19 +115,6 @@ sometimes needed: git submodule update --init --recursive -CMake Structure ---------------- - -The CMake build lives under ``transformer_engine/common/CMakeLists.txt`` and produces -``libtransformer_engine.so``. Key configuration: - -- **Architecture flags**: ``NVTE_CUDA_ARCHS`` maps to CMake's - ``CMAKE_CUDA_ARCHITECTURES``. Only specified architectures are compiled, which - significantly affects build time. -- **Conditional compilation**: Some features are gated on CUDA version (e.g., FP4 support - requires CUDA 12.8+). -- **Header paths**: cuDNN headers are found via ``CUDNN_PATH`` or system paths. - Developer Build Tips -------------------- @@ -88,11 +124,14 @@ Developer Build Tips cmake --build build/cmake --parallel 4 -**Minimal architecture for fast iteration**: +**Enable ccache** for fast rebuilds when iterating on C++ code: .. code-block:: bash - NVTE_CUDA_ARCHS="90" pip install -e . -v --no-build-isolation + NVTE_USE_CCACHE=1 pip install -e . -v --no-build-isolation + + # Or with sccache (common in CI): + NVTE_USE_CCACHE=1 NVTE_CCACHE_BIN=sccache pip install -e . -v --no-build-isolation **Verbose CMake output** (for debugging build issues): @@ -100,10 +139,22 @@ Developer Build Tips cmake --build build/cmake --verbose --parallel 4 -**Common build issues:** - -- ``fatal error: cudnn.h: No such file or directory`` → Set ``CUDNN_PATH`` -- ``nvcc fatal: Unsupported gpu architecture`` → Check ``NVTE_CUDA_ARCHS`` matches - your GPU -- Submodule errors → Run ``git submodule update --init --recursive`` -- Out of memory during compilation → Reduce ``MAX_JOBS`` (try ``MAX_JOBS=1``) +Common Build Failures +--------------------- + +Build failures are common for new developers. Here are the most frequent issues: + +- ``fatal error: cudnn.h: No such file or directory`` → Set ``CUDNN_PATH`` to your + cuDNN installation directory. +- ``nvcc fatal: Unsupported gpu architecture`` → Your ``NVTE_CUDA_ARCHS`` includes an + architecture not supported by your CUDA toolkit version. Either upgrade CUDA or remove + the unsupported architecture. +- Submodule errors (``CMake Error ... CUTLASS``) → Run + ``git submodule update --init --recursive``. +- Out of memory during compilation → Reduce ``MAX_JOBS`` (try ``MAX_JOBS=1``). Each + compilation unit for templated CUDA kernels can consume several GB of memory. +- ``ModuleNotFoundError: No module named 'torch'`` → You need ``--no-build-isolation`` + so that the build can find your installed PyTorch. +- Build succeeds but import fails with undefined symbols → This usually means the core + library and the framework extension were built against different CUDA versions. Clean + the build (``rm -rf build/``) and rebuild from scratch. diff --git a/docs/developer/cpp_core/kernel_areas.rst b/docs/developer/cpp_core/kernel_areas.rst index e8267294a1..d59abd665e 100644 --- a/docs/developer/cpp_core/kernel_areas.rst +++ b/docs/developer/cpp_core/kernel_areas.rst @@ -12,177 +12,87 @@ The C++ core organizes CUDA kernels into functional areas, each in its own subdi under ``transformer_engine/common/``. Every area exposes a C API header in ``include/transformer_engine/`` and implements one or more CUDA kernels. -GEMM (``gemm/``) ------------------ +The major areas include GEMM (``gemm/``), normalization (``normalization/``), +quantization/cast (``cast/``), activation (``activation/``), fused attention +(``fused_attn/``), communication-GEMM overlap (``comm_gemm/``), and several others. +See the directory listing for the full set. -Matrix multiplication via cuBLASLt with FP8/MXFP8/NVFP4 support. +Rather than cataloging each area, this page describes the general architectural patterns +that a developer writing new kernels should understand. -- **Header**: ``include/transformer_engine/gemm.h`` -- **Key file**: ``gemm/cublaslt_gemm.cu`` -- **Entry point**: ``nvte_general_gemm()`` -- Handles: scale/scale-inverse application, bias addition, pre-GeLU fusion, grouped GEMM -- Dispatches to cuBLASLt with appropriate compute types based on input precision and - scaling mode. +Output-Driven Computation +-------------------------- -Normalization (``normalization/``) ----------------------------------- +A key design principle: the **output tensor** dictates what computation the kernel +performs. Kernels inspect the output ``NVTETensor`` to determine: -LayerNorm and RMSNorm with optional FP8 output quantization. +- Whether to produce rowwise data (``output.has_data()``), columnwise data + (``output.has_columnwise_data()``), or both. +- The scaling mode (``output.scaling_mode``), which determines scale granularity. +- The output dtype, which determines the quantization target. -- **Header**: ``include/transformer_engine/normalization.h`` -- **Key files**: ``normalization/layernorm/``, ``normalization/rmsnorm/`` -- Variants: forward, backward, fused with quantization (cast output to FP8 in the same - kernel to save memory bandwidth) -- Architecture dispatch: separate kernel implementations for different GPU architectures - (e.g., Hopper uses warp-specialized kernels). +The input tensor is then validated to confirm it provides the fields needed for the +requested computation (e.g., the input must have data and the correct dtype). This +output-driven design makes sense because, for example, in a quantize kernel the input is +high-precision and knows nothing about rowwise/columnwise layouts — only the output +carries that information. -Activation (``activation/``) ------------------------------ +Scaling Mode Dispatch +---------------------- -Element-wise activation functions with optional FP8 output. +Kernels that handle multiple scaling modes use a switch on the output's +``scaling_mode`` to dispatch to the appropriate implementation. For example, in +``cast/dispatch/quantize.cuh``: -- **Header**: ``include/transformer_engine/activation.h`` -- Supports: GeLU, SiLU/Swish, ReLU, QuickGeLU, SReLU, and their gated variants -- Fused quantization: activation + cast to FP8 in a single kernel pass. - -Cast (``cast/``) ------------------ - -Type-casting and quantization kernels. - -- **Header**: ``include/transformer_engine/cast.h`` -- ``nvte_quantize()`` — cast high-precision data to quantized format with scale computation -- Handles all scaling modes (tensor, block, MXFP8, NVFP4) with mode-specific kernels. - -Transpose (``transpose/``) ---------------------------- - -Transpose operations, often fused with casting. - -- **Header**: ``include/transformer_engine/transpose.h`` -- ``nvte_transpose()`` — standalone transpose -- Fused variants: cast + transpose in a single kernel (critical for producing columnwise - data efficiently during forward pass). - -Fused Attention (``fused_attn/``) ---------------------------------- - -The largest and most complex kernel area. See :doc:`/developer/attention/fused_attn_kernels` -for detailed coverage. - -- **Header**: ``include/transformer_engine/fused_attn.h`` -- Multiple backends: cuDNN-based (F16 and FP8), custom CUDA kernels -- Supports: MHA, GQA, MQA, arbitrary head dims, causal/padding masks, dropout, - sliding window, FP8 quantization. - -Fused RoPE (``fused_rope/``) ------------------------------ - -Rotary positional embedding applied as a fused CUDA kernel. - -- **Header**: ``include/transformer_engine/fused_rope.h`` -- Applies rotary embedding in-place during forward and backward passes. - -Fused Softmax (``fused_softmax/``) ------------------------------------ - -Scaled softmax variants optimized for attention score computation. - -- **Header**: ``include/transformer_engine/softmax.h`` -- Variants: scaled softmax, scaled masked softmax, scaled upper-triangular masked softmax. - -Fused Router (``fused_router/``) ---------------------------------- - -Mixture-of-Experts (MoE) routing kernels. - -- Permutation and unpermutation of tokens across experts. -- Includes its own type-switch macros (note: these need DType casting, see - :doc:`type_system`). - -Communication-GEMM Overlap (``comm_gemm/``) --------------------------------------------- - -Overlapping NCCL collectives with GEMM computation using CUDA multi-stream. - -- See :doc:`/developer/distributed/comm_gemm_overlap` for the design. -- Uses CUDA events and user-allocated buffers (``UserBuffers``) to pipeline communication - and computation. - -Multi-Tensor Operations (``multi_tensor/``) -------------------------------------------- - -Fused operations over lists of tensors (e.g., multi-tensor scale for optimizer steps). - -- **Header**: ``include/transformer_engine/multi_tensor.h`` - -Hadamard Transform (``hadamard_transform/``) --------------------------------------------- - -Hadamard and group Hadamard transforms, with optional fused quantization. - -- **Header**: ``include/transformer_engine/hadamard_transform.h`` -- Variants: standard, group, fused with cast, graph-safe implementations. - -Recipe (``recipe/``) --------------------- - -Scaling recipe kernels that compute quantization scales for each recipe type. - -- **Header**: ``include/transformer_engine/recipe.h`` -- Per-recipe files: ``current_scaling.cu``, ``delayed_scaling.cu``, - ``fp8_block_scaling.cu``, ``mxfp8_scaling.cu``, ``nvfp4.cu``. - -Dropout (``dropout/``) ------------------------ - -Dropout kernel. - -- **Header**: ``include/transformer_engine/dropout.h`` +.. code-block:: text -Swizzle (``swizzle/``) ------------------------ + switch (output_tensor->scaling_mode): + NVTE_DELAYED_TENSOR_SCALING → fp8::quantize() or cast_transpose() + NVTE_MXFP8_1D_SCALING → mxfp8::quantize() + NVTE_BLOCK_SCALING_1D/2D → quantize_transpose_square_blockwise() + NVTE_NVFP4_1D_SCALING → nvfp4::quantize_transpose() -Scale swizzling for GEMM-compatible layouts. +Each scaling mode may have separate code paths for rowwise-only, columnwise-only, or +both layouts, again driven by which fields are populated on the output tensor. -- **Header**: ``include/transformer_engine/swizzle.h`` -- Includes block scaling variant (``swizzle_block_scaling.cu``). +Architecture Dispatch +---------------------- -Permutation (``permutation/``) -------------------------------- +Many kernel areas provide architecture-specific implementations for different GPU +generations (Ampere, Hopper, Blackwell). The dispatch mechanism varies by area: -Token permutation/unpermutation for MoE expert routing. +- **Normalization** uses a ``KernelRegistry`` (in ``normalization/common.h``) that + maps (hidden_size, dtype, warp_config) tuples to kernel function pointers. Kernels are + heavily templated and the appropriate specialization is selected at runtime from the + registry. -- **Header**: ``include/transformer_engine/permutation.h`` +- **Other areas** (cast, GEMM, activation, etc.) use direct dispatch — typically switch + statements or if/else chains on the compute capability. Some kernels only support + certain GPU architectures (e.g., FP4 kernels require Blackwell). -Padding (``util/padding.cu``) ------------------------------- +There is no rigid pattern — the dispatch strategy is chosen per area based on what +makes sense for that area's complexity. The ``NVTE_CUDA_ARCHS`` CMake flag controls +which architectures are compiled. This affects which template instantiations are +generated, which in turn determines build time and binary size. See :doc:`build_system` +for details. -Utility kernel for padding tensors to alignment requirements. +Fused Kernels +-------------- -- **Header**: ``include/transformer_engine/padding.h`` +Many areas support fusing quantization with the primary computation: -Architecture Dispatch ---------------------- +- **Normalization + quantize**: The normalization kernel can produce FP8 output directly, + avoiding a separate cast kernel launch. +- **Activation + quantize**: Similarly, activation functions can cast their output to FP8 + in the same kernel pass. +- **Cast + transpose**: A single kernel produces both rowwise and transposed (columnwise) + output, critical for efficiently preparing data for the backward pass. -Many kernel areas provide architecture-specific implementations. Unlike some CUDA -libraries that use separate source files per architecture (e.g., ``kernel_sm80.cu``, -``kernel_sm90.cu``), TE uses **template specialization and compile-time dispatch** within -shared source files: +These fusions are driven by the same output-driven pattern: if the output tensor's dtype +is an FP8 type and has the appropriate scale fields, the kernel produces quantized output +directly. -.. code-block:: text +For fused attention, see :doc:`/developer/attention/fused_attn_kernels` for dedicated +coverage. - normalization/ - ├── common.h # Shared types and helpers - ├── common.cpp # KernelRegistry, runtime dispatch - ├── layernorm/ - │ ├── ln_fwd_cuda_kernel.cu # Forward kernels (all archs) - │ ├── ln_bwd_semi_cuda_kernel.cu # Backward kernels - │ ├── ln_fwd_kernels.cuh # Templated kernel implementations - │ └── ln_bwd_kernels.cuh - -Kernels are heavily templated (data types, hidden sizes, warp configurations) and the -appropriate specialization is selected at runtime via a ``KernelRegistry``. CMake compile -flags (``NVTE_CUDA_ARCHS``) control which GPU architectures are compiled — this affects -which template instantiations are generated, not which source files are included. See -:doc:`build_system` for details. +For communication-GEMM overlap, see :doc:`/developer/distributed/comm_gemm_overlap`. diff --git a/docs/developer/cpp_core/scaling_modes.rst b/docs/developer/cpp_core/scaling_modes.rst index c521cdd337..3ffa2cc6c3 100644 --- a/docs/developer/cpp_core/scaling_modes.rst +++ b/docs/developer/cpp_core/scaling_modes.rst @@ -13,6 +13,12 @@ different trade-offs between precision, performance, and hardware requirements. scaling mode is selected via the ``NVTEScalingMode`` enum and propagated through the ``Tensor.scaling_mode`` field. +For user-facing documentation on recipes and how to select scaling modes, see +:doc:`/features/low_precision_training/index`. For state management details (amax +history, scale updates, recipe-to-quantizer mapping), see +:doc:`/developer/quantization/scaling_recipes`. This page focuses on what the C++ core +needs to know: data layouts, scale tensor shapes, and GEMM constraints. + .. figure:: ./img/scaling_modes_comparison.svg :align: center :width: 80% @@ -23,136 +29,125 @@ scaling mode is selected via the ``NVTEScalingMode`` enum and propagated through Diagram description for ``scaling_modes_comparison.svg``: A table/matrix with 5 rows (one per scaling mode) and columns: Mode Name | Enum Value | Scale Granularity | Block Size | Data Type | Min Arch. - DELAYED_TENSOR_SCALING | 0 | 1 scale per tensor | N/A | FP8 E4M3/E5M2 | Hopper + DELAYED_TENSOR_SCALING | 0 | 1 scale per tensor | N/A | FP8 E4M3/E5M2 | Ada MXFP8_1D_SCALING | 1 | 1 scale per 32 elements | 32 | FP8 + E8M0 scales | Blackwell - BLOCK_SCALING_1D | 2 | 1 scale per 1×N block | configurable | FP8 E4M3/E5M2 | Blackwell - BLOCK_SCALING_2D | 3 | 1 scale per N×N block | configurable | FP8 E4M3/E5M2 | Blackwell + BLOCK_SCALING_1D | 2 | 1 scale per 1×128 block | 128 | FP8 E4M3/E5M2 | Blackwell + BLOCK_SCALING_2D | 3 | 1 scale per 128×128 block | 128 | FP8 E4M3/E5M2 | Blackwell NVFP4_1D_SCALING | 4 | 1 scale per 16 elements | 16 | FP4 E2M1 + E8M0 | Blackwell -Overview --------- - -.. list-table:: - :header-rows: 1 - :widths: 25 10 25 15 25 - - * - Mode - - Enum - - Scale Granularity - - Block Size - - Data Type - * - Delayed Tensor Scaling - - ``0`` - - 1 scale per tensor - - N/A - - FP8 E4M3 / E5M2 - * - MXFP8 1D Scaling - - ``1`` - - 1 scale per 32 elements - - 32 - - FP8 + E8M0 scales - * - Block Scaling 1D - - ``2`` - - 1 scale per 1×N block - - Configurable - - FP8 E4M3 / E5M2 - * - Block Scaling 2D - - ``3`` - - 1 scale per N×N block - - Configurable - - FP8 E4M3 / E5M2 - * - NVFP4 1D Scaling - - ``4`` - - 1 scale per 16 elements - - 16 - - FP4 E2M1 + E8M0 - -Delayed Tensor Scaling (``NVTE_DELAYED_TENSOR_SCALING``) --------------------------------------------------------- - -The original FP8 scaling mode. A single scale factor applies to the entire tensor, -computed from the *previous iteration's* amax (absolute maximum) value. - -**How it works:** - -1. After each forward/backward pass, the kernel records the amax of the output. -2. The ``FP8GlobalStateManager`` maintains an amax history window (typically 1024 values). -3. Before the next iteration, the scale is computed as: - ``scale = fp8_max / amax_history.max()`` -4. This scale is applied uniformly to all elements during quantization. - -**Trade-offs:** - -- Simple and fast (no per-element scale overhead in GEMM). -- Scale is one iteration stale — can cause overflow/underflow on loss spikes. -- Requires the amax feedback loop between iterations. - -**C++ helpers:** +Non-TN GEMM Support +--------------------- + +A cross-cutting concern that affects all scaling modes is non-TN GEMM support. +cuBLASLt traditionally required FP8 GEMMs to use TN (transpose-normal) layout, which +is why the forward pass needs rowwise data for one operand and columnwise data for +the other. On architectures that support non-TN FP8 GEMM (compute capability 10.x and +13.0+, i.e. Blackwell), this constraint is relaxed — cuBLASLt can consume rowwise data +for both operands. + +This has a significant impact on data layout: when non-TN GEMM is supported, +``data`` and ``columnwise_data`` may point to the same buffer (since a separate +transposed copy is no longer needed for the GEMM). When non-TN is not supported, they +must be separate buffers with physically transposed data. + +Tensor Scaling (``NVTE_DELAYED_TENSOR_SCALING``) +-------------------------------------------------- + +A single scale factor applies to the entire tensor. Minimum architecture: Ada (sm89). + +.. note:: + + The enum name ``NVTE_DELAYED_TENSOR_SCALING`` is a legacy misnomer. This mode is + used for all per-tensor scaling (delayed and current) as well as for high-precision + (unquantized) tensors. The C++ core does not distinguish between delayed and current + scaling — both use the same code paths. The distinction only matters at the framework + level (see :doc:`/developer/quantization/scaling_recipes`). + +Data layout: + +- ``data``: row-major, one scale for the entire tensor +- ``columnwise_data``: column-major (physically transposed), same single scale +- ``scale_inv``: shape ``[1]`` (one element) for FP8 tensors, empty for high-precision +- ``amax``: populated by delayed scaling (framework responsibility), unused by current + scaling + +C++ helpers: .. code-block:: cpp + // True for NVTE_DELAYED_TENSOR_SCALING (per-tensor scaling, including current) bool is_tensor_scaling(const NVTEScalingMode &mode); + // Alias — identical behavior bool is_delayed_tensor_scaling(const NVTEScalingMode &mode); MXFP8 1D Scaling (``NVTE_MXFP8_1D_SCALING``) ---------------------------------------------- Microscaling FP8 format per the OCP MX specification. Each block of 32 contiguous -elements shares a single E8M0 (8-bit exponent-only) scale factor. +elements shares a single E8M0 (8-bit exponent-only) scale factor. Minimum architecture: +Blackwell. -**Key details:** +Data layout: -- Block size is fixed at 32 elements. -- Scales are E8M0 format (power-of-two only, no mantissa bits). -- Scales may need "GEMM swizzling" — a specific memory layout that cuBLASLt expects. - The ``with_gemm_swizzled_scales`` flag on ``Tensor`` tracks this. -- Current scaling (no amax history needed). +- ``data``: row-major FP8, shape ``[M, N]`` +- ``columnwise_data``: the physically transposed data, shape ``[N, M]`` — this is a + separate buffer with independently quantized blocks along the transposed dimension +- ``scale_inv`` (rowwise): shape ``[ceil(M), ceil(N/32)]``, padded to ``[128, 4]`` + multiples for GEMM alignment +- ``columnwise_scale_inv``: shape ``[ceil(N), ceil(M/32)]``, padded to ``[128, 4]`` + multiples -**C++ helpers:** +Scales may need "GEMM swizzling" — a specific memory layout that cuBLASLt expects for +block-scaled operands. See the `cuBLAS documentation on block scaling factors layout +`_. +The ``with_gemm_swizzled_scales`` flag on ``Tensor`` tracks whether scales have been +swizzled. + +C++ helper: .. code-block:: cpp bool is_mxfp8_scaling(const NVTEScalingMode &mode); - bool is_mxfp_scaling(const NVTEScalingMode &mode); // alias Block Scaling (``NVTE_BLOCK_SCALING_1D`` / ``NVTE_BLOCK_SCALING_2D``) ---------------------------------------------------------------------- -Per-block FP8 scaling with configurable block dimensions. ``1D`` uses 1×N blocks -(one scale per row-slice of N elements), while ``2D`` uses N×N blocks. - -**Key details:** +Per-block FP8 scaling with a fixed block size of 128 (matching the DeepSeek v3 recipe). +``1D`` uses 1×128 blocks (one scale per row-slice of 128 elements), while ``2D`` uses +128×128 blocks. The block size itself is not configurable — only the 1D vs 2D aspect is +selectable via the recipe. Minimum architecture: Blackwell. -- Block size is configurable via the quantizer's ``block_scaling_dim`` property. -- Scales are FP32, not E8M0 (more precise than MXFP8). -- Current scaling (computed just-in-time, no history). -- Available on Blackwell and later architectures. +Scales are FP32 (more precise than MXFP8's E8M0). Like MXFP8, scales may require GEMM +swizzling. -**C++ helpers:** +C++ helper: .. code-block:: cpp bool is_block_scaling(const NVTEScalingMode &mode); - // Returns true for all non-tensor-scaling modes + // Returns true for both 1D and 2D block scaling NVFP4 1D Scaling (``NVTE_NVFP4_1D_SCALING``) ---------------------------------------------- -4-bit floating point with E8M0 block scales, targeting maximum compression. +4-bit floating point with E8M0 block scales, targeting maximum compression. Minimum +architecture: Blackwell. + +Data layout: -**Key details:** +- ``data``: FP4 E2M1 format, packed 2 elements per byte (uint8 storage) +- ``columnwise_data``: physically transposed FP4 data +- ``scale_inv`` (rowwise): E8M0 format, shape ``[ceil(M/16), ceil(N/16)]``, padded to + ``[4, 128]`` multiples for GEMM alignment +- ``columnwise_scale_inv``: shape ``[ceil(N/16), ceil(M/16)]``, padded to ``[128, 4]`` -- Data is FP4 E2M1 format (2 exponent bits, 1 mantissa bit). -- Each block of 16 elements shares one E8M0 scale. -- Like MXFP8, may require GEMM-swizzled scales. -- Blackwell and later only. +Like MXFP8, scales may require GEMM swizzling. -**C++ helpers:** +C++ helper: .. code-block:: cpp bool is_nvfp4_scaling(const NVTEScalingMode &mode); - bool is_nvfp_scaling(const NVTEScalingMode &mode); // alias Scale Storage and Layout ------------------------ diff --git a/docs/developer/cpp_core/type_system.rst b/docs/developer/cpp_core/type_system.rst index afa0373424..844483af68 100644 --- a/docs/developer/cpp_core/type_system.rst +++ b/docs/developer/cpp_core/type_system.rst @@ -16,10 +16,11 @@ with opaque ``NVTETensor`` handles rather than C++ objects directly. The C++ ``T and ``SimpleTensor`` structs in ``common.h`` are internal implementation details — they are never exposed across the API boundary. -This design means there are **two dtype enums** (``NVTEDType`` in C, ``DType`` in C++) -and **two tensor abstractions** (``NVTEBasicTensor`` in C, ``SimpleTensor`` in C++). They -mirror each other but are not implicitly convertible. The rest of this page documents both -systems and the conversion patterns between them. +This follows the standard pattern of having a pure C API with corresponding C++ wrappers +that add type safety and convenience. There are two dtype enums (``NVTEDType`` in C, +``DType`` in C++) and two tensor abstractions (``NVTEBasicTensor`` in C, ``SimpleTensor`` +in C++). They have identical numeric values and ``SimpleTensor`` provides implicit +conversion operators, so crossing the boundary is straightforward. C Types (Public API) -------------------- @@ -65,8 +66,14 @@ NVTETensor ^^^^^^^^^^ An opaque handle (``typedef void *NVTETensor``) to the internal ``Tensor`` object. -Created with ``nvte_create_tensor()`` and populated via ``nvte_tensor_set()`` with -``NVTETensorParam`` keys: +Despite the ``void *`` typedef, ``NVTETensor`` is not a direct pointer to a ``Tensor`` +struct. Instead, it is a 1-based integer index into a pre-allocated pool of ``Tensor`` +objects (see ``transformer_engine.cpp``). The pool has a fixed capacity of approximately +26,000 tensors (20 MB / sizeof(Tensor)). Freed tensors are returned to a free-list for +reuse. This means ``nvte_create_tensor()`` does not allocate memory — it retrieves an +entry from the pool. If the pool is exhausted, an error is raised. + +``NVTETensor`` is populated via ``nvte_tensor_set()`` with ``NVTETensorParam`` keys: .. code-block:: c @@ -84,7 +91,11 @@ Created with ``nvte_create_tensor()`` and populated via ``nvte_tensor_set()`` wi NVTEScalingMode ^^^^^^^^^^^^^^^ -Controls the quantization granularity — see :doc:`scaling_modes` for details. +Controls the quantization granularity — see :doc:`scaling_modes` for details. Note that +``NVTE_DELAYED_TENSOR_SCALING`` (value 0) is used not only for delayed scaling but also +for current scaling and for high-precision (unquantized) tensors. The C++ core +distinguishes these cases by checking the data dtype (FP8 vs high-precision) and whether +the amax field is populated. C++ Types (Internal) -------------------- @@ -185,9 +196,13 @@ The main internal tensor type, composed of multiple ``SimpleTensor`` members: bool with_gemm_swizzled_scales; // MXFP8/NVFP4 scale format flag }; -This structure holds all the metadata needed for a quantized tensor: the data itself, -scaling factors for both rowwise and columnwise layouts, and amax values for delayed -scaling. +This structure holds all the metadata needed for a quantized tensor. Not all fields are +valid at the same time — which fields are populated depends on the ``scaling_mode`` and +the data dtype. For example, ``amax`` is only used with delayed tensor scaling and FP8 +dtypes. ``scale_inv`` has shape ``[1]`` for tensor scaling but a multi-element shape for +block scaling modes. ``columnwise_data`` and ``columnwise_scale_inv`` are only populated +when the tensor was quantized with columnwise usage enabled. Validation functions in +``transformer_engine.cpp`` (e.g., ``CheckScaleTensorShape``) enforce these constraints. DType / NVTEDType Conversion Patterns -------------------------------------- @@ -241,6 +256,9 @@ Common Pitfalls won't be found by ADL when called with ``NVTEDType`` arguments (which are in global scope). Use the namespace qualifier: ``transformer_engine::to_string(nvte_type)``. -4. **Custom switch macros**: Macros like ``TRANSFORMER_ENGINE_TYPE_SWITCH_ALL`` cast to - ``DType`` internally, but custom switch macros (e.g., in ``fused_router/utils.h``) - may need the same treatment. +4. **Domain-specific switch macros**: Some kernel areas define their own switch macros + that restrict the set of valid types. For example, ``fused_router/utils.h`` defines + ``TE_ROUTER_PROBS_TYPE_SWITCH_ALL`` (only FP32/FP16/BF16) and + ``TE_ROUTER_INDEX_TYPE_SWITCH_ALL`` (only Int32/Int64/BF16/FP32). These use + ``DType`` directly in their switch cases, so the same ``static_cast`` rules apply + when calling them with ``NVTEDType`` values. diff --git a/docs/developer/pytorch_frontend/ops_framework.rst b/docs/developer/pytorch_frontend/ops_framework.rst index 19ee5e2de3..248355367f 100644 --- a/docs/developer/pytorch_frontend/ops_framework.rst +++ b/docs/developer/pytorch_frontend/ops_framework.rst @@ -181,121 +181,17 @@ This allows an op to **quantize its output eagerly** using the next op's quantiz enabling fused cast kernels. For example, a LayerNorm op receiving the next Linear's input quantizer can fuse the normalization and FP8 cast into a single kernel. -Catalog of Basic Operations ----------------------------- +Available Operations +--------------------- -.. list-table:: - :header-rows: 1 - :widths: 25 40 35 - - * - Operation - - File - - Description - * - ``BasicLinear`` - - ``basic/basic_linear.py`` - - Core GEMM (no bias). Supports FP8, TP, SP. - * - ``Bias`` - - ``basic/bias.py`` - - Additive bias parameter. - * - ``LayerNorm`` - - ``basic/layer_norm.py`` - - Layer normalization with learnable scale/bias. - * - ``RMSNorm`` - - ``basic/rmsnorm.py`` - - Root mean square normalization. - * - ``GELU``, ``SiLU``, ``ReLU``, ``QGELU`` - - ``basic/activation.py`` - - Element-wise activation functions. - * - ``GEGLU``, ``SwiGLU``, ``ReGLU``, ``SReGLU``, ``QGEGLU``, ``GLU``, - ``ClampedSwiGLU``, ``ScaledSwiGLU`` - - ``basic/activation.py``, ``basic/swiglu.py`` - - Gated activation functions (2× input split). - * - ``Quantize`` - - ``basic/quantize.py`` - - Explicit FP8 cast (identity if FP8 disabled). - * - ``AllGather`` - - ``basic/all_gather.py`` - - Distributed all-gather along first dim. - * - ``AllReduce`` - - ``basic/all_reduce.py`` - - Distributed all-reduce. - * - ``ReduceScatter`` - - ``basic/reduce_scatter.py`` - - Distributed reduce-scatter along first dim. - * - ``AddExtraInput`` - - ``basic/add_extra_input.py`` - - Add an extra tensor input (for residual connections). - * - ``MakeExtraOutput`` - - ``basic/make_extra_output.py`` - - Export intermediate tensor as extra output. - * - ``Dropout`` - - ``basic/dropout.py`` - - Random zeroing. - * - ``Reshape`` - - ``basic/reshape.py`` - - Tensor reshape. - * - ``Identity`` - - ``basic/identity.py`` - - Pass-through (useful as fusion anchor). - * - ``ConstantScale`` - - ``basic/constant_scale.py`` - - Multiply by constant factor. - * - ``L2Normalization`` - - ``basic/l2normalization.py`` - - L2 normalization. - * - ``GroupedLinear`` - - ``basic/grouped_linear.py`` - - Multiple parallel GEMMs (MoE). - -Catalog of Fused Operations ----------------------------- - -**Forward fusions** (registered in ``ops/fused/__init__.py``): +Basic operations are located in ``transformer_engine/pytorch/ops/basic/``. Each file +implements a single atomic operation (GEMM, normalization, activation, communication, +etc.). -.. list-table:: - :header-rows: 1 - :widths: 30 35 35 - - * - Fused Operation - - Pattern - - Description - * - ``ForwardLinearBiasActivation`` - - ``BasicLinear`` + ``Bias`` + activation - - Fused GEMM + bias + activation - * - ``ForwardLinearBiasAdd`` - - ``BasicLinear`` + ``Bias`` + ``AddExtraInput`` - - Fused GEMM + bias + residual add - * - ``ForwardLinearScaleAdd`` - - ``BasicLinear`` + ``ConstantScale`` + ``AddExtraInput`` - - Fused GEMM + scale + add - * - ``UserbuffersForwardLinear`` - - ``BasicLinear`` + ``Bias`` + ``ReduceScatter`` - - GEMM overlapped with TP communication - -**Backward fusions**: - -.. list-table:: - :header-rows: 1 - :widths: 30 35 35 - - * - Fused Operation - - Pattern - - Description - * - ``BackwardActivationBias`` - - Backward of activation + bias - - Fused activation and bias gradients - * - ``BackwardAddRMSNorm`` - - Backward of add + RMSNorm - - Fused residual add and norm gradients - * - ``BackwardLinearAdd`` - - ``MakeExtraOutput`` + ``BasicLinear`` - - Fused dgrad GEMM + residual add - * - ``BackwardLinearScale`` - - Backward GEMM + scale - - Fused backward GEMM and scaling - * - ``UserbuffersBackwardLinear`` - - Backward with Userbuffers - - Backward GEMM overlapped with TP communication +Fused operations are located in ``transformer_engine/pytorch/ops/fused/``. Forward and +backward fusions are registered separately in ``ops/fused/__init__.py``. Each fused +operation matches a specific pattern of basic operations (e.g., GEMM + bias + activation) +and replaces them with a single fused kernel call. Extra Inputs and Outputs ------------------------- @@ -376,9 +272,10 @@ functionality: - Production - Experimental -Both approaches call the same C++ extensions and CUDA kernels. The ops framework is -intended to eventually replace the monolithic modules as the primary implementation, -once it reaches feature parity. +Both approaches call the same C++ extensions and CUDA kernels. The two implementations +will eventually be merged, with the ops framework being the more future-proof API. New +developers should be aware that the ops framework is the preferred direction for new +work, even though the monolithic modules are currently more mature. See Also -------- diff --git a/docs/developer/pytorch_frontend/transformer_layer.rst b/docs/developer/pytorch_frontend/transformer_layer.rst index 3dbc580415..e1a2302cd0 100644 --- a/docs/developer/pytorch_frontend/transformer_layer.rst +++ b/docs/developer/pytorch_frontend/transformer_layer.rst @@ -9,76 +9,95 @@ TransformerLayer ================ ``TransformerLayer`` (``transformer_engine/pytorch/transformer.py``) composes TE modules -into a complete Transformer block. Unlike other TE modules, it does **not** inherit from +into a complete Transformer block. Unlike other TE modules, it does not inherit from ``TransformerEngineBaseModule`` — it is a pure composition of TE sub-modules. -.. figure:: ./img/transformer_layer.svg - :align: center - :width: 40% - - Data flow through a TransformerLayer. - -.. - Diagram description for ``transformer_layer.svg``: - Vertical flow diagram: - "Input" → "LayerNorm" → "Self-Attention (MHA)" → "+" (residual add) → - "LayerNorm" → "MLP (LayerNormMLP)" → "+" (residual add) → "Output" - Residual connections shown as arrows bypassing the attention and MLP blocks. - -Architecture ------------- - -A standard ``TransformerLayer`` contains: - -.. code-block:: text - - TransformerLayer - ├── self_attention: MultiheadAttention - │ ├── qkv_projection: LayerNormLinear # Fused LN + QKV projection - │ ├── core_attention: DotProductAttention # Attention computation - │ └── output_projection: Linear # Output projection - ├── layernorm_mlp: LayerNormMLP # Fused LN + FC1 + Act + FC2 - └── (optional) cross_attention: MultiheadAttention - -Configuration -------------- - -Key constructor parameters: - -.. list-table:: - :header-rows: 1 - :widths: 30 70 - - * - Parameter - - Description - * - ``hidden_size`` - - Model hidden dimension - * - ``ffn_hidden_size`` - - Feed-forward network intermediate size - * - ``num_attention_heads`` - - Number of attention heads - * - ``num_gqa_groups`` - - Number of GQA groups (for grouped-query attention) - * - ``layer_type`` - - ``"encoder"`` or ``"decoder"`` (controls cross-attention) - * - ``self_attn_mask_type`` - - ``"causal"``, ``"padding"``, ``"no_mask"``, etc. - * - ``normalization`` - - ``"LayerNorm"`` or ``"RMSNorm"`` - * - ``activation`` - - ``"gelu"``, ``"swiglu"``, ``"geglu"``, etc. - -Distributed Training --------------------- - -When ``set_tensor_parallel_group()`` is called, the sub-modules automatically configure: - -- QKV projection as **column-parallel** (split heads across ranks). -- Attention output projection as **row-parallel** (each rank has partial output, reduced - via all-reduce or reduce-scatter). -- MLP FC1 as **column-parallel**, FC2 as **row-parallel**. - -With sequence parallelism enabled, the LayerNorm and residual additions operate on -sequence-partitioned data. - -See :doc:`/developer/distributed/tensor_parallel` for details on the parallelism patterns. +For the full constructor API and parameter descriptions, see the class docstring. This +page covers implementation details that are non-obvious or differ from a standard +Transformer implementation. + +QKV Weight Handling +------------------- + +The ``fuse_qkv_params`` parameter (default ``True``) controls whether Q, K, and V +projections share a single fused weight tensor or use separate weights. When +``fuse_qkv_params=True``, the QKV projection is a single ``LayerNormLinear`` with output +size ``3 * hidden_size`` (or adjusted for GQA), and the weight is split into Q, K, V +slices internally. + +When ``fuse_qkv_params=False``: + +- The ``qkv_weight_interleaved`` parameter is silently overridden to ``False``, + regardless of what the user passed. This is because interleaved layout only makes + sense with a fused weight tensor. +- ``fuse_wgrad_accumulation`` is not supported and raises an error. +- Despite using separate Q, K, V weight tensors, the code still performs a single GEMM. + This is achieved through the ``parameters_split`` feature of the ``Linear`` / + ``LayerNormLinear`` module: the three weight matrices are registered as separate + parameters but are concatenated (via a no-op cat that avoids actual copying when the + tensors are already contiguous in memory) into a single weight for the GEMM. The + output is then split back into Q, K, V slices. This means the per-parameter API is + preserved for checkpoint compatibility while the compute path remains a single fused + GEMM. + +The ``qkv_weight_interleaved`` parameter controls the memory layout of the fused QKV +weight. When ``True``, the Q, K, V weights for each attention head are interleaved +(``[Q_head0, K_head0, V_head0, Q_head1, ...]``). When ``False``, they are contiguous +blocks (``[Q_all_heads, K_all_heads, V_all_heads]``). The interleaved layout can be +more efficient for certain GEMM configurations. + +Output Structure Varies by Configuration +----------------------------------------- + +The return values from ``self_attention()`` and ``layernorm_mlp()`` change structure +depending on configuration flags: + +- With ``apply_residual_connection_post_layernorm=False`` (default): returns + ``(output, bias)`` +- With ``apply_residual_connection_post_layernorm=True``: returns + ``(output, bias, residual)`` +- With ``parallel_attention_mlp=True``: ``return_bias`` is forced to ``False``, so the + attention and MLP outputs are plain tensors that get combined in a single + bias-dropout-add call. + +The unpacking logic in ``forward()`` relies on the control flow to determine the tuple +shape — there is no explicit type checking. + +Bias-Dropout-Add Executor Context +----------------------------------- + +The bias-dropout-add fusion uses a context manager that varies by PyTorch version. +The minimum supported PyTorch version is 2.1, so in practice the handler is +``torch.enable_grad`` for PyTorch >= 2.2 (the current code path). The +``torch.enable_grad()`` context ensures that gradients are properly tracked through the +fused bias + dropout + residual addition operation, even when called from within a +``torch.no_grad()`` region during certain compilation or warmup phases. + +.. note:: + + The code still contains a legacy branch for PyTorch < 2.2 that uses ``nullcontext`` + (relying on NVFuser). Since TE requires PyTorch >= 2.1, this branch is only reachable + on PyTorch 2.1.x. The NVFuser path may be removed in a future cleanup. + +Activation Checkpointing +------------------------- + +``TransformerLayer`` itself contains no checkpointing logic. The +``checkpoint_core_attention`` parameter is forwarded directly to +``DotProductAttention``, which handles the actual recomputation. For full-layer +checkpointing, use the TE-provided ``checkpoint()`` function from +``transformer_engine.pytorch.distributed`` (see +:doc:`autograd_integration`). + +Parallel Attention and MLP +--------------------------- + +When ``parallel_attention_mlp=True``, the attention and MLP blocks run on the same input +(rather than sequentially) and their outputs are summed. This changes the data flow: + +- Both blocks receive the same ``hidden_states`` input. +- The ``return_bias`` flag is forced to ``False`` for both sub-modules. +- A single ``_bias_dropout_add`` call combines both outputs with the residual. + +This mode corresponds to the parallel formulation used in some model architectures +(e.g., PaLM). diff --git a/docs/developer/quantization/rowwise_columnwise.rst b/docs/developer/quantization/rowwise_columnwise.rst index 5b2e7a781a..3d13faa924 100644 --- a/docs/developer/quantization/rowwise_columnwise.rst +++ b/docs/developer/quantization/rowwise_columnwise.rst @@ -88,11 +88,12 @@ In the C++ ``Tensor`` struct, the layout is reflected in separate fields: // ... }; -For **tensor scaling** (delayed/current), ``data`` and ``columnwise_data`` point to -different memory buffers (the data is physically transposed). - -For **block scaling** modes, the data may be the same buffer but with different scale -tensors (row-oriented vs. column-oriented block scales). +Whether ``data`` and ``columnwise_data`` point to the same or different memory buffers +depends on non-TN GEMM support (see :doc:`/developer/cpp_core/scaling_modes`). On +architectures that support non-TN FP8 GEMM (Blackwell and later), cuBLASLt can consume +rowwise data directly for both operands, so ``data`` and ``columnwise_data`` may share +the same buffer. On architectures without non-TN support, the data must be physically +transposed into a separate buffer for the columnwise layout. Performance Implications ------------------------ diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index 41a8fd1112..91d73841c3 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -60,10 +60,6 @@ inline bool is_nvfp4_scaling(const NVTEScalingMode &mode) { return mode == NVTE_ inline bool is_mxfp8_scaling(const NVTEScalingMode &mode) { return mode == NVTE_MXFP8_1D_SCALING; } -inline bool is_mxfp_scaling(const NVTEScalingMode &mode) { return mode == NVTE_MXFP8_1D_SCALING; } - -inline bool is_nvfp_scaling(const NVTEScalingMode &mode) { return mode == NVTE_NVFP4_1D_SCALING; } - inline size_t product(const std::vector &shape, const size_t begin, const size_t end) { NVTE_CHECK(begin <= end && end <= shape.size(), "Attempted to access entries ", begin, " to ", end, " in a vector with ", shape.size(), " entries"); diff --git a/transformer_engine/common/gemm/cublaslt_gemm.cu b/transformer_engine/common/gemm/cublaslt_gemm.cu index 144aea1a07..4fefadebab 100644 --- a/transformer_engine/common/gemm/cublaslt_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_gemm.cu @@ -118,8 +118,8 @@ GemmParam CanonicalizeGemmInput(const transformer_engine::Tensor &A, const cubla bool is_B_transposed = transB == CUBLAS_OP_T; // Set conditions for MXFP8 and NVFP4 gemm execution. - const auto nvfp4 = is_nvfp_scaling(A.scaling_mode) && is_nvfp_scaling(B.scaling_mode); - const auto mxfp8 = !nvfp4 && is_mxfp_scaling(A.scaling_mode) && is_mxfp_scaling(B.scaling_mode); + const auto nvfp4 = is_nvfp4_scaling(A.scaling_mode) && is_nvfp4_scaling(B.scaling_mode); + const auto mxfp8 = !nvfp4 && is_mxfp8_scaling(A.scaling_mode) && is_mxfp8_scaling(B.scaling_mode); int is_nvte_non_tn_fp8_gemm_supported = 0; // needed only for per tensor scaling if (is_tensor_scaling(A.scaling_mode) || is_tensor_scaling(B.scaling_mode)) { is_nvte_non_tn_fp8_gemm_supported = nvte_is_non_tn_fp8_gemm_supported(); @@ -820,7 +820,7 @@ void nvte_cublas_gemm(const NVTETensor A, const NVTETensor B, NVTETensor D, cons // Check for NVFP4 // TODO Remove once alpha scale logic is moved into cublas_gemm function - if (is_nvfp_scaling(inputA->scaling_mode) || is_nvfp_scaling(inputB->scaling_mode)) { + if (is_nvfp4_scaling(inputA->scaling_mode) || is_nvfp4_scaling(inputB->scaling_mode)) { NVTE_ERROR("nvte_cublas_gemm does not support NVFP4 data. Use nvte_cublas_gemm_v2 instead."); } @@ -908,7 +908,7 @@ void nvte_cublas_gemm_scaled(const NVTETensor A, const NVTETensor B, NVTETensor // Check for NVFP4 // TODO Remove once alpha scale logic is moved into cublas_gemm function - if (is_nvfp_scaling(inputA->scaling_mode) || is_nvfp_scaling(inputB->scaling_mode)) { + if (is_nvfp4_scaling(inputA->scaling_mode) || is_nvfp4_scaling(inputB->scaling_mode)) { NVTE_ERROR("nvte_cublas_gemm does not support NVFP4 data. Use nvte_cublas_gemm_v2 instead."); } diff --git a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu index 7069debc56..2e3dd1e2d2 100644 --- a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu @@ -244,12 +244,12 @@ inline GroupedGemmInputProperties validate_grouped_gemm_inputs( "Grouped GEMM inputs must be FP8 or BF16."); NVTE_CHECK(is_fp8_dtype(inputA->dtype()) == is_fp8_dtype(inputB->dtype()), "Grouped GEMM: A and B must both be FP8 or both be non-FP8."); - NVTE_CHECK(transformer_engine::is_mxfp_scaling(inputA->scaling_mode) == - transformer_engine::is_mxfp_scaling(inputB->scaling_mode), + NVTE_CHECK(transformer_engine::is_mxfp8_scaling(inputA->scaling_mode) == + transformer_engine::is_mxfp8_scaling(inputB->scaling_mode), "Grouped GEMM: A and B must both use MXFP8 scaling or both use tensor scaling, " "mixed configurations are not supported."); const bool is_fp8 = is_fp8_dtype(inputA->dtype()); - const bool is_mxfp8 = transformer_engine::is_mxfp_scaling(inputA->scaling_mode); + const bool is_mxfp8 = transformer_engine::is_mxfp8_scaling(inputA->scaling_mode); if (is_mxfp8) { NVTE_CHECK(inputA->with_gemm_swizzled_scales, "MXFP8 grouped GEMM: A scales must be swizzled for GEMM"); @@ -318,7 +318,7 @@ inline GroupedOperandSelection select_grouped_operand(const transformer_engine:: "Grouped GEMM operand is missing both row-wise and column-wise data"); const auto sm = t->scaling_mode; - const bool mxfp8 = is_mxfp_scaling(sm); + const bool mxfp8 = is_mxfp8_scaling(sm); // Validate scaling mode NVTE_CHECK(sm == NVTE_DELAYED_TENSOR_SCALING || mxfp8, From a1930a538f5057cb9735a0cba1d904473dcae450 Mon Sep 17 00:00:00 2001 From: Przemek Tredak Date: Tue, 24 Mar 2026 16:49:54 -0700 Subject: [PATCH 08/20] More fixes from review and removing dead code Signed-off-by: Przemek Tredak --- .../developer/attention/backend_selection.rst | 36 ++-- docs/developer/attention/backends.rst | 157 +++++++++--------- docs/developer/cpp_core/scaling_modes.rst | 3 +- docs/developer/cpp_core/type_system.rst | 77 +++++++-- docs/developer/index.rst | 6 +- docs/developer/linear_walkthrough.rst | 8 +- .../pytorch_frontend/module_hierarchy.rst | 20 ++- docs/developer/testing.rst | 110 ++++++++++++ docs/envvars.rst | 6 - tests/pytorch/attention/test_attention.py | 14 -- .../dot_product_attention.py | 26 +-- 11 files changed, 289 insertions(+), 174 deletions(-) create mode 100644 docs/developer/testing.rst diff --git a/docs/developer/attention/backend_selection.rst b/docs/developer/attention/backend_selection.rst index 006d3c48fd..4bae7d9b49 100644 --- a/docs/developer/attention/backend_selection.rst +++ b/docs/developer/attention/backend_selection.rst @@ -16,19 +16,22 @@ Selection Flow -------------- The backend is selected at the start of each ``DotProductAttention.forward()`` call. -The logic (in ``transformer_engine/pytorch/attention/dot_product_attention/``) follows -a priority order: +The selection result is **cached** — if the attention parameters haven't changed since the +last call, the cached backend is reused without re-running the selection logic (see +``_attention_backends`` in ``dot_product_attention.py``). + +The selection logic lives in ``get_attention_backend()`` in +``transformer_engine/pytorch/attention/dot_product_attention/utils.py``. It filters +backends by compatibility (head dimension, sequence length, mask type, dropout, GQA +groups, etc.) and then applies a priority order: .. code-block:: text - 1. Check user override (env vars, constructor args) - 2. Check FP8 — if FP8 enabled, try cuDNN FP8 fused - 3. Check cuDNN F16 fused — if supported by config - 4. Check FlashAttention — if installed and supported - 5. Fall back to unfused attention + On Hopper+ (sm90): FusedAttention > FlashAttention > Unfused + On pre-Hopper: FlashAttention > FusedAttention > Unfused -At each step, the logic checks whether the backend supports the current configuration: -head dimension, sequence length, mask type, dropout, GQA groups, etc. +FusedAttention is preferred on Hopper+ for performance reasons. On pre-Hopper hardware, +FlashAttention takes priority when both are available. Environment Variables --------------------- @@ -44,9 +47,9 @@ Environment Variables * - ``NVTE_FLASH_ATTN`` - ``0`` to disable FlashAttention, ``1`` to enable (default: ``1``) * - ``NVTE_FUSED_ATTN_BACKEND`` - - Force a specific cuDNN fused attention sub-backend (integer ID) - * - ``NVTE_FUSED_ATTN_FORCE_WORKSPACE_OPT`` - - Force workspace optimization for fused attention + - Force a specific cuDNN fused attention sub-backend by integer ID: + ``0`` = F16 max512, ``1`` = F16 arbitrary seqlen, ``2`` = FP8. + These correspond to the ``NVTE_Fused_Attn_Backend`` enum in ``fused_attn.h``. Constructor Override -------------------- @@ -87,7 +90,7 @@ Not all backends support all features. Key restrictions: - Yes - Yes * - Sliding window - - No + - Yes (via arbitrary mask) - Yes - Yes - Yes @@ -124,9 +127,12 @@ To see which backend was selected, set: .. code-block:: bash - NVTE_DEBUG=1 + NVTE_DEBUG=1 NVTE_DEBUG_LEVEL=1 -This logs the backend selection decision and the reason for any fallbacks. +``NVTE_DEBUG_LEVEL`` controls verbosity: ``1`` logs the selected backend (INFO level), +``2`` logs the full selection process including each filter step and why backends were +disabled (DEBUG level). Output goes through Python's ``logging`` module under the +``DotProductAttention`` logger name. See Also -------- diff --git a/docs/developer/attention/backends.rst b/docs/developer/attention/backends.rst index ffa411b9ac..94fe6668d1 100644 --- a/docs/developer/attention/backends.rst +++ b/docs/developer/attention/backends.rst @@ -21,60 +21,18 @@ hardware, precision, and feature combinations. Diagram description for ``attention_backends.svg``: Tree/taxonomy diagram: "Attention Backends" - ├── "Unfused" (FlashAttention via torch) - │ └── Any GPU, no FP8, fallback - ├── "Flash Attention" (Tri Dao's FlashAttention) - │ └── Ampere+, BF16/FP16, fastest for non-FP8 - └── "Fused Attention" (cuDNN-based) - ├── "F16 Fused" — cuDNN with BF16/FP16 - │ └── Ampere+, supports arbitrary seq len - ├── "FP8 Fused" — cuDNN with FP8 - │ └── Hopper+, best perf with FP8 training - └── "Custom" — TE's own CUDA kernels - └── Legacy, limited feature support - -Overview --------- - -.. list-table:: - :header-rows: 1 - :widths: 20 15 15 15 35 - - * - Backend - - Precision - - Min Arch - - cuDNN Required - - Notes - * - Unfused - - BF16/FP16 - - Any - - No - - Fallback; uses PyTorch's native attention or manual QK^TV - * - FlashAttention - - BF16/FP16 - - Ampere - - No - - Tri Dao's FlashAttention; fast, memory-efficient - * - Fused F16 (max512) - - BF16/FP16 - - Ampere - - Yes - - cuDNN graph-based; sequence length ≤ 512 - * - Fused F16 (arbitrary) - - BF16/FP16 - - Ampere - - Yes - - cuDNN graph-based; arbitrary sequence length, supports all mask types - * - Fused (FP8) - - FP8 - - Hopper - - Yes - - FP8 Q/K/V with cuDNN; sequence length ≤ 512; best throughput for FP8 training - * - Custom/THD - - BF16/FP16 - - Hopper - - No - - TE custom kernels; used for specific context-parallel patterns + ├── "Unfused" — PyTorch native ops (manual QK^TV) + │ └── Any GPU, BF16/FP16, no cuDNN required, fallback + ├── "FlashAttention" — Tri Dao's flash-attn package (external) + │ └── Ampere+, BF16/FP16, no cuDNN required + └── "FusedAttention" — cuDNN graph API + ├── Sub-backend 0: "F16 max512" — seqlen ≤ 512, Ampere+ + ├── Sub-backend 1: "F16 arbitrary" — any seqlen, Ampere+ + └── Sub-backend 2: "FP8" — FP8 Q/K/V, Hopper+ + +There are exactly three backends (Unfused, FlashAttention, FusedAttention). +FusedAttention has three cuDNN sub-backends selected automatically based on precision +and sequence length. Backend Details --------------- @@ -82,50 +40,85 @@ Backend Details Unfused Attention ^^^^^^^^^^^^^^^^^ -The simplest backend. Computes attention as separate matrix multiplications: - -.. code-block:: text - - scores = Q @ K^T / sqrt(d) - scores = mask(scores) - weights = softmax(scores) - output = weights @ V +**Class**: ``UnfusedDotProductAttention`` in +``transformer_engine/pytorch/attention/dot_product_attention/backends.py`` +Computes attention as separate PyTorch operations (``torch.baddbmm`` for QK^T, +``FusedScaleMaskSoftmax`` for masking + softmax, ``torch.bmm`` for the V multiply). Used as a fallback when no fused backend supports the requested configuration. Also -useful for debugging (produces identical numerics to textbook attention). +useful for debugging since it produces identical numerics to textbook attention. + +Sliding window is supported by converting the window specification into an ``arbitrary`` +attention mask (see ``get_full_mask()`` in ``utils.py``). FlashAttention ^^^^^^^^^^^^^^ -Uses Tri Dao's FlashAttention algorithm, which tiles the computation to avoid -materializing the full attention matrix. Key properties: +**Class**: ``FlashAttention`` in +``transformer_engine/pytorch/attention/dot_product_attention/backends.py`` + +Integrates Tri Dao's `flash-attn `_ +package as an external dependency (not bundled). The integration imports from +``flash_attn.flash_attn_interface`` (FA2) and ``flash_attn_3.flash_attn_interface`` +(FA3). Version constraints are managed in ``FlashAttentionUtils`` in ``utils.py``. + +Code flow: -- O(N) memory instead of O(N²) for the attention matrix. -- IO-aware: optimized for GPU memory hierarchy. -- Supports causal masks, variable-length sequences. -- No FP8 support (BF16/FP16 only). +1. ``FlashAttention.forward()`` converts TE's tensor layouts (sbhd/thd) to FlashAttention's + expected bshd format. +2. Selects between ``flash_attn_func`` (dense), ``flash_attn_varlen_func`` + (variable-length), or ``flash_attn_with_kvcache`` (inference) based on the + configuration. +3. For context parallelism, wraps the call via ``attn_forward_func_with_cp()``. + +To modify FlashAttention integration: start with ``FlashAttention.forward()`` in +``backends.py`` for the Python wrapper, and ``FlashAttentionUtils`` in ``utils.py`` for +version and feature compatibility checks. cuDNN Fused Attention (F16 and FP8) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -Uses cuDNN's graph API to fuse the entire attention computation (QK^T, scaling, masking, -softmax, dropout, V multiply) into a single cuDNN operation. Key properties: +**Class**: ``FusedAttention`` in +``transformer_engine/pytorch/attention/dot_product_attention/backends.py`` + +Uses cuDNN's **cudnn-frontend graph API** to fuse the entire attention computation +(QK^T, scaling, masking, softmax, dropout, V multiply) into a single cuDNN operation. + +Code flow: + +1. **Python** (``backends.py``): ``FusedAttention.forward()`` → ``FusedAttnFunc.apply()`` + (custom autograd function). +2. **pybind11** (``pytorch/csrc/extensions/attention.cpp``): ``fused_attn_fwd()`` converts + Python tensors to ``TensorWrapper`` objects and calls the C++ layer. +3. **C++ dispatcher** (``common/fused_attn/fused_attn.cpp``): Routes to the appropriate + sub-backend implementation based on ``nvte_get_fused_attn_backend()``. +4. **cuDNN graph builders** — one file per sub-backend: + + - ``fused_attn_f16_max512_seqlen.cu`` — sub-backend 0 + - ``fused_attn_f16_arbitrary_seqlen.cu`` — sub-backend 1 + - ``fused_attn_fp8.cu`` — sub-backend 2 + + Each builds a ``cudnn_frontend::graph::Graph``, configures SDPA attributes (masking, + dropout, scaling), and executes it. **Graphs are cached** per thread in a thread-local + cache keyed by ``FADescriptor`` to avoid rebuilding identical configurations. -- **F16 variant**: BF16/FP16 inputs, Ampere and later. -- **FP8 variant**: FP8 Q/K/V inputs with per-tensor or block scales, Hopper and later. -- Supports: arbitrary masks (causal, padding, sliding window), GQA/MQA, dropout. -- Requires cuDNN 9.3+ for FP8. +To modify cuDNN fused attention: start with the relevant ``.cu`` file for the sub-backend, +specifically the ``fused_attn_*_fwd_impl()`` / ``fused_attn_*_bwd_impl()`` functions +where the cuDNN graph is built. For backend selection logic, see +``nvte_get_fused_attn_backend()`` in ``fused_attn.cpp``. -The cuDNN backend is the recommended choice for production FP8 training. +Helper CUDA Kernels +^^^^^^^^^^^^^^^^^^^ -Custom/THD Kernels -^^^^^^^^^^^^^^^^^^ +The ``transformer_engine/common/fused_attn/`` directory also contains CUDA helper kernels +that support the attention backends but do **not** perform the full attention operation +themselves: -TE's own CUDA attention kernels, used for specific context-parallel patterns -(THD = Token-Head-Dimension layout). These handle cases where cuDNN doesn't support the -required distributed communication pattern. +- ``utils.cu`` — stride calculation, auxiliary tensor setup +- ``context_parallel.cu`` — KV communication for context parallelism +- ``kv_cache.cu`` — KV cache management for inference -See :doc:`fused_attn_kernels` for the C++ kernel organization. +See :doc:`fused_attn_kernels` for the full C++ kernel organization. See Also -------- diff --git a/docs/developer/cpp_core/scaling_modes.rst b/docs/developer/cpp_core/scaling_modes.rst index 3ffa2cc6c3..6270cea28f 100644 --- a/docs/developer/cpp_core/scaling_modes.rst +++ b/docs/developer/cpp_core/scaling_modes.rst @@ -125,7 +125,8 @@ C++ helper: .. code-block:: cpp bool is_block_scaling(const NVTEScalingMode &mode); - // Returns true for both 1D and 2D block scaling + // Returns true for all non-tensor-scaling modes (MXFP8, block 1D/2D, NVFP4). + // For finer discrimination, use is_mxfp8_scaling() or is_nvfp4_scaling(). NVFP4 1D Scaling (``NVTE_NVFP4_1D_SCALING``) ---------------------------------------------- diff --git a/docs/developer/cpp_core/type_system.rst b/docs/developer/cpp_core/type_system.rst index 844483af68..660373d877 100644 --- a/docs/developer/cpp_core/type_system.rst +++ b/docs/developer/cpp_core/type_system.rst @@ -8,19 +8,18 @@ Type System =========== -Transformer Engine maintains two parallel type systems: a C type system for the public -API and a C++ type system for internal use. +Transformer Engine maintains three layers of types: a C API for ABI stability, C++ +wrappers for external consumers, and C++ internal types for the core implementation. The C++ core exposes a **C API** (not C++) for ABI stability. Framework bindings interact -with opaque ``NVTETensor`` handles rather than C++ objects directly. The C++ ``Tensor`` -and ``SimpleTensor`` structs in ``common.h`` are internal implementation details — -they are never exposed across the API boundary. +with opaque ``NVTETensor`` handles rather than C++ objects directly. The same header +(``transformer_engine.h``) also provides C++ wrapper types (``DType``, +``TensorWrapper``) that add type safety and convenience around the C API. These wrappers +are intended for use by framework bindings (pybind11 code, tests) — they wrap the C API +but do not expose internal implementation details. -This follows the standard pattern of having a pure C API with corresponding C++ wrappers -that add type safety and convenience. There are two dtype enums (``NVTEDType`` in C, -``DType`` in C++) and two tensor abstractions (``NVTEBasicTensor`` in C, ``SimpleTensor`` -in C++). They have identical numeric values and ``SimpleTensor`` provides implicit -conversion operators, so crossing the boundary is straightforward. +Separately, ``common.h`` defines internal C++ types (``SimpleTensor``, ``Tensor``) used +only within the core library. These are never exposed across the API boundary. C Types (Public API) -------------------- @@ -97,15 +96,17 @@ for current scaling and for high-precision (unquantized) tensors. The C++ core distinguishes these cases by checking the data dtype (FP8 vs high-precision) and whether the amax field is populated. -C++ Types (Internal) --------------------- +C++ Wrappers (External) +----------------------- -Defined in ``transformer_engine/common/common.h``. +Defined in ``transformer_engine/common/include/transformer_engine/transformer_engine.h``, +within the ``transformer_engine`` C++ namespace. These wrap the C API with type safety and +RAII semantics, and are intended for use by framework bindings (pybind11 code) and tests. DType ^^^^^ -A C++ ``enum class`` in the ``transformer_engine`` namespace, mirroring ``NVTEDType``: +A C++ ``enum class`` mirroring ``NVTEDType``: .. code-block:: cpp @@ -121,13 +122,53 @@ A C++ ``enum class`` in the ``transformer_engine`` namespace, mirroring ``NVTEDT .. warning:: ``DType`` and ``NVTEDType`` have the same numeric values but are **not** implicitly - convertible. Use ``static_cast`` for conversions. + convertible to each other (no implicit construction or assignment). Use ``static_cast`` + for conversions. However, comparison operators (``==``, ``!=``) are overloaded to work + across the two types — see :ref:`dtype-conversion-patterns` below. + +The header also provides inline helpers that accept ``DType``: ``is_fp8_dtype()``, +``is_fp4_dtype()``, ``is_high_precision_dtype()``. + +TensorWrapper +^^^^^^^^^^^^^ + +A C++ RAII wrapper around the opaque ``NVTETensor`` handle. It manages the lifecycle of +the underlying C tensor (``nvte_create_tensor`` / ``nvte_destroy_tensor``) and provides +constructors that populate the tensor's fields via the C API: + +.. code-block:: cpp + + class TensorWrapper { + public: + // Construct with data, shape, dtype, and optional scale/amax pointers + TensorWrapper(void *dptr, const NVTEShape &shape, const DType dtype, + float *amax_dptr = nullptr, float *scale_dptr = nullptr, + float *scale_inv_dptr = nullptr, + NVTEShape scale_inv_shape = defaultShape, + const NVTEScalingMode scaling_mode = NVTE_DELAYED_TENSOR_SCALING); + + // Construct empty tensor + explicit TensorWrapper(const NVTEScalingMode scaling_mode = NVTE_DELAYED_TENSOR_SCALING); + + ~TensorWrapper(); // Calls nvte_destroy_tensor + + NVTETensor data(); // Returns the underlying opaque handle + }; + +``TensorWrapper`` is the primary way framework bindings create ``NVTETensor`` handles to +pass into C API functions. It does not own the data memory — it only owns the pool entry. + +C++ Types (Internal) +-------------------- + +Defined in ``transformer_engine/common/common.h``. These types are used exclusively within +the core library and are never exposed across the API boundary. SimpleTensor ^^^^^^^^^^^^ -A C++ wrapper around ``NVTEBasicTensor`` that provides constructors, implicit conversions, -and utility methods: +A lightweight internal tensor descriptor with implicit conversion to/from +``NVTEBasicTensor``: .. code-block:: cpp @@ -204,6 +245,8 @@ block scaling modes. ``columnwise_data`` and ``columnwise_scale_inv`` are only p when the tensor was quantized with columnwise usage enabled. Validation functions in ``transformer_engine.cpp`` (e.g., ``CheckScaleTensorShape``) enforce these constraints. +.. _dtype-conversion-patterns: + DType / NVTEDType Conversion Patterns -------------------------------------- diff --git a/docs/developer/index.rst b/docs/developer/index.rst index ffc1133cf1..7ae3276b17 100644 --- a/docs/developer/index.rst +++ b/docs/developer/index.rst @@ -20,8 +20,9 @@ How to Use This Guide - **New contributors**: Start with :doc:`architecture_overview` for the big picture, then read :doc:`linear_walkthrough` for an end-to-end trace through a concrete PyTorch - module. For building from source, see :doc:`/installation`. For running tests and - contributing guidelines, see ``CONTRIBUTING.rst`` at the repository root. + module. For building from source, see :doc:`/installation`. For running tests, see + :doc:`testing`. For contributing guidelines, see ``CONTRIBUTING.rst`` at the repository + root. - **Working on new quantization recipe**: See :doc:`quantization/index` for the Quantizer/Storage/Tensor design and :doc:`cpp_core/type_system` for the underlying C/C++ types. - **Working on attention**: See :doc:`attention/index` for backend selection and kernel @@ -37,6 +38,7 @@ How to Use This Guide architecture_overview linear_walkthrough + testing cpp_core/index quantization/index pytorch_frontend/index diff --git a/docs/developer/linear_walkthrough.rst b/docs/developer/linear_walkthrough.rst index 04dbc5a8f8..67c5078e45 100644 --- a/docs/developer/linear_walkthrough.rst +++ b/docs/developer/linear_walkthrough.rst @@ -32,7 +32,7 @@ quantization, GEMM, distributed communication, and autograd. 2. "grad_output_quantizer(grad)" → 3. "Dgrad GEMM: general_gemm(grad, weight)" producing "grad_input" → 4. "Wgrad GEMM: general_gemm(input_columnwise, grad_columnwise)" producing "grad_weight" → - 5. "Reduce-scatter (if row-parallel + SP) or All-reduce (if row-parallel)grad_input" + 5. "Reduce-scatter (if row-parallel + SP) or All-reduce (if row-parallel) grad_input" Dotted arrows from forward boxes 3→backward box 4 labeled "saved columnwise input" and forward weight→backward box 3 labeled "saved weight". @@ -113,7 +113,8 @@ that actually perform casts. - ``output_quantizer`` — for the forward output (optional, when ``fp8_output=True``) - ``grad_output_quantizer`` — for the backward gradient - ``grad_input_quantizer`` — for the backward gradient input (optional, when ``fp8_grad=True``) - - ``grad_weight_quantizer`` — for the backward weight gradient (currently unused) + - ``grad_weight_quantizer`` — for the backward weight gradient (placeholder for future + use; allocated for API completeness but not consumed by any current code path) All six are created here because the backward pass executes outside the ``autocast`` context and therefore no longer has access to the recipe @@ -276,7 +277,8 @@ The call traverses four layers, matching the :doc:`architecture_overview`: accumulators, trading some performance for numerical accuracy. This is controlled by the recipe and defaults to ``False`` for the forward pass (where some accumulation error is acceptable) and ``True`` for backward (where gradient accuracy matters more). This setting -is Hopper-specific and is a no-op on Blackwell. +is Hopper-specific and is a no-op on Blackwell (see the compute capability checks in +``transformer_engine/common/gemm/`` for the architecture dispatch logic). Phase 5: _Linear.forward() — Output Communication ---------------------------------------------------- diff --git a/docs/developer/pytorch_frontend/module_hierarchy.rst b/docs/developer/pytorch_frontend/module_hierarchy.rst index ccdb164a34..9baf5f3220 100644 --- a/docs/developer/pytorch_frontend/module_hierarchy.rst +++ b/docs/developer/pytorch_frontend/module_hierarchy.rst @@ -23,13 +23,14 @@ FP8 state management, quantizer lifecycle, and distributed training hooks. Below it, two branches: Branch 1: TransformerEngineBaseModule, with children: ├── Linear - ├── LayerNorm - ├── RMSNorm ├── LayerNormLinear ├── LayerNormMLP ├── GroupedLinear └── DotProductAttention - Branch 2: Directly from torch.nn.Module: + Branch 2: BasicOperation (ops framework), with children: + ├── LayerNorm (via _LayerNormOp) + └── RMSNorm (via _RMSNormOp) + Branch 3: Directly from torch.nn.Module: ├── MultiheadAttention (composes DotProductAttention + Linear projections) ├── TransformerLayer (composes the above modules) ├── Fp8Padding @@ -47,9 +48,9 @@ TE modules that support FP8 quantization. Key responsibilities: - ``init_fp8_metadata()``: Creates quantizer instances and amax history buffers based on the active recipe. -- ``pre_forward()``: Called at the start of each forward pass to update FP8 state (refresh - scales from amax history, check if FP8 is enabled). -- ``post_forward()``: Called after forward to record amax values. +- ``prepare_forward()``: Called at the start of each forward pass to update FP8 state + (refresh scales from amax history, check if FP8 is enabled). +- ``end_forward()``: Called after forward to record amax values. - FP8 metadata is registered as module buffers for serialization. **Quantizer Access** @@ -61,7 +62,8 @@ TE modules that support FP8 quantization. Key responsibilities: **Distributed Hooks** - ``set_tensor_parallel_group()``: Configure TP process group. -- ``set_sequence_parallel()``: Enable sequence parallelism. +- ``sequence_parallel``: A plain boolean attribute (set in ``__init__``, no dedicated + setter method). - Manages all-reduce of amax values across distributed ranks. **Weight Caching** @@ -82,13 +84,13 @@ A typical forward pass through a TE module: 1. autocast() sets global FP8 state 2. module.forward() called - a. pre_forward() — refresh FP8 scales, create quantizers + a. prepare_forward() — refresh FP8 scales, create quantizers b. Quantize input via input_quantizer(input) c. Get quantized weight (cached or quantize via weight_quantizer) d. Call the module's custom autograd function i. general_gemm(quantized_input, quantized_weight) ii. Save tensors for backward - e. post_forward() — record amax values + e. end_forward() — record amax values 3. Return output See :doc:`autograd_integration` for details on step (d). diff --git a/docs/developer/testing.rst b/docs/developer/testing.rst new file mode 100644 index 0000000000..0a7f6cac48 --- /dev/null +++ b/docs/developer/testing.rst @@ -0,0 +1,110 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +.. _testing: + +Testing Infrastructure +====================== + +This page describes how tests are organized, how to run them, and how to add new tests. + +Test Hierarchy +-------------- + +Tests are organized into four levels by resource usage and execution frequency: + +- **L0** — Fast, single-GPU unit tests. Run on every commit. All L0 tests must pass before + submitting PRs. +- **L1** — Multi-GPU and integration tests. Run nightly. Includes distributed tests, + Megatron-Core integration, ONNX export, and Thunder integration. +- **L2** — Extended tests. Run weekly. Broader coverage of edge cases and configurations. +- **L3** — Compatibility tests. Run per-release. Tests against specific external library + versions (e.g., Flash Attention version compatibility). + +Test scripts live in ``qa/L_/test.sh``. Each script sets up the environment +and invokes ``pytest`` or framework-specific test runners. + +Test Locations +-------------- + +- ``tests/pytorch/`` — PyTorch unit tests (sanity, numerics, operators, etc.) +- ``tests/pytorch/distributed/`` — PyTorch distributed/multi-GPU tests +- ``tests/jax/`` — JAX unit tests +- ``tests/cpp/operator/`` — C++ CUDA kernel tests (cast, activation, transpose, softmax, + normalization, GEMM, etc.) + +Running Tests +------------- + +**PyTorch — full L0 suite**: + +.. code-block:: bash + + TE_PATH=/path/to/TransformerEngine bash qa/L0_pytorch_unittest/test.sh + +**PyTorch — single test file**: + +.. code-block:: bash + + python3 -m pytest tests/pytorch/test_sanity.py -v + +**JAX — full L0 suite**: + +.. code-block:: bash + + TE_PATH=/path/to/TransformerEngine bash qa/L0_jax_unittest/test.sh + +**C++ kernel tests**: + +C++ tests are built separately with CMake/Ninja and run via ``ctest``. These test CUDA +kernels directly (cast, activation, transpose, softmax, normalization, etc.) without +framework overhead. + +.. code-block:: bash + + cd tests/cpp + cmake -GNinja -Bbuild . + cmake --build build + ctest --test-dir build -j4 + +``LD_LIBRARY_PATH`` must include the path to the installed ``libtransformer_engine.so``. +The ``qa/L0_cppunittest/test.sh`` script handles this automatically by querying ``pip3 +show transformer-engine``. + +Key Environment Variables +------------------------- + +Some tests require specific environment variables for reproducibility: + +- ``PYTORCH_JIT=0`` and ``NVTE_TORCH_COMPILE=0`` — Disable JIT compilation and Torch + Compile. Used for tests that need deterministic behavior, such as CUDA graph tests + where capture correctness must be verified. +- ``TE_PATH`` — Root of the Transformer Engine source tree. Used by ``qa/`` test scripts + to locate test files and configuration. + +Adding a New Test +----------------- + +1. Create the test file in the appropriate directory (``tests/pytorch/``, ``tests/jax/``, + or ``tests/cpp/operator/``). +2. For Python tests, use ``pytest`` conventions (functions or classes prefixed with + ``test_``). +3. For C++ tests, add a ``.cu`` file in ``tests/cpp/operator/`` and register it in the + ``CMakeLists.txt`` in that directory. +4. Register the test in the corresponding ``qa/L0_*/test.sh`` script (or the appropriate + level), setting any needed environment variables. +5. Verify the test passes locally before submitting. + +Linting +------- + +Pre-commit linting runs automatically on every commit via pre-commit.ci. You can also +run it locally: + +.. code-block:: bash + + bash qa/format.sh + +See ``CONTRIBUTING.rst`` at the repository root for full formatting and style guidelines. diff --git a/docs/envvars.rst b/docs/envvars.rst index 85445430f8..babe48898d 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -146,12 +146,6 @@ Attention Backend Selection :Default: Auto-selected :Description: Force a specific FusedAttention backend. ``0`` = F16_max512_seqlen (cuDNN, ≤512 seq len), ``1`` = F16_arbitrary_seqlen (cuDNN, any seq len), ``2`` = FP8 backend. If not set, the backend is automatically selected based on the input configuration. -.. envvar:: NVTE_FUSED_ATTN_FORCE_WORKSPACE_OPT - - :Type: ``int`` (0 or 1) - :Default: Auto-determined - :Description: Control workspace-related optimizations in FusedAttention. ``0`` disables optimizations, ``1`` enables them. These optimizations trade memory for performance. When unset, Transformer Engine determines the code path based on internal logic. For deterministic behavior with cuDNN ≥8.9.5 and <9.0.0, this is automatically set to ``1``. - .. envvar:: NVTE_FUSED_ATTN_USE_FAv2_BWD :Type: ``int`` (0 or 1) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 60ade522e3..91484ea74b 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -121,7 +121,6 @@ def reset_global_fp8_state(): @pytest.mark.parametrize("model_configs", [model_configs_base]) @pytest.mark.parametrize("model", model_configs_base.keys()) @pytest.mark.parametrize("ckpt_attn", [False]) -@pytest.mark.parametrize("workspace_opt", [True, False]) @pytest.mark.parametrize("qkv_layout", [None]) @pytest.mark.parametrize("swa", [False]) @pytest.mark.parametrize("pad_between_seqs", [False]) @@ -130,7 +129,6 @@ def test_dot_product_attention( model_configs, model, ckpt_attn, - workspace_opt, qkv_layout, swa, pad_between_seqs, @@ -220,7 +218,6 @@ def test_dot_product_attention( "UnfusedDotProductAttention", ckpt_attn, qkv_layout, - workspace_opt, pad_between_seqs, is_training, ) @@ -234,7 +231,6 @@ def test_dot_product_attention( "FusedAttention", ckpt_attn, qkv_layout, - workspace_opt, pad_between_seqs, is_training, ) @@ -246,7 +242,6 @@ def test_dot_product_attention( "FusedAttention", ckpt_attn, qkv_layout, - workspace_opt, pad_between_seqs, is_training, ) @@ -257,7 +252,6 @@ def test_dot_product_attention( "FusedAttention", ckpt_attn, qkv_layout, - workspace_opt, pad_between_seqs, is_training, ) @@ -270,7 +264,6 @@ def test_dot_product_attention( "FlashAttention", ckpt_attn, qkv_layout, - workspace_opt, pad_between_seqs, is_training, ) @@ -910,7 +903,6 @@ def _run_dot_product_attention( backend: str, ckpt_attn: bool, qkv_layout: str, - workspace_opt: bool, pad_between_seqs: bool, is_training: bool, ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: @@ -924,7 +916,6 @@ def _run_dot_product_attention( os.environ["NVTE_FLASH_ATTN"] = "1" if backend == "FusedAttention": os.environ["NVTE_FUSED_ATTN"] = "1" - os.environ["NVTE_FUSED_ATTN_FORCE_WORKSPACE_OPT"] = "1" if workspace_opt else "0" if backend == "UnfusedDotProductAttention": os.environ["NVTE_UNFUSED_ATTN"] = "1" _attention_backends["backend_selection_requires_update"] = True @@ -1323,7 +1314,6 @@ def test_transformer_layer( # Get configs config = model_configs[model] tols = dict(atol=5e-2, rtol=5e-2) - workspace_opt = True # Test backend availability is_training = True @@ -1367,7 +1357,6 @@ def test_transformer_layer( "UnfusedDotProductAttention", ckpt_attn, qkv_format, - workspace_opt, fused_qkv_params, RoPE, is_training, @@ -1381,7 +1370,6 @@ def test_transformer_layer( "FusedAttention", ckpt_attn, qkv_format, - workspace_opt, fused_qkv_params, RoPE, is_training, @@ -1395,7 +1383,6 @@ def test_transformer_layer( "FlashAttention", ckpt_attn, qkv_format, - workspace_opt, fused_qkv_params, RoPE, is_training, @@ -1465,7 +1452,6 @@ def _run_transformer_layer( backend: str, ckpt_attn: bool, qkv_format: str, - workspace_opt: bool, fused_qkv_params: bool, RoPE: bool, is_training: bool, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index 2dc42be18a..5546313328 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -20,7 +20,6 @@ DelayedScaling, Float8CurrentScaling, ) -from transformer_engine.pytorch.utils import get_cudnn_version from transformer_engine.pytorch.quantization import ( get_fp8_te_dtype, FP8GlobalStateManager, @@ -403,25 +402,6 @@ def __init__( not bool(int(os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1"))) or torch.are_deterministic_algorithms_enabled() ) - # To use the workspace optimization path for determinism, please - # set NVTE_FUSED_ATTN_FORCE_WORKSPACE_OPT=1 for cuDNN >=8.9.5 and <9.0.0, - # and set NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 for cuDNN >=9.0.0. - cudnn_version = get_cudnn_version() - if (8, 9, 5) <= cudnn_version < (9, 0, 0): - if self.deterministic: - os.environ["NVTE_FUSED_ATTN_FORCE_WORKSPACE_OPT"] = "1" - - # CUDNN_FRONTEND_ATTN_DP_WORKSPACE_LIMIT - # - unset: enables workspace optimization when required workspace is <= 256MB - # or when bias gradient needs to be computed - # - n: enables workspace optimization when required workspace is <= n bytes - # - -1: enables workspace optimization always - # - 0: disables workspace optimization always - if "NVTE_FUSED_ATTN_FORCE_WORKSPACE_OPT" in os.environ: - if os.environ["NVTE_FUSED_ATTN_FORCE_WORKSPACE_OPT"] == "0": - os.environ["CUDNN_FRONTEND_ATTN_DP_WORKSPACE_LIMIT"] = "0" - if os.environ["NVTE_FUSED_ATTN_FORCE_WORKSPACE_OPT"] == "1": - os.environ["CUDNN_FRONTEND_ATTN_DP_WORKSPACE_LIMIT"] = "-1" assert attention_type in AttnTypes, f"attention_type {attention_type} not supported" @@ -857,11 +837,7 @@ def forward( If FusedAttention is being used, users can also choose to switch to flash-attn's implementation for backward by setting :attr:`NVTE_FUSED_ATTN_USE_FAv2_BWD=1` (default: 0), because of the performance differences between various versions of - flash-attn and FusedAttention. Further, :attr:`NVTE_FUSED_ATTN_FORCE_WORKSPACE_OPT` - can be used to enable (:attr:`1`) or disable (:attr:`0`) the workspace related - optimizations in FusedAttention. When unset, Transformer Engine determines the code path - based on its internal logic. These optimizations trade memory for performance - and should be used with care. + flash-attn and FusedAttention. .. note:: .. _cu_seqlens note: From b993fe6946e739a990f68344f9b60353ba1fc128 Mon Sep 17 00:00:00 2001 From: Przemek Tredak Date: Tue, 24 Mar 2026 17:06:31 -0700 Subject: [PATCH 09/20] More fixes from review Signed-off-by: Przemek Tredak --- docs/developer/attention/backends.rst | 26 ++- docs/developer/attention/context_parallel.rst | 198 +++++++++++++++--- .../attention/fused_attn_kernels.rst | 34 ++- docs/developer/cpp_core/kernel_areas.rst | 8 +- 4 files changed, 216 insertions(+), 50 deletions(-) diff --git a/docs/developer/attention/backends.rst b/docs/developer/attention/backends.rst index 94fe6668d1..c4dd034b5a 100644 --- a/docs/developer/attention/backends.rst +++ b/docs/developer/attention/backends.rst @@ -64,11 +64,12 @@ package as an external dependency (not bundled). The integration imports from Code flow: -1. ``FlashAttention.forward()`` converts TE's tensor layouts (sbhd/thd) to FlashAttention's - expected bshd format. -2. Selects between ``flash_attn_func`` (dense), ``flash_attn_varlen_func`` - (variable-length), or ``flash_attn_with_kvcache`` (inference) based on the - configuration. +1. ``FlashAttention.forward()`` converts TE's tensor layouts as needed — sbhd tensors + are transposed to bshd, while THD (variable-length) inputs use FlashAttention's + native ``flash_attn_varlen_func`` without converting to dense bshd. +2. Selects between ``flash_attn_func`` (dense bshd/sbhd without padding), + ``flash_attn_varlen_func`` (THD or padding masks), or ``flash_attn_with_kvcache`` + (inference) based on the configuration. 3. For context parallelism, wraps the call via ``attn_forward_func_with_cp()``. To modify FlashAttention integration: start with ``FlashAttention.forward()`` in @@ -94,13 +95,16 @@ Code flow: sub-backend implementation based on ``nvte_get_fused_attn_backend()``. 4. **cuDNN graph builders** — one file per sub-backend: - - ``fused_attn_f16_max512_seqlen.cu`` — sub-backend 0 - - ``fused_attn_f16_arbitrary_seqlen.cu`` — sub-backend 1 + - ``fused_attn_f16_max512_seqlen.cu`` — sub-backend 0 (legacy; only selected when + sub-backend 1 is unavailable due to older cuDNN or restrictive parameters) + - ``fused_attn_f16_arbitrary_seqlen.cu`` — sub-backend 1 (preferred for all F16 cases) - ``fused_attn_fp8.cu`` — sub-backend 2 Each builds a ``cudnn_frontend::graph::Graph``, configures SDPA attributes (masking, - dropout, scaling), and executes it. **Graphs are cached** per thread in a thread-local - cache keyed by ``FADescriptor`` to avoid rebuilding identical configurations. + dropout, scaling), and executes it. **Graphs are cached** per thread in thread-local + caches (e.g., ``sdpa_f16_fprop_cache`` keyed by ``FADescriptor_v1`` for sub-backend 1, + ``fa_fprop_cache`` keyed by ``FADescriptor`` for FP8) to avoid rebuilding identical + configurations. See ``utils.h`` for the descriptor definitions. To modify cuDNN fused attention: start with the relevant ``.cu`` file for the sub-backend, specifically the ``fused_attn_*_fwd_impl()`` / ``fused_attn_*_bwd_impl()`` functions @@ -111,8 +115,7 @@ Helper CUDA Kernels ^^^^^^^^^^^^^^^^^^^ The ``transformer_engine/common/fused_attn/`` directory also contains CUDA helper kernels -that support the attention backends but do **not** perform the full attention operation -themselves: +that support the attention backends: - ``utils.cu`` — stride calculation, auxiliary tensor setup - ``context_parallel.cu`` — KV communication for context parallelism @@ -126,3 +129,4 @@ See Also - :doc:`backend_selection` — How the backend is chosen at runtime - :doc:`pytorch_attention` — PyTorch DotProductAttention and MultiheadAttention modules - :doc:`fused_attn_kernels` — C++ kernel organization and cuDNN integration +- :doc:`context_parallel` — Context parallelism strategies and implementation diff --git a/docs/developer/attention/context_parallel.rst b/docs/developer/attention/context_parallel.rst index c435a2fda6..3216e20f0a 100644 --- a/docs/developer/attention/context_parallel.rst +++ b/docs/developer/attention/context_parallel.rst @@ -9,64 +9,181 @@ Context Parallelism =================== Context Parallelism (CP) distributes the sequence dimension across multiple GPUs, -enabling training with very long sequences that don't fit in a single GPU's memory. +enabling training with very long sequences that don't fit in a single GPU's memory. It is +the most complex part of the attention subsystem. + +**Implementation**: ``transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py`` +(~4300 lines). Entry point: ``flash_attn_with_cp()``, which dispatches to the appropriate +strategy class based on ``cp_comm_type``. + +Overview +-------- + +In CP, each GPU holds a shard of the sequence for Q, K, and V. To compute full attention, +each GPU needs access to the K and V from *all* sequence shards. CP achieves this by +exchanging KV (or QKV) chunks between GPUs using one of several communication strategies. + +The key challenge is **load balancing**: with causal attention, GPUs holding later sequence +positions must attend to more KV positions than GPUs holding earlier ones. Without +balancing, later GPUs would be idle while earlier ones compute. TE addresses this with +**dual chunk** distribution — see :ref:`cp-load-balancing` below. + +Communication Strategies +------------------------ + +TE implements four CP communication strategies, selected via the ``cp_comm_type`` +parameter on ``DotProductAttention``: + +**P2P (Point-to-Point Ring)** + +- **Class**: ``AttnFuncWithCPAndKVP2P`` +- KV chunks are passed around a ring of GPUs using asynchronous ``isend``/``irecv``. + Each GPU computes attention for its local Q against the currently-held KV chunk, then + sends its KV to the next GPU while receiving the next chunk. This overlaps communication + with computation. +- For causal attention, the computation is divided into **sections** based on the + relationship between Q and KV positions: + + - **Diagonal** (step 0): Q attends to its own KV shard (causal mask applies). + - **Lower-triangle** (steps where KV is from earlier positions): Q's second half + attends to full KV. + - **Upper-triangle** (steps where KV is from later positions): Q's first half attends + to full KV. + + Each section computes a partial attention output and softmax LSE (log-sum-exp). + **Output correction** merges these partial results: + ``out = out + out_per_step * exp(lse_per_step - lse_combined)``. The correction + functions (``flash_attn_fwd_out_correction()``, + ``flash_attn_fwd_softmax_lse_correction()``) handle the numerically stable + accumulation. +- Works with both FusedAttention and FlashAttention backends. + +**All-Gather** + +- **Class**: ``AttnFuncWithCPAndKVAllGather`` +- All GPUs perform an all-gather to collect the full K and V tensors, then compute + attention locally. Simpler than P2P but requires more memory (full K and V on each GPU) + and cannot overlap communication with computation. +- Reference: Section 3.3.2 of `The Llama 3 Herd of Models `_. + +**All-to-All (A2A)** + +- **Class**: ``AttnFuncWithCPAndQKVOA2A`` +- Follows the `DeepSpeed Ulysses `_ approach: scatters + attention heads across the CP group and gathers to get the full sequence of QKV. Unlike + P2P and all-gather, this exchanges Q in addition to KV — each GPU ends up with a + subset of heads but the full sequence. + +**A2A + P2P (Hierarchical)** + +- Uses a two-level CP group: ``cp_group = [cp_group_a2a, cp_group_p2p]``. A2A is applied + within subgroups (e.g., via NVLink), then P2P between subgroups (e.g., via IB). +- Reference: `USP `_ and + `LongVILA `_. + +.. _cp-load-balancing: + +Load Balancing +-------------- + +With causal attention, naively assigning contiguous sequence shards would leave GPUs with +early positions idle while GPUs with late positions compute more. TE uses two load +balancing strategies — **dual chunk** (PyTorch) and **striped** (JAX) — that both ensure +each GPU processes a balanced mix of early and late sequence positions. .. figure:: ./img/context_parallel_ring.svg :align: center :width: 70% - Ring of 4 GPUs passing KV chunks in context parallelism. + Dual chunk distribution for P2P ring with 4 GPUs. .. Diagram description for ``context_parallel_ring.svg``: - Four GPU boxes arranged in a ring (square layout). - GPU 0 has "seq[0:L/4]", GPU 1 has "seq[L/4:L/2]", GPU 2 has "seq[L/2:3L/4]", - GPU 3 has "seq[3L/4:L]". - Arrows form a ring: GPU0 → GPU1 → GPU2 → GPU3 → GPU0. - Each arrow is labeled "KV chunk". - Center text: "Each GPU computes local Q × all KV via ring passes" + Horizontal bar representing a sequence of length L, divided into 8 chunks + (2 × cp_size, where cp_size=4), numbered 0-7 left to right. + Below, 4 GPU boxes showing which chunks each GPU holds: + GPU 0: chunks 0 and 7 (start and end) + GPU 1: chunks 2 and 5 + GPU 2: chunks 4 and 3 + GPU 3: chunks 6 and 1 + Arrows between GPUs form a P2P ring: GPU0 → GPU1 → GPU2 → GPU3 → GPU0. + Each arrow labeled "KV chunks". -Overview --------- +**Dual Chunk (PyTorch)** -In CP, each GPU holds a shard of the sequence for Q, K, and V. To compute full attention, -each GPU needs access to the K and V from *all* sequence shards. CP achieves this by -passing KV chunks between GPUs in a ring or all-gather pattern. +The dual chunk strategy splits the sequence into ``2 × cp_size`` chunks and assigns each +GPU two discontiguous chunks — one from the start and one from the end of the sequence: -Strategies ----------- +- GPU rank ``r`` receives chunks ``[2r, 2×cp_size - 2r - 1]``. +- For ``cp_size=4``: GPU 0 gets chunks {0, 7}, GPU 1 gets {2, 5}, GPU 2 gets {4, 3}, + GPU 3 gets {6, 1}. -**Ring Attention** +The reordering logic is in ``get_seq_chunk_ids_for_reordering_before_attn()`` and +``get_seq_chunk_ids_for_reordering_after_attn()`` in ``context_parallel.py``. -KV chunks are passed around a ring of GPUs. Each GPU: +**Striped (JAX)** -1. Computes attention for its local Q against the currently-held KV chunk. -2. Sends its KV chunk to the next GPU in the ring. -3. Receives a new KV chunk from the previous GPU. -4. Repeats until all KV chunks have been seen. +Distributes tokens in a round-robin (interleaved) pattern across GPUs. For ``cp_size=4`` +and a 16-token sequence: -This overlaps communication with computation: while computing attention on the current -KV chunk, the next chunk is being transferred. +- GPU 0 gets tokens {0, 4, 8, 12} +- GPU 1 gets tokens {1, 5, 9, 13} +- GPU 2 gets tokens {2, 6, 10, 14} +- GPU 3 gets tokens {3, 7, 11, 15} -**All-Gather KV** +GPU rank ``r`` holds tokens at original positions ``[r, r + cp_size, r + 2×cp_size, ...]``. +The reordering reshapes the sequence dimension into +``[S/(cp_size × stripe_size), cp_size, stripe_size]`` and swaps the middle two axes. The +``stripe_size`` parameter (default 1) controls the granularity of the interleaving. -All GPUs perform an all-gather to collect the full K and V tensors, then compute -attention locally. Simpler than ring attention but requires more memory (full K and V -on each GPU). +Dual chunk is used in PyTorch (with both BSHD and THD layouts), while the striped +strategy is used in JAX. Implementation: +``reorder_causal_striped()`` in ``transformer_engine/jax/cpp_extensions/attention.py``, +dispatched via ``reorder_causal_load_balancing()`` in ``transformer_engine/jax/attention.py``. -**THD (Token-Head-Dimension) Layout** +Differences from Regular Attention Flow +--------------------------------------- -A specialized layout where tokens from different sequence positions are interleaved -across GPUs in a way that enables efficient attention computation with minimal -communication. Used with TE's custom CUDA attention kernels. +CP attention differs from the non-CP flow (described in :doc:`/developer/linear_walkthrough`) +in several key ways: + +1. **Multi-step computation**: Instead of a single attention call, P2P CP iterates over + ``cp_size`` steps, computing partial attention at each step against a different KV + chunk. + +2. **Output correction**: Partial outputs from each step are merged using softmax LSE + correction. This is mathematically equivalent to computing attention over the full + sequence but requires tracking per-step LSE statistics. + +3. **Overlapped communication**: P2P uses separate CUDA streams to overlap KV + communication with attention computation. While step ``i`` computes, step ``i+1``'s + KV is being transferred. + +4. **Section-based computation (causal)**: For causal attention, each P2P step computes + only the relevant portion of the attention matrix (diagonal, lower-triangle, or + upper-triangle section), not the full Q×KV product. + +THD Layout Considerations +------------------------- + +THD (Token-Head-Dimension) is a variable-length sequence layout, not a CP strategy. Each +of the strategies above must handle THD inputs specially because: + +- Sequence metadata (``cu_seqlens``) is a GPU tensor, not part of the tensor shape. + Chunking logic must use index operations rather than view/reshape. +- Token partitioning must respect sequence boundaries (cannot split within a sequence). +- Softmax LSE tensors have different shapes (``TH1`` packed vs ``BHS1`` dense). +- C++ helper kernels in ``context_parallel.cu`` (``thd_partition_indices_kernel``, + ``thd_read_half_tensor_kernel``, ``thd_lse_kernel``, ``thd_out_correction_kernel``) + handle the THD-specific index and correction operations. Integration with Attention Backends ------------------------------------ CP support varies by attention backend: -- **cuDNN Fused**: Full CP support via both ring and all-gather strategies. -- **FlashAttention**: Limited CP support. +- **cuDNN Fused**: Full CP support via P2P, all-gather, and A2A strategies. +- **FlashAttention**: CP support via P2P and A2A. - **Unfused**: No CP support. The CP strategy is configured via ``DotProductAttention``: @@ -80,6 +197,21 @@ The CP strategy is configured via ``DotProductAttention``: cp_stream=cuda_stream, ) +Key Entry Points for Developers +--------------------------------- + +- **Strategy dispatch**: ``flash_attn_with_cp()`` in ``context_parallel.py`` — selects + P2P, all-gather, A2A, or A2A+P2P based on ``cp_comm_type``. +- **P2P forward**: ``AttnFuncWithCPAndKVP2P.forward()`` — the step loop, section logic, + and output correction. +- **Communication**: ``flash_attn_p2p_communicate()`` (P2P isend/irecv), + ``flash_attn_a2a_communicate()`` (all-to-all with reordering). +- **Output correction**: ``flash_attn_fwd_out_correction()`` and + ``flash_attn_fwd_softmax_lse_correction()``. +- **Load balancing**: ``get_seq_chunk_ids_for_reordering_before_attn()`` and + ``get_seq_chunk_ids_for_reordering_after_attn()``. +- **THD C++ kernels**: ``transformer_engine/common/fused_attn/context_parallel.cu``. + See Also -------- diff --git a/docs/developer/attention/fused_attn_kernels.rst b/docs/developer/attention/fused_attn_kernels.rst index 9c69f535e2..29cb453672 100644 --- a/docs/developer/attention/fused_attn_kernels.rst +++ b/docs/developer/attention/fused_attn_kernels.rst @@ -8,8 +8,7 @@ Fused Attention Kernels ======================= -The C++ fused attention implementation lives in ``transformer_engine/common/fused_attn/`` -and represents the most complex kernel area in Transformer Engine. +The C++ fused attention implementation lives in ``transformer_engine/common/fused_attn/``. C API ----- @@ -87,6 +86,37 @@ before execution: This two-pass pattern (query size, then execute) is common across TE's C API. +Auxiliary Tensors (NVTETensorPack) +---------------------------------- + +The forward pass produces auxiliary tensors needed for the backward pass — softmax +statistics, RNG state, and optionally bias gradients. Because different sub-backends +produce different auxiliary tensors, they are passed through an ``NVTETensorPack`` +(defined in ``transformer_engine.h``): + +.. code-block:: c + + struct NVTETensorPack { + static const int MAX_SIZE = 10; + NVTETensor tensors[MAX_SIZE]; + size_t size = 0; + }; + +Each sub-backend populates the pack differently: + +- **Sub-backend 0** (F16 max512): 1 tensor — ``S`` (full softmax intermediate). +- **Sub-backend 1** (F16 arbitrary): 2+ tensors — softmax stats (``S`` or + ``Max``/``Sum_Exp`` depending on ``return_max_logit``), ``rng_state``, and optionally + ``Bias`` and ``SoftmaxOffset``. +- **Sub-backend 2** (FP8): 3 tensors — ``M`` (row max), ``ZInv`` (inverse softmax + denominator), ``rng_state``. + +On the Python/pybind11 side (``pytorch/csrc/extensions/attention.cpp``), the forward call +uses the two-pass pattern: the first call with empty tensors discovers the required shapes +and dtypes; PyTorch tensors are then allocated; the second call executes with the +memory-backed pack. The backward call receives the same pack to reuse the saved +statistics. + See Also -------- diff --git a/docs/developer/cpp_core/kernel_areas.rst b/docs/developer/cpp_core/kernel_areas.rst index d59abd665e..e76f5fd6cd 100644 --- a/docs/developer/cpp_core/kernel_areas.rst +++ b/docs/developer/cpp_core/kernel_areas.rst @@ -41,8 +41,10 @@ Scaling Mode Dispatch ---------------------- Kernels that handle multiple scaling modes use a switch on the output's -``scaling_mode`` to dispatch to the appropriate implementation. For example, in -``cast/dispatch/quantize.cuh``: +``scaling_mode`` to dispatch to the appropriate implementation. The scaling mode +effectively acts as a tag in a tagged union — the ``Tensor`` struct holds the same set of +fields, but which fields are valid and how they are interpreted depends on the scaling +mode. For example, in ``cast/dispatch/quantize.cuh``: .. code-block:: text @@ -85,8 +87,6 @@ Many areas support fusing quantization with the primary computation: avoiding a separate cast kernel launch. - **Activation + quantize**: Similarly, activation functions can cast their output to FP8 in the same kernel pass. -- **Cast + transpose**: A single kernel produces both rowwise and transposed (columnwise) - output, critical for efficiently preparing data for the backward pass. These fusions are driven by the same output-driven pattern: if the output tensor's dtype is an FP8 type and has the appropriate scale fields, the kernel produces quantized output From e052e04b8236886692b4a1dd5a8332fb8b3dcd6a Mon Sep 17 00:00:00 2001 From: Przemek Tredak Date: Mon, 30 Mar 2026 14:03:47 -0700 Subject: [PATCH 10/20] After more fixes Signed-off-by: Przemek Tredak --- docs/developer/attention/backends.rst | 10 ++++--- docs/developer/attention/context_parallel.rst | 8 +++--- docs/developer/cpp_core/scaling_modes.rst | 4 ++- docs/developer/index.rst | 2 +- docs/developer/jax_frontend/quantization.rst | 2 +- .../quantization/class_hierarchy.rst | 9 +------ .../quantization/rowwise_columnwise.rst | 27 ++++++++++++------- .../quantization/scaling_recipes.rst | 2 +- 8 files changed, 35 insertions(+), 29 deletions(-) diff --git a/docs/developer/attention/backends.rst b/docs/developer/attention/backends.rst index c4dd034b5a..d0b51f0c0b 100644 --- a/docs/developer/attention/backends.rst +++ b/docs/developer/attention/backends.rst @@ -58,9 +58,13 @@ FlashAttention ``transformer_engine/pytorch/attention/dot_product_attention/backends.py`` Integrates Tri Dao's `flash-attn `_ -package as an external dependency (not bundled). The integration imports from -``flash_attn.flash_attn_interface`` (FA2) and ``flash_attn_3.flash_attn_interface`` -(FA3). Version constraints are managed in ``FlashAttentionUtils`` in ``utils.py``. +package as an external dependency (not bundled). Two versions are supported: + +- **FA2** (``flash_attn.flash_attn_interface``): Ampere and later (sm80+). +- **FA3** (``flash_attn_3.flash_attn_interface``): Hopper only (sm90). Selected + automatically when installed and running on sm90; disabled on other architectures. + +Version constraints are managed in ``FlashAttentionUtils`` in ``utils.py``. Code flow: diff --git a/docs/developer/attention/context_parallel.rst b/docs/developer/attention/context_parallel.rst index 3216e20f0a..9632b8229d 100644 --- a/docs/developer/attention/context_parallel.rst +++ b/docs/developer/attention/context_parallel.rst @@ -13,8 +13,8 @@ enabling training with very long sequences that don't fit in a single GPU's memo the most complex part of the attention subsystem. **Implementation**: ``transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py`` -(~4300 lines). Entry point: ``flash_attn_with_cp()``, which dispatches to the appropriate -strategy class based on ``cp_comm_type``. +(~4300 lines). Entry point: ``attn_forward_func_with_cp()``, which dispatches to the +appropriate strategy class based on ``cp_comm_type``. Overview -------- @@ -200,8 +200,8 @@ The CP strategy is configured via ``DotProductAttention``: Key Entry Points for Developers --------------------------------- -- **Strategy dispatch**: ``flash_attn_with_cp()`` in ``context_parallel.py`` — selects - P2P, all-gather, A2A, or A2A+P2P based on ``cp_comm_type``. +- **Strategy dispatch**: ``attn_forward_func_with_cp()`` in ``context_parallel.py`` — + selects P2P, all-gather, A2A, or A2A+P2P based on ``cp_comm_type``. - **P2P forward**: ``AttnFuncWithCPAndKVP2P.forward()`` — the step loop, section logic, and output correction. - **Communication**: ``flash_attn_p2p_communicate()`` (P2P isend/irecv), diff --git a/docs/developer/cpp_core/scaling_modes.rst b/docs/developer/cpp_core/scaling_modes.rst index 6270cea28f..04559f0954 100644 --- a/docs/developer/cpp_core/scaling_modes.rst +++ b/docs/developer/cpp_core/scaling_modes.rst @@ -53,7 +53,9 @@ must be separate buffers with physically transposed data. Tensor Scaling (``NVTE_DELAYED_TENSOR_SCALING``) -------------------------------------------------- -A single scale factor applies to the entire tensor. Minimum architecture: Ada (sm89). +A single scale factor applies to the entire tensor. Minimum architecture: Ada (sm89) for +FP8 data types, though practical FP8 training with cuBLASLt GEMMs is typically done on +Hopper (sm90) and later. .. note:: diff --git a/docs/developer/index.rst b/docs/developer/index.rst index 7ae3276b17..b90cc53782 100644 --- a/docs/developer/index.rst +++ b/docs/developer/index.rst @@ -38,10 +38,10 @@ How to Use This Guide architecture_overview linear_walkthrough - testing cpp_core/index quantization/index pytorch_frontend/index jax_frontend/index attention/index distributed/index + testing diff --git a/docs/developer/jax_frontend/quantization.rst b/docs/developer/jax_frontend/quantization.rst index e68af1bdda..59e7805772 100644 --- a/docs/developer/jax_frontend/quantization.rst +++ b/docs/developer/jax_frontend/quantization.rst @@ -31,7 +31,7 @@ The JAX quantizer hierarchy mirrors PyTorch's: - Current tensor * - ``BlockScaleQuantizer`` - ``MXFP8Quantizer`` - - MXFP8 block + - MXFP8 block (JAX does not yet support generic block scaling 1D/2D) * - ``NVFP4Quantizer`` - ``NVFP4Quantizer`` - NVFP4 diff --git a/docs/developer/quantization/class_hierarchy.rst b/docs/developer/quantization/class_hierarchy.rst index 718e4565f7..273a001646 100644 --- a/docs/developer/quantization/class_hierarchy.rst +++ b/docs/developer/quantization/class_hierarchy.rst @@ -77,12 +77,6 @@ responsibilities: # Usage control def set_usage(self, *, rowwise=False, columnwise=False): ... - # Properties - @property - def dtype(self) -> torch.dtype: ... - @property - def scaling_mode(self) -> str: ... - The ``internal`` Flag ^^^^^^^^^^^^^^^^^^^^^ @@ -171,8 +165,7 @@ QuantizedTensor (Base: ``transformer_engine/pytorch/quantized_tensor.py``) # torch.Tensor subclass machinery def __torch_dispatch__(cls, func, types, args, kwargs): ... - # Data access - def get_data_tensors(self) -> dict: ... # Extract raw data for kernel calls + # Data access (get_data_tensors() is on concrete Storage subclasses, not here) def get_metadata(self) -> dict: ... # Scaling metadata # Dequantize diff --git a/docs/developer/quantization/rowwise_columnwise.rst b/docs/developer/quantization/rowwise_columnwise.rst index 3d13faa924..329097f5d7 100644 --- a/docs/developer/quantization/rowwise_columnwise.rst +++ b/docs/developer/quantization/rowwise_columnwise.rst @@ -40,15 +40,22 @@ cuBLASLt's FP8 GEMM computes ``C = A × B`` where: In the **forward pass**: -- ``output = activation × weight^T`` -- Activation is the ``A`` operand → needs **rowwise** quantization. -- Weight is the ``B`` operand → needs **columnwise** quantization. +- ``output = weight × activation^T`` (via ``general_gemm(weightmat, inputmat)``) +- Weight is the ``A`` operand → needs **rowwise** quantization. +- Activation is the ``B`` operand → needs **rowwise** quantization. + +In the **backward dgrad pass**: + +- ``grad_input = weight^T × grad_output`` (via ``general_gemm(weight, grad_output)``) +- Weight transposed uses **columnwise** data saved from forward. +- Grad output needs **rowwise** quantization. In the **backward wgrad pass**: -- ``weight_grad = activation^T × grad_output`` -- Activation transposed is the ``A`` operand → needs **columnwise** data from forward. -- Grad output is the ``B`` operand → needs **columnwise** quantization. +- ``weight_grad = input^T × grad_output`` + (via ``general_gemm(inputmat_columnwise, grad_output_columnwise)``) +- Input transposed uses **columnwise** data saved from forward. +- Grad output uses **columnwise** data. This is why the forward pass must produce *both* rowwise and columnwise quantized data for activations: the rowwise data is consumed immediately by the forward GEMM, while the @@ -61,13 +68,13 @@ The ``Quantizer.set_usage()`` method controls which layouts are produced: .. code-block:: python - # Forward activation quantizer: needs both layouts + # Forward activation quantizer: rowwise for forward GEMM, columnwise for backward wgrad fwd_quantizer.set_usage(rowwise=True, columnwise=True) - # Weight quantizer: only columnwise for forward GEMM - weight_quantizer.set_usage(rowwise=False, columnwise=True) + # Weight quantizer: rowwise for forward GEMM, columnwise for backward dgrad + weight_quantizer.set_usage(rowwise=True, columnwise=True) - # Backward grad_output quantizer: needs both for dgrad and wgrad + # Backward grad_output quantizer: rowwise for dgrad, columnwise for wgrad bwd_quantizer.set_usage(rowwise=True, columnwise=True) When ``columnwise=True``, the quantization kernel produces an additional transposed diff --git a/docs/developer/quantization/scaling_recipes.rst b/docs/developer/quantization/scaling_recipes.rst index 7eacc550b5..ab4022046d 100644 --- a/docs/developer/quantization/scaling_recipes.rst +++ b/docs/developer/quantization/scaling_recipes.rst @@ -85,7 +85,7 @@ Each ``TransformerEngineBaseModule`` maintains its own quantization state: - **Scale tensors**: Computed from amax history before each forward pass. The module's ``init_fp8_metadata()`` method creates this state, and -``pre_forward()`` / ``post_forward()`` manage the per-iteration lifecycle. +``prepare_forward()`` / ``end_forward()`` manage the per-iteration lifecycle. Amax Update Flow (Delayed Scaling) ----------------------------------- From 9b22bfe8e0a569040b91e1c7105bcff0eb282f23 Mon Sep 17 00:00:00 2001 From: Przemek Tredak Date: Tue, 31 Mar 2026 11:31:01 -0700 Subject: [PATCH 11/20] And more fixes Signed-off-by: Przemek Tredak --- docs/developer/architecture_overview.rst | 2 +- docs/developer/cpp_core/scaling_modes.rst | 10 +++++++--- docs/developer/linear_walkthrough.rst | 19 ++++++++++--------- .../pytorch_frontend/cpp_extensions.rst | 2 +- .../quantization/class_hierarchy.rst | 8 ++++---- .../quantization/rowwise_columnwise.rst | 11 +++++++---- 6 files changed, 30 insertions(+), 22 deletions(-) diff --git a/docs/developer/architecture_overview.rst b/docs/developer/architecture_overview.rst index 8d4d1c389d..ad83b33d62 100644 --- a/docs/developer/architecture_overview.rst +++ b/docs/developer/architecture_overview.rst @@ -88,7 +88,7 @@ Directory Map ├── quantize/ # JAX quantizer and ScaledTensor ├── attention.py # JAX attention entry points ├── sharding.py # Mesh and sharding utilities - └── distributed.py # JAX distributed collectives + └── dense.py # JAX dense (linear) layer Data Flow: Forward Pass ----------------------- diff --git a/docs/developer/cpp_core/scaling_modes.rst b/docs/developer/cpp_core/scaling_modes.rst index 04559f0954..d9012da67e 100644 --- a/docs/developer/cpp_core/scaling_modes.rst +++ b/docs/developer/cpp_core/scaling_modes.rst @@ -68,7 +68,9 @@ Hopper (sm90) and later. Data layout: - ``data``: row-major, one scale for the entire tensor -- ``columnwise_data``: column-major (physically transposed), same single scale +- ``columnwise_data``: column-major layout, same single scale. On architectures without + non-TN FP8 GEMM support this is a physically transposed copy; on architectures with + non-TN support (Blackwell+) it may alias ``data`` - ``scale_inv``: shape ``[1]`` (one element) for FP8 tensors, empty for high-precision - ``amax``: populated by delayed scaling (framework responsibility), unused by current scaling @@ -92,8 +94,10 @@ Blackwell. Data layout: - ``data``: row-major FP8, shape ``[M, N]`` -- ``columnwise_data``: the physically transposed data, shape ``[N, M]`` — this is a - separate buffer with independently quantized blocks along the transposed dimension +- ``columnwise_data``: the columnwise layout, shape ``[N, M]`` — on architectures + without non-TN FP8 GEMM support this is a physically transposed separate buffer + with independently quantized blocks along the transposed dimension; on architectures + with non-TN support (Blackwell+) it may share the rowwise buffer - ``scale_inv`` (rowwise): shape ``[ceil(M), ceil(N/32)]``, padded to ``[128, 4]`` multiples for GEMM alignment - ``columnwise_scale_inv``: shape ``[ceil(N), ceil(M/32)]``, padded to ``[128, 4]`` diff --git a/docs/developer/linear_walkthrough.rst b/docs/developer/linear_walkthrough.rst index 67c5078e45..c3f126028a 100644 --- a/docs/developer/linear_walkthrough.rst +++ b/docs/developer/linear_walkthrough.rst @@ -105,21 +105,22 @@ that actually perform casts. ``amax_epsilon`` and ``power_2_scale`` on quantizers, while MXFP8 needs no such configuration. -2. Six quantizers are prepared — three for the forward pass and three for the backward - pass: +2. Six quantizer slots are prepared — three for the forward pass and three for the + backward pass: - ``input_quantizer`` — for the activation tensor (forward GEMM input) - ``weight_quantizer`` — for the weight parameter (forward GEMM input) - ``output_quantizer`` — for the forward output (optional, when ``fp8_output=True``) - ``grad_output_quantizer`` — for the backward gradient - ``grad_input_quantizer`` — for the backward gradient input (optional, when ``fp8_grad=True``) - - ``grad_weight_quantizer`` — for the backward weight gradient (placeholder for future - use; allocated for API completeness but not consumed by any current code path) - - All six are created here because the backward pass executes outside the - ``autocast`` context and therefore no longer has access to the recipe - configuration. By creating the backward quantizers during forward setup, we - capture the recipe information while it is still available. + - ``grad_weight_quantizer`` — reserved for the backward weight gradient (currently + left as ``None``; no code path uses it yet, but the slot is carried through + the API for future use) + + The backward quantizers that are populated are created here because the backward + pass executes outside the ``autocast`` context and therefore no longer has access + to the recipe configuration. By creating the backward quantizers during forward + setup, we capture the recipe information while it is still available. Each tensor gets its own quantizer because activations, weights, and gradients have different value distributions — sharing a single scale would waste dynamic range. diff --git a/docs/developer/pytorch_frontend/cpp_extensions.rst b/docs/developer/pytorch_frontend/cpp_extensions.rst index 1cddfad259..e4672960bf 100644 --- a/docs/developer/pytorch_frontend/cpp_extensions.rst +++ b/docs/developer/pytorch_frontend/cpp_extensions.rst @@ -144,7 +144,7 @@ accepts ``py::handle`` inputs: This function performs runtime type detection to handle all Python tensor types: 1. It iterates over a registry of known quantized types (defined in - ``transformer_engine/pytorch/csrc/extensions/pybind.h`` as + ``transformer_engine/pytorch/csrc/pybind.h`` as ``custom_types_converters``). Each entry maps a ``PyTypeObject*`` to an extraction function. This registry is populated during extension initialization (``init_extension()``), which is why the extensions must be initialized before any diff --git a/docs/developer/quantization/class_hierarchy.rst b/docs/developer/quantization/class_hierarchy.rst index 273a001646..3bfd626676 100644 --- a/docs/developer/quantization/class_hierarchy.rst +++ b/docs/developer/quantization/class_hierarchy.rst @@ -71,8 +71,8 @@ responsibilities: class Quantizer: # Core interface - def __call__(self, tensor, *, noop=False) -> QuantizedTensor: ... - def quantize(self, tensor, *, noop=False) -> QuantizedTensorStorage: ... + def __call__(self, tensor) -> QuantizedTensor: ... + def quantize(self, tensor, *, out=None, dtype=None) -> QuantizedTensor: ... # Usage control def set_usage(self, *, rowwise=False, columnwise=False): ... @@ -113,7 +113,7 @@ Concrete Quantizers - MXFP8 with 32-element E8M0 block scales * - ``Float8BlockQuantizer`` - ``tensor/float8_blockwise_tensor.py`` - - 1D or 2D block scaling with configurable block size + - 1D or 2D block scaling with a fixed block length of 128 (dimensionality is configurable via ``block_scaling_dim``) * - ``NVFP4Quantizer`` - ``tensor/nvfp4_tensor.py`` - NVFP4 with 16-element E8M0 block scales @@ -134,7 +134,7 @@ support and minimal overhead. # Public interface def update_usage(self, *, rowwise_usage=None, columnwise_usage=None): ... - def get_usages(self) -> tuple[bool, bool]: ... + def get_usages(self) -> Dict[str, bool]: ... def prepare_for_saving(self) -> tuple: ... # For autograd save_for_backward def restore_from_saved(cls, saved) -> "QuantizedTensorStorage": ... def quantize_(self, tensor, quantizer): ... # In-place quantization diff --git a/docs/developer/quantization/rowwise_columnwise.rst b/docs/developer/quantization/rowwise_columnwise.rst index 329097f5d7..2dc99f8412 100644 --- a/docs/developer/quantization/rowwise_columnwise.rst +++ b/docs/developer/quantization/rowwise_columnwise.rst @@ -77,8 +77,11 @@ The ``Quantizer.set_usage()`` method controls which layouts are produced: # Backward grad_output quantizer: rowwise for dgrad, columnwise for wgrad bwd_quantizer.set_usage(rowwise=True, columnwise=True) -When ``columnwise=True``, the quantization kernel produces an additional transposed -copy of the data (or, for block-scaling modes, computes column-oriented block scales). +When ``columnwise=True``, the quantizer prepares a columnwise layout of the data. +On architectures without non-TN FP8 GEMM support, this means producing an additional +physically transposed copy (or, for block-scaling modes, computing column-oriented +block scales). On architectures that support non-TN FP8 GEMM (Blackwell and later), +the rowwise data can be reused directly, so no extra buffer is allocated. Impact on the Tensor Struct ---------------------------- @@ -105,8 +108,8 @@ transposed into a separate buffer for the columnwise layout. Performance Implications ------------------------ -Producing both layouts doubles the quantization work and memory for activations. This is -a deliberate trade-off: +On architectures without non-TN FP8 GEMM support, producing both layouts doubles the +quantization work and memory for activations. This is a deliberate trade-off: - **Without dual layout**: The wgrad GEMM would need to transpose and requantize at backward time. Beyond the performance cost (extra kernel launch and temporary memory), From 6ba9c5107effacffc05075af903f360519a07718 Mon Sep 17 00:00:00 2001 From: Przemek Tredak Date: Fri, 3 Apr 2026 09:38:12 -0700 Subject: [PATCH 12/20] Changes after merging main Signed-off-by: Przemek Tredak --- docs/developer/attention/backend_selection.rst | 16 ++++++++++++++++ docs/developer/attention/backends.rst | 6 ++++++ docs/developer/attention/context_parallel.rst | 6 ++++++ docs/developer/attention/fused_attn_kernels.rst | 6 +++--- docs/developer/jax_frontend/quantization.rst | 6 +++++- docs/developer/linear_walkthrough.rst | 8 +++++--- .../pytorch_frontend/autograd_integration.rst | 13 +++++-------- .../pytorch_frontend/cpp_extensions.rst | 3 ++- docs/developer/quantization/class_hierarchy.rst | 3 +++ 9 files changed, 51 insertions(+), 16 deletions(-) diff --git a/docs/developer/attention/backend_selection.rst b/docs/developer/attention/backend_selection.rst index 4bae7d9b49..e479f34e8b 100644 --- a/docs/developer/attention/backend_selection.rst +++ b/docs/developer/attention/backend_selection.rst @@ -33,6 +33,22 @@ groups, etc.) and then applies a priority order: FusedAttention is preferred on Hopper+ for performance reasons. On pre-Hopper hardware, FlashAttention takes priority when both are available. +SM120 (Blackwell) Architecture Notes +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +SM120 GPUs have additional restrictions that affect backend selection: + +- **KV caching**: Both FusedAttention and FlashAttention are disabled when KV caching + (inference mode) is active. Only the Unfused backend is available for KV-cached + inference on SM120. +- **cuDNN version**: FusedAttention requires cuDNN >= 9.18.1 for THD layout on SM120. + With older cuDNN versions, FusedAttention is disabled for THD. +- **Layout restrictions**: Even with cuDNN >= 9.18.1, the T3HD and TH3D layouts + (interleaved Q/K/V) are not supported on SM120. +- **Deterministic mode**: Deterministic attention is not supported on SM120. +- **Softmax LSE packed format**: The packed softmax LSE format for THD is excluded on + SM120 (requires cuDNN >= 9.6.0 and a non-SM120 architecture). + Environment Variables --------------------- diff --git a/docs/developer/attention/backends.rst b/docs/developer/attention/backends.rst index d0b51f0c0b..114d4ced1d 100644 --- a/docs/developer/attention/backends.rst +++ b/docs/developer/attention/backends.rst @@ -115,6 +115,12 @@ specifically the ``fused_attn_*_fwd_impl()`` / ``fused_attn_*_bwd_impl()`` funct where the cuDNN graph is built. For backend selection logic, see ``nvte_get_fused_attn_backend()`` in ``fused_attn.cpp``. +.. note:: + + SM120 (Blackwell) has additional validation in the C++ dispatcher: cuDNN >= 9.18.1 is + required, deterministic mode is not supported, and T3HD/TH3D interleaved layouts are + not available. See :doc:`backend_selection` for the full set of SM120 restrictions. + Helper CUDA Kernels ^^^^^^^^^^^^^^^^^^^ diff --git a/docs/developer/attention/context_parallel.rst b/docs/developer/attention/context_parallel.rst index 9632b8229d..127a7557d9 100644 --- a/docs/developer/attention/context_parallel.rst +++ b/docs/developer/attention/context_parallel.rst @@ -186,6 +186,12 @@ CP support varies by attention backend: - **FlashAttention**: CP support via P2P and A2A. - **Unfused**: No CP support. +.. note:: + + On SM120 (Blackwell), the ragged softmax stats optimization used for THD layout in + context parallelism is not available (cuDNN limitation). The implementation falls back + to the dense stats format on SM120. + The CP strategy is configured via ``DotProductAttention``: .. code-block:: python diff --git a/docs/developer/attention/fused_attn_kernels.rst b/docs/developer/attention/fused_attn_kernels.rst index 29cb453672..bdc11a1db2 100644 --- a/docs/developer/attention/fused_attn_kernels.rst +++ b/docs/developer/attention/fused_attn_kernels.rst @@ -105,9 +105,9 @@ produce different auxiliary tensors, they are passed through an ``NVTETensorPack Each sub-backend populates the pack differently: - **Sub-backend 0** (F16 max512): 1 tensor — ``S`` (full softmax intermediate). -- **Sub-backend 1** (F16 arbitrary): 2+ tensors — softmax stats (``S`` or - ``Max``/``Sum_Exp`` depending on ``return_max_logit``), ``rng_state``, and optionally - ``Bias`` and ``SoftmaxOffset``. +- **Sub-backend 1** (F16 arbitrary): 2+ tensors — ``Stats`` (log-sum-exp, always + present), optionally ``Max`` (when ``return_max_logit=True``), ``rng_state``, and + optionally ``Bias`` and ``SoftmaxOffset``. - **Sub-backend 2** (FP8): 3 tensors — ``M`` (row max), ``ZInv`` (inverse softmax denominator), ``rng_state``. diff --git a/docs/developer/jax_frontend/quantization.rst b/docs/developer/jax_frontend/quantization.rst index 59e7805772..3922fdb474 100644 --- a/docs/developer/jax_frontend/quantization.rst +++ b/docs/developer/jax_frontend/quantization.rst @@ -51,7 +51,11 @@ Instead of PyTorch's ``QuantizedTensor`` (which is a ``torch.Tensor`` subclass), - ``ScaledTensor2x`` — Dual-layout tensor wrapping both a rowwise and columnwise ``ScaledTensor1x``. -A ``NoScaleTensor`` type wraps unquantized data with optional amax tracking. +A ``NoScaleTensor`` type wraps unquantized data with optional amax tracking. For grouped +operations, ``GroupedScaledTensor1x`` extends the single-layout type with support for +ragged grouping (specified via ``first_dims`` and ``last_dims`` arrays that describe +per-group sizes along each dimension), and ``GroupedNoScaleTensor`` provides an +unquantized grouped tensor for use with grouped GEMM. Since JAX arrays are immutable, these are simple data classes (not JAX array subclasses), registered as pytree nodes so JAX can trace through them: diff --git a/docs/developer/linear_walkthrough.rst b/docs/developer/linear_walkthrough.rst index c3f126028a..b164a56e76 100644 --- a/docs/developer/linear_walkthrough.rst +++ b/docs/developer/linear_walkthrough.rst @@ -328,10 +328,12 @@ layouts to keep. ctx.save_for_backward(*tensors_to_save) ctx.tensor_objects = tensor_objects -The ``prepare_for_saving`` / ``restore_from_saved`` pair handles the fact that +The ``prepare_for_saving`` / ``restore_from_func_ctx`` pair handles the fact that ``QuantizedTensorStorage`` is not a ``torch.Tensor`` and therefore cannot be passed -directly to ``ctx.save_for_backward()``. It splits each storage into metadata and raw -tensors so that PyTorch can manage their lifetime correctly. See +directly to ``ctx.save_for_backward()``. ``prepare_for_saving`` splits each storage into +metadata and raw tensors so that PyTorch can manage their lifetime correctly. +``restore_from_func_ctx`` reassembles the storages and automatically cleans up +``ctx.tensor_objects`` references to avoid holding stale tensor references. See :doc:`pytorch_frontend/autograd_integration` for the full explanation of why this is necessary and how it works. diff --git a/docs/developer/pytorch_frontend/autograd_integration.rst b/docs/developer/pytorch_frontend/autograd_integration.rst index 7a6d57c767..117d87fae8 100644 --- a/docs/developer/pytorch_frontend/autograd_integration.rst +++ b/docs/developer/pytorch_frontend/autograd_integration.rst @@ -89,9 +89,9 @@ Since ``QuantizedTensorStorage`` is not a ``torch.Tensor`` (see - Its **metadata** (with all tensor fields set to ``None``) — stored on ``ctx.tensor_objects``. - A list of raw **``torch.Tensor`` objects** — passed through ``ctx.save_for_backward()``. -In the backward pass, ``restore_from_saved()`` reassembles the original -``QuantizedTensorStorage`` objects from the saved tensors and metadata. After -reassembling, ``ctx.tensor_objects`` must be deleted to avoid keeping references to the +In the backward pass, ``restore_from_func_ctx()`` reassembles the original +``QuantizedTensorStorage`` objects from the saved tensors and metadata. It also +automatically deletes ``ctx.tensor_objects`` to avoid keeping references to the reassembled tensors on ``ctx`` (which would defeat the purpose of using ``save_for_backward`` in the first place). @@ -102,11 +102,8 @@ reassembled tensors on ``ctx`` (which would defeat the purpose of using ctx.save_for_backward(*tensors_to_save) ctx.tensor_objects = tensor_objects - # Backward: reassemble and release ctx references - inputmat, weightmat, weight, bias = restore_from_saved( - ctx.saved_tensors, ctx.tensor_objects, - ) - del ctx.tensor_objects # Avoid holding tensor references on ctx + # Backward: reassemble and release ctx references in one call + inputmat, weightmat, weight, bias = restore_from_func_ctx(ctx) Activation Recomputation ------------------------ diff --git a/docs/developer/pytorch_frontend/cpp_extensions.rst b/docs/developer/pytorch_frontend/cpp_extensions.rst index e4672960bf..0ab2027a1c 100644 --- a/docs/developer/pytorch_frontend/cpp_extensions.rst +++ b/docs/developer/pytorch_frontend/cpp_extensions.rst @@ -36,7 +36,8 @@ Python Side Only a subset of C++ extensions have dedicated Python wrapper modules: -- ``gemm.py`` — ``general_gemm()``, ``grouped_gemm()`` +- ``gemm.py`` — ``general_gemm()``, ``grouped_gemm()``, + ``general_grouped_gemm_for_grouped_tensor()`` - ``fused_attn.py`` — ``fused_attn_fwd()``, ``fused_attn_bwd()`` Most other C++ extensions (normalization, activation, cast, transpose, etc.) are exposed diff --git a/docs/developer/quantization/class_hierarchy.rst b/docs/developer/quantization/class_hierarchy.rst index 3bfd626676..d2bc533d8c 100644 --- a/docs/developer/quantization/class_hierarchy.rst +++ b/docs/developer/quantization/class_hierarchy.rst @@ -140,6 +140,9 @@ support and minimal overhead. def quantize_(self, tensor, quantizer): ... # In-place quantization def copy_from_storage(self, other): ... + # Module-level helper (preferred over calling restore_from_saved directly) + def restore_from_func_ctx(func_ctx) -> list: ... # Reassemble + cleanup ctx + Storage objects are used internally by: - C++ extension calls (converted to ``NVTETensor`` for kernel dispatch) From c59e41b97e85aab8568e28731debaeb8d3876134 Mon Sep 17 00:00:00 2001 From: Przemek Tredak Date: Wed, 8 Apr 2026 17:13:13 -0700 Subject: [PATCH 13/20] Fixes to scaling factor overview Signed-off-by: Przemek Tredak --- docs/developer/cpp_core/scaling_modes.rst | 87 ++++++++----------- transformer_engine/common/common.h | 3 - .../common/gemm/cublaslt_gemm.cu | 4 +- .../common/normalization/layernorm/ln_api.cpp | 4 +- .../normalization/rmsnorm/rmsnorm_api.cpp | 4 +- 5 files changed, 41 insertions(+), 61 deletions(-) diff --git a/docs/developer/cpp_core/scaling_modes.rst b/docs/developer/cpp_core/scaling_modes.rst index d9012da67e..a5a4708af7 100644 --- a/docs/developer/cpp_core/scaling_modes.rst +++ b/docs/developer/cpp_core/scaling_modes.rst @@ -39,23 +39,21 @@ Non-TN GEMM Support --------------------- A cross-cutting concern that affects all scaling modes is non-TN GEMM support. -cuBLASLt traditionally required FP8 GEMMs to use TN (transpose-normal) layout, which -is why the forward pass needs rowwise data for one operand and columnwise data for -the other. On architectures that support non-TN FP8 GEMM (compute capability 10.x and -13.0+, i.e. Blackwell), this constraint is relaxed — cuBLASLt can consume rowwise data -for both operands. +On Hopper and Ada, hardware only supports TN layout for FP8 GEMMs, +which is why the forward pass needs rowwise data for one operand and columnwise data for +the other. On Blackwell (compute capability 10.x), non-TN FP8 GEMM is supported for FP8 +data formats (but not FP4), so cuBLASLt can consume rowwise data for both operands. This has a significant impact on data layout: when non-TN GEMM is supported, -``data`` and ``columnwise_data`` may point to the same buffer (since a separate -transposed copy is no longer needed for the GEMM). When non-TN is not supported, they -must be separate buffers with physically transposed data. +``columnwise_data`` may point to the same buffer as ``data`` or have the same rowwise +layout (since the transposed version is no longer needed for the GEMM). When non-TN +is not supported, they must be separate buffers with physically transposed data. Tensor Scaling (``NVTE_DELAYED_TENSOR_SCALING``) -------------------------------------------------- A single scale factor applies to the entire tensor. Minimum architecture: Ada (sm89) for -FP8 data types, though practical FP8 training with cuBLASLt GEMMs is typically done on -Hopper (sm90) and later. +FP8 data types. .. note:: @@ -70,7 +68,7 @@ Data layout: - ``data``: row-major, one scale for the entire tensor - ``columnwise_data``: column-major layout, same single scale. On architectures without non-TN FP8 GEMM support this is a physically transposed copy; on architectures with - non-TN support (Blackwell+) it may alias ``data`` + non-TN support (Blackwell+) it is empty and the GEMM uses ``data`` directly - ``scale_inv``: shape ``[1]`` (one element) for FP8 tensors, empty for high-precision - ``amax``: populated by delayed scaling (framework responsibility), unused by current scaling @@ -81,31 +79,30 @@ C++ helpers: // True for NVTE_DELAYED_TENSOR_SCALING (per-tensor scaling, including current) bool is_tensor_scaling(const NVTEScalingMode &mode); - // Alias — identical behavior - bool is_delayed_tensor_scaling(const NVTEScalingMode &mode); MXFP8 1D Scaling (``NVTE_MXFP8_1D_SCALING``) ---------------------------------------------- -Microscaling FP8 format per the OCP MX specification. Each block of 32 contiguous -elements shares a single E8M0 (8-bit exponent-only) scale factor. Minimum architecture: -Blackwell. +Microscaling FP8 format inspired by the OCP MX specification (differs in the rounding +mode used to compute scaling factors). Each block of 32 consecutive elements along the +K dimension (contracting dimension) of the GEMM shares a single E8M0 (8-bit +exponent-only) scale factor. Minimum architecture: Blackwell. Data layout: - ``data``: row-major FP8, shape ``[M, N]`` -- ``columnwise_data``: the columnwise layout, shape ``[N, M]`` — on architectures - without non-TN FP8 GEMM support this is a physically transposed separate buffer - with independently quantized blocks along the transposed dimension; on architectures - with non-TN support (Blackwell+) it may share the rowwise buffer -- ``scale_inv`` (rowwise): shape ``[ceil(M), ceil(N/32)]``, padded to ``[128, 4]`` +- ``columnwise_data``: rowwise in memory with shape ``[M, N]``, but quantized along the + column direction (i.e. blocks of 32 elements are formed along columns). This is a + separate buffer from ``data`` since the quantization direction differs +- ``scale_inv``: shape ``[ceil(M), ceil(N/32)]``, padded to ``[128, 4]`` multiples for GEMM alignment - ``columnwise_scale_inv``: shape ``[ceil(N), ceil(M/32)]``, padded to ``[128, 4]`` - multiples + multiples. Note that this scale layout differs from the ``columnwise_data`` tensor + layout -Scales may need "GEMM swizzling" — a specific memory layout that cuBLASLt expects for -block-scaled operands. See the `cuBLAS documentation on block scaling factors layout -`_. +Scales, before being passed to the GEMM, may need "swizzling" — a specific memory layout that +cuBLASLt expects for block-scaled operands. See the `cuBLAS documentation on block scaling factors +layout `_. The ``with_gemm_swizzled_scales`` flag on ``Tensor`` tracks whether scales have been swizzled. @@ -116,15 +113,17 @@ C++ helper: bool is_mxfp8_scaling(const NVTEScalingMode &mode); Block Scaling (``NVTE_BLOCK_SCALING_1D`` / ``NVTE_BLOCK_SCALING_2D``) ----------------------------------------------------------------------- +--------------------------------------------------------------------- Per-block FP8 scaling with a fixed block size of 128 (matching the DeepSeek v3 recipe). ``1D`` uses 1×128 blocks (one scale per row-slice of 128 elements), while ``2D`` uses 128×128 blocks. The block size itself is not configurable — only the 1D vs 2D aspect is selectable via the recipe. Minimum architecture: Blackwell. -Scales are FP32 (more precise than MXFP8's E8M0). Like MXFP8, scales may require GEMM -swizzling. +Scales are FP32 (more precise than MXFP8's E8M0), but can be limited to +powers of 2 only, which is the default behavior. This is needed to enable emulation of the block +scaling recipe with MXFP8 GEMMs on Blackwell, which have higher throughput but do not support +non-power-of-2 scaling factors. Like MXFP8, scales require GEMM swizzling. C++ helper: @@ -137,18 +136,18 @@ C++ helper: NVFP4 1D Scaling (``NVTE_NVFP4_1D_SCALING``) ---------------------------------------------- -4-bit floating point with E8M0 block scales, targeting maximum compression. Minimum -architecture: Blackwell. +4-bit floating point with E4M3 block scales and a second-level per-tensor FP32 scale factor. +Minimum architecture: Blackwell. Data layout: -- ``data``: FP4 E2M1 format, packed 2 elements per byte (uint8 storage) -- ``columnwise_data``: physically transposed FP4 data -- ``scale_inv`` (rowwise): E8M0 format, shape ``[ceil(M/16), ceil(N/16)]``, padded to +- ``data``: FP4 E2M1 format, packed 2 elements per byte (uint8 storage), shape ``[M, N]`` +- ``columnwise_data``: physically transposed FP4 data, shape ``[N, M]`` +- ``scale_inv`` (rowwise): E4M3 format, shape ``[M, ceil(N/16)]``, padded to ``[4, 128]`` multiples for GEMM alignment -- ``columnwise_scale_inv``: shape ``[ceil(N/16), ceil(M/16)]``, padded to ``[128, 4]`` +- ``columnwise_scale_inv``: shape ``[N, ceil(M/16)]``, padded to ``[128, 4]`` -Like MXFP8, scales may require GEMM swizzling. +Like MXFP8, scales require GEMM swizzling. C++ helper: @@ -165,27 +164,11 @@ layouts: - ``scale_inv`` — scale inverses for decoding rowwise data (used in forward GEMM). - ``columnwise_scale_inv`` — scale inverses for decoding columnwise data (used in wgrad - GEMM). + and dgrad GEMMs). See :doc:`/developer/quantization/rowwise_columnwise` for how these layouts interact with GEMM execution. -Checking Scaling Mode in Code ------------------------------ - -The ``common.h`` header provides inline helpers that should be preferred over direct -enum comparison: - -.. code-block:: cpp - - // Prefer this: - if (is_mxfp8_scaling(tensor.scaling_mode)) { ... } - - // Over this: - if (tensor.scaling_mode == NVTE_MXFP8_1D_SCALING) { ... } - -This makes code resilient to future additions to the ``NVTEScalingMode`` enum. - Invalid Scaling Sentinel ^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index 71cfe6dfca..7e97a51d78 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -79,9 +79,6 @@ inline bool is_tensor_scaling(const NVTEScalingMode &mode) { inline bool is_block_scaling(const NVTEScalingMode &mode) { return !is_tensor_scaling(mode); } -inline bool is_delayed_tensor_scaling(const NVTEScalingMode &mode) { - return mode == NVTE_DELAYED_TENSOR_SCALING; -} inline bool is_nvfp4_scaling(const NVTEScalingMode &mode) { return mode == NVTE_NVFP4_1D_SCALING; } diff --git a/transformer_engine/common/gemm/cublaslt_gemm.cu b/transformer_engine/common/gemm/cublaslt_gemm.cu index 4fefadebab..5747b68919 100644 --- a/transformer_engine/common/gemm/cublaslt_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_gemm.cu @@ -956,8 +956,8 @@ void nvte_cublas_atomic_gemm(const NVTETensor A, const NVTETensor B, NVTETensor const void *alpha_ptr = GetScalarOne(); const void *beta_ptr = accumulate ? GetScalarOne() : GetScalarZero(); - NVTE_CHECK(is_delayed_tensor_scaling(inputA->scaling_mode) && - is_delayed_tensor_scaling(inputB->scaling_mode), + NVTE_CHECK(is_tensor_scaling(inputA->scaling_mode) && + is_tensor_scaling(inputB->scaling_mode), "Atomic GEMM only supports delayed scaling."); cublas_gemm(inputA, inputB, outputD, biasTensor, outputGelu, (transa) ? CUBLAS_OP_T : CUBLAS_OP_N, (transb) ? CUBLAS_OP_T : CUBLAS_OP_N, grad, wspace->data.dptr, wspace->data.shape[0], diff --git a/transformer_engine/common/normalization/layernorm/ln_api.cpp b/transformer_engine/common/normalization/layernorm/ln_api.cpp index 7bd5a1bbd0..2ac162618c 100644 --- a/transformer_engine/common/normalization/layernorm/ln_api.cpp +++ b/transformer_engine/common/normalization/layernorm/ln_api.cpp @@ -28,7 +28,7 @@ void layernorm_fwd(const Tensor& x, // BxSxhidden_size const int multiprocessorCount, const bool zero_centered_gamma, cudaStream_t stream) { // Check for unsupported configurations - if (is_fp8_dtype(z->data.dtype) && !is_delayed_tensor_scaling(z->scaling_mode) && + if (is_fp8_dtype(z->data.dtype) && !is_tensor_scaling(z->scaling_mode) && !is_mxfp8_scaling(z->scaling_mode)) { NVTE_ERROR("Not implemented scaling mode: " + to_string(z->scaling_mode) + "."); } @@ -87,7 +87,7 @@ void layernorm_fwd(const Tensor& x, // BxSxhidden_size } bool training = - is_delayed_tensor_scaling(z->scaling_mode) || (z->columnwise_data).dptr != nullptr; + is_tensor_scaling(z->scaling_mode) || (z->columnwise_data).dptr != nullptr; auto plan = NormalizationPlanRegistry::getInstance().getNormalizationPlan( norm_backend, NVTE_Norm_Type::LayerNorm, NVTE_Norm_Stage::Forward, diff --git a/transformer_engine/common/normalization/rmsnorm/rmsnorm_api.cpp b/transformer_engine/common/normalization/rmsnorm/rmsnorm_api.cpp index adf2ccee04..67664cabd7 100644 --- a/transformer_engine/common/normalization/rmsnorm/rmsnorm_api.cpp +++ b/transformer_engine/common/normalization/rmsnorm/rmsnorm_api.cpp @@ -24,7 +24,7 @@ void rmsnorm_fwd(const Tensor &x, const Tensor &gamma, const float epsilon, Tens Tensor *rsigma, Tensor *workspace, const int multiprocessorCount, const bool zero_centered_gamma, cudaStream_t stream) { // Check for unsupported configurations - if (is_fp8_dtype(z->data.dtype) && !is_delayed_tensor_scaling(z->scaling_mode) && + if (is_fp8_dtype(z->data.dtype) && !is_tensor_scaling(z->scaling_mode) && !is_mxfp8_scaling(z->scaling_mode)) { NVTE_ERROR("Not implemented scaling mode: " + to_string(z->scaling_mode) + "."); } @@ -62,7 +62,7 @@ void rmsnorm_fwd(const Tensor &x, const Tensor &gamma, const float epsilon, Tens } bool training = - is_delayed_tensor_scaling(z->scaling_mode) || (z->columnwise_data).dptr != nullptr; + is_tensor_scaling(z->scaling_mode) || (z->columnwise_data).dptr != nullptr; bool gamma_in_weight_dtype = false; if (cudnn_backend) { From 4be5a50f5c673a9c582d73500841c29e4565c946 Mon Sep 17 00:00:00 2001 From: Przemek Tredak Date: Tue, 14 Apr 2026 10:49:33 -0700 Subject: [PATCH 14/20] More fixes Signed-off-by: Przemek Tredak --- .../pytorch_frontend/autograd_integration.rst | 48 ++++++++----------- 1 file changed, 20 insertions(+), 28 deletions(-) diff --git a/docs/developer/pytorch_frontend/autograd_integration.rst b/docs/developer/pytorch_frontend/autograd_integration.rst index 117d87fae8..c52ab37b60 100644 --- a/docs/developer/pytorch_frontend/autograd_integration.rst +++ b/docs/developer/pytorch_frontend/autograd_integration.rst @@ -35,8 +35,12 @@ This autograd function wraps the forward GEMM and defines the corresponding back # 3. Forward GEMM: output = qinput @ qweight^T output = general_gemm(qinput, qweight, bias=bias) - # 4. Save for backward - ctx.save_for_backward(qinput, qweight, ...) + # 4. Save for backward — prepare_for_saving splits QuantizedTensor + # and QuantizedTensorStorage objects into raw torch.Tensors + # (for save_for_backward) and metadata (for ctx.tensor_objects) + tensors_to_save, tensor_objects = prepare_for_saving(qinput, qweight, weight, bias) + ctx.save_for_backward(*tensors_to_save) + ctx.tensor_objects = tensor_objects return output @@ -46,7 +50,9 @@ This autograd function wraps the forward GEMM and defines the corresponding back @staticmethod def backward(ctx, grad_output): - qinput, qweight, ... = ctx.saved_tensors + # restore_from_func_ctx reassembles the QuantizedTensorStorage + # objects from saved tensors + metadata, and deletes ctx.tensor_objects + qinput, qweight, weight, bias = restore_from_func_ctx(ctx) # 1. Quantize grad_output qgrad = grad_quantizer(grad_output) @@ -67,12 +73,9 @@ What gets saved for backward depends on the configuration: **FP8 disabled**: Standard PyTorch behavior — save full-precision input and weight. -**FP8 enabled (no recompute)**: Save the ``QuantizedTensor`` objects from forward. These -contain both rowwise data (used by dgrad GEMM) and columnwise data (used by wgrad GEMM). -Memory cost: ~2× the FP8 data size (rowwise + columnwise). - -Saving QuantizedTensorStorage via prepare_for_saving -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +**FP8 enabled (no recompute)**: Save the ``QuantizedTensor`` or ``QuantizedTensorStorage`` +objects from forward. These contain both rowwise data (used by dgrad GEMM) and columnwise +data (used by wgrad GEMM). Memory cost: ~2× the FP8 data size (rowwise + columnwise). PyTorch offers two ways to pass data from forward to backward: direct attributes on ``ctx``, and ``ctx.save_for_backward()``. Direct ``ctx`` attributes are not released after @@ -84,26 +87,15 @@ promptly, but it only accepts ``torch.Tensor`` objects. Since ``QuantizedTensorStorage`` is not a ``torch.Tensor`` (see :doc:`/developer/quantization/class_hierarchy` for the distinction between ``QuantizedTensorStorage`` and ``QuantizedTensor``), the helper function -``prepare_for_saving()`` (in ``quantized_tensor.py``) splits each storage into: +``prepare_for_saving()`` (in ``quantized_tensor.py``) splits each object into: - Its **metadata** (with all tensor fields set to ``None``) — stored on ``ctx.tensor_objects``. - A list of raw **``torch.Tensor`` objects** — passed through ``ctx.save_for_backward()``. -In the backward pass, ``restore_from_func_ctx()`` reassembles the original -``QuantizedTensorStorage`` objects from the saved tensors and metadata. It also -automatically deletes ``ctx.tensor_objects`` to avoid keeping references to the -reassembled tensors on ``ctx`` (which would defeat the purpose of using -``save_for_backward`` in the first place). - -.. code-block:: python - - # Forward: split and save - tensors_to_save, tensor_objects = prepare_for_saving(inputmat, weightmat, weight, bias) - ctx.save_for_backward(*tensors_to_save) - ctx.tensor_objects = tensor_objects - - # Backward: reassemble and release ctx references in one call - inputmat, weightmat, weight, bias = restore_from_func_ctx(ctx) +In the backward pass, ``restore_from_func_ctx()`` reassembles the original objects from +the saved tensors and metadata, as shown in the pseudocode above. It also automatically +deletes ``ctx.tensor_objects`` to avoid keeping references to the reassembled tensors on +``ctx`` (which would defeat the purpose of using ``save_for_backward`` in the first place). Activation Recomputation ------------------------ @@ -136,9 +128,9 @@ __torch_dispatch__ and Automatic Dequantization This is the most complex part of the ``QuantizedTensor`` class. The dispatch logic handles three cases: -1. **TE-optimized operations** (e.g., GEMM): These recognize ``QuantizedTensor`` inputs - and extract the raw FP8 data + scales directly via ``get_data_tensors()``, avoiding - any dequantization. +1. **TE-optimized operations** (e.g., GEMM): ``QuantizedTensor`` and + ``QuantizedTensorStorage`` objects are handled natively by the C++ extensions, + so no dequantization occurs. 2. **Non-mutable operations**: For standard PyTorch ops that don't modify their inputs (e.g., ``torch.add``, ``torch.matmul``), the dispatch automatically dequantizes all From dc1020b07f59b7b2743eb69d070dee201d24c3e3 Mon Sep 17 00:00:00 2001 From: Przemek Tredak Date: Wed, 1 Jul 2026 07:17:50 -0700 Subject: [PATCH 15/20] Updating the docs with the latest main Signed-off-by: Przemek Tredak --- .../developer/attention/backend_selection.rst | 28 ++++++------ docs/developer/attention/backends.rst | 13 +++--- .../attention/fused_attn_kernels.rst | 13 +++--- .../attention/img/attention_backends.svg | 16 ++----- docs/developer/attention/jax_attention.rst | 44 ++++++++++++++----- docs/developer/cpp_core/build_system.rst | 12 +++-- docs/developer/cpp_core/img/tensor_struct.svg | 19 +++++--- docs/developer/cpp_core/type_system.rst | 16 +++++-- .../distributed/fsdp_integration.rst | 11 +++-- .../distributed/img/column_parallel.svg | 18 ++++---- .../distributed/img/row_parallel.svg | 4 +- .../developer/distributed/tensor_parallel.rst | 23 ++++++---- docs/developer/index.rst | 2 +- docs/developer/jax_frontend/module_system.rst | 36 +++++++++++---- .../jax_frontend/xla_ffi_primitives.rst | 11 +++-- .../pytorch_frontend/autograd_integration.rst | 6 +-- .../pytorch_frontend/transformer_layer.rst | 2 +- .../quantization/class_hierarchy.rst | 14 +++++- docs/envvars.rst | 6 --- .../dot_product_attention.py | 6 +-- 20 files changed, 187 insertions(+), 113 deletions(-) diff --git a/docs/developer/attention/backend_selection.rst b/docs/developer/attention/backend_selection.rst index e479f34e8b..c91f0cdbbd 100644 --- a/docs/developer/attention/backend_selection.rst +++ b/docs/developer/attention/backend_selection.rst @@ -62,24 +62,26 @@ Environment Variables - ``0`` to disable cuDNN fused attention, ``1`` to enable (default: ``1``) * - ``NVTE_FLASH_ATTN`` - ``0`` to disable FlashAttention, ``1`` to enable (default: ``1``) - * - ``NVTE_FUSED_ATTN_BACKEND`` - - Force a specific cuDNN fused attention sub-backend by integer ID: - ``0`` = F16 max512, ``1`` = F16 arbitrary seqlen, ``2`` = FP8. - These correspond to the ``NVTE_Fused_Attn_Backend`` enum in ``fused_attn.h``. + * - ``NVTE_UNFUSED_ATTN`` + - ``0`` to disable native PyTorch attention, ``1`` to enable (default: ``1``) -Constructor Override --------------------- +Selecting a Backend +------------------- -``DotProductAttention`` accepts a ``backend`` parameter: +``DotProductAttention`` does not expose a constructor argument that forces one backend. +To select a backend for debugging or testing, enable the desired backend and disable the +other two before the module is first called. For example, to select FlashAttention: .. code-block:: python - attn = te.DotProductAttention( - num_attention_heads=32, - kv_channels=128, - attention_type="self", - backend="flash_attention", # Force FlashAttention - ) + import os + + os.environ["NVTE_FLASH_ATTN"] = "1" + os.environ["NVTE_FUSED_ATTN"] = "0" + os.environ["NVTE_UNFUSED_ATTN"] = "0" + +The environment variables are read during backend selection. Set them before the first +``DotProductAttention.forward()`` call because the selection result is cached. Feature Compatibility Matrix ----------------------------- diff --git a/docs/developer/attention/backends.rst b/docs/developer/attention/backends.rst index 114d4ced1d..8beb23ad04 100644 --- a/docs/developer/attention/backends.rst +++ b/docs/developer/attention/backends.rst @@ -26,13 +26,11 @@ hardware, precision, and feature combinations. ├── "FlashAttention" — Tri Dao's flash-attn package (external) │ └── Ampere+, BF16/FP16, no cuDNN required └── "FusedAttention" — cuDNN graph API - ├── Sub-backend 0: "F16 max512" — seqlen ≤ 512, Ampere+ ├── Sub-backend 1: "F16 arbitrary" — any seqlen, Ampere+ └── Sub-backend 2: "FP8" — FP8 Q/K/V, Hopper+ There are exactly three backends (Unfused, FlashAttention, FusedAttention). -FusedAttention has three cuDNN sub-backends selected automatically based on precision -and sequence length. +FusedAttention has two cuDNN sub-backends selected automatically based on precision. Backend Details --------------- @@ -58,11 +56,14 @@ FlashAttention ``transformer_engine/pytorch/attention/dot_product_attention/backends.py`` Integrates Tri Dao's `flash-attn `_ -package as an external dependency (not bundled). Two versions are supported: +package as an external dependency (not bundled). Three versions are supported: - **FA2** (``flash_attn.flash_attn_interface``): Ampere and later (sm80+). - **FA3** (``flash_attn_3.flash_attn_interface``): Hopper only (sm90). Selected automatically when installed and running on sm90; disabled on other architectures. +- **FA4** (``flash_attn.cute``): Ampere, Hopper, Blackwell datacenter, and Blackwell + GeForce architectures (sm80, sm90, sm100, and sm120). On Hopper, FA3 is preferred + over FA4 when both are installed. Version constraints are managed in ``FlashAttentionUtils`` in ``utils.py``. @@ -99,9 +100,7 @@ Code flow: sub-backend implementation based on ``nvte_get_fused_attn_backend()``. 4. **cuDNN graph builders** — one file per sub-backend: - - ``fused_attn_f16_max512_seqlen.cu`` — sub-backend 0 (legacy; only selected when - sub-backend 1 is unavailable due to older cuDNN or restrictive parameters) - - ``fused_attn_f16_arbitrary_seqlen.cu`` — sub-backend 1 (preferred for all F16 cases) + - ``fused_attn_f16_arbitrary_seqlen.cu`` — sub-backend 1 (F16/BF16) - ``fused_attn_fp8.cu`` — sub-backend 2 Each builds a ``cudnn_frontend::graph::Graph``, configures SDPA attributes (masking, diff --git a/docs/developer/attention/fused_attn_kernels.rst b/docs/developer/attention/fused_attn_kernels.rst index bdc11a1db2..bc4da4cef6 100644 --- a/docs/developer/attention/fused_attn_kernels.rst +++ b/docs/developer/attention/fused_attn_kernels.rst @@ -38,13 +38,13 @@ Directory Structure .. code-block:: text fused_attn/ - ├── fused_attn.h # Internal C++ API ├── fused_attn.cpp # C API implementation, backend dispatch - ├── fused_attn_f16_max512_seqlen.h/.cu # cuDNN F16 (short sequences) ├── fused_attn_f16_arbitrary_seqlen.h/.cu # cuDNN F16 (any sequence length) ├── fused_attn_fp8.h/.cu # cuDNN FP8 - ├── thd_utils.h # THD layout utilities for CP - └── utils.h # Shared utilities + ├── flash_attn.cu # FlashAttention-related helpers + ├── context_parallel.cu # Context-parallel helper kernels + ├── kv_cache.cu # KV-cache helper kernels + └── utils.cu/.h # Shared utilities cuDNN Integration ----------------- @@ -104,12 +104,11 @@ produce different auxiliary tensors, they are passed through an ``NVTETensorPack Each sub-backend populates the pack differently: -- **Sub-backend 0** (F16 max512): 1 tensor — ``S`` (full softmax intermediate). - **Sub-backend 1** (F16 arbitrary): 2+ tensors — ``Stats`` (log-sum-exp, always present), optionally ``Max`` (when ``return_max_logit=True``), ``rng_state``, and optionally ``Bias`` and ``SoftmaxOffset``. -- **Sub-backend 2** (FP8): 3 tensors — ``M`` (row max), ``ZInv`` (inverse softmax - denominator), ``rng_state``. +- **Sub-backend 2** (FP8): ``M`` (softmax statistics), ``rng_state``, and optionally + ``SoftmaxOffset`` for non-vanilla softmax. On the Python/pybind11 side (``pytorch/csrc/extensions/attention.cpp``), the forward call uses the two-pass pattern: the first call with empty tensors discovers the required shapes diff --git a/docs/developer/attention/img/attention_backends.svg b/docs/developer/attention/img/attention_backends.svg index de3013407d..817c48a4ad 100644 --- a/docs/developer/attention/img/attention_backends.svg +++ b/docs/developer/attention/img/attention_backends.svg @@ -64,7 +64,7 @@ Tiled, IO-aware algorithm Ampere+, BF16 / FP16 - O(N) memory, no materialization + FA2 / FA3 / FA4 @@ -78,7 +78,7 @@ - + @@ -100,16 +100,6 @@ Hopper+ | SM90+ Best FP8 performance - - - - - Custom / THD - - - - Hopper+ | SM90+ - CP-specific kernels @@ -120,5 +110,5 @@ FlashAttention cuDNN Fused - SM80 = Ampere (A100) | SM90 = Hopper (H100) | SM70 = Volta (V100) + SM80 = Ampere | SM90 = Hopper | SM100/SM120 = Blackwell \ No newline at end of file diff --git a/docs/developer/attention/jax_attention.rst b/docs/developer/attention/jax_attention.rst index 9309320440..ca48c82d95 100644 --- a/docs/developer/attention/jax_attention.rst +++ b/docs/developer/attention/jax_attention.rst @@ -14,14 +14,15 @@ Entry Points **Location**: ``transformer_engine/jax/attention.py`` -The primary functions are: +The primary public function is: - ``fused_attn()`` — Fused attention using cuDNN (equivalent to PyTorch's cuDNN fused backend). -- ``fused_attn_fwd()`` / ``fused_attn_bwd()`` — Forward and backward primitives. -These are JAX primitives registered via the XLA FFI mechanism -(see :doc:`/developer/jax_frontend/xla_ffi_primitives`). +It is a ``jax.custom_vjp`` function backed by the low-level ``fused_attn_fwd()`` and +``fused_attn_bwd()`` wrappers in ``transformer_engine/jax/cpp_extensions/attention.py``. +Those wrappers invoke JAX primitives registered through XLA FFI (see +:doc:`/developer/jax_frontend/xla_ffi_primitives`). Differences from PyTorch ------------------------- @@ -54,14 +55,37 @@ Usage Example .. code-block:: python - from transformer_engine.jax.attention import fused_attn + import jax.numpy as jnp + from transformer_engine.jax.attention import ( + AttnBiasType, + AttnMaskType, + AttnSoftmaxType, + QKVLayout, + SequenceDescriptor, + fused_attn, + ) + + batch, seqlen, heads, head_dim = 2, 128, 16, 64 + qkv = jnp.zeros((batch, seqlen, 3, heads, head_dim), dtype=jnp.bfloat16) + sequence_descriptor = SequenceDescriptor.from_seqlens( + seqlens=( + jnp.full((batch,), seqlen), + jnp.full((batch,), seqlen), + ) + ) output = fused_attn( - query, key, value, - bias=attn_bias, - mask=attn_mask, - scaling_factor=1.0 / math.sqrt(head_dim), - is_training=True, + (qkv,), + None, # bias + sequence_descriptor, + None, # dropout seed; no dropout in this example + AttnBiasType.NO_BIAS, + AttnMaskType.CAUSAL_MASK, + QKVLayout.BS3HD, + AttnSoftmaxType.VANILLA_SOFTMAX, + head_dim**-0.5, + 0.0, # dropout probability + True, # is_training ) See Also diff --git a/docs/developer/cpp_core/build_system.rst b/docs/developer/cpp_core/build_system.rst index ac99526f83..5967193419 100644 --- a/docs/developer/cpp_core/build_system.rst +++ b/docs/developer/cpp_core/build_system.rst @@ -78,6 +78,8 @@ Key Environment Variables - Path to ccache binary (default: ``ccache``). E.g. ``sccache`` for CI. * - ``NVTE_BUILD_DEBUG`` - Set to ``1`` to build C++ extensions with debug symbols + * - ``NVTE_WITH_NCCL_EP`` + - Build NCCL-based expert-parallel support (default: ``1``; requires an sm90+ target) NVTE_CUDA_ARCHS ^^^^^^^^^^^^^^^^ @@ -101,15 +103,19 @@ For development, set this to only your target GPU to minimize build time: Git Submodules -------------- -The build requires two submodules in ``3rdparty/``: +The repository uses four submodules in ``3rdparty/``: - **CUTLASS** (``3rdparty/cutlass``) — NVIDIA's CUDA Templates for Linear Algebra Subroutines. Used for custom GEMM kernels and attention kernels. - **cuDNN Frontend** (``3rdparty/cudnn-frontend``) — C++ frontend for cuDNN graph API. Used for fused attention and normalization. +- **NCCL** (``3rdparty/nccl``) — Builds the static NCCL expert-parallel component when + ``NVTE_WITH_NCCL_EP=1`` and at least one target architecture is sm90 or newer. Set + ``NVTE_WITH_NCCL_EP=0`` to omit it. +- **GoogleTest** (``3rdparty/googletest``) — Used by the C++ test suite. -``setup.py`` automatically initializes these if they are missing, but manual setup is -sometimes needed: +``setup.py`` attempts to initialize every registered submodule unless +``NVTE_SKIP_SUBMODULE_CHECKS_DURING_BUILD`` is set. Manual setup is sometimes needed: .. code-block:: bash diff --git a/docs/developer/cpp_core/img/tensor_struct.svg b/docs/developer/cpp_core/img/tensor_struct.svg index d4adb519d1..187e310ff5 100644 --- a/docs/developer/cpp_core/img/tensor_struct.svg +++ b/docs/developer/cpp_core/img/tensor_struct.svg @@ -22,9 +22,9 @@ - + - + @@ -64,6 +64,11 @@ + with_gemm_swizzled_scales : bool + + row_scaled_nvfp4 + : bool + + + nvfp4_e4m3_max + : int @@ -86,7 +91,7 @@ : void* + shape - : vector<size_t> + : Shape + dtype : DType @@ -96,14 +101,14 @@ - + - + - + - + diff --git a/docs/developer/cpp_core/type_system.rst b/docs/developer/cpp_core/type_system.rst index 660373d877..cddd54210b 100644 --- a/docs/developer/cpp_core/type_system.rst +++ b/docs/developer/cpp_core/type_system.rst @@ -72,7 +72,7 @@ objects (see ``transformer_engine.cpp``). The pool has a fixed capacity of appro reuse. This means ``nvte_create_tensor()`` does not allocate memory — it retrieves an entry from the pool. If the pool is exhausted, an error is raised. -``NVTETensor`` is populated via ``nvte_tensor_set()`` with ``NVTETensorParam`` keys: +``NVTETensor`` is populated via ``nvte_set_tensor_param_v2()`` with ``NVTETensorParam`` keys: .. code-block:: c @@ -85,6 +85,9 @@ entry from the pool. If the pool is exhausted, an error is raised. kNVTEColumnwiseScaleInv = 5, kNVTEColumnwiseAmax = 6, kNVTEWithGEMMSwizzledScales = 7, + kNVTERowScaledNVFP4 = 8, + kNVTENVFP4E4M3Max = 9, + kNVTENumTensorParams, }; NVTEScalingMode @@ -174,7 +177,7 @@ A lightweight internal tensor descriptor with implicit conversion to/from struct SimpleTensor { void *dptr; - std::vector shape; + Shape shape; DType dtype; // Implicit conversion from NVTEBasicTensor @@ -215,9 +218,11 @@ The main internal tensor type, composed of multiple ``SimpleTensor`` members: + scaling_mode : NVTEScalingMode + nvte_tensor : NVTETensor (opaque handle) + with_gemm_swizzled_scales : bool + + row_scaled_nvfp4 : bool + + nvfp4_e4m3_max : int Below, a smaller UML box for "SimpleTensor": + dptr : void* - + shape : vector + + shape : Shape + dtype : DType Arrow from each SimpleTensor field in Tensor to the SimpleTensor box (composition). @@ -235,6 +240,8 @@ The main internal tensor type, composed of multiple ``SimpleTensor`` members: NVTEScalingMode scaling_mode; NVTETensor nvte_tensor; // Opaque handle for C API bool with_gemm_swizzled_scales; // MXFP8/NVFP4 scale format flag + bool row_scaled_nvfp4; // NVFP4 uses one amax per row + int nvfp4_e4m3_max; // Global E4M3 scale bound (448 or 256) }; This structure holds all the metadata needed for a quantized tensor. Not all fields are @@ -244,6 +251,9 @@ dtypes. ``scale_inv`` has shape ``[1]`` for tensor scaling but a multi-element s block scaling modes. ``columnwise_data`` and ``columnwise_scale_inv`` are only populated when the tensor was quantized with columnwise usage enabled. Validation functions in ``transformer_engine.cpp`` (e.g., ``CheckScaleTensorShape``) enforce these constraints. +For NVFP4, ``row_scaled_nvfp4`` changes ``amax`` from a per-tensor value to a per-row +array, while ``nvfp4_e4m3_max`` records the scale bound used by both quantization and +downstream consumers. .. _dtype-conversion-patterns: diff --git a/docs/developer/distributed/fsdp_integration.rst b/docs/developer/distributed/fsdp_integration.rst index c25b5a134b..40698ccb57 100644 --- a/docs/developer/distributed/fsdp_integration.rst +++ b/docs/developer/distributed/fsdp_integration.rst @@ -68,6 +68,10 @@ during weight all-gather. No explicit setup or registration is needed. - Gathers uint8 data + scale_inv. Handles amax reduction for current scaling. * - ``MXFP8Tensor`` - Gathers data + scale_inv with padding/unpadding for block-aligned scales. + * - ``Float8BlockwiseQTensor`` + - Supports 2D block scaling; 1D block scaling is not compatible with dim-0 all-gather. + * - ``NVFP4Tensor`` + - Gathers rowwise data and scales, then derives columnwise data locally when needed. .. note:: @@ -174,9 +178,10 @@ Limitations - FP8 weight caching interacts with FSDP gather/scatter — weights may be re-quantized on each gather. - Mixed FSDP + TP configurations require careful process group setup. -- ``fsdp_pre_all_gather`` / ``fsdp_post_all_gather`` are currently implemented only for - ``Float8Tensor`` and ``MXFP8Tensor``, not for ``Float8BlockwiseQTensor`` or - ``NVFP4Tensor``. +- ``Float8BlockwiseQTensor`` FSDP2 support is limited to 2D block scaling. +- ``NVFP4Tensor`` does not support FSDP2 with GEMM-swizzled scales. Columnwise usage + requires 2D quantization, and the flattened row count of each shard must be a multiple + of the NVFP4 block size. See Also -------- diff --git a/docs/developer/distributed/img/column_parallel.svg b/docs/developer/distributed/img/column_parallel.svg index 37ccdefacf..422f74d681 100644 --- a/docs/developer/distributed/img/column_parallel.svg +++ b/docs/developer/distributed/img/column_parallel.svg @@ -35,24 +35,24 @@ @ - + Weight W (split along output dim) - - W₀ (GPU 0) + + W₀ (GPU 0) - - W₁ (GPU 1) + + W₁ (GPU 1) - - W₂ (GPU 2) + + W₂ (GPU 2) - - W₃ (GPU 3) + + W₃ (GPU 3) diff --git a/docs/developer/distributed/img/row_parallel.svg b/docs/developer/distributed/img/row_parallel.svg index 429e0f95e5..674ecbded8 100644 --- a/docs/developer/distributed/img/row_parallel.svg +++ b/docs/developer/distributed/img/row_parallel.svg @@ -26,7 +26,7 @@ Row-Parallel Linear - + Input X (partitioned) @@ -48,7 +48,7 @@ - + Weight W (split along input dim) diff --git a/docs/developer/distributed/tensor_parallel.rst b/docs/developer/distributed/tensor_parallel.rst index 50c1fe9bef..760c3ef9d6 100644 --- a/docs/developer/distributed/tensor_parallel.rst +++ b/docs/developer/distributed/tensor_parallel.rst @@ -16,12 +16,12 @@ Megatron-style column-parallel and row-parallel linear layers. :align: center :width: 70% - Column-parallel linear: input is broadcast, weight columns are split across GPUs. + Column-parallel linear: input is replicated, weight output-feature rows are split. .. Diagram description for ``column_parallel.svg``: Left: "Input X" (full tensor, same on all GPUs). - Center: Weight matrix W split vertically into W_0, W_1, W_2, W_3 (4 GPUs). + Center: Weight matrix W split along its output-feature rows across 4 GPUs. Each GPU computes Y_i = X @ W_i^T (partial output). Right: Partial outputs Y_0..Y_3 are concatenated (or kept split for next layer). Label: "All-gather input if needed; output is split along hidden dim" @@ -30,12 +30,12 @@ Megatron-style column-parallel and row-parallel linear layers. :align: center :width: 70% - Row-parallel linear: input is split, weight rows are split, output is reduced. + Row-parallel linear: input and weight input-feature columns are split; output is reduced. .. Diagram description for ``row_parallel.svg``: Left: "Input X" split into X_0, X_1, X_2, X_3 (one per GPU). - Center: Weight matrix W split horizontally into W_0, W_1, W_2, W_3. + Center: Weight matrix W split along its input-feature columns across 4 GPUs. Each GPU computes Y_i = X_i @ W_i^T (partial result). Right: "All-reduce (or reduce-scatter)" to produce final output Y = sum(Y_i). Label: "Input is partitioned; output requires reduction" @@ -54,8 +54,10 @@ Used for the **first** linear layer in a pair (e.g., QKV projection, MLP FC1): tp_group=tp_process_group, ) -**Forward**: Each GPU holds columns ``W[:, start:end]`` and computes partial output. -If the input is not already partitioned, an all-gather collects the full input first. +**Forward**: TE stores a linear weight as ``[out_features, in_features]``. Each GPU holds +a shard ``W[start:end, :]`` along the output-feature dimension and computes the +corresponding output shard. With sequence parallelism, the sequence-sharded input is +all-gathered before the GEMM; otherwise the full input is already available on each rank. **Backward**: Gradient of the output is split; weight gradient is computed locally. Input gradient requires an all-reduce (or reduce-scatter with sequence parallelism). @@ -75,10 +77,13 @@ MLP FC2): tp_group=tp_process_group, ) -**Forward**: Each GPU holds rows ``W[start:end, :]`` and computes a partial result. -The partial results are summed via all-reduce (or reduce-scatter). +**Forward**: Each GPU holds ``W[:, start:end]``, a shard along the input-feature +dimension, and consumes the matching input shard. The partial results are summed via +all-reduce (or reduce-scatter with sequence parallelism). -**Backward**: Gradient is broadcast (or all-gathered); weight gradient is computed locally. +**Backward**: The output gradient is already present on every rank in the non-SP case. +With sequence parallelism it is all-gathered along the sequence dimension. Each rank +then computes its local weight-gradient and input-gradient shards. Column + Row Pairing -------------------- diff --git a/docs/developer/index.rst b/docs/developer/index.rst index b90cc53782..af08f5e37b 100644 --- a/docs/developer/index.rst +++ b/docs/developer/index.rst @@ -23,7 +23,7 @@ How to Use This Guide module. For building from source, see :doc:`/installation`. For running tests, see :doc:`testing`. For contributing guidelines, see ``CONTRIBUTING.rst`` at the repository root. -- **Working on new quantization recipe**: See :doc:`quantization/index` for the +- **Working on new quantization recipe**: See :doc:`quantization/index` for the Quantizer/Storage/Tensor design and :doc:`cpp_core/type_system` for the underlying C/C++ types. - **Working on attention**: See :doc:`attention/index` for backend selection and kernel organization. diff --git a/docs/developer/jax_frontend/module_system.rst b/docs/developer/jax_frontend/module_system.rst index b25527ed7f..54ce96fadc 100644 --- a/docs/developer/jax_frontend/module_system.rst +++ b/docs/developer/jax_frontend/module_system.rst @@ -86,24 +86,44 @@ stored as attributes: output = layer.apply(params, input) **Sharding via annotations**: Instead of ``parallel_mode`` constructor args, JAX modules -use ``MeshResource`` and rely on XLA SPMD for communication: +use logical axis annotations and a ``MeshResource`` set by ``global_shard_guard()`` or +``autocast()``. Within an active JAX ``Mesh``, logical axis rules map those annotations +to physical mesh axes, and XLA SPMD uses the mapping to insert communication: .. code-block:: python + import flax.linen as nn + import transformer_engine.jax.flax as te_jax + from transformer_engine.jax.sharding import MeshResource, global_shard_guard + layer = te_jax.DenseGeneral( features=3072, - mesh_resource=MeshResource(tp_resource="tp"), + kernel_axes=("input", "tp"), ) -**Quantization context**: Instead of ``autocast()``, JAX uses explicit quantizer -arguments: + mesh_resource = MeshResource(tp_resource="tp") + axis_rules = te_jax.extend_logical_axis_rules((("input", None), ("tp", "tp"))) + with global_shard_guard(mesh_resource), nn.logical_axis_rules(axis_rules): + params = layer.init(rng, input) + output = layer.apply(params, input) + +**Quantization context**: As in the PyTorch frontend, low-precision execution is selected +with ``autocast()``. JAX modules create their quantizer sets from the active recipe; +quantizers are not passed to ``Module.apply()`` explicitly: .. code-block:: python - output = layer.apply( - params, input, - quantizer=fp8_quantizer, - ) + from transformer_engine.common.recipe import MXFP8BlockScaling + from transformer_engine.jax import autocast + from transformer_engine.jax.sharding import MeshResource + + with autocast( + enabled=True, + recipe=MXFP8BlockScaling(), + mesh_resource=MeshResource(), + ): + params = layer.init(rng, input) + output = layer.apply(params, input) See Also -------- diff --git a/docs/developer/jax_frontend/xla_ffi_primitives.rst b/docs/developer/jax_frontend/xla_ffi_primitives.rst index edbe0cc5bf..84fc000e63 100644 --- a/docs/developer/jax_frontend/xla_ffi_primitives.rst +++ b/docs/developer/jax_frontend/xla_ffi_primitives.rst @@ -18,11 +18,14 @@ BasePrimitive ``BasePrimitive`` is the base class for all JAX TE primitives. It provides the machinery to: -1. **Register** a custom call with XLA (both forward and backward). +1. **Register** inner and outer JAX primitives and lower them to XLA FFI calls. 2. **Define abstract evaluation** (shape/dtype inference without running the kernel). -3. **Define the custom VJP** (JAX's equivalent of PyTorch's autograd backward). +3. **Register eager and ``vmap`` behavior** for each primitive. 4. **Handle sharding** (propagate partition specs through the operation). +``BasePrimitive`` does not define autodiff rules. Higher-level JAX functions compose +forward and backward primitives and expose the pair through ``jax.custom_vjp``. + .. code-block:: python class BasePrimitive(metaclass=ABCMeta): @@ -69,7 +72,9 @@ When the JAX TE module is imported, each primitive: 3. Registers ``lowering`` as the MLIR lowering rule (maps to XLA custom call via FFI). 4. Registers ``batcher`` for ``jax.vmap`` support. 5. Registers ``partition`` and ``shardy_sharding_rule`` for XLA SPMD. -6. Backward is registered via ``jax.custom_vjp`` on the outer primitive. + +Higher-level functions such as ``dense.py::_dense`` and ``attention.py::fused_attn`` +register custom VJPs and call the corresponding primitives. Primitives can be selectively enabled/disabled via the ``NVTE_JAX_CUSTOM_CALLS`` environment variable. diff --git a/docs/developer/pytorch_frontend/autograd_integration.rst b/docs/developer/pytorch_frontend/autograd_integration.rst index c52ab37b60..e186971b15 100644 --- a/docs/developer/pytorch_frontend/autograd_integration.rst +++ b/docs/developer/pytorch_frontend/autograd_integration.rst @@ -110,9 +110,9 @@ TE supports activation recomputation (gradient checkpointing) at multiple granul function re-enters the ``autocast`` context during recomputation so that quantizers and FP8 state are correctly restored. - **Selective recompute**: Only recompute specific operations (e.g., the core attention - computation) while keeping other activations saved. The ``TransformerLayer`` module - provides an ``activation_checkpointing`` parameter that controls selective - recomputation. + computation) while keeping other activations saved. Pass + ``checkpoint_core_attention=True`` to ``TransformerLayer.forward()`` to recompute the + core attention operation during backward. FP8 Tensors in Autograd ------------------------ diff --git a/docs/developer/pytorch_frontend/transformer_layer.rst b/docs/developer/pytorch_frontend/transformer_layer.rst index e1a2302cd0..c2a950b71a 100644 --- a/docs/developer/pytorch_frontend/transformer_layer.rst +++ b/docs/developer/pytorch_frontend/transformer_layer.rst @@ -19,7 +19,7 @@ Transformer implementation. QKV Weight Handling ------------------- -The ``fuse_qkv_params`` parameter (default ``True``) controls whether Q, K, and V +The ``fuse_qkv_params`` parameter (default ``False``) controls whether Q, K, and V projections share a single fused weight tensor or use separate weights. When ``fuse_qkv_params=True``, the QKV projection is a single ``LayerNormLinear`` with output size ``3 * hidden_size`` (or adjusted for GQA), and the weight is split into Q, K, V diff --git a/docs/developer/quantization/class_hierarchy.rst b/docs/developer/quantization/class_hierarchy.rst index d2bc533d8c..ff180903f9 100644 --- a/docs/developer/quantization/class_hierarchy.rst +++ b/docs/developer/quantization/class_hierarchy.rst @@ -197,8 +197,18 @@ A typical quantization lifecycle: .. code-block:: python - # 1. Create quantizer (usually done once per module) - quantizer = Float8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) + import torch + import transformer_engine.pytorch as te + from transformer_engine.pytorch.tensor.float8_tensor import Float8Quantizer + + # 1. Create quantizer state (usually created by a module's recipe machinery) + scale = torch.ones(1, dtype=torch.float32, device="cuda") + amax = torch.zeros(1, dtype=torch.float32, device="cuda") + quantizer = Float8Quantizer( + scale=scale, + amax=amax, + fp8_dtype=te.DType.kFloat8E4M3, + ) quantizer.set_usage(rowwise=True, columnwise=True) # 2. Quantize a tensor (returns QuantizedTensor for autograd) diff --git a/docs/envvars.rst b/docs/envvars.rst index ba8335d7c1..4715dfb958 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -140,12 +140,6 @@ Attention Backend Selection :Default: ``1`` :Description: Enable or disable UnfusedDotProductAttention backend (native PyTorch). When set to ``0``, UnfusedDotProductAttention will not be used. -.. envvar:: NVTE_FUSED_ATTN_BACKEND - - :Type: ``int`` (1 or 2) - :Default: Auto-selected - :Description: Force a specific FusedAttention backend. ``1`` = F16_arbitrary_seqlen (cuDNN, any seq len), ``2`` = FP8 backend. If not set, the backend is automatically selected based on the input configuration. - .. envvar:: NVTE_FUSED_ATTN_USE_FAv2_BWD :Type: ``int`` (0 or 1) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index e4eb883af8..0c1946c2dc 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -1038,9 +1038,9 @@ def forward( .. note:: Users can use environment variables :attr:`NVTE_FLASH_ATTN`, :attr:`NVTE_FUSED_ATTN`, - and :attr:`NVTE_FUSED_ATTN_BACKEND` to control which DotProductAttention backend, - and FusedAttention backend if applicable, to use. Transformer Engine prioritizes - FlashAttention over FusedAttention and over UnfusedDotProductAttention. + and :attr:`NVTE_UNFUSED_ATTN` to enable or disable each DotProductAttention backend. + When multiple backends are available, Transformer Engine prefers FusedAttention on + Hopper and newer architectures and FlashAttention on pre-Hopper architectures. If FusedAttention is being used, users can also choose to switch to flash-attn's implementation for backward by setting :attr:`NVTE_FUSED_ATTN_USE_FAv2_BWD=1` (default: 0), because of the performance differences between various versions of From 55445a7c281b81429f0279ca8c99adbedef9e5f1 Mon Sep 17 00:00:00 2001 From: Przemek Tredak Date: Wed, 1 Jul 2026 12:25:24 -0700 Subject: [PATCH 16/20] Beginning of take 2 Signed-off-by: Przemek Tredak --- docs/developer2/design_decisions.rst | 15 +++ docs/developer2/development_workflow.rst | 48 ++++++++++ .../extending_transformer_engine.rst | 18 ++++ docs/developer2/guide_overview.rst | 91 +++++++++++++++++++ .../implementation_architecture.rst | 20 ++++ docs/developer2/index.rst | 27 ++++++ docs/developer2/maintainer_operations.rst | 18 ++++ docs/developer2/project_and_architecture.rst | 18 ++++ docs/developer2/reference.rst | 18 ++++ docs/developer2/setup_build_and_run.rst | 18 ++++ .../testing_and_engineering_quality.rst | 19 ++++ 11 files changed, 310 insertions(+) create mode 100644 docs/developer2/design_decisions.rst create mode 100644 docs/developer2/development_workflow.rst create mode 100644 docs/developer2/extending_transformer_engine.rst create mode 100644 docs/developer2/guide_overview.rst create mode 100644 docs/developer2/implementation_architecture.rst create mode 100644 docs/developer2/index.rst create mode 100644 docs/developer2/maintainer_operations.rst create mode 100644 docs/developer2/project_and_architecture.rst create mode 100644 docs/developer2/reference.rst create mode 100644 docs/developer2/setup_build_and_run.rst create mode 100644 docs/developer2/testing_and_engineering_quality.rst diff --git a/docs/developer2/design_decisions.rst b/docs/developer2/design_decisions.rst new file mode 100644 index 0000000000..ff87d1a75d --- /dev/null +++ b/docs/developer2/design_decisions.rst @@ -0,0 +1,15 @@ +Design Decisions +================ + +Durable context explaining why the architecture has its current shape. + +Planned coverage +---------------- + +* Architecture decision records. +* Significant migrations and their compatibility constraints. +* Rejected alternatives when they remain relevant to future work. +* Known architectural constraints and technical debt. + +.. TODO: Define an architecture-decision-record format and migrate rationale + that is currently scattered across subsystem documentation. diff --git a/docs/developer2/development_workflow.rst b/docs/developer2/development_workflow.rst new file mode 100644 index 0000000000..9b1dad30d6 --- /dev/null +++ b/docs/developer2/development_workflow.rst @@ -0,0 +1,48 @@ +Development Workflow +==================== + +The expected path from identifying a change through implementation, validation, +review, and merge. + +Planned coverage +---------------- + +* Preparing and scoping a change. +* Coding conventions for each implementation language. +* Formatting, linting, commits, and pull requests. +* Contribution, ownership, review, and security expectations. +* Assessing change impact across APIs, frontends, tests, packaging, and + documentation. + +Documentation conventions +------------------------- + +Technical chapters should use the following structure where it is useful: + +#. Purpose and responsibilities. +#. Relevant files and symbols. +#. Concepts and terminology. +#. Execution or data flow. +#. Invariants and constraints. +#. Common change points. +#. Validation procedures. +#. Common pitfalls. +#. Related documentation. + +Not every chapter needs every part. Conceptual, procedural, and reference +material should remain distinct enough that readers can find an explanation, +perform a task, or look up a fact without reading unrelated content. + +Command examples should state: + +* the directory in which they are run; +* prerequisites and significant environment assumptions; +* placeholders that the reader must replace; and +* the expected result or success criterion. + +Pages should distinguish mandatory requirements from recommendations and +examples. Platform, framework, hardware, and version restrictions should be +stated next to the guidance they qualify. + +.. TODO: Describe the shortest reliable path from a clean checkout to a + reviewable change, then identify any workflows that need dedicated pages. diff --git a/docs/developer2/extending_transformer_engine.rst b/docs/developer2/extending_transformer_engine.rst new file mode 100644 index 0000000000..47f01d5a5a --- /dev/null +++ b/docs/developer2/extending_transformer_engine.rst @@ -0,0 +1,18 @@ +Extending Transformer Engine +============================ + +Principles and repeatable procedures for adding supported functionality. + +Planned coverage +---------------- + +* General extension principles and cross-layer impact. +* Adding kernels, operators, modules, and fused operations. +* Extending the C API and native bindings. +* Adding PyTorch and JAX integrations. +* Adding quantization types, recipes, and backend support. +* Task playbooks with prerequisites, decision points, validation, and + completion criteria. + +.. TODO: Seed the playbooks with the existing quantization-type, extension, + and operation-fuser procedures. diff --git a/docs/developer2/guide_overview.rst b/docs/developer2/guide_overview.rst new file mode 100644 index 0000000000..09f44ca3b6 --- /dev/null +++ b/docs/developer2/guide_overview.rst @@ -0,0 +1,91 @@ +Developer Guide Overview +======================== + +Transformer Engine spans several implementation layers, framework +integrations, and hardware-specific paths. This guide helps contributors +understand how those pieces fit together, make changes safely, and determine +how those changes should be validated. + +It complements the user guide and API reference. Those documents explain how +to use Transformer Engine and describe its public interfaces; this guide +focuses on the implementation, its design decisions, and the workflows used to +develop and maintain it. + +Where to start +-------------- + +You do not need to read the guide from beginning to end. Start with the section +that best matches what you are trying to do. + +.. list-table:: + :header-rows: 1 + :widths: 38 62 + + * - Goal + - Start with + * - Prepare a first contribution + - :doc:`Setup, Build, and Run `, followed by + :doc:`Development Workflow ` + * - Understand the codebase + - :doc:`Project and Architecture ` + * - Modify an existing subsystem + - :doc:`Implementation Architecture ` + * - Add a new capability + - :doc:`Extending Transformer Engine ` + * - Test a change or investigate a failure + - :doc:`Testing and Engineering Quality + ` + * - Understand why a durable design choice was made + - :doc:`Design Decisions ` + * - Perform packaging, documentation, or release work + - :doc:`Maintainer Operations ` + * - Find a command, term, support boundary, or external reference + - :doc:`Reference ` + +How the documentation fits together +----------------------------------- + +The `nightly documentation `_ +follows the in-development version of Transformer Engine. The `released +documentation +`_ +covers supported releases. Both contain the user guide and API reference. + +Contribution requirements are maintained in `CONTRIBUTING.rst +`_ at +the repository root. Changes between releases are recorded on `GitHub Releases +`_ and in the published +`release notes +`_. + +Keeping the guide useful +------------------------ + +Reference material should have one maintained home. Developer pages should +link to public APIs, configuration options, compatibility tables, and release +information instead of maintaining separate copies. A small amount of context +may still be repeated when it is needed to explain an internal design. + +The implementation and its tests are the primary references for internal +behavior. Build, test, and automation behavior is defined by the corresponding +project configuration and scripts. If these sources disagree with the +documentation, verify the intended behavior and treat the disagreement as a +defect rather than documenting both versions. + +When pointing into the source tree, use a repository-relative file path and a +stable symbol such as a function, class, module, or build target. Avoid exact +line numbers, which drift as the implementation changes. For example, prefer +``path/to/file.py:ClassName.method`` over a bare filename or line reference. + +Documentation status +-------------------- + +This replacement developer guide is being assembled incrementally. Skeleton +pages and explicit TODO comments describe planned coverage; they are not +normative project guidance. + +The guide is versioned with the rest of the documentation and should describe +the corresponding repository revision. The documentation landing page displays +a warning when the nightly, in-development version is shown and links to the +released documentation. Stale guidance should be corrected or reported like +any other project defect. diff --git a/docs/developer2/implementation_architecture.rst b/docs/developer2/implementation_architecture.rst new file mode 100644 index 0000000000..507cb4a67c --- /dev/null +++ b/docs/developer2/implementation_architecture.rst @@ -0,0 +1,20 @@ +Implementation Architecture +=========================== + +Detailed guides to the major subsystems, their interfaces, execution flows, +invariants, common change points, validation, and pitfalls. + +Planned coverage +---------------- + +* C and C++ core, type system, build pipeline, and kernels. +* PyTorch frontend, autograd integration, and native bindings. +* JAX frontend, module system, primitives, and FFI. +* Quantization types, recipes, layouts, and scaling modes. +* Attention dispatch, backends, kernels, and context parallelism. +* Distributed training and communication overlap. +* Operation fusing. +* Other major subsystems identified during migration. + +.. TODO: Migrate the strongest material from the existing developer guide and + split this section into subsystem pages as part of that migration. diff --git a/docs/developer2/index.rst b/docs/developer2/index.rst new file mode 100644 index 0000000000..075723211e --- /dev/null +++ b/docs/developer2/index.rst @@ -0,0 +1,27 @@ +.. + Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +Developer Guide +=============== + +This is the provisional structure for the Transformer Engine developer guide. +The ten sections below consolidate related concerns; each can be split into +focused pages when its content and navigation requirements are better +understood. + +.. toctree:: + :maxdepth: 2 + :caption: Developer Guide + + guide_overview + project_and_architecture + setup_build_and_run + development_workflow + testing_and_engineering_quality + implementation_architecture + extending_transformer_engine + maintainer_operations + design_decisions + reference diff --git a/docs/developer2/maintainer_operations.rst b/docs/developer2/maintainer_operations.rst new file mode 100644 index 0000000000..8dd3fd77a1 --- /dev/null +++ b/docs/developer2/maintainer_operations.rst @@ -0,0 +1,18 @@ +Maintainer Operations +===================== + +Procedures and policies needed to maintain and publish the project rather than +to implement an ordinary contributor change. + +Planned coverage +---------------- + +* Dependencies, submodules, and third-party code. +* Generated code, documentation, build artifacts, and caches. +* Documentation authoring, building, and validation. +* Packaging and supported distribution artifacts. +* Versioning, release preparation, release validation, and deprecation. +* Licensing and security-sensitive maintenance processes. + +.. TODO: Identify which operations are public contributor knowledge and which + require maintainer-only documentation. diff --git a/docs/developer2/project_and_architecture.rst b/docs/developer2/project_and_architecture.rst new file mode 100644 index 0000000000..5b19e5d8b8 --- /dev/null +++ b/docs/developer2/project_and_architecture.rst @@ -0,0 +1,18 @@ +Project and Architecture +======================== + +The concepts and structural context needed to understand changes across +Transformer Engine. + +Planned coverage +---------------- + +* Project goals, major capabilities, boundaries, and terminology. +* Repository structure and ownership boundaries. +* Layered system architecture and component relationships. +* An end-to-end execution walkthrough. +* Interfaces between languages and frameworks. +* Guiding design principles and important invariants. + +.. TODO: Migrate the existing architecture overview and execution walkthrough, + expanding the repository map beyond the main package directory. diff --git a/docs/developer2/reference.rst b/docs/developer2/reference.rst new file mode 100644 index 0000000000..04b78e4fc0 --- /dev/null +++ b/docs/developer2/reference.rst @@ -0,0 +1,18 @@ +Reference +========= + +Concise contributor-oriented reference material and links to authoritative +project resources. + +Planned coverage +---------------- + +* Command reference. +* Glossary. +* Supported hardware, software, and feature combinations. +* Environment-variable and public-API references. +* Repository source and ownership map. +* Index of important invariants and constraints. + +.. TODO: Keep reference material concise, precise, and linked from conceptual + and procedural chapters. diff --git a/docs/developer2/setup_build_and_run.rst b/docs/developer2/setup_build_and_run.rst new file mode 100644 index 0000000000..d50cb99259 --- /dev/null +++ b/docs/developer2/setup_build_and_run.rst @@ -0,0 +1,18 @@ +Setup, Build, and Run +===================== + +A reproducible path from a clean checkout to a usable development build. + +Planned coverage +---------------- + +* Prerequisites and supported development environments. +* Dependency and toolchain setup. +* Standard, optional, and framework-specific builds. +* Editable builds and selecting a development build at runtime. +* Build outputs and incremental rebuilds. +* Import checks and representative smoke tests. +* Build and setup troubleshooting. + +.. TODO: Establish canonical contributor commands while linking to the + authoritative installation and environment-variable references. diff --git a/docs/developer2/testing_and_engineering_quality.rst b/docs/developer2/testing_and_engineering_quality.rst new file mode 100644 index 0000000000..3739c86fad --- /dev/null +++ b/docs/developer2/testing_and_engineering_quality.rst @@ -0,0 +1,19 @@ +Testing and Engineering Quality +=============================== + +The practices used to establish correctness, compatibility, reliability, and +performance before a change is accepted. + +Planned coverage +---------------- + +* Test organization, selection, and canonical commands. +* Writing tests and choosing appropriate test levels. +* Numerical references, tolerances, stochastic behavior, and determinism. +* Continuous-integration topology and local reproduction. +* Debugging tools and troubleshooting workflows. +* Profiling, benchmarking, and performance-regression evaluation. +* API, ABI, checkpoint, hardware, and software compatibility policies. + +.. TODO: Consolidate existing testing guidance and identify which quality + policies need to be established rather than merely documented. From c1a7e6b7a2fdb671450eab5c132fa7c3fd7a3e73 Mon Sep 17 00:00:00 2001 From: Przemek Tredak Date: Sat, 4 Jul 2026 07:43:23 -0700 Subject: [PATCH 17/20] More decuments Signed-off-by: Przemek Tredak --- docs/developer2/build_configuration.rst | 40 +++ docs/developer2/build_troubleshooting.rst | 39 ++ docs/developer2/development_environment.rst | 35 ++ docs/developer2/img/layered_architecture.svg | 87 +++++ .../img/pytorch_linear_walkthrough.svg | 119 +++++++ docs/developer2/layered_architecture.rst | 164 +++++++++ docs/developer2/project_and_architecture.rst | 212 ++++++++++- .../developer2/pytorch_linear_walkthrough.rst | 332 ++++++++++++++++++ docs/developer2/setup_build_and_run.rst | 27 +- docs/developer2/split_distribution_builds.rst | 53 +++ docs/developer2/unified_editable_build.rst | 48 +++ .../using_and_validating_builds.rst | 41 +++ 12 files changed, 1170 insertions(+), 27 deletions(-) create mode 100644 docs/developer2/build_configuration.rst create mode 100644 docs/developer2/build_troubleshooting.rst create mode 100644 docs/developer2/development_environment.rst create mode 100644 docs/developer2/img/layered_architecture.svg create mode 100644 docs/developer2/img/pytorch_linear_walkthrough.svg create mode 100644 docs/developer2/layered_architecture.rst create mode 100644 docs/developer2/pytorch_linear_walkthrough.rst create mode 100644 docs/developer2/split_distribution_builds.rst create mode 100644 docs/developer2/unified_editable_build.rst create mode 100644 docs/developer2/using_and_validating_builds.rst diff --git a/docs/developer2/build_configuration.rst b/docs/developer2/build_configuration.rst new file mode 100644 index 0000000000..38f79a431c --- /dev/null +++ b/docs/developer2/build_configuration.rst @@ -0,0 +1,40 @@ +Build Configuration +=================== + +Configuration principles +------------------------ + +.. TODO: Explain which settings are intended for contributors, which are + internal to release packaging, and which authoritative reference defines + each environment variable. + +Framework selection +------------------- + +.. TODO: Cover framework auto-detection and explicit selection without + duplicating the unified-build procedure. + +CUDA toolkit and target architectures +------------------------------------- + +.. TODO: Cover toolkit discovery, include paths, CUDA-major selection, and GPU + architecture targets, including their effect on build time and portability. + +Build type, caching, and parallelism +------------------------------------ + +.. TODO: Cover debug builds, CMake build directories, ccache or sccache, and + controls for parallel compilation. + +Optional components +------------------- + +.. TODO: Group options for MPI, NVSHMEM, cuBLASMp, NCCL expert parallelism, + fast-math kernels, and other optional native features by dependency and + effect. + +Release-only controls +--------------------- + +.. TODO: Identify release-build, metapackage, package-layout, and wheel-build + settings that ordinary development builds should not set. diff --git a/docs/developer2/build_troubleshooting.rst b/docs/developer2/build_troubleshooting.rst new file mode 100644 index 0000000000..2b169effa1 --- /dev/null +++ b/docs/developer2/build_troubleshooting.rst @@ -0,0 +1,39 @@ +Build and Setup Troubleshooting +=============================== + +Collect build diagnostics +------------------------- + +.. TODO: Define the command output, environment details, tool versions, and + generated logs needed to investigate a build failure. + +Toolchain and configuration failures +------------------------------------ + +.. TODO: Organize failures involving CUDA discovery, unsupported target + architectures, CMake, compilers, missing headers, and optional libraries. + +Submodule and source-dependency failures +---------------------------------------- + +.. TODO: Cover uninitialized, stale, or locally modified submodules and the + difference between correcting them and intentionally bypassing checks. + +Resource and incremental-build failures +--------------------------------------- + +.. TODO: Cover compiler memory pressure, parallel-job limits, stale CMake + state, compiler caches, and when a clean rebuild is appropriate. + +Import, linking, and ABI failures +--------------------------------- + +.. TODO: Cover imports resolving to the wrong installation, missing shared + libraries, mismatched CUDA or framework builds, undefined symbols, and + conflicts between editable and distribution-package layouts. + +Distribution-build failures +--------------------------- + +.. TODO: Cover container setup, platform tags, package metadata, wheel or + source-distribution dependencies, cached wheel lookup, and wheelhouse logs. diff --git a/docs/developer2/development_environment.rst b/docs/developer2/development_environment.rst new file mode 100644 index 0000000000..d6a3ec687b --- /dev/null +++ b/docs/developer2/development_environment.rst @@ -0,0 +1,35 @@ +Development Environment +======================= + +Prerequisites +------------- + +.. TODO: Document supported host platforms, Python versions, compilers, CUDA, + CUDA libraries, and build tools. Link to the authoritative installation and + support documentation instead of duplicating version matrices. + +Python and framework environment +-------------------------------- + +.. TODO: Document the expected virtual environment and the PyTorch and/or JAX + dependencies required before a no-build-isolation source build. + +Repository dependencies +----------------------- + +.. TODO: Document required Git submodules and when they are initialized or + checked automatically by the build. + +Optional native dependencies +---------------------------- + +.. TODO: Identify optional MPI, NVSHMEM, cuBLASMp, NCCL expert-parallel, and + other native dependencies without turning this page into a build-variable + reference. + +Verify the environment +---------------------- + +.. TODO: Provide a short pre-build checklist that verifies the selected Python + interpreter, framework installations, CUDA toolkit, compiler, CMake, Ninja, + and relevant library paths. diff --git a/docs/developer2/img/layered_architecture.svg b/docs/developer2/img/layered_architecture.svg new file mode 100644 index 0000000000..562458096a --- /dev/null +++ b/docs/developer2/img/layered_architecture.svg @@ -0,0 +1,87 @@ + + Transformer Engine layered architecture + PyTorch and JAX users enter framework-aware integrations that own public Python APIs, composite operations, framework state and memory, and internal C++ bindings. Those integrations call the public C API and import the shared Python API in the framework-independent common layer. Standalone clients can call the C API directly. The common layer owns focused operations, kernel selection, C++ and CUDA kernels, CUDA library integrations, and Triton or other framework-independent kernel implementations. + + + + + + + + + Transformer Engine layered architecture + + + + PyTorch users + + + JAX users + + + Standalone + C API users + + + + Framework-aware layer + + + + + + PyTorch integration + Public Python API + Modules and composite operations + Framework state and memory + Internal C++ bindings (pybind11) + + + JAX integration + Public Python API + Primitives and composite operations + Framework state and memory + Internal C++ bindings (XLA FFI) + + + + Framework-independent common layer + + + + Shared Python imports + + + C API calls + + + + + Shared Python API + Recipes and shared utilities + + + Focused operations and kernel selection + Quantize, GEMM, normalization, attention, ... + + + Public C API + Internal and standalone use + + + + + + Framework-independent kernel implementations + C++ / CUDA · cuBLASLt / cuDNN · Triton and other kernel technologies + diff --git a/docs/developer2/img/pytorch_linear_walkthrough.svg b/docs/developer2/img/pytorch_linear_walkthrough.svg new file mode 100644 index 0000000000..ce51bc5bca --- /dev/null +++ b/docs/developer2/img/pytorch_linear_walkthrough.svg @@ -0,0 +1,119 @@ + + PyTorch Linear forward and backward lifecycle + The forward path runs from Linear.forward through recipe and quantizer setup, common quantization operations, the common GEMM operation, and framework-owned output and saved backward state. The backward path starts with the output gradient, prepares its required representations, runs dgrad and wgrad GEMMs in the common layer, and returns gradients after framework cleanup. Saved input and weight representations connect the forward and backward paths. + + + + + + + + + + + + + + + PyTorch Linear: representative training path + + + + Framework orchestration + + Common operation + + + + Forward + + + Public entry + Linear.forward + + + + + Configure execution + Recipe → quantizers + LinearFwdArgs + _Linear.apply + + + + + Prepare operands + Quantize input + Quantize or reuse weight + Required layouts + + + + + Forward GEMM + y = xWᵀ + b + + + + + Finish forward + Return output + Save backward state + Refresh weight cache + + + + Backward + + + Backward entry + grad_output + + + + + Prepare gradients + Restore saved state + Prepare grad_output + Required layouts + + + + + Dgrad GEMM + dx = dyW + + + + + Wgrad GEMM + dW = dyᵀx + + + + + Finish backward + Return dx, dW, and db + Wait for pending communication + Update recipe state if required + + + + Saved input and weight representations + + + Autograd context + + Tensor- and sequence-parallel communication is added around these phases; it is not shown in this baseline path. + diff --git a/docs/developer2/layered_architecture.rst b/docs/developer2/layered_architecture.rst new file mode 100644 index 0000000000..e58fe1bc99 --- /dev/null +++ b/docs/developer2/layered_architecture.rst @@ -0,0 +1,164 @@ +Layered Architecture and Component Relationships +================================================ + +Transformer Engine has two ownership layers: a framework-independent common +layer and framework-aware integrations for PyTorch and JAX. Calls cross several +interfaces within those layers, but the ownership boundary is consistent: +framework-specific behavior stays in the framework integration, while reusable +operations and kernel implementations stay in the common layer. + +.. figure:: img/layered_architecture.svg + :align: center + :width: 100% + :alt: Transformer Engine framework-aware and common architecture layers + + Framework-aware integrations compose operations and manage framework state, + while the common layer provides public interfaces, focused operations, and + framework-independent kernel implementations. + +.. + Diagram description for ``layered_architecture.svg``: + PyTorch and JAX users enter separate framework-aware integrations. Each + integration contains its public Python API, modules or primitives, composite + operations, framework state and memory, and internal C++ bindings. The + integrations import the shared common Python API and call the public C API. + Standalone clients can call the C API directly without a framework. The + framework-independent common layer contains the shared Python API, the public + C API, focused operations and kernel selection, and kernel implementations + written with C++, CUDA, CUDA libraries such as cuBLASLt and cuDNN, Triton, or + other framework-independent kernel technologies. + +The framework-aware layer +------------------------- + +The PyTorch and JAX integrations expose public Python APIs using the constructs +that are natural to each framework. They are intended to provide broadly +symmetrical functionality, although their API shapes differ and feature gaps +may exist between them. + +This layer owns composite operations. A module such as ``Linear`` interprets +the recipe and options selected by the user, then coordinates the sequence of +framework-native and Transformer Engine operations needed to produce the +result. That sequence may include quantization, communication, GEMM, and other +steps. + +Both the Python and C++ code that interacts with a framework belong to this +layer. The C++ bindings translate framework tensors and execution conventions +into calls to the common C API. They do not contain GPU kernels. In particular, +the framework integrations are designed not to require NVCC; functionality +that requires kernel compilation belongs in the common layer. + +The common layer +---------------- + +The common layer exposes focused operations with well-defined inputs and +outputs. The caller chooses what operation to perform and supplies the relevant +policy, such as the quantization type. The common layer selects and launches an +appropriate implementation of that operation. + +Kernel technology does not change this ownership. C++ and CUDA kernels, calls +to CUDA libraries such as cuBLASLt and cuDNN, Triton implementations, and other +framework-independent kernel technologies all belong to the common layer. Its +shared Python API also defines concepts used by both framework integrations, +including the common definition of a quantization recipe. + +The C API is public. Transformer Engine's framework integrations use it +internally, and external software can use it directly without PyTorch or JAX. + +Policy and dispatch +------------------- + +Responsibility for a decision depends on its level: + +.. list-table:: + :header-rows: 1 + :widths: 48 24 28 + + * - Decision + - Owner + - Example + * - User-visible policy and feature configuration + - Framework-aware layer + - Selecting a quantization recipe + * - Composition of multiple operations + - Framework-aware layer + - Combining quantization, communication, and GEMM in ``Linear`` + * - Selection between framework-level implementations + - Framework-aware layer + - Choosing fused, FlashAttention, or unfused attention + * - Parameters of a focused common operation + - Caller + - Passing the quantization type to a quantize operation + * - Selection of the low-level kernel implementation + - Common layer + - Choosing a kernel for the requested quantize operation + +Some features have more than one dispatch level. Attention is an important +example: the framework integration first chooses a high-level implementation. +If it selects the cuDNN-backed implementation, it calls the common fused +attention API, which then selects among the available cuDNN kernels. The +attention architecture documents that hierarchy in detail. + +Memory and state ownership +-------------------------- + +The framework-aware layer owns nearly all device memory. Outputs, quantization +scales, amax values, workspaces, and other persistent or temporary data are +allocated and stored as ordinary framework tensors, then passed to common +operations. + +The common layer allocates device memory only in rare cases where the required +allocation cannot be obtained through a framework API. Symmetric memory needed +by some JAX operations is one example. This is an exception rather than a +general allocation mechanism. + +Execution plans and library handles are different: the common layer creates +and caches them because they are implementation details of the kernels and +libraries it manages. + +Cross-cutting responsibilities +------------------------------ + +Several features span the two layers: + +**Quantization recipes** + The framework-independent recipe concepts are defined in the common Python + API. The framework integrations interpret the selected recipe and realize + it by composing framework-native behavior with common operations. + +**Numerical debugging** + Debugging is primarily a framework-layer concern and may call kernels in the + common layer when needed. The design is intended to be cross-framework, but + the current implementation supports PyTorch only. + +**Distributed execution** + Framework-native communication normally remains in the framework-aware + layer. The common layer participates when communication is intertwined with + computation, as in expert-parallel dispatch and combine operations or + communication and GEMM overlap. + +Distributed scope +----------------- + +Distributed functionality is in scope when it changes computation within an +individual Transformer layer. This includes tensor, sequence, expert, and +context parallelism. Model- and optimizer-level orchestration, including +pipeline parallelism and distributed optimizers, belongs to higher-level +training toolkits. + +Fully Sharded Data Parallel is also an external facility. Transformer Engine +does not implement FSDP, but it must understand FSDP behavior and provide the +integration required for Transformer Engine modules to work with it. + +API boundaries and stability +---------------------------- + +The public C and Python APIs are compatibility boundaries. Existing C APIs are +not changed incompatibly; when a change is necessary, a new API version is +introduced and the previous version remains available until the next major +release. Public Python APIs evolve compatibly, normally through additions such +as new optional keyword arguments. + +Private interfaces between a framework integration and its bindings are not +compatibility boundaries. PyTorch C++ extensions and similar internal +interfaces may change together with their callers. diff --git a/docs/developer2/project_and_architecture.rst b/docs/developer2/project_and_architecture.rst index 5b19e5d8b8..291e0cd1e5 100644 --- a/docs/developer2/project_and_architecture.rst +++ b/docs/developer2/project_and_architecture.rst @@ -1,18 +1,206 @@ Project and Architecture ======================== -The concepts and structural context needed to understand changes across -Transformer Engine. +This section describes where Transformer Engine fits in the software stack, +how its major pieces work together, and how those boundaries shape the code. -Planned coverage ----------------- +Role in the software stack +-------------------------- -* Project goals, major capabilities, boundaries, and terminology. -* Repository structure and ownership boundaries. -* Layered system architecture and component relationships. -* An end-to-end execution walkthrough. -* Interfaces between languages and frameworks. -* Guiding design principles and important invariants. +Transformer Engine provides optimized building blocks for Transformer models +on NVIDIA GPUs. It is not a complete LLM training toolkit. Higher-level +frameworks and training toolkits compose its operations and modules into full +models and own the orchestration of the training system. -.. TODO: Migrate the existing architecture overview and execution walkthrough, - expanding the repository map beyond the main package directory. +At a high level, Transformer Engine has two pieces: + +.. list-table:: + :header-rows: 1 + :widths: 16 24 44 16 + + * - Piece + - Location + - Responsibilities + - API surface + * - Common layer + - ``transformer_engine/common/`` + - Implements framework-independent functionality. Performance-critical + operations are written in C++ and CUDA and use CUDA libraries such as + cuBLASLt and cuDNN where appropriate. Shared Python code provides + functionality used by both framework integrations. The layer exposes + focused operations with well-defined inputs and outputs, such as + quantization, GEMM, and RMSNorm forward and backward passes. + - C API and shared Python API + * - Framework-aware layer + - ``transformer_engine/pytorch/`` and ``transformer_engine/jax/`` + - Integrates the common operations with PyTorch and JAX. It adapts + framework tensors and execution models, manages framework-visible state, + and composes common operations into larger components. This layer also + presents features such as quantization recipes. + - Framework-specific Python APIs + +The binding code sits at the boundary between the two pieces. It translates +framework objects and execution conventions into calls to the C API while +keeping the common layer independent of PyTorch and JAX. + +Scope boundary +~~~~~~~~~~~~~~ + +Most Transformer Engine functionality operates at or below the level of an +individual Transformer layer. The project provides optimized operations and +modules that a larger training system can combine, rather than owning the +entire model or training architecture. + +This boundary is especially important for distributed execution. Transformer +Engine may implement or integrate parallel techniques that affect its +operations and layers, but model-level orchestration belongs to higher-level +toolkits. Pipeline parallelism, for example, is outside Transformer Engine's +scope and is the responsibility of a toolkit such as Megatron. + +Choosing where a change belongs +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +This division provides a useful starting point when locating or designing a +change: + +* Framework-independent computation and GPU execution belong in the common + layer and should be exposed through the C API when needed by a frontend. + Framework-independent Python functionality shared by both frontends also + belongs in the common layer. +* Framework tensor handling, execution semantics, state management, recipes, + and composition into user-facing modules belong in the framework-aware + layer. +* Features that span both pieces should keep the reusable computation in the + common layer and place framework-specific policy and integration in the + corresponding frontend. + +Repository structure +-------------------- + +The following map focuses on the durable areas that help contributors decide +where to look or make a change. Detailed subsystem maps belong in the +implementation chapters. + +.. list-table:: + :header-rows: 1 + :widths: 32 68 + + * - Location + - Purpose + * - ``transformer_engine/common/`` + - Framework-independent C++, CUDA, and shared Python functionality, as + described above. + * - ``transformer_engine/pytorch/`` + - PyTorch integration and framework-specific Python API. + * - ``transformer_engine/jax/`` + - JAX integration and framework-specific Python API. + * - ``transformer_engine/debug/`` + - Numerical debugging tools. The design is intended to be + cross-framework, although the current implementation supports PyTorch + only. + * - ``tests/`` + - The test implementations, organized into C++, distributed C++, PyTorch, + and JAX areas. + * - ``qa/`` + - Launcher scripts that select, configure, and run validation in the CI + test classes. L0 jobs, such as ``L0_pytorch_unittest``, run whenever CI + is launched. L1, L2, and L3 classes are reserved for less frequent + nightly or weekly validation. + * - ``benchmarks/`` + - Performance benchmarks and profiling utilities, kept separate from + correctness tests. + * - ``docs/`` + - Sphinx documentation sources, including the user guide, API reference, + developer documentation, and the tutorials and notebooks under ``docs/examples/``. + * - ``examples/`` + - Runnable standalone PyTorch and JAX examples. These focus on executable + code with little accompanying explanation, unlike the tutorials under + ``docs/examples/``. + * - ``build_tools/`` and top-level packaging files + - Helpers for building the common library and framework extensions, + version handling, and wheel creation. ``setup.py``, ``pyproject.toml``, + and ``MANIFEST.in`` provide the top-level build and packaging entry + points. + * - ``3rdparty/`` + - External source dependencies managed as Git submodules, including cuDNN + Frontend, CUTLASS, GoogleTest, and NCCL. + +Architecture topics +------------------- + +.. toctree:: + :maxdepth: 1 + + layered_architecture + pytorch_linear_walkthrough + +Working architectural principles +-------------------------------- + +.. important:: + + The statements below are working proposals, not settled project policy. + They summarize the architecture described so far and must be revisited as + the remaining subsystem documentation is developed. Before they become + normative, each statement should be checked against multiple subsystems and + any necessary exceptions should be documented. + +Candidate invariants +~~~~~~~~~~~~~~~~~~~~ + +* **The common layer is framework-independent.** It may contain C++, CUDA, + shared Python, Triton, and other kernel technologies, but it must not depend + on PyTorch or JAX. +* **Framework-aware layers do not contain GPU kernels.** In particular, they + should not require NVCC. Kernel implementations belong in the common layer. +* **Framework integrations use the public C API for compiled common + functionality.** The same API can be used without PyTorch or JAX. +* **Common operations are focused and caller-directed.** The caller selects an + operation and supplies policy such as the quantization type; the common layer + selects the concrete kernel implementation. +* **Composite operations belong in the framework-aware layer.** Modules such + as ``Linear`` interpret recipes and options, then coordinate framework-native + operations, common operations, communication, and state. +* **Tensor memory is normally owned by the framework.** Outputs, workspaces, + scales, amax values, and saved state should be regular framework tensors + passed to common operations. +* **Common-layer device allocation is exceptional.** It is appropriate only + when the required allocation cannot be obtained through a framework API, + such as certain symmetric-memory allocations. +* **Execution plans and library handles belong to the common layer.** They are + implementation details of the kernels and libraries managed there and may + be cached by that layer. +* **Public APIs are compatibility boundaries.** Existing C APIs remain + available until the next major version when a replacement is necessary. + Public Python APIs evolve compatibly, normally through optional keyword + arguments. +* **Private framework bindings are not compatibility boundaries.** PyTorch C++ + extensions, JAX FFI bindings, and their internal callers may evolve together. + +Candidate design guidance +~~~~~~~~~~~~~~~~~~~~~~~~~ + +* **Expose equivalent intent through framework-native constructs.** PyTorch + and JAX should provide broadly symmetrical functionality without being + forced into identical APIs. +* **Place shared concepts in the common layer.** A concept such as a + quantization recipe should have one framework-independent definition, while + each frontend determines how to realize it during execution. +* **Separate policy from mechanism.** User-visible choices and multi-operation + composition belong in the framework-aware layer; low-level dispatch belongs + in the common layer. +* **Prefer framework-native functionality when it is sufficient.** Common + implementations are most appropriate when specialized kernels are needed or + when communication and computation must be designed together. +* **Optimize representations across the full operation lifecycle.** Forward + execution should account for what backward will need while avoiding + unnecessary storage, communication, transposes, and requantization. +* **Keep Transformer Engine focused on composable building blocks.** + Functionality generally belongs at or below the Transformer-layer level + rather than implementing a complete model-training system. +* **Distributed ownership follows the computation boundary.** Framework-native + communication stays in the frontend; the common layer participates when + communication is inseparable from kernel execution. +* **External facilities may still require explicit integration.** Transformer + Engine does not implement FSDP, for example, but its modules must understand + enough about FSDP to behave correctly with it. diff --git a/docs/developer2/pytorch_linear_walkthrough.rst b/docs/developer2/pytorch_linear_walkthrough.rst new file mode 100644 index 0000000000..90ed879d8f --- /dev/null +++ b/docs/developer2/pytorch_linear_walkthrough.rst @@ -0,0 +1,332 @@ +PyTorch Linear: End-to-End Walkthrough +====================================== + +``transformer_engine.pytorch.Linear`` is a useful guide through Transformer +Engine because one training step crosses nearly every important boundary: a +public framework API interprets a quantization recipe, framework code prepares +state and tensor representations, common operations quantize data and execute +GEMMs, and PyTorch autograd connects the forward and backward passes. + +This walkthrough explains those relationships. It intentionally describes the +main execution path rather than every branch in the implementation. + +Path traced here +---------------- + +The baseline path uses: + +* the PyTorch ``Linear`` module in training mode; +* an MXFP8 block-scaling recipe; +* a platform for which ``te.is_mxfp8_available()`` reports support; +* BF16 inputs, parameters, and unquantized outputs; +* gradients for the input, weight, and bias; and +* a single process, without tensor or sequence parallelism. + +FSDP, CPU offloading, communication overlap, output quantization, and delayed +weight-gradient computation are discussed later as variations. Establishing a +baseline keeps configuration-dependent branches from looking like universal +behavior. + +The following example exercises the baseline path: + +.. code-block:: python + + import torch + import transformer_engine.pytorch as te + from transformer_engine.common import recipe + + linear = te.Linear( + 4096, + 16384, + bias=True, + params_dtype=torch.bfloat16, + device="cuda", + ) + inp = torch.randn( + 32, + 4096, + dtype=torch.bfloat16, + device="cuda", + requires_grad=True, + ) + + fp8_recipe = recipe.MXFP8BlockScaling() + with te.autocast(enabled=True, recipe=fp8_recipe): + out = linear(inp) + + out.sum().backward() + +The computation is conceptually: + +.. math:: + + Y = XW^T + b, \qquad + \frac{\partial L}{\partial X} = \frac{\partial L}{\partial Y} W, \qquad + \frac{\partial L}{\partial W} = \left(\frac{\partial L}{\partial Y}\right)^T X. + +The implementation adds quantized representations, saved state, and optional +communication around these three matrix multiplications. + +.. figure:: img/pytorch_linear_walkthrough.svg + :align: center + :width: 100% + :alt: PyTorch Linear forward and backward lifecycle + + A representative local training path. Blue boxes are framework-layer + orchestration; green boxes invoke focused operations from the common layer. + +.. + Diagram description for ``pytorch_linear_walkthrough.svg``: + The forward lane proceeds from ``Linear.forward`` to recipe and quantizer + setup, construction of ``LinearFwdArgs``, and ``_Linear.apply``. Common + operations prepare quantized input and weight representations and execute + the forward GEMM. The framework layer returns the output, saves the state + required by backward, and refreshes the cached weight workspace. Dashed + arrows carry saved input and weight representations and the autograd context + into the backward lane. Backward starts from ``grad_output``, restores saved + state, and prepares the gradient representations. Common GEMMs compute dgrad + and wgrad. The framework layer returns input, weight, and bias gradients, + waits for pending communication, and updates recipe state when required. + Tensor- and sequence-parallel communication is omitted from this baseline + diagram and described separately below. + +1. Public entry and execution configuration +------------------------------------------- + +The public entry point is +``transformer_engine/pytorch/module/linear.py:Linear.forward``. It first calls +``TransformerEngineBaseModule.prepare_forward`` to establish the module's +execution context and activation dtype. It then obtains the current weight and +bias tensors and selects the quantizers required for this invocation through +``Linear._get_quantizers``. + +The active recipe determines the concrete quantizer implementations and their +scaling behavior. ``Linear`` assigns those quantizers roles for the input, +weight, optional output, output gradient, and optional input and weight +gradients. Not every role is populated for every call; inference, output +quantization, debugging, custom recipes, and requested gradients all affect the +selection. + +The module packages execution policy in ``LinearFwdArgs``. Differentiable +tensors remain explicit arguments to ``_Linear.apply`` so that PyTorch can +track them, while quantizers, parallel configuration, numerical options, and +cached workspaces travel in the structured argument object. + +When gradients are enabled, ``_Linear.apply`` enters the custom autograd +function. With gradients disabled, ``Linear.forward`` calls ``_Linear.forward`` +directly and does not construct backward state. + +2. Input and weight preparation +------------------------------- + +``transformer_engine/pytorch/module/linear.py:_linear_forward_impl`` performs +the forward computation. + +For the local baseline, the input is already present in full and no collective +communication is needed. The input quantizer creates the representation needed +by the forward GEMM. When the weight requires a gradient, it also preserves or +can later recover the representation required by wgrad. For MXFP8, rowwise and columnwise forms are not numerically +equivalent transposes. When both are needed, the quantizer produces both directly from the +high-precision tensor to avoid a second quantization step. Other recipes may +use a different storage strategy; the invariant is that each GEMM receives a +representation compatible with its layout. + +The weight follows a similar process through +``transformer_engine/pytorch/module/linear.py:quantize_weight``. A quantized +weight workspace may be cached because a parameter is normally unchanged +across the microbatches of a gradient-accumulation step. ``is_first_microbatch`` +controls when the workspace is refreshed. The newly produced workspace is +returned through ``_Linear`` and stored by ``Linear.forward`` in +``Linear._fp8_workspaces``; it remains framework-owned state rather than hidden +allocation in the common layer. + +Quantized storage can contain data, scales, and metadata for more than one +layout. Consequently, statements such as "FP8 uses half the memory" apply only +to the raw element payload and should not be interpreted as an exact total +memory ratio. + +3. Forward GEMM and the common-layer boundary +--------------------------------------------- + +Once both operands are ready, ``_linear_forward_impl`` calls +``transformer_engine/pytorch/cpp_extensions/gemm.py:general_gemm``. In the +baseline it computes :math:`Y = XW^T + b`, returning BF16 output. The bias can +be fused into the GEMM. If the caller requests a quantized output, an output +quantizer is passed through ``quantization_params`` instead. + +This call illustrates the boundary between the two architectural layers: + +#. ``general_gemm`` prepares framework-owned output and workspace tensors and + selects the appropriate framework binding. +#. ``transformer_engine/pytorch/csrc/extensions/gemm.cpp`` converts PyTorch and + quantized-storage objects into the tensor handles expected by the public C + API. These handles are views; the framework tensors retain ownership of the + underlying memory. +#. A public C GEMM entry point declared in + ``transformer_engine/common/include/transformer_engine/gemm.h`` receives the + operation description. +#. The common implementation under ``transformer_engine/common/gemm/`` selects + and launches the appropriate implementation, such as a cuBLASLt GEMM. + +The PyTorch C++ extension is a private binding and can evolve with its Python +caller. The C API is a public compatibility boundary. + +4. Output and saved backward state +---------------------------------- + +The local baseline returns the GEMM result directly. Parallel modes may insert +collective communication before the output is ready, as described below. + +If gradients are required, ``_linear_forward_impl`` retains only the state that +backward needs. Depending on the quantizer, it may discard a rowwise input +representation after the forward GEMM while retaining a columnwise +representation for wgrad. Other configurations save the original input or +reconstruct the needed layout later, so this is an optimization rather than a +universal storage rule. + +``transformer_engine/pytorch/module/linear.py:_linear_setup_ctx`` transfers the +required configuration into ``LinearBwdArgs``. The +``prepare_for_saving``/``restore_from_func_ctx`` mechanism separates +quantized-storage objects into PyTorch tensors and metadata so PyTorch can own +their lifetime through the autograd context. + +5. Backward entry and gradient preparation +------------------------------------------ + +PyTorch invokes ``transformer_engine/pytorch/module/linear.py:_Linear.backward`` +with the output gradient. It restores the saved tensors, attaches the gradient +to ``LinearBwdArgs``, and calls +``transformer_engine/pytorch/module/linear.py:_linear_backward``. + +``TransformerEngineBaseModule.grad_output_preprocess`` prepares +``grad_output`` for the requested gradients. Dgrad consumes a rowwise +representation, while wgrad consumes a columnwise representation. The +quantizer may create both together or create them at different points, +depending on communication and storage capabilities. Bias-gradient computation +can also be fused into this preparation or a GEMM path. + +6. Dgrad and wgrad +------------------ + +When the input requires a gradient, dgrad computes +:math:`\partial L / \partial X`. ``_linear_backward`` calls ``general_gemm`` +with ``layout="NN"``, using the output gradient and the weight representation +saved or reconstructed from forward. In the local baseline, the result needs no +collective communication. + +When the weight requires a gradient, wgrad computes +:math:`\partial L / \partial W`. Before the GEMM, the implementation makes sure +that both the input and output gradient have representations compatible with +the transposed access pattern. It then calls ``general_gemm`` with +``layout="NT"``. + +Wgrad may write directly into ``weight.main_grad`` when fused gradient +accumulation is enabled. It may also be deferred through the weight-gradient +store so that a higher-level schedule can execute it later. Those modes change +the destination or timing of the GEMM, not the ownership boundary: framework +code schedules the work and the common layer executes the focused operation. + +After the required GEMMs finish, the implementation completes outstanding +communication, releases temporary references, and returns the requested input, +weight, and bias gradients. The MXFP8 baseline computes its scales during +quantization and has no delayed amax-history update. Recipes with persistent +state may add recipe-specific bookkeeping after backward. + +Parallel execution variants +--------------------------- + +Tensor and sequence parallelism insert communication around the same local +phases. A ``Linear`` instance is either column-parallel or row-parallel, so +these paths should not be combined into one unconditional flow. + +.. list-table:: + :header-rows: 1 + :widths: 20 38 42 + + * - Mode + - Forward + - Backward + * - Column parallel + - Each rank owns a shard of output features. With sequence parallelism, + the input is gathered before the forward GEMM. No output reduction is + required. + - Dgrad partial results are all-reduced, or reduce-scattered with sequence + parallelism. The input required by wgrad may be gathered again. + * - Row parallel + - Each rank owns a shard along the GEMM reduction dimension. Partial + outputs are all-reduced, or reduce-scattered with sequence parallelism. + - With sequence parallelism, ``grad_output`` is gathered before the + backward GEMMs. Each rank produces its local input-gradient shard. + +Quantizer capabilities determine whether communication operates on rowwise or +columnwise quantized data and whether a representation is produced before or +after a collective. Communication-overlap implementations may fuse or pipeline +the same logical operations, but they preserve the mathematical dependencies +shown above. + +Other variations +---------------- + +The implementation contains additional branches for important features: + +**Output and gradient quantization** + ``fp8_output`` and ``fp8_grad`` request quantized boundaries between adjacent + operations and add output quantizers to the corresponding GEMMs. + +**FSDP and FSDP2** + FSDP remains an external framework facility. Transformer Engine scatters or + gathers saved state where needed and may reconstruct a weight workspace from + an FSDP2-managed weight during backward. + +**CPU offloading** + Saved activations may be marked for offload while computation continues. + +**Communication overlap** + Userbuffers or cuBLASMp paths can overlap all-gather or reduce-scatter with a + GEMM. They preserve the high-level data dependencies but change buffer and + scheduling details. + +**Backward overrides and saved inputs** + A recipe may request high-precision or dequantized backward computation. + ``save_original_input`` retains the original activation instead of relying + on a saved quantized representation. + +**Delayed scaling** + Delayed scaling maintains amax history across iterations and reduces and + updates that state after backward. This bookkeeping is not part of the + MXFP8 baseline. + +**Weight-gradient scheduling** + Fused accumulation changes the wgrad destination, and delayed wgrad changes + when the GEMM executes. + +Architectural invariants +------------------------ + +The walkthrough illustrates several rules that apply beyond ``Linear``: + +* The framework layer interprets user policy and composes framework-native and + common operations. +* Kernel implementations remain in the framework-independent common layer. +* Framework tensors own outputs, workspaces, scales, amax values, and other + operation state; common tensor handles are non-owning views. +* Forward prepares or preserves the information required by backward, but the + optimal representation depends on the recipe and execution configuration. +* Public Python and C APIs are compatibility boundaries; internal framework + bindings are not. + +How the path is tested +---------------------- + +``tests/pytorch/test_numerics.py:test_linear_accuracy`` demonstrates the +preferred correctness pattern. It constructs equivalent Transformer Engine and +``torch.nn.Linear`` modules, copies identical parameters into the native +PyTorch reference, and runs both forward and backward. + +``tests/pytorch/test_numerics.py:_test_granular_accuracy`` collects the output, +input gradient, and parameter gradients. The test compares each result with +dtype-appropriate tolerances. This is stronger than a smoke test: it establishes +the expected numerical relationship with the framework-native implementation. + +Recipe-, layout-, distributed-, and hardware-specific tests extend this +baseline for paths that do not have an exact native equivalent. diff --git a/docs/developer2/setup_build_and_run.rst b/docs/developer2/setup_build_and_run.rst index d50cb99259..e44dc2353e 100644 --- a/docs/developer2/setup_build_and_run.rst +++ b/docs/developer2/setup_build_and_run.rst @@ -1,18 +1,15 @@ Setup, Build, and Run ===================== -A reproducible path from a clean checkout to a usable development build. - -Planned coverage ----------------- - -* Prerequisites and supported development environments. -* Dependency and toolchain setup. -* Standard, optional, and framework-specific builds. -* Editable builds and selecting a development build at runtime. -* Build outputs and incremental rebuilds. -* Import checks and representative smoke tests. -* Build and setup troubleshooting. - -.. TODO: Establish canonical contributor commands while linking to the - authoritative installation and environment-variable references. +This section will describe the path from a clean checkout to a development +build that can be imported, exercised, and rebuilt efficiently. + +.. toctree:: + :maxdepth: 1 + + development_environment + unified_editable_build + split_distribution_builds + build_configuration + using_and_validating_builds + build_troubleshooting diff --git a/docs/developer2/split_distribution_builds.rst b/docs/developer2/split_distribution_builds.rst new file mode 100644 index 0000000000..3a34dbec2f --- /dev/null +++ b/docs/developer2/split_distribution_builds.rst @@ -0,0 +1,53 @@ +Split Distribution Builds +========================= + +Purpose and package boundaries +------------------------------ + +.. TODO: Explain why release packaging separates the metapackage, CUDA-major + common package, PyTorch package, and JAX package, and show their dependency + relationships. + +Distribution entry points +------------------------- + +.. TODO: Document the responsibilities of ``build_tools/wheel_utils/``, the + architecture-specific launch scripts, Dockerfiles, and + ``build_wheels.sh``. + +Common package +-------------- + +.. TODO: Outline how the common library is built, repackaged for a CUDA-major + package name, and made independent of a particular Python minor version. + +PyTorch package +--------------- + +.. TODO: Outline the PyTorch source-distribution and wheel flow, dependency on + the matching common package, framework and CUDA compatibility tags, cached + wheel lookup, and source-build fallback. + +JAX package +----------- + +.. TODO: Outline the JAX source-distribution and extension-build flow and its + dependency on the matching common package. + +Platform and CUDA matrix +------------------------ + +.. TODO: Document how x86-64 and AArch64 manylinux builds and CUDA-major + variants are selected. Keep the current matrix in one authoritative place. + +Artifacts and logs +------------------ + +.. TODO: Document wheelhouse contents, build logs, package naming, and how to + distinguish wheels from source distributions and metapackages. + +Validate distribution artifacts +------------------------------- + +.. TODO: Define metadata, dependency, installation, import, ABI, and runtime + checks for each distribution artifact. diff --git a/docs/developer2/unified_editable_build.rst b/docs/developer2/unified_editable_build.rst new file mode 100644 index 0000000000..10fed55388 --- /dev/null +++ b/docs/developer2/unified_editable_build.rst @@ -0,0 +1,48 @@ +Unified Editable Build +====================== + +Purpose and scope +----------------- + +.. TODO: Explain that the root source build compiles the common library and + all selected framework extensions together for development. + +Canonical command +----------------- + +.. TODO: Document ``pip install -e . -v --no-build-isolation`` from the + repository root, including its prerequisites and expected success criteria. + +Select framework integrations +----------------------------- + +.. TODO: Document framework auto-detection and explicit selection through + ``NVTE_FRAMEWORK``, including common-only, PyTorch-only, JAX-only, and + combined builds. + +Build pipeline +-------------- + +.. TODO: Trace the root ``setup.py`` flow through the common CMake build, + PyTorch ``CppExtension``, and JAX pybind11 extension. Identify the relevant + functions and files without relying on line numbers. + +Build outputs +------------- + +.. TODO: Document the editable installation layout, common shared library, + framework extension modules, CMake build tree, and framework-owned Python + package files. + +Incremental rebuilds +-------------------- + +.. TODO: Explain which edits require reinstalling or rebuilding, how + ``NVTE_CMAKE_BUILD_DIR`` affects reuse, and how compiler caching and build + parallelism fit into the development loop. + +Clean rebuilds +-------------- + +.. TODO: Define a safe clean-build procedure and distinguish generated build + artifacts from source files. diff --git a/docs/developer2/using_and_validating_builds.rst b/docs/developer2/using_and_validating_builds.rst new file mode 100644 index 0000000000..8b97004347 --- /dev/null +++ b/docs/developer2/using_and_validating_builds.rst @@ -0,0 +1,41 @@ +Using and Validating Development Builds +======================================= + +Confirm the selected installation +--------------------------------- + +.. TODO: Show how to verify that Python imports the current checkout rather + than a previously installed release, and how to locate the loaded common + and framework extension libraries. + +Run common functionality +------------------------ + +.. TODO: Define a minimal validation path for a common-only build and the + public C API. + +Run the PyTorch integration +--------------------------- + +.. TODO: Provide a minimal import and execution check that verifies both the + Python package and compiled PyTorch binding. + +Run the JAX integration +----------------------- + +.. TODO: Provide a minimal import and execution check that verifies both the + Python package and compiled JAX binding. + +Choose validation after a change +-------------------------------- + +.. TODO: Map Python-only, binding, common C++, CUDA, packaging, and + framework-specific changes to the smallest useful smoke checks before the + full testing guidance is applied. + +Use runnable examples +--------------------- + +.. TODO: Explain when the standalone programs under ``examples/`` are useful + for validating a development build and how they differ from test and + tutorial coverage. From c4c04b3833497642e53c8d7a8b39264d7c33de6f Mon Sep 17 00:00:00 2001 From: Przemek Tredak Date: Thu, 9 Jul 2026 08:56:45 -0700 Subject: [PATCH 18/20] Build instructions Signed-off-by: Przemek Tredak --- docs/developer2/build_configuration.rst | 40 -- docs/developer2/build_troubleshooting.rst | 39 -- .../index.rst} | 0 docs/developer2/development_environment.rst | 35 -- .../index.rst} | 0 .../index.rst} | 0 .../index.rst} | 0 docs/developer2/index.rst | 20 +- .../index.rst} | 0 .../index.rst} | 18 +- .../img/layered_architecture.svg | 0 .../img/pytorch_linear_walkthrough.svg | 0 .../index.rst} | 0 .../layered_architecture.rst | 0 .../pytorch_linear_walkthrough.rst | 0 .../{reference.rst => reference/index.rst} | 0 docs/developer2/setup_build_and_run.rst | 15 - .../build_configuration.rst | 226 +++++++++++ .../build_troubleshooting.rst | 221 +++++++++++ .../development_environment.rst | 365 ++++++++++++++++++ docs/developer2/setup_build_and_run/index.rst | 49 +++ .../split_distribution_builds.rst | 207 ++++++++++ .../unified_editable_build.rst | 183 +++++++++ .../using_and_validating_builds.rst | 174 +++++++++ docs/developer2/split_distribution_builds.rst | 53 --- .../index.rst} | 0 docs/developer2/unified_editable_build.rst | 48 --- .../using_and_validating_builds.rst | 41 -- 28 files changed, 1444 insertions(+), 290 deletions(-) delete mode 100644 docs/developer2/build_configuration.rst delete mode 100644 docs/developer2/build_troubleshooting.rst rename docs/developer2/{design_decisions.rst => design_decisions/index.rst} (100%) delete mode 100644 docs/developer2/development_environment.rst rename docs/developer2/{development_workflow.rst => development_workflow/index.rst} (100%) rename docs/developer2/{extending_transformer_engine.rst => extending_transformer_engine/index.rst} (100%) rename docs/developer2/{implementation_architecture.rst => implementation_architecture/index.rst} (100%) rename docs/developer2/{maintainer_operations.rst => maintainer_operations/index.rst} (100%) rename docs/developer2/{guide_overview.rst => overview/index.rst} (85%) rename docs/developer2/{ => project_and_architecture}/img/layered_architecture.svg (100%) rename docs/developer2/{ => project_and_architecture}/img/pytorch_linear_walkthrough.svg (100%) rename docs/developer2/{project_and_architecture.rst => project_and_architecture/index.rst} (100%) rename docs/developer2/{ => project_and_architecture}/layered_architecture.rst (100%) rename docs/developer2/{ => project_and_architecture}/pytorch_linear_walkthrough.rst (100%) rename docs/developer2/{reference.rst => reference/index.rst} (100%) delete mode 100644 docs/developer2/setup_build_and_run.rst create mode 100644 docs/developer2/setup_build_and_run/build_configuration.rst create mode 100644 docs/developer2/setup_build_and_run/build_troubleshooting.rst create mode 100644 docs/developer2/setup_build_and_run/development_environment.rst create mode 100644 docs/developer2/setup_build_and_run/index.rst create mode 100644 docs/developer2/setup_build_and_run/split_distribution_builds.rst create mode 100644 docs/developer2/setup_build_and_run/unified_editable_build.rst create mode 100644 docs/developer2/setup_build_and_run/using_and_validating_builds.rst delete mode 100644 docs/developer2/split_distribution_builds.rst rename docs/developer2/{testing_and_engineering_quality.rst => testing_and_engineering_quality/index.rst} (100%) delete mode 100644 docs/developer2/unified_editable_build.rst delete mode 100644 docs/developer2/using_and_validating_builds.rst diff --git a/docs/developer2/build_configuration.rst b/docs/developer2/build_configuration.rst deleted file mode 100644 index 38f79a431c..0000000000 --- a/docs/developer2/build_configuration.rst +++ /dev/null @@ -1,40 +0,0 @@ -Build Configuration -=================== - -Configuration principles ------------------------- - -.. TODO: Explain which settings are intended for contributors, which are - internal to release packaging, and which authoritative reference defines - each environment variable. - -Framework selection -------------------- - -.. TODO: Cover framework auto-detection and explicit selection without - duplicating the unified-build procedure. - -CUDA toolkit and target architectures -------------------------------------- - -.. TODO: Cover toolkit discovery, include paths, CUDA-major selection, and GPU - architecture targets, including their effect on build time and portability. - -Build type, caching, and parallelism ------------------------------------- - -.. TODO: Cover debug builds, CMake build directories, ccache or sccache, and - controls for parallel compilation. - -Optional components -------------------- - -.. TODO: Group options for MPI, NVSHMEM, cuBLASMp, NCCL expert parallelism, - fast-math kernels, and other optional native features by dependency and - effect. - -Release-only controls ---------------------- - -.. TODO: Identify release-build, metapackage, package-layout, and wheel-build - settings that ordinary development builds should not set. diff --git a/docs/developer2/build_troubleshooting.rst b/docs/developer2/build_troubleshooting.rst deleted file mode 100644 index 2b169effa1..0000000000 --- a/docs/developer2/build_troubleshooting.rst +++ /dev/null @@ -1,39 +0,0 @@ -Build and Setup Troubleshooting -=============================== - -Collect build diagnostics -------------------------- - -.. TODO: Define the command output, environment details, tool versions, and - generated logs needed to investigate a build failure. - -Toolchain and configuration failures ------------------------------------- - -.. TODO: Organize failures involving CUDA discovery, unsupported target - architectures, CMake, compilers, missing headers, and optional libraries. - -Submodule and source-dependency failures ----------------------------------------- - -.. TODO: Cover uninitialized, stale, or locally modified submodules and the - difference between correcting them and intentionally bypassing checks. - -Resource and incremental-build failures ---------------------------------------- - -.. TODO: Cover compiler memory pressure, parallel-job limits, stale CMake - state, compiler caches, and when a clean rebuild is appropriate. - -Import, linking, and ABI failures ---------------------------------- - -.. TODO: Cover imports resolving to the wrong installation, missing shared - libraries, mismatched CUDA or framework builds, undefined symbols, and - conflicts between editable and distribution-package layouts. - -Distribution-build failures ---------------------------- - -.. TODO: Cover container setup, platform tags, package metadata, wheel or - source-distribution dependencies, cached wheel lookup, and wheelhouse logs. diff --git a/docs/developer2/design_decisions.rst b/docs/developer2/design_decisions/index.rst similarity index 100% rename from docs/developer2/design_decisions.rst rename to docs/developer2/design_decisions/index.rst diff --git a/docs/developer2/development_environment.rst b/docs/developer2/development_environment.rst deleted file mode 100644 index d6a3ec687b..0000000000 --- a/docs/developer2/development_environment.rst +++ /dev/null @@ -1,35 +0,0 @@ -Development Environment -======================= - -Prerequisites -------------- - -.. TODO: Document supported host platforms, Python versions, compilers, CUDA, - CUDA libraries, and build tools. Link to the authoritative installation and - support documentation instead of duplicating version matrices. - -Python and framework environment --------------------------------- - -.. TODO: Document the expected virtual environment and the PyTorch and/or JAX - dependencies required before a no-build-isolation source build. - -Repository dependencies ------------------------ - -.. TODO: Document required Git submodules and when they are initialized or - checked automatically by the build. - -Optional native dependencies ----------------------------- - -.. TODO: Identify optional MPI, NVSHMEM, cuBLASMp, NCCL expert-parallel, and - other native dependencies without turning this page into a build-variable - reference. - -Verify the environment ----------------------- - -.. TODO: Provide a short pre-build checklist that verifies the selected Python - interpreter, framework installations, CUDA toolkit, compiler, CMake, Ninja, - and relevant library paths. diff --git a/docs/developer2/development_workflow.rst b/docs/developer2/development_workflow/index.rst similarity index 100% rename from docs/developer2/development_workflow.rst rename to docs/developer2/development_workflow/index.rst diff --git a/docs/developer2/extending_transformer_engine.rst b/docs/developer2/extending_transformer_engine/index.rst similarity index 100% rename from docs/developer2/extending_transformer_engine.rst rename to docs/developer2/extending_transformer_engine/index.rst diff --git a/docs/developer2/implementation_architecture.rst b/docs/developer2/implementation_architecture/index.rst similarity index 100% rename from docs/developer2/implementation_architecture.rst rename to docs/developer2/implementation_architecture/index.rst diff --git a/docs/developer2/index.rst b/docs/developer2/index.rst index 075723211e..b59a1503ff 100644 --- a/docs/developer2/index.rst +++ b/docs/developer2/index.rst @@ -15,13 +15,13 @@ understood. :maxdepth: 2 :caption: Developer Guide - guide_overview - project_and_architecture - setup_build_and_run - development_workflow - testing_and_engineering_quality - implementation_architecture - extending_transformer_engine - maintainer_operations - design_decisions - reference + overview/index + project_and_architecture/index + setup_build_and_run/index + development_workflow/index + testing_and_engineering_quality/index + implementation_architecture/index + extending_transformer_engine/index + maintainer_operations/index + design_decisions/index + reference/index diff --git a/docs/developer2/maintainer_operations.rst b/docs/developer2/maintainer_operations/index.rst similarity index 100% rename from docs/developer2/maintainer_operations.rst rename to docs/developer2/maintainer_operations/index.rst diff --git a/docs/developer2/guide_overview.rst b/docs/developer2/overview/index.rst similarity index 85% rename from docs/developer2/guide_overview.rst rename to docs/developer2/overview/index.rst index 09f44ca3b6..6dd55a5747 100644 --- a/docs/developer2/guide_overview.rst +++ b/docs/developer2/overview/index.rst @@ -24,23 +24,23 @@ that best matches what you are trying to do. * - Goal - Start with * - Prepare a first contribution - - :doc:`Setup, Build, and Run `, followed by - :doc:`Development Workflow ` + - :doc:`Setup, Build, and Run <../setup_build_and_run/index>`, followed by + :doc:`Development Workflow <../development_workflow/index>` * - Understand the codebase - - :doc:`Project and Architecture ` + - :doc:`Project and Architecture <../project_and_architecture/index>` * - Modify an existing subsystem - - :doc:`Implementation Architecture ` + - :doc:`Implementation Architecture <../implementation_architecture/index>` * - Add a new capability - - :doc:`Extending Transformer Engine ` + - :doc:`Extending Transformer Engine <../extending_transformer_engine/index>` * - Test a change or investigate a failure - :doc:`Testing and Engineering Quality - ` + <../testing_and_engineering_quality/index>` * - Understand why a durable design choice was made - - :doc:`Design Decisions ` + - :doc:`Design Decisions <../design_decisions/index>` * - Perform packaging, documentation, or release work - - :doc:`Maintainer Operations ` + - :doc:`Maintainer Operations <../maintainer_operations/index>` * - Find a command, term, support boundary, or external reference - - :doc:`Reference ` + - :doc:`Reference <../reference/index>` How the documentation fits together ----------------------------------- diff --git a/docs/developer2/img/layered_architecture.svg b/docs/developer2/project_and_architecture/img/layered_architecture.svg similarity index 100% rename from docs/developer2/img/layered_architecture.svg rename to docs/developer2/project_and_architecture/img/layered_architecture.svg diff --git a/docs/developer2/img/pytorch_linear_walkthrough.svg b/docs/developer2/project_and_architecture/img/pytorch_linear_walkthrough.svg similarity index 100% rename from docs/developer2/img/pytorch_linear_walkthrough.svg rename to docs/developer2/project_and_architecture/img/pytorch_linear_walkthrough.svg diff --git a/docs/developer2/project_and_architecture.rst b/docs/developer2/project_and_architecture/index.rst similarity index 100% rename from docs/developer2/project_and_architecture.rst rename to docs/developer2/project_and_architecture/index.rst diff --git a/docs/developer2/layered_architecture.rst b/docs/developer2/project_and_architecture/layered_architecture.rst similarity index 100% rename from docs/developer2/layered_architecture.rst rename to docs/developer2/project_and_architecture/layered_architecture.rst diff --git a/docs/developer2/pytorch_linear_walkthrough.rst b/docs/developer2/project_and_architecture/pytorch_linear_walkthrough.rst similarity index 100% rename from docs/developer2/pytorch_linear_walkthrough.rst rename to docs/developer2/project_and_architecture/pytorch_linear_walkthrough.rst diff --git a/docs/developer2/reference.rst b/docs/developer2/reference/index.rst similarity index 100% rename from docs/developer2/reference.rst rename to docs/developer2/reference/index.rst diff --git a/docs/developer2/setup_build_and_run.rst b/docs/developer2/setup_build_and_run.rst deleted file mode 100644 index e44dc2353e..0000000000 --- a/docs/developer2/setup_build_and_run.rst +++ /dev/null @@ -1,15 +0,0 @@ -Setup, Build, and Run -===================== - -This section will describe the path from a clean checkout to a development -build that can be imported, exercised, and rebuilt efficiently. - -.. toctree:: - :maxdepth: 1 - - development_environment - unified_editable_build - split_distribution_builds - build_configuration - using_and_validating_builds - build_troubleshooting diff --git a/docs/developer2/setup_build_and_run/build_configuration.rst b/docs/developer2/setup_build_and_run/build_configuration.rst new file mode 100644 index 0000000000..3ade5b6113 --- /dev/null +++ b/docs/developer2/setup_build_and_run/build_configuration.rst @@ -0,0 +1,226 @@ +Build Configuration +=================== + +Configuration principles +------------------------ + +Transformer Engine's source build is configured primarily through environment +variables. Set them before invoking ``pip``; ``setup.py``, CMake, and the +framework extension builders read the environment while the build is being +configured. + +Use three rules when changing build configuration: + +* Set only the controls needed for the current change. The default development + build should remain the baseline. +* Record non-default values in bug reports and test results. Architecture, + optional-library, and numerical flags can materially change the compiled + code. +* Reconfigure after changing a native option. If CMake reuse is questionable, + select a new ``NVTE_CMAKE_BUILD_DIR`` rather than trusting cached state. + +This page groups the important controls by purpose. ``docs/envvars.rst`` +provides the broader environment-variable reference, while the implementation +in ``setup.py``, ``build_tools/``, and +``transformer_engine/common/CMakeLists.txt`` remains authoritative. + +Framework selection +------------------- + +``NVTE_FRAMEWORK`` is a comma-separated framework selection understood by +``build_tools/utils.py` (``get_frameworks``). + +.. list-table:: + :header-rows: 1 + :widths: 28 26 46 + + * - Value + - Common library + - Framework binding + * - ``pytorch`` + - Built + - PyTorch + * - ``jax`` + - Built + - JAX + * - ``pytorch,jax`` or ``all`` + - Built once + - PyTorch and JAX + * - ``none`` + - Built + - None + * - Unset + - Built + - Every installed framework that can be imported + +An explicit value is recommended for development and CI. Auto-detection is +sensitive to unrelated packages installed in the environment. + +CUDA toolkit and target architectures +------------------------------------- + +Toolkit discovery starts with ``CUDA_HOME``, then searches for ``nvcc`` on +``PATH``, and finally checks the conventional ``/usr/local/cuda`` +location. The selected NVCC determines the CUDA toolkit version used to choose +default architecture targets. + +.. list-table:: + :header-rows: 1 + :widths: 32 26 42 + + * - Setting + - Typical value + - Effect + * - ``CUDA_HOME`` + - ``/usr/local/cuda`` + - Selects the CUDA toolkit used for NVCC and headers when several + toolkits are installed. + * - ``CUDA_PATH`` + - Toolkit prefix + - Additional CUDA location honored by CUDA-dependent components and + library discovery. + * - ``CUDNN_PATH`` or ``CUDNN_HOME`` + - cuDNN installation prefix + - Identifies a nonstandard cuDNN installation. Runtime library paths must + resolve the same installation used for the build. + * - ``NVTE_CUDA_ARCHS`` + - ``"90"`` or ``"80;90"`` + - Selects compiled GPU architectures. More targets increase build time + and binary size. + * - ``NVTE_BUILD_USE_NVIDIA_WHEELS=1`` + - Disabled by default + - Forces CUDA include discovery through installed NVIDIA Python packages + instead of a toolkit include directory. + * - ``NVTE_CMAKE_EXTRA_ARGS`` + - Space-separated CMake arguments + - Appends advanced arguments to the common CMake configure command. + +For local development, set ``NVTE_CUDA_ARCHS`` to the compute capabilities +that will actually execute the build. See +:ref:`developer-ngc-environment` for a discovery example. A target unsupported +by the selected NVCC fails during compilation; a target omitted from the list +is not expected to run on that GPU. + +``NVTE_CMAKE_EXTRA_ARGS`` is an escape hatch rather than a stable public +interface. Prefer a named build option when one exists, and include the exact +value in any reproduction. + +Build type, caching, and parallelism +------------------------------------ + +.. list-table:: + :header-rows: 1 + :widths: 32 22 46 + + * - Setting + - Default + - Effect + * - ``MAX_JOBS`` + - Maximum available + - Standard cap on concurrent native compilation. + * - ``NVTE_BUILD_MAX_JOBS`` + - Unset + - TE-specific job cap. When set, the Python build helper gives it + precedence over ``MAX_JOBS``. + * - ``NVTE_BUILD_THREADS_PER_JOB`` + - ``1`` + - Threads used internally by each NVCC process. + * - ``NVTE_BUILD_DEBUG=1`` + - ``0`` + - Selects a Debug common-library configuration and adds debug information + and assertions to framework bindings. + * - ``NVTE_CMAKE_BUILD_DIR`` + - ``build/cmake`` + - Chooses the common CMake state and object directory. + * - ``NVTE_USE_CCACHE=1`` + - ``0`` + - Uses a compiler launcher for common C++ and CUDA compilation. + * - ``NVTE_CCACHE_BIN`` + - ``ccache`` + - Selects the launcher executable; for example, an environment may use + ``sccache``. + +Parallelism and compiler threads multiply resource use. If compilation runs +out of memory, reduce ``MAX_JOBS`` first and leave +``NVTE_BUILD_THREADS_PER_JOB=1``. Caching affects build time but not the +resulting library; changing compilers, CUDA toolkits, or important flags should +use a separate cache namespace or a cleared cache. + +Optional components +------------------- + +These options change compiled capabilities and require additional libraries: + +.. list-table:: + :header-rows: 1 + :widths: 34 20 46 + + * - Setting + - Default + - Requirement and effect + * - ``NVTE_WITH_NCCL_EP`` + - Enabled when applicable + - Builds NCCL expert-parallel support for compute capability 9.0 or newer. + ``NCCL_HOME`` can select the NCCL installation. Setting ``0`` + builds stubs and omits the NCCL EP implementation. + * - ``NVTE_UB_WITH_MPI=1`` + - ``0`` + - Enables MPI bootstrap for userbuffers and requires ``MPI_HOME``. + * - ``NVTE_ENABLE_NVSHMEM=1`` + - ``0`` + - Enables NVSHMEM support and requires ``NVSHMEM_HOME``. + * - ``NVTE_WITH_CUBLASMP=1`` + - ``0`` + - Enables cuBLASMp paths and uses ``CUBLASMP_HOME`` or the matching + installed NVIDIA package. + * - ``NVTE_WITH_CUSOLVERMP=1`` + - ``0`` + - Enables distributed Newton-Schulz support and uses + ``CUSOLVERMP_HOME``. + +Enable the same capability in the common library and selected framework +binding by using one environment for the complete root build. Do not assemble +native outputs produced with conflicting feature flags. + +Specialized semantic controls +----------------------------- + +Some build options change numerical behavior rather than just build mechanics. +``NVTE_BUILD_ACTIVATION_WITH_FAST_MATH=1`` compiles selected activation +kernels with CUDA fast math, trading some numerical precision for performance. +``NVTE_BUILD_NUM_PHILOX_ROUNDS`` sets the number of Philox rounds compiled +into stochastic-rounding kernels, and must be a positive integer. Such +options should be changed only for targeted development or +experimentation. Record their value and rebuild every affected binary before +comparing numerical results. + +Release-only controls +--------------------- + +The following controls belong to packaging and release workflows, not ordinary +editable builds: + +.. list-table:: + :header-rows: 1 + :widths: 34 66 + + * - Setting + - Purpose + * - ``NVTE_RELEASE_BUILD=1`` + - Switches the root build to the split-distribution layout and suppresses + the local Git suffix in the package version. + * - ``NVTE_BUILD_METAPACKAGE=1`` + - Builds the dependency-selecting ``transformer-engine`` metapackage + without compiled extensions. It requires release-build mode. + * - ``NVTE_NO_LOCAL_VERSION=1`` + - Suppresses the Git commit local-version suffix outside release mode. + * - ``NVTE_PYTORCH_FORCE_BUILD=TRUE`` + - Prevents the PyTorch package builder from using a compatible cached + release wheel and forces local compilation. + * - ``NVTE_PYTORCH_FORCE_CXX11_ABI=TRUE`` + - Overrides the ABI reported by PyTorch for specialized packaging. An + incorrect value produces an unloadable extension. + +``NVTE_PROJECT_BUILDING`` is set internally by the setup scripts and should +not be set by contributors. The distribution workflow and its package-specific +controls are described in :doc:`split_distribution_builds`. diff --git a/docs/developer2/setup_build_and_run/build_troubleshooting.rst b/docs/developer2/setup_build_and_run/build_troubleshooting.rst new file mode 100644 index 0000000000..4f3cb41784 --- /dev/null +++ b/docs/developer2/setup_build_and_run/build_troubleshooting.rst @@ -0,0 +1,221 @@ +Build and Setup Troubleshooting +=============================== + +Troubleshoot the first failing stage. A later import or runtime error is often +a consequence of an earlier environment or linking mismatch, so preserve the +original verbose output instead of rebuilding repeatedly with different +settings. + +Collect build diagnostics +------------------------- + +Reproduce the failure from the repository root and retain the complete log: + +.. code-block:: bash + + pip install -e . -vvv --no-build-isolation 2>&1 | tee /tmp/te-build.log + +Record the settings that affect native compilation: + +.. code-block:: bash + + env | sort | grep -E '^(CUDA|CUDNN|NCCL|MPI|NVSHMEM|CUBLASMP|CUSOLVERMP|NVTE|MAX_JOBS|CC|CXX|LD_LIBRARY_PATH)=' + +Also capture the selected tools: + +.. code-block:: bash + + python --version + python -m pip --version + nvcc --version + c++ --version + cmake --version + ninja --version + git submodule status --recursive + +Record framework and backend information with the framework that failed: + +.. tabs:: + + .. tab:: PyTorch + + .. code-block:: python + + import torch + + print(torch.__version__) + print(torch.version.cuda) + print(torch._C._GLIBCXX_USE_CXX11_ABI) + print(torch.cuda.is_available()) + + .. tab:: JAX + + .. code-block:: python + + import jax + + print(jax.__version__) + print(jax.devices()) + +For an import or runtime failure, also print the package and shared-library +paths from :doc:`using_and_validating_builds`. A useful report states the +container image or host OS, GPU model, driver, exact Git revision, framework +selection, CUDA architectures, and whether a new CMake build directory changes +the result. + +Toolchain and configuration failures +------------------------------------ + +Use the last command printed before the error to identify the failing layer: + +.. list-table:: + :header-rows: 1 + :widths: 28 34 38 + + * - Stage + - Typical symptom + - First checks + * - Python build setup + - Framework or pybind11 cannot be imported. + - Confirm the intended Python and ``pip``, install the framework first, + and retain ``--no-build-isolation``. + * - CMake configure + - CUDA compiler, headers, cuDNN, or another package is not found. + - Check ``CUDA_HOME``, compiler paths, library prefixes, and enabled + optional components. + * - Common compilation + - NVCC rejects an architecture or a header is missing. + - Compare ``NVTE_CUDA_ARCHS`` with the selected NVCC and inspect the + first compiler error. + * - Framework binding compilation + - XLA, PyTorch, CUDA, MPI, or NVSHMEM headers are missing. + - Confirm the selected framework installation and that optional flags are + consistent with the common build. + * - Install or import + - Shared library is absent or a symbol cannot be resolved. + - Inspect loaded artifact paths and compare the build and runtime library + environments. + +If ``nvcc`` is missing or the wrong version is selected, set +``CUDA_HOME`` and ensure ``$CUDA_HOME/bin`` precedes other toolkits on +``PATH``. An ``unsupported gpu architecture`` error means at least one +``NVTE_CUDA_ARCHS`` target is not supported by that NVCC; remove the target +or select a newer compatible toolkit. + +Missing optional-library headers should not be worked around by adding +unrelated include paths. Either provide the installation required by the +feature and its documented ``*_HOME`` variable, or disable the feature. For +JAX, ``build_tools/jax.py`` (``xla_path``) normally obtains XLA headers +from the installed JAX FFI package; ``XLA_HOME`` is the fallback for a +deliberate external XLA tree. + +Submodule and source-dependency failures +---------------------------------------- + +Inspect the recorded state: + +.. code-block:: bash + + git submodule status --recursive + +A leading ``-`` indicates an uninitialized submodule. Initialize all recorded +revisions with: + +.. code-block:: bash + + git submodule update --init --recursive + +A leading ``+`` indicates that the submodule is checked out at a commit +other than the one recorded by the parent repository. If that was not +intentional, restore the recorded revision with the same update command. If it +was intentional development, use +``NVTE_SKIP_SUBMODULE_CHECKS_DURING_BUILD=1`` and disclose the override in +test results. + +Do not use the bypass merely to make a stale checkout compile. The build uses +submodule headers and sources directly, so an unintended revision can produce +compile errors or behavior that does not correspond to the parent commit. + +Resource and incremental-build failures +--------------------------------------- + +CUDA template compilation may use substantial memory. If the system becomes +unresponsive, NVCC is killed, or the build fails with an out-of-memory error: + +1. set ``MAX_JOBS=1``; +2. keep ``NVTE_BUILD_THREADS_PER_JOB=1``; +3. target only the needed GPU with ``NVTE_CUDA_ARCHS``; +4. retry from the existing build tree. + +If the error appears only in an incremental build, point +``NVTE_CMAKE_BUILD_DIR`` at a new empty directory and rerun the same command. +Success with the same compiler, environment, and sources isolates the old +CMake state as the cause. Changing the cache directory or temporarily +disabling ``NVTE_USE_CCACHE`` can similarly isolate a bad compiler-cache +entry. + +Do not begin troubleshooting by deleting the complete checkout or submodules. +Preserve the failing build log and user changes, and remove only identified +generated artifacts. + +Import, linking, and ABI failures +--------------------------------- + +First confirm that ``transformer_engine.__file__`` points to the intended +checkout. An NGC container already contains a released Transformer Engine +installation; the editable package must take precedence during development. + +``transformer_engine/common/__init__.py`` +(``_get_shared_object_file``) raises an error when it finds either no +matching shared object or multiple candidates. Multiple candidates usually +mean that editable, source-install, and split-wheel artifacts have been mixed +in one environment. Use a clean environment or uninstall the unrelated +distribution rather than copying libraries until the import succeeds. + +Undefined symbols and loader errors commonly indicate one of these mismatches: + +* the framework was upgraded after its private TE binding was compiled; +* common and framework libraries came from different builds or versions; +* the build and runtime resolve different CUDA, cuDNN, NCCL, or C++ ABI + libraries; +* a native optional component was enabled in only part of the build. + +Rebuild the common library and selected binding together after a framework, +compiler, or CUDA change. For PyTorch, record +``torch._C._GLIBCXX_USE_CXX11_ABI``; the package builder normally mirrors +that value. For cuDNN loading failures in a virtual environment, ensure the +build and ``LD_LIBRARY_PATH`` resolve the same cuDNN installation rather +than mixing a system library with a Python-distributed library. + +A JAX error reporting no registered CUDA implementation for a TE custom call +usually means the JAX extension was not built for the active environment. +Reinstall with ``NVTE_FRAMEWORK=jax`` and +``--no-build-isolation``, then confirm the loaded +``transformer_engine_jax`` path. + +Distribution-build failures +--------------------------- + +Release artifacts add package metadata and compatibility dimensions to the +native build. Diagnose them in the container and wheelhouse produced by +``build_tools/wheel_utils/``: + +* inspect ``wheelhouse/logs/metapackage.txt``, ``common.txt``, + ``torch.txt``, or ``jax.txt`` for the first failure; +* verify that every artifact has the same Transformer Engine version; +* confirm the requested CUDA major and manylinux platform tag; +* distinguish framework source distributions from compiled wheels; +* install into a clean environment with the checkout absent from + ``PYTHONPATH``; +* run ``python -m pip check`` before import and runtime tests. + +The PyTorch package first looks for a cached wheel matching TE, CUDA major, +PyTorch, Python, platform, and C++ ABI. A missing cached wheel is not itself an +error: the builder falls back to source compilation. Set +``NVTE_PYTORCH_FORCE_BUILD=TRUE`` when the source-build path is the behavior +under test. + +Use ``qa/L0_pytorch_wheel/test.sh`` and +``qa/L0_jax_wheel/test.sh`` as the executable references for package +assembly and clean import checks. See :doc:`split_distribution_builds` for +the artifact relationships. diff --git a/docs/developer2/setup_build_and_run/development_environment.rst b/docs/developer2/setup_build_and_run/development_environment.rst new file mode 100644 index 0000000000..ba51d55134 --- /dev/null +++ b/docs/developer2/setup_build_and_run/development_environment.rst @@ -0,0 +1,365 @@ +Development Environment +======================= + +The recommended way to prepare a Transformer Engine development environment +is to start from an NVIDIA GPU Cloud (NGC) framework container. The PyTorch +and JAX containers provide compatible framework, CUDA, cuDNN, compiler, and +Python environments, which avoids reproducing that compatibility setup on the +host. + +.. _developer-ngc-environment: + +Recommended: NGC container +-------------------------- + +Choose the container for the framework being modified: + +.. list-table:: + :header-rows: 1 + :widths: 20 40 20 + + * - Integration + - Image + - Build selection + * - PyTorch + - ``nvcr.io/nvidia/pytorch:`` + - ``NVTE_FRAMEWORK=pytorch`` + * - JAX + - ``nvcr.io/nvidia/jax:`` + - ``NVTE_FRAMEWORK=jax`` + +For example, ``26.06-py3`` identifies the June 2026 container release. Treat +the tag as an example rather than a permanent recommendation. Select a +container compatible with the hardware and the version against which the +change will be tested. Current tags are listed in the `PyTorch NGC catalog +`__ and the +`JAX NGC catalog +`__. + +Start the container +^^^^^^^^^^^^^^^^^^^ + +From the repository root on the host, initialize the Git submodules: + +.. code-block:: bash + + git submodule update --init --recursive + +Then start the container for the framework being developed: + +.. tabs:: + + .. tab:: PyTorch + + .. code-block:: bash + + export TE_IMAGE=nvcr.io/nvidia/pytorch:26.06-py3 + + docker run --gpus all --rm -it --ipc=host \ + --volume "$PWD:/workspace/TransformerEngine" \ + --workdir /workspace/TransformerEngine \ + "$TE_IMAGE" + + .. tab:: JAX + + .. code-block:: bash + + export TE_IMAGE=nvcr.io/nvidia/jax:26.06-py3 + + docker run --gpus all --rm -it --ipc=host \ + --volume "$PWD:/workspace/TransformerEngine" \ + --workdir /workspace/TransformerEngine \ + "$TE_IMAGE" + +These commands assume that the current directory is the Transformer Engine +checkout. Each provides all visible host GPUs, mounts the checkout at +``/workspace/TransformerEngine``, and makes that directory the container's +working directory. + +``--ipc=host`` is not required for compilation, but it avoids the small +default shared-memory limit when the same container is subsequently used for +framework and distributed tests. Use a more restrictive GPU or IPC setup +when the local execution environment requires it. + +Prepare the package environment +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +NGC containers set ``PIP_CONSTRAINT`` to keep Python packages compatible with +the versions shipped in the image. That constraint can prevent an editable +installation of a different Transformer Engine version. Unset it inside the +container before building the checkout: + +.. code-block:: bash + + unset PIP_CONSTRAINT + +Customize the development build +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +The default build targets several CUDA architectures and uses all available +parallel build jobs. Those defaults produce broadly usable binaries, but they +can make a local development build unnecessarily slow or exhaust the +container's memory. Configure both settings before invoking ``pip``. + +Target GPU architectures +~~~~~~~~~~~~~~~~~~~~~~~~ + +``NVTE_CUDA_ARCHS`` is a semicolon-separated list of CUDA compute +capabilities. Each additional architecture causes another set of CUDA kernels +to be compiled, so a development build should normally include only the GPUs +on which it will run. + +Use ``nvidia-smi`` to inspect the compute capabilities visible in the +container: + +.. code-block:: bash + + nvidia-smi --query-gpu=compute_cap --format=csv,noheader + +Remove the decimal point when setting the build target. For example, a reported +compute capability of ``9.0`` corresponds to: + +.. code-block:: bash + + export NVTE_CUDA_ARCHS="90" + +If the build must run on more than one GPU architecture, include each required +target, for example ``"80;90"``. Quote multi-architecture values so the +shell does not interpret the semicolon as a command separator. A binary built +for only one architecture is not intended to run on GPUs with a different +compute capability. + +Limit parallel compilation +~~~~~~~~~~~~~~~~~~~~~~~~~~ + +``MAX_JOBS`` limits the number of compiler processes that may run in +parallel. Native and CUDA compilation can consume substantial memory per job; +using every CPU core may cause memory pressure or an out-of-memory failure +instead of reducing the build time. + +Choose a value appropriate for the CPU and memory available to the container. +For example: + +.. code-block:: bash + + export MAX_JOBS=4 + +Reduce the value further if compilation exhausts memory; ``MAX_JOBS=1`` is +the lowest-resource fallback. A larger value may reduce wall-clock time on a +machine with enough memory, but does not change the resulting binaries. + +Reuse compilation results +~~~~~~~~~~~~~~~~~~~~~~~~~ + +The default ``build/cmake`` directory is inside the mounted checkout, so its +CMake state persists when the container exits. Enable compiler caching to +reuse unchanged C++ and CUDA compilation results during repeated builds: + +.. code-block:: bash + + export NVTE_USE_CCACHE=1 + +The cache benefits repeated builds in the same container. To reuse it across +disposable container sessions, mount the ccache directory from the host. See +:doc:`build_configuration` for the cache location and build-directory +controls. + +Build the selected integration +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Set the framework explicitly and install the checkout in editable mode: + +.. tabs:: + + .. tab:: PyTorch + + .. code-block:: bash + + export NVTE_FRAMEWORK=pytorch + pip install -e . -v --no-build-isolation + + .. tab:: JAX + + .. code-block:: bash + + export NVTE_FRAMEWORK=jax + pip install -e . -v --no-build-isolation + +Explicit selection makes the result independent of any other framework that +happens to be installed in the image. The editable build compiles the common +library and the selected framework extension from the mounted checkout. See +:doc:`unified_editable_build` for the build stages, outputs, and incremental +rebuild behavior. + +Confirm the environment +^^^^^^^^^^^^^^^^^^^^^^^ + +Before running a larger test, check that the selected framework sees a GPU and +that Transformer Engine resolves to the mounted checkout: + +.. tabs:: + + .. tab:: PyTorch + + .. code-block:: python + + import torch + import transformer_engine + import transformer_engine.pytorch + + print(f"PyTorch: {torch.__version__}") + print(f"CUDA available: {torch.cuda.is_available()}") + print(f"Transformer Engine: {transformer_engine.__file__}") + + .. tab:: JAX + + .. code-block:: python + + import jax + import transformer_engine + import transformer_engine.jax + + print(f"JAX: {jax.__version__}") + print(f"JAX devices: {jax.devices()}") + print(f"Transformer Engine: {transformer_engine.__file__}") + +The reported Transformer Engine path should point into +``/workspace/TransformerEngine``. The PyTorch check should report CUDA as +available, and the JAX device list should contain at least one GPU. These +checks catch accidental use of the Transformer Engine version already shipped +in the container as well as failures to expose the host GPUs. + +Existing compatible environment +-------------------------------- + +A container is not required if the host already provides a compatible +framework and native toolchain. This path is useful for persistent development +machines and CI images, but the contributor is responsible for keeping the +components compatible. + +The authoritative supported-version information belongs in the `installation +guide `__. In +addition to a supported GPU and driver, a source build needs: + +* a supported Python environment; +* PyTorch, JAX, or both, installed with GPU support; +* a CUDA toolkit that includes NVCC and development headers; +* cuDNN and the other CUDA libraries required by the selected features; +* a C++17-capable host compiler; +* CMake, Ninja, Git, pybind11, setuptools, and wheel. + +The framework must be installed before Transformer Engine. The root build +imports the selected framework while configuring its C++ extension, which is +why the installation uses ``--no-build-isolation``. + +Verify the native tools from the same shell that will run ``pip``: + +.. code-block:: bash + + python --version + python -m pip --version + nvcc --version + c++ --version + cmake --version + ninja --version + +Also confirm that pybind11 is visible to that Python interpreter: + +.. code-block:: python + + import pybind11 + + print(pybind11.__version__) + print(pybind11.get_cmake_dir()) + +Finally, verify the selected framework and its GPU backend: + +.. tabs:: + + .. tab:: PyTorch + + .. code-block:: python + + import torch + + print(torch.__version__) + print(torch.version.cuda) + print(torch.cuda.is_available()) + + .. tab:: JAX + + .. code-block:: python + + import jax + + print(jax.__version__) + print(jax.devices()) + +If multiple CUDA installations are present, set ``CUDA_HOME`` to the toolkit +that should provide NVCC and headers. Library search paths must resolve a +compatible cuDNN, NCCL, and CUDA runtime at both build and execution time. See +:doc:`build_configuration` for discovery controls and +:doc:`build_troubleshooting` for mismatch symptoms. + +Repository dependencies +----------------------- + +Transformer Engine records its source dependencies as Git submodules. +Initialize the revisions recorded by the checkout before the first build: + +.. code-block:: bash + + git submodule update --init --recursive + +The root ``setup.py`` function ``git_check_submodules`` performs the same +initialization when Git and ``.gitmodules`` are available. It accepts an +uninitialized submodule or the commit recorded by the parent repository. A +submodule checked out at a different commit is treated as a configuration +error so that a stale dependency is not used silently. + +``NVTE_SKIP_SUBMODULE_CHECKS_DURING_BUILD=1`` bypasses both the revision +check and automatic update. Use it only when intentionally developing against +modified submodules, and record that fact when reporting test results or build +failures. + +Optional native dependencies +---------------------------- + +Most contributors should begin with the default feature set. Enable an +optional component only when the change or test requires it because each +component adds discovery, compatibility, and linking requirements. + +.. list-table:: + :header-rows: 1 + :widths: 24 30 46 + + * - Component + - Build selection + - Additional requirement + * - NCCL expert parallelism + - ``NVTE_WITH_NCCL_EP`` + - Enabled by default for targets with compute capability 9.0 or newer. + Requires NCCL headers and a runtime-compatible NCCL library, and builds + the expert-parallel code from ``3rdparty/nccl``. Set the variable to + ``0`` when this capability is not needed. + * - MPI userbuffers bootstrap + - ``NVTE_UB_WITH_MPI=1`` + - Requires ``MPI_HOME`` to identify the MPI installation prefix. + * - NVSHMEM + - ``NVTE_ENABLE_NVSHMEM=1`` + - Requires ``NVSHMEM_HOME``; the framework binding and common library + must be built consistently. + * - cuBLASMp + - ``NVTE_WITH_CUBLASMP=1`` + - Uses ``CUBLASMP_HOME`` when set, otherwise the matching NVIDIA Python + distribution must be installed. + * - cuSolverMp + - ``NVTE_WITH_CUSOLVERMP=1`` + - Uses ``CUSOLVERMP_HOME``, with ``/usr`` as the build default. + +The exact option handling is defined by ``setup.py`` (``setup_common_extension``), +``build_tools/pytorch.py`` (``setup_pytorch_extension``), +``build_tools/jax.py`` (``setup_jax_extension``), and +``transformer_engine/common/CMakeLists.txt``. The configuration summary in +:doc:`build_configuration` explains which settings belong in normal +development and which are reserved for distribution builds. diff --git a/docs/developer2/setup_build_and_run/index.rst b/docs/developer2/setup_build_and_run/index.rst new file mode 100644 index 0000000000..d2ad6f4474 --- /dev/null +++ b/docs/developer2/setup_build_and_run/index.rst @@ -0,0 +1,49 @@ +Setup, Build, and Run +===================== + +This section takes a contributor from a source checkout to a development build +that can be imported, exercised, and rebuilt efficiently. The recommended path +uses a framework-specific NGC container; an existing compatible host +environment is also supported. + +Start with these pages in order: + +1. :doc:`development_environment` prepares PyTorch or JAX, the native + toolchain, submodules, and optional dependencies. +2. :doc:`unified_editable_build` explains the root editable build and its + generated artifacts. +3. :doc:`using_and_validating_builds` confirms that the checkout and compiled + libraries are the ones actually executing. +4. :doc:`build_troubleshooting` organizes failures by build stage. + +Transformer Engine has two build models: + +.. list-table:: + :header-rows: 1 + :widths: 26 34 40 + + * - Build model + - Intended use + - Entry point + * - Unified editable build + - Day-to-day implementation and testing + - ``pip install -e . -v --no-build-isolation`` from the repository + root + * - Split distribution build + - Release package construction and package-boundary validation + - Containerized scripts under ``build_tools/wheel_utils/`` + +Use :doc:`build_configuration` when a change needs controls beyond the +recommended architecture, job-limit, and compiler-cache settings introduced +in the environment page. Use :doc:`split_distribution_builds` only for +packaging work or when validating the installed-package layout. + +.. toctree:: + :maxdepth: 1 + + development_environment + unified_editable_build + build_configuration + using_and_validating_builds + build_troubleshooting + split_distribution_builds diff --git a/docs/developer2/setup_build_and_run/split_distribution_builds.rst b/docs/developer2/setup_build_and_run/split_distribution_builds.rst new file mode 100644 index 0000000000..984888a527 --- /dev/null +++ b/docs/developer2/setup_build_and_run/split_distribution_builds.rst @@ -0,0 +1,207 @@ +Split Distribution Builds +========================= + +Purpose and package boundaries +------------------------------ + +Transformer Engine's release installation is split so that the large common +library is selected by CUDA major version while each framework binding can be +built against the framework installed by the user. This is a packaging +workflow, not the recommended contributor build. + +The distributions have distinct responsibilities: + +.. list-table:: + :header-rows: 1 + :widths: 27 28 45 + + * - Distribution + - Artifact role + - Relationship + * - ``transformer-engine`` + - Dependency-selecting metapackage + - Provides extras such as ``core``, ``core_cu12``, + ``core_cu13``, ``pytorch``, and ``jax`` that select the + concrete packages. + * - ``transformer-engine-cu12`` or + ``transformer-engine-cu13`` + - CUDA-major common package + - Contains the Transformer Engine Python package and common shared + library for one CUDA major. + * - ``transformer-engine-torch`` + - PyTorch binding package + - Depends on the exactly matching common package selected from the CUDA + major reported by PyTorch. + * - ``transformer-engine-jax`` + - JAX binding package + - Depends on the exactly matching common package selected from the JAX + CUDA backend. + +At import time, ``transformer_engine/common/__init__.py`` +(``sanity_checks_for_pypi_installation`` and +``load_framework_extension``) checks that the metapackage, common package, +and installed framework package have matching Transformer Engine versions. +The private bindings may vary with framework and ABI even though the public +Python package is shared. + +Distribution entry points +------------------------- + +Release packaging is coordinated by ``build_tools/wheel_utils/``: + +.. list-table:: + :header-rows: 1 + :widths: 34 66 + + * - Entry point + - Responsibility + * - ``Dockerfile.x86`` + - Creates a ``manylinux_2_28_x86_64`` environment with the requested + CUDA toolkit, cuDNN, NCCL, and build tools. + * - ``Dockerfile.aarch`` + - Creates the corresponding ``manylinux_2_28_aarch64`` environment. + * - ``launch_x86.sh`` + - Builds and runs the x86 container for the CUDA and artifact + combinations currently produced on x86. + * - ``launch_aarch.sh`` + - Builds and runs the AArch64 container for the combinations currently + produced on AArch64. + * - ``build_wheels.sh`` + - Builds the requested metapackage, common wheel, and framework source + distributions inside the manylinux container, and writes artifacts and + logs to ``/wheelhouse``. + +The launch scripts are the authoritative description of the current platform, +CUDA-major, and artifact matrix. Do not duplicate that matrix in prose; review +the build arguments passed by the relevant launcher when changing release +coverage. + +.. warning:: + + The launch scripts remove their existing named wheelhouse directories and + perform clean container builds. Run them from an appropriate release + checkout after preserving any output needed from an earlier run. They are + substantially more expensive than the root editable build. + +Common package +-------------- + +With ``NVTE_RELEASE_BUILD=1``, the root ``setup.py`` builds the common +CMake extension but omits the PyTorch and JAX extensions. The wheel script then: + +1. builds a platform wheel from the root package; +2. unpacks the wheel; +3. changes the distribution name to + ``transformer-engine-cu``; +4. updates the distribution metadata and wheel tag; +5. repacks it with a ``py3-none-`` tag. + +The common native library does not use the Python C API, so the finished +artifact is independent of a particular Python minor version while remaining +platform-specific. The CUDA major is part of the distribution name because +the native dependencies are not interchangeable across CUDA majors. + +The empty metapackage is built separately with both +``NVTE_RELEASE_BUILD=1`` and ``NVTE_BUILD_METAPACKAGE=1``. Its extras +select common and framework packages at the exact same TE version. + +PyTorch package +--------------- + +``build_wheels.sh`` creates a source distribution from +``transformer_engine/pytorch/setup.py``. That source distribution contains +the binding sources, copied common headers, build helpers, and package +metadata, but not the common library itself. + +When a wheel is requested from the PyTorch package, its +``CachedWheelsCommand`` first constructs a release URL from: + +* Transformer Engine version; +* CUDA major reported by PyTorch; +* PyTorch version; +* Python ABI; +* operating-system platform; +* PyTorch's C++11 ABI setting. + +If that exact cached wheel exists, the command uses it. Otherwise it compiles +``transformer_engine_torch`` from source against the installed PyTorch and +depends on the exactly matching CUDA-major common package. +``NVTE_PYTORCH_FORCE_BUILD=TRUE`` bypasses the cached-wheel attempt when +testing source compilation. + +JAX package +----------- + +``build_wheels.sh`` creates a source distribution from +``transformer_engine/jax/setup.py``. During installation, the setup script: + +1. imports the installed JAX package and obtains its XLA FFI headers; +2. determines the CUDA major from the JAX GPU backend; +3. copies the common C and C++ headers required by the binding; +4. compiles ``transformer_engine_jax`` with pybind11; +5. depends on the exactly matching CUDA-major common package. + +A GPU-enabled JAX environment is therefore required when building the JAX +extension package. Unlike the PyTorch setup path, the current JAX setup does +not implement a cached release-wheel lookup. + +Platform and CUDA matrix +------------------------ + +The release matrix has three independent dimensions: + +* manylinux platform: x86-64 or AArch64; +* CUDA major and the toolkit minor used to produce that artifact; +* artifact selection: metapackage, common package, PyTorch, and JAX. + +The Dockerfiles define how one platform/toolkit environment is constructed. +The ``BUILD_METAPACKAGE``, ``BUILD_COMMON``, ``BUILD_PYTORCH``, +``BUILD_JAX``, and ``CUDA_MAJOR`` build arguments select its outputs. +The launch scripts provide the checked-in combinations. A matrix change should +update those scripts and their CI coverage together rather than only changing +this page. + +Artifacts and logs +------------------ + +Each container writes products under ``/wheelhouse``. The launcher copies +that directory to a platform- and CUDA-specific directory on the host. + +Expected contents include: + +* a ``transformer_engine`` metapackage wheel when requested; +* a ``transformer_engine_cu`` common wheel; +* a ``transformer_engine_torch`` source distribution when requested; +* a ``transformer_engine_jax`` source distribution when requested; +* ``logs/metapackage.txt``, ``common.txt``, ``torch.txt``, and + ``jax.txt`` for the corresponding requested builds. + +Artifact presence alone is not success. Inspect the log for the command that +created it, and confirm the filename, internal metadata, version, CUDA major, +Python tag, and platform tag agree. + +Validate distribution artifacts +------------------------------- + +Validate release products outside the source checkout so editable imports +cannot hide a packaging error: + +1. create a clean environment with the intended framework and GPU backend; +2. install artifacts only from the candidate wheelhouse or source + distributions; +3. run ``python -m pip check``; +4. inspect installed distribution names and versions; +5. import the common package and selected framework integration; +6. print the loaded common and framework shared-library paths; +7. execute a small forward and backward operation on a GPU; +8. repeat for every supported framework, CUDA major, Python, platform, and ABI + combination represented by the artifact. + +The executable repository references are +``qa/L0_pytorch_wheel/test.sh`` and +``qa/L0_jax_wheel/test.sh``. They assemble the common and metapackage +artifacts, build the framework package, install the candidates without +dependency substitution, and run +``tests/pytorch/test_sanity_import.py`` or +``tests/jax/test_sanity_import.py``. Release validation should add runtime +coverage beyond those import checks for the combinations being published. diff --git a/docs/developer2/setup_build_and_run/unified_editable_build.rst b/docs/developer2/setup_build_and_run/unified_editable_build.rst new file mode 100644 index 0000000000..a33f9cb71f --- /dev/null +++ b/docs/developer2/setup_build_and_run/unified_editable_build.rst @@ -0,0 +1,183 @@ +Unified Editable Build +====================== + +Purpose and scope +----------------- + +The root editable build is the normal development build. One invocation builds +the framework-independent common library, builds each selected framework +binding, and exposes the Python sources directly from the checkout. This keeps +changes to Python files immediately visible while native changes can be +recompiled in place. + +This path is distinct from :doc:`split_distribution_builds`, which creates +separate release artifacts for the common, PyTorch, and JAX packages. + +Canonical command +----------------- + +Run the build from the repository root after preparing the environment and +selecting the desired framework: + +.. code-block:: bash + + pip install -e . -v --no-build-isolation + +The options are intentional: + +* ``-e`` installs the Python package in editable mode, so imports resolve to + the checkout. +* ``-v`` exposes the CMake and compiler commands needed to diagnose a + failure. +* ``--no-build-isolation`` makes the build use the PyTorch or JAX + installation already present in the development environment. An isolated + build environment would not necessarily contain the selected framework or + the matching CUDA packages. + +A successful command must finish without a compiler or installer error. It is +not sufficient on its own: confirm the import path and execute a compiled +operation as described in :doc:`using_and_validating_builds`. + +Select framework integrations +----------------------------- + +Set ``NVTE_FRAMEWORK`` explicitly for reproducible development builds: + +.. list-table:: + :header-rows: 1 + :widths: 30 70 + + * - Value + - Result + * - ``pytorch`` + - Build the common library and PyTorch binding. + * - ``jax`` + - Build the common library and JAX binding. + * - ``pytorch,jax`` or ``all`` + - Build the common library and both bindings. + * - ``none`` + - Build only the common library and framework-independent Python package. + * - Unset + - Import PyTorch and JAX if available and build every framework that is + detected. + +Auto-detection is convenient for an interactive experiment, but it can change +when packages are added to the environment. Contributor instructions, bug +reproductions, and CI jobs should normally set the value explicitly. The +selection logic is implemented by ``build_tools/utils.py`` +(``get_frameworks``). + +Build pipeline +-------------- + +The root ``setup.py`` coordinates these build stages: + +1. ``git_check_submodules`` verifies and initializes the repository + submodules unless the check was intentionally disabled. +2. ``get_frameworks`` resolves the explicit or detected framework list. +3. ``setup_common_extension`` creates a ``CMakeExtension`` for + ``transformer_engine/common``. The build command in + ``build_tools/build_ext.py`` configures, builds, and installs + ``libtransformer_engine``. +4. For PyTorch, ``build_tools/pytorch.py`` + (``setup_pytorch_extension``) creates a + ``torch.utils.cpp_extension.CppExtension`` named + ``transformer_engine_torch``. +5. For JAX, ``build_tools/jax.py`` (``setup_jax_extension``) creates a + pybind11 extension named ``transformer_engine_jax``. + +The framework bindings contain C++ integration code but do not own CUDA +kernels. Kernel compilation remains in the common build, consistent with the +layering described in +:doc:`../project_and_architecture/layered_architecture`. + +Build outputs +------------- + +An editable installation combines files in the checkout with generated native +artifacts: + +.. list-table:: + :header-rows: 1 + :widths: 28 72 + + * - Output + - Purpose + * - Python source under ``transformer_engine/`` + - Imported directly from the checkout through editable-install metadata. + * - ``build/cmake/`` + - Default CMake configure and object-file tree for the common library. + ``NVTE_CMAKE_BUILD_DIR`` can place this tree elsewhere. + * - ``libtransformer_engine`` shared library + - Framework-independent compiled implementation. + * - ``transformer_engine_torch`` shared module + - Private PyTorch C++ binding, present only when PyTorch was selected. + * - ``transformer_engine_jax`` shared module + - Private JAX C++ binding, present only when JAX was selected. + * - Editable distribution metadata + - Allows ``pip`` and ``importlib.metadata`` to identify the source + checkout as the installed ``transformer-engine`` distribution. + +Shared-object placement varies between editable, regular source, and split +wheel installations. Runtime discovery is centralized in +``transformer_engine/common/__init__.py`` +(``_get_shared_object_file``); callers should not hard-code a generated +filename or directory. + +Incremental rebuilds +-------------------- + +Use the smallest rebuild that matches the changed files: + +.. list-table:: + :header-rows: 1 + :widths: 31 34 35 + + * - Change + - Rebuild + - Follow-up + * - Python-only code + - None for an existing editable install. + - Restart the Python process if it already imported the changed module. + * - Common C++ or CUDA + - Rerun the editable-install command so CMake builds and installs the + changed targets. + - Run a common test and each affected framework test. + * - PyTorch binding C++ + - Rerun with ``NVTE_FRAMEWORK=pytorch``. + - Import PyTorch TE and run the affected operation. + * - JAX binding C++ + - Rerun with ``NVTE_FRAMEWORK=jax``. + - Import JAX TE and run the affected operation. + * - Build scripts, compiler flags, framework selection, or optional feature + - Rerun the complete editable build. Use a fresh build tree if the old + configuration may be incompatible. + - Recheck loaded library paths and enabled features. + * - Git submodule revision or submodule sources + - Update the submodule and rerun the native build. + - Run tests covering the dependent component. + +CMake reuses ``build/cmake`` by default. ``NVTE_USE_CCACHE=1`` adds +compiler-level reuse, while ``MAX_JOBS`` controls concurrent compilation. +See :doc:`build_configuration` for their precedence and scope. + +Clean rebuilds +-------------- + +Start with an incremental rebuild. A clean build is appropriate after changing +the CUDA toolkit, host compiler, framework ABI, target architectures, or +native feature set, or when generated artifacts no longer agree with the +configuration shown in the build log. + +The safest diagnostic is to point CMake at a new empty directory rather than +deleting the existing tree: + +.. code-block:: bash + + export NVTE_CMAKE_BUILD_DIR="$(mktemp -d /tmp/te-cmake-build.XXXXXX)" + pip install -e . -v --no-build-isolation + +If that succeeds, the previous CMake tree was stale. Before removing generated +files from the checkout, inspect ``git status --ignored`` and remove only +known build artifacts. Do not remove or reset source files, submodule changes, +or unrelated user work as part of a clean build. diff --git a/docs/developer2/setup_build_and_run/using_and_validating_builds.rst b/docs/developer2/setup_build_and_run/using_and_validating_builds.rst new file mode 100644 index 0000000000..130cb80704 --- /dev/null +++ b/docs/developer2/setup_build_and_run/using_and_validating_builds.rst @@ -0,0 +1,174 @@ +Using and Validating Development Builds +======================================= + +A successful compiler invocation proves only that artifacts were produced. +Before testing a change, verify that the current Python process imports the +checkout, loads the newly built native libraries, and can execute an operation +through the selected framework. + +Confirm the selected installation +--------------------------------- + +Inspect the imported package and installed distribution: + +.. code-block:: python + + from importlib.metadata import distribution + from pathlib import Path + + import transformer_engine + + package_path = Path(transformer_engine.__file__).resolve() + dist = distribution("transformer-engine") + + print(f"Imported package: {package_path}") + print(f"Installed distribution: {dist.locate_file('')}") + +For the NGC workflow, the imported package path should be under +``/workspace/TransformerEngine``. For another editable installation, it +should be under that checkout rather than only under ``site-packages``. + +The runtime loader can also report the native artifacts it selected: + +.. code-block:: python + + from transformer_engine.common import _get_shared_object_file + + print(f"Common library: {_get_shared_object_file('core')}") + +``_get_shared_object_file`` is a private diagnostic helper, not a public +application API. It searches the editable checkout, a regular source +installation, and the split-wheel library directory. Use it to diagnose a +development environment, but do not depend on it in user code. + +Run common functionality +------------------------ + +A common-only build can be checked without importing PyTorch or JAX: + +.. code-block:: python + + import ctypes + + from transformer_engine.common import _get_shared_object_file + + common_library = _get_shared_object_file("core") + ctypes.CDLL(str(common_library)) + print(common_library) + +This confirms that the framework-independent shared library exists and that +the dynamic loader can resolve its dependencies. It does not validate an +operation's numerical behavior. Changes to the C API or common kernels should +also run the focused C++ tests under ``tests/cpp/``; the testing section +describes that workflow. + +Run a framework integration +--------------------------- + +Use a small forward and backward computation to exercise Python code, the +private framework binding, and the common library together. + +.. tabs:: + + .. tab:: PyTorch + + .. code-block:: python + + import torch + import transformer_engine.pytorch as te + + x = torch.randn(8, 16, device="cuda", requires_grad=True) + layer = te.Linear(16, 16).cuda() + + y = layer(x) + y.sum().backward() + torch.cuda.synchronize() + + print(y.shape) + print(x.grad.shape) + + .. tab:: JAX + + .. code-block:: python + + import jax + import jax.numpy as jnp + import transformer_engine.jax.flax as te_flax + + x = jnp.ones((8, 16), dtype=jnp.float32) + layer = te_flax.DenseGeneral(features=16) + variables = layer.init(jax.random.PRNGKey(0), x) + + y = jax.jit(layer.apply)(variables, x) + y.block_until_ready() + + print(y.shape) + print(jax.devices()) + +The snippets should complete on a GPU and report an output shape of +``(8, 16)``. They are smoke checks, not substitutes for comparison with a +native-framework reference. + +Choose validation after a change +-------------------------------- + +Start with the smallest check that can fail for the modified boundary, then +expand to the relevant unit and integration tests. + +.. list-table:: + :header-rows: 1 + :widths: 31 32 37 + + * - Change + - First validation + - Follow-up + * - Framework-independent Python + - Import the affected module and execute its focused behavior. + - Test through both frameworks when the code is shared. + * - PyTorch Python + - PyTorch smoke computation above. + - Focused test under ``tests/pytorch/``. + * - JAX Python + - JAX smoke computation above. + - Focused test under ``tests/jax/``. + * - Framework binding C++ + - Print the selected framework library path and execute the affected + operation. + - Framework-specific numerical or integration test. + * - Common C++ or CUDA + - Load the common library and execute the affected operation through one + frontend. + - Focused ``tests/cpp/`` test plus every affected frontend test. + * - Build or packaging logic + - Build in a clean directory or environment and inspect artifact paths. + - Matching wheel test under ``qa/L0_pytorch_wheel/`` or + ``qa/L0_jax_wheel/``. + * - Optional distributed component + - Import with the component enabled and inspect the build log. + - Focused multi-GPU launcher from the appropriate ``qa/`` class. + +For numerical changes, compare outputs and gradients with the native framework +reference rather than checking only that execution completed. For example, +``tests/pytorch/test_numerics.py`` +(``test_linear_accuracy``) is the reference pattern for a PyTorch linear +layer. + +Use runnable examples +--------------------- + +Programs under ``examples/`` are useful after the focused smoke and unit +tests pass. They exercise realistic compositions such as MNIST models, +Transformer encoders, FSDP integration, expert parallelism, and communication +overlap. + +Examples are not the primary correctness oracle: + +* they may require extra packages, datasets, or multiple GPUs; +* they cover an end-to-end configuration rather than every boundary case; +* many are intended to demonstrate execution rather than compare against a + high-precision or native-framework reference. + +Use ``examples/pytorch/`` or ``examples/jax/`` for runnable standalone +programs. The material under ``docs/examples/`` is tutorial-oriented and +may be embedded in documentation or notebooks. Test commands and CI class +selection belong in :doc:`../testing_and_engineering_quality/index`. diff --git a/docs/developer2/split_distribution_builds.rst b/docs/developer2/split_distribution_builds.rst deleted file mode 100644 index 3a34dbec2f..0000000000 --- a/docs/developer2/split_distribution_builds.rst +++ /dev/null @@ -1,53 +0,0 @@ -Split Distribution Builds -========================= - -Purpose and package boundaries ------------------------------- - -.. TODO: Explain why release packaging separates the metapackage, CUDA-major - common package, PyTorch package, and JAX package, and show their dependency - relationships. - -Distribution entry points -------------------------- - -.. TODO: Document the responsibilities of ``build_tools/wheel_utils/``, the - architecture-specific launch scripts, Dockerfiles, and - ``build_wheels.sh``. - -Common package --------------- - -.. TODO: Outline how the common library is built, repackaged for a CUDA-major - package name, and made independent of a particular Python minor version. - -PyTorch package ---------------- - -.. TODO: Outline the PyTorch source-distribution and wheel flow, dependency on - the matching common package, framework and CUDA compatibility tags, cached - wheel lookup, and source-build fallback. - -JAX package ------------ - -.. TODO: Outline the JAX source-distribution and extension-build flow and its - dependency on the matching common package. - -Platform and CUDA matrix ------------------------- - -.. TODO: Document how x86-64 and AArch64 manylinux builds and CUDA-major - variants are selected. Keep the current matrix in one authoritative place. - -Artifacts and logs ------------------- - -.. TODO: Document wheelhouse contents, build logs, package naming, and how to - distinguish wheels from source distributions and metapackages. - -Validate distribution artifacts -------------------------------- - -.. TODO: Define metadata, dependency, installation, import, ABI, and runtime - checks for each distribution artifact. diff --git a/docs/developer2/testing_and_engineering_quality.rst b/docs/developer2/testing_and_engineering_quality/index.rst similarity index 100% rename from docs/developer2/testing_and_engineering_quality.rst rename to docs/developer2/testing_and_engineering_quality/index.rst diff --git a/docs/developer2/unified_editable_build.rst b/docs/developer2/unified_editable_build.rst deleted file mode 100644 index 10fed55388..0000000000 --- a/docs/developer2/unified_editable_build.rst +++ /dev/null @@ -1,48 +0,0 @@ -Unified Editable Build -====================== - -Purpose and scope ------------------ - -.. TODO: Explain that the root source build compiles the common library and - all selected framework extensions together for development. - -Canonical command ------------------ - -.. TODO: Document ``pip install -e . -v --no-build-isolation`` from the - repository root, including its prerequisites and expected success criteria. - -Select framework integrations ------------------------------ - -.. TODO: Document framework auto-detection and explicit selection through - ``NVTE_FRAMEWORK``, including common-only, PyTorch-only, JAX-only, and - combined builds. - -Build pipeline --------------- - -.. TODO: Trace the root ``setup.py`` flow through the common CMake build, - PyTorch ``CppExtension``, and JAX pybind11 extension. Identify the relevant - functions and files without relying on line numbers. - -Build outputs -------------- - -.. TODO: Document the editable installation layout, common shared library, - framework extension modules, CMake build tree, and framework-owned Python - package files. - -Incremental rebuilds --------------------- - -.. TODO: Explain which edits require reinstalling or rebuilding, how - ``NVTE_CMAKE_BUILD_DIR`` affects reuse, and how compiler caching and build - parallelism fit into the development loop. - -Clean rebuilds --------------- - -.. TODO: Define a safe clean-build procedure and distinguish generated build - artifacts from source files. diff --git a/docs/developer2/using_and_validating_builds.rst b/docs/developer2/using_and_validating_builds.rst deleted file mode 100644 index 8b97004347..0000000000 --- a/docs/developer2/using_and_validating_builds.rst +++ /dev/null @@ -1,41 +0,0 @@ -Using and Validating Development Builds -======================================= - -Confirm the selected installation ---------------------------------- - -.. TODO: Show how to verify that Python imports the current checkout rather - than a previously installed release, and how to locate the loaded common - and framework extension libraries. - -Run common functionality ------------------------- - -.. TODO: Define a minimal validation path for a common-only build and the - public C API. - -Run the PyTorch integration ---------------------------- - -.. TODO: Provide a minimal import and execution check that verifies both the - Python package and compiled PyTorch binding. - -Run the JAX integration ------------------------ - -.. TODO: Provide a minimal import and execution check that verifies both the - Python package and compiled JAX binding. - -Choose validation after a change --------------------------------- - -.. TODO: Map Python-only, binding, common C++, CUDA, packaging, and - framework-specific changes to the smallest useful smoke checks before the - full testing guidance is applied. - -Use runnable examples ---------------------- - -.. TODO: Explain when the standalone programs under ``examples/`` are useful - for validating a development build and how they differ from test and - tutorial coverage. From 2875f467e6581583b2bea7219c9b210170194735 Mon Sep 17 00:00:00 2001 From: Przemek Tredak Date: Thu, 9 Jul 2026 10:41:09 -0700 Subject: [PATCH 19/20] More changes to build Signed-off-by: Przemek Tredak --- .../unified_editable_build.rst | 37 +++++++++++++------ 1 file changed, 25 insertions(+), 12 deletions(-) diff --git a/docs/developer2/setup_build_and_run/unified_editable_build.rst b/docs/developer2/setup_build_and_run/unified_editable_build.rst index a33f9cb71f..5176b7251c 100644 --- a/docs/developer2/setup_build_and_run/unified_editable_build.rst +++ b/docs/developer2/setup_build_and_run/unified_editable_build.rst @@ -164,20 +164,33 @@ See :doc:`build_configuration` for their precedence and scope. Clean rebuilds -------------- -Start with an incremental rebuild. A clean build is appropriate after changing -the CUDA toolkit, host compiler, framework ABI, target architectures, or -native feature set, or when generated artifacts no longer agree with the -configuration shown in the build log. - -The safest diagnostic is to point CMake at a new empty directory rather than -deleting the existing tree: +CMake and Ninja normally handle incremental changes. However, a clean rebuild is +needed when the generated build state refers to a toolchain or dependency that no +longer exists. + +Typical symptoms include: + +* Ninja reports that it cannot build a target because an old CUDA library, + such as a particular ``libcublas`` version, is missing and there is no rule + to produce it. This commonly happens after the CUDA installation is updated + from version X to version Y while the old absolute library path remains in + the build graph. +* CMake continues to use a compiler or CUDA toolkit path that no longer + matches ``CUDA_HOME`` or ``PATH``. +* CMake reports that the configured generator differs from the requested + generator. +* An optional library was upgraded, removed, or moved, but Ninja still refers + to its previous headers or shared libraries. +* Architecture or native feature settings changed and the generated targets + no longer agree with the current environment. + +In the case of build in the default location: .. code-block:: bash - export NVTE_CMAKE_BUILD_DIR="$(mktemp -d /tmp/te-cmake-build.XXXXXX)" + rm -rf build pip install -e . -v --no-build-isolation -If that succeeds, the previous CMake tree was stale. Before removing generated -files from the checkout, inspect ``git status --ignored`` and remove only -known build artifacts. Do not remove or reset source files, submodule changes, -or unrelated user work as part of a clean build. +If ``NVTE_CMAKE_BUILD_DIR`` points outside ``build/``, verify its value and +remove that directory instead. Removing the build directory does not clear the +compiler cache, so ccache can still reuse compatible compilation results. From 5cbcfb3f65d65b80bc89971def9e3f62b38dfded Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 9 Jul 2026 17:49:03 +0000 Subject: [PATCH 20/20] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/attention/test_attention.py | 8 ++------ transformer_engine/common/common.h | 1 - transformer_engine/common/gemm/cublaslt_gemm.cu | 3 +-- .../common/normalization/layernorm/ln_api.cpp | 3 +-- .../common/normalization/rmsnorm/rmsnorm_api.cpp | 3 +-- 5 files changed, 5 insertions(+), 13 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 6a78d1dd5e..f2ab4a8495 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -392,9 +392,7 @@ def test_dpa_fa4_hdim256(dtype, model_configs, model): @pytest.mark.parametrize("model", model_configs_fa4_mla.keys()) def test_dpa_fa4_mla(dtype, model_configs, model): """Test DotProductAttention with FA4: MLA (head_dim_qk != head_dim_v)""" - test_dot_product_attention( - dtype, model_configs, model, False, "bshd_bshd_bshd", False, False - ) + test_dot_product_attention(dtype, model_configs, model, False, "bshd_bshd_bshd", False, False) model_configs_fa4_swa = { @@ -543,9 +541,7 @@ def test_dpa_fa4_mask(dtype, model_configs, model): @pytest.mark.parametrize("model", model_configs_softmax.keys()) def test_dpa_softmax(dtype, model_configs, model): """Test DotProductAttention module with different softmax types""" - test_dot_product_attention( - dtype, model_configs, model, True, "bshd_bshd_bshd", False, False - ) + test_dot_product_attention(dtype, model_configs, model, True, "bshd_bshd_bshd", False, False) @pytest.mark.skipif(get_cudnn_version() < (9, 18, 0), reason="cuDNN 9.18.0+ is required.") diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index f5640a4c40..56d994a4a1 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -83,7 +83,6 @@ inline bool is_tensor_scaling(const NVTEScalingMode &mode) { inline bool is_block_scaling(const NVTEScalingMode &mode) { return !is_tensor_scaling(mode); } - inline bool is_nvfp4_scaling(const NVTEScalingMode &mode) { return mode == NVTE_NVFP4_1D_SCALING; } inline bool is_mxfp8_scaling(const NVTEScalingMode &mode) { return mode == NVTE_MXFP8_1D_SCALING; } diff --git a/transformer_engine/common/gemm/cublaslt_gemm.cu b/transformer_engine/common/gemm/cublaslt_gemm.cu index 6e2b0a1497..1033dd93e0 100644 --- a/transformer_engine/common/gemm/cublaslt_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_gemm.cu @@ -957,8 +957,7 @@ void nvte_cublas_atomic_gemm(const NVTETensor A, const NVTETensor B, NVTETensor const void *alpha_ptr = GetScalarOne(); const void *beta_ptr = accumulate ? GetScalarOne() : GetScalarZero(); - NVTE_CHECK(is_tensor_scaling(inputA->scaling_mode) && - is_tensor_scaling(inputB->scaling_mode), + NVTE_CHECK(is_tensor_scaling(inputA->scaling_mode) && is_tensor_scaling(inputB->scaling_mode), "Atomic GEMM only supports delayed scaling."); cublas_gemm(inputA, inputB, outputD, biasTensor, outputGelu, (transa) ? CUBLAS_OP_T : CUBLAS_OP_N, (transb) ? CUBLAS_OP_T : CUBLAS_OP_N, grad, wspace->data.dptr, wspace->data.shape[0], diff --git a/transformer_engine/common/normalization/layernorm/ln_api.cpp b/transformer_engine/common/normalization/layernorm/ln_api.cpp index 922080235f..21f034985b 100644 --- a/transformer_engine/common/normalization/layernorm/ln_api.cpp +++ b/transformer_engine/common/normalization/layernorm/ln_api.cpp @@ -85,8 +85,7 @@ void layernorm_fwd(const Tensor& x, // BxSxhidden_size mu->data.dptr, rsigma->data.dptr); } - bool training = - is_tensor_scaling(z->scaling_mode) || (z->columnwise_data).dptr != nullptr; + bool training = is_tensor_scaling(z->scaling_mode) || (z->columnwise_data).dptr != nullptr; auto plan = NormalizationPlanRegistry::getInstance().getNormalizationPlan( norm_backend, NVTE_Norm_Type::LayerNorm, NVTE_Norm_Stage::Forward, diff --git a/transformer_engine/common/normalization/rmsnorm/rmsnorm_api.cpp b/transformer_engine/common/normalization/rmsnorm/rmsnorm_api.cpp index 963fd8cee3..db799e758a 100644 --- a/transformer_engine/common/normalization/rmsnorm/rmsnorm_api.cpp +++ b/transformer_engine/common/normalization/rmsnorm/rmsnorm_api.cpp @@ -60,8 +60,7 @@ void rmsnorm_fwd(const Tensor &x, const Tensor &gamma, const float epsilon, Tens "cuDNN does not currently support amax output for non quantized output"); } - bool training = - is_tensor_scaling(z->scaling_mode) || (z->columnwise_data).dptr != nullptr; + bool training = is_tensor_scaling(z->scaling_mode) || (z->columnwise_data).dptr != nullptr; bool gamma_in_weight_dtype = false; if (cudnn_backend) {