diff --git a/docs/developer/architecture_overview.rst b/docs/developer/architecture_overview.rst new file mode 100644 index 0000000000..ad83b33d62 --- /dev/null +++ b/docs/developer/architecture_overview.rst @@ -0,0 +1,178 @@ +.. + 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, + Quantize, Activation, Attention, Other. + 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``, etc.) 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 + └── dense.py # JAX dense (linear) layer + +Data Flow: Forward Pass +----------------------- + +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 + :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). 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 +------------------- + +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. 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 ``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 +---------- + +- :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..c91f0cdbbd --- /dev/null +++ b/docs/developer/attention/backend_selection.rst @@ -0,0 +1,159 @@ +.. + 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 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 + + On Hopper+ (sm90): FusedAttention > FlashAttention > Unfused + On pre-Hopper: FlashAttention > FusedAttention > Unfused + +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 +--------------------- + +.. 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_UNFUSED_ATTN`` + - ``0`` to disable native PyTorch attention, ``1`` to enable (default: ``1``) + +Selecting a Backend +------------------- + +``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 + + 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 +----------------------------- + +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 + - Yes (via arbitrary mask) + - 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 NVTE_DEBUG_LEVEL=1 + +``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 +-------- + +- :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..8beb23ad04 --- /dev/null +++ b/docs/developer/attention/backends.rst @@ -0,0 +1,141 @@ +.. + 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" — 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 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 two cuDNN sub-backends selected automatically based on precision. + +Backend Details +--------------- + +Unfused Attention +^^^^^^^^^^^^^^^^^ + +**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 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 +^^^^^^^^^^^^^^ + +**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). 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``. + +Code flow: + +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 +``backends.py`` for the Python wrapper, and ``FlashAttentionUtils`` in ``utils.py`` for +version and feature compatibility checks. + +cuDNN Fused Attention (F16 and FP8) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +**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_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, + 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 +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 +^^^^^^^^^^^^^^^^^^^ + +The ``transformer_engine/common/fused_attn/`` directory also contains CUDA helper kernels +that support the attention backends: + +- ``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 full 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 +- :doc:`context_parallel` — Context parallelism strategies and implementation diff --git a/docs/developer/attention/context_parallel.rst b/docs/developer/attention/context_parallel.rst new file mode 100644 index 0000000000..127a7557d9 --- /dev/null +++ b/docs/developer/attention/context_parallel.rst @@ -0,0 +1,225 @@ +.. + 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. 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: ``attn_forward_func_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% + + Dual chunk distribution for P2P ring with 4 GPUs. + +.. + Diagram description for ``context_parallel_ring.svg``: + 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". + +**Dual Chunk (PyTorch)** + +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: + +- 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}. + +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``. + +**Striped (JAX)** + +Distributes tokens in a round-robin (interleaved) pattern across GPUs. For ``cp_size=4`` +and a 16-token sequence: + +- 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} + +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. + +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``. + +Differences from Regular Attention Flow +--------------------------------------- + +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 P2P, all-gather, and A2A strategies. +- **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 + + attn = te.DotProductAttention( + num_attention_heads=32, + kv_channels=128, + cp_group=cp_process_group, + cp_stream=cuda_stream, + ) + +Key Entry Points for Developers +--------------------------------- + +- **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), + ``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 +-------- + +- :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..bc4da4cef6 --- /dev/null +++ b/docs/developer/attention/fused_attn_kernels.rst @@ -0,0 +1,123 @@ +.. + 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/``. + +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.cpp # C API implementation, backend dispatch + ├── fused_attn_f16_arbitrary_seqlen.h/.cu # cuDNN F16 (any sequence length) + ├── fused_attn_fp8.h/.cu # cuDNN FP8 + ├── 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 +----------------- + +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. + +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 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): ``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 +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 +-------- + +- :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..817c48a4ad --- /dev/null +++ b/docs/developer/attention/img/attention_backends.svg @@ -0,0 +1,114 @@ + + + + + + + + + + + + + + + + + + + + + + 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 + FA2 / FA3 / FA4 + + + + SM80+ (Ampere) + + + + Fused (cuDNN) + + + + + + + + + + + F16 Fused + + + + Ampere+ | SM80+ + cuDNN graph API + + + + + FP8 Fused + + + + Hopper+ | SM90+ + Best FP8 performance + + + + + Legend: + + Unfused (fallback) + + FlashAttention + + cuDNN Fused + SM80 = Ampere | SM90 = Hopper | SM100/SM120 = Blackwell + \ 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..ca48c82d95 --- /dev/null +++ b/docs/developer/attention/jax_attention.rst @@ -0,0 +1,95 @@ +.. + 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 public function is: + +- ``fused_attn()`` — Fused attention using cuDNN (equivalent to PyTorch's cuDNN fused + backend). + +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 +------------------------- + +.. 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 + + 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( + (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 +-------- + +- :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..09edfcbb29 --- /dev/null +++ b/docs/developer/attention/pytorch_attention.rst @@ -0,0 +1,106 @@ +.. + 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) +- ``ScaledAlignedCausalMaskedSoftmax`` — Aligned causal mask variant + +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..5967193419 --- /dev/null +++ b/docs/developer/cpp_core/build_system.rst @@ -0,0 +1,166 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +.. _build-system: + +Build System +============ + +Transformer Engine has two build paths: a monolithic developer build and a split +pip-wheel build for distribution. + +Monolithic Developer Build +-------------------------- + +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 libtransformer_engine.so (C++ core) + └── Builds framework-specific 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 +------------------------- + +.. list-table:: + :header-rows: 1 + :widths: 30 70 + + * - Variable + - Description + * - ``NVTE_FRAMEWORK`` + - Select framework: ``pytorch``, ``jax``, ``all``, ``none``, 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`` + - 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_WITH_NCCL_EP`` + - Build NCCL-based expert-parallel support (default: ``1``; requires an sm90+ target) + +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 +-------------- + +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`` attempts to initialize every registered submodule unless +``NVTE_SKIP_SUBMODULE_CHECKS_DURING_BUILD`` is set. Manual setup is sometimes needed: + +.. code-block:: bash + + git submodule update --init --recursive + +Developer Build Tips +-------------------- + +**Fast incremental rebuild** (C++ only, skip Python): + +.. code-block:: bash + + cmake --build build/cmake --parallel 4 + +**Enable ccache** for fast rebuilds when iterating on C++ code: + +.. code-block:: bash + + 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): + +.. code-block:: bash + + cmake --build build/cmake --verbose --parallel 4 + +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/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..187e310ff5 --- /dev/null +++ b/docs/developer/cpp_core/img/tensor_struct.svg @@ -0,0 +1,119 @@ + + + + + + + + + + + + + + + + + + + + + + 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 + + row_scaled_nvfp4 + : bool + + + nvfp4_e4m3_max + : int + + + + + + + + + + + + + + + + SimpleTensor + + + + dptr + : void* + + + shape + : Shape + + + 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..e76f5fd6cd --- /dev/null +++ b/docs/developer/cpp_core/kernel_areas.rst @@ -0,0 +1,98 @@ +.. + 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. + +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. + +Rather than cataloging each area, this page describes the general architectural patterns +that a developer writing new kernels should understand. + +Output-Driven Computation +-------------------------- + +A key design principle: the **output tensor** dictates what computation the kernel +performs. Kernels inspect the output ``NVTETensor`` to determine: + +- 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. + +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. + +Scaling Mode Dispatch +---------------------- + +Kernels that handle multiple scaling modes use a switch on the output's +``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 + + 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() + +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. + +Architecture Dispatch +---------------------- + +Many kernel areas provide architecture-specific implementations for different GPU +generations (Ampere, Hopper, Blackwell). The dispatch mechanism varies by area: + +- **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. + +- **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). + +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. + +Fused Kernels +-------------- + +Many areas support fusing quantization with the primary computation: + +- **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. + +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. + +For fused attention, see :doc:`/developer/attention/fused_attn_kernels` for dedicated +coverage. + +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 new file mode 100644 index 0000000000..a5a4708af7 --- /dev/null +++ b/docs/developer/cpp_core/scaling_modes.rst @@ -0,0 +1,177 @@ +.. + 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. + +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% + + 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 | Ada + MXFP8_1D_SCALING | 1 | 1 scale per 32 elements | 32 | FP8 + E8M0 scales | 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 + +Non-TN GEMM Support +--------------------- + +A cross-cutting concern that affects all scaling modes is non-TN GEMM support. +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, +``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. + +.. 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 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 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 + +C++ helpers: + +.. code-block:: cpp + + // True for NVTE_DELAYED_TENSOR_SCALING (per-tensor scaling, including current) + bool is_tensor_scaling(const NVTEScalingMode &mode); + +MXFP8 1D Scaling (``NVTE_MXFP8_1D_SCALING``) +---------------------------------------------- + +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``: 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. Note that this scale layout differs from the ``columnwise_data`` tensor + 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. + +C++ helper: + +.. code-block:: cpp + + 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), 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: + +.. code-block:: cpp + + bool is_block_scaling(const NVTEScalingMode &mode); + // 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``) +---------------------------------------------- + +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), 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 ``[N, ceil(M/16)]``, padded to ``[128, 4]`` + +Like MXFP8, scales require GEMM swizzling. + +C++ helper: + +.. code-block:: cpp + + bool is_nvfp4_scaling(const NVTEScalingMode &mode); + +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 + and dgrad GEMMs). + +See :doc:`/developer/quantization/rowwise_columnwise` for how these layouts interact +with GEMM execution. + +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/cpp_core/type_system.rst b/docs/developer/cpp_core/type_system.rst new file mode 100644 index 0000000000..cddd54210b --- /dev/null +++ b/docs/developer/cpp_core/type_system.rst @@ -0,0 +1,317 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +.. _type-system: + +Type System +=========== + +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 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. + +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) +-------------------- + +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. +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_set_tensor_param_v2()`` with ``NVTETensorParam`` keys: + +.. code-block:: c + + enum NVTETensorParam { + kNVTERowwiseData = 0, + kNVTEColumnwiseData = 1, + kNVTEScale = 2, + kNVTEAmax = 3, + kNVTERowwiseScaleInv = 4, + kNVTEColumnwiseScaleInv = 5, + kNVTEColumnwiseAmax = 6, + kNVTEWithGEMMSwizzledScales = 7, + kNVTERowScaledNVFP4 = 8, + kNVTENVFP4E4M3Max = 9, + kNVTENumTensorParams, + }; + +NVTEScalingMode +^^^^^^^^^^^^^^^ + +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++ Wrappers (External) +----------------------- + +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`` 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 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 lightweight internal tensor descriptor with implicit conversion to/from +``NVTEBasicTensor``: + +.. code-block:: cpp + + struct SimpleTensor { + void *dptr; + Shape 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 + + row_scaled_nvfp4 : bool + + nvfp4_e4m3_max : int + Below, a smaller UML box for "SimpleTensor": + + dptr : void* + + shape : Shape + + 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 + 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 +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. +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: + +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. **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/distributed/comm_gemm_overlap.rst b/docs/developer/distributed/comm_gemm_overlap.rst new file mode 100644 index 0000000000..1906edea36 --- /dev/null +++ b/docs/developer/distributed/comm_gemm_overlap.rst @@ -0,0 +1,137 @@ +.. + 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 + * - ``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( + hidden_size=4096, + ffn_hidden_size=16384, + num_attention_heads=32, + tp_group=tp_group, + ub_overlap_rs=True, # Overlap reduce-scatter in forward (row-parallel) + ub_overlap_ag=True, # Overlap all-gather in forward (column-parallel) + ) + +At the ``_Linear`` autograd level, finer-grained control is available via parameters that +specify overlap for each pass direction: + +- ``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 +-------------------------- + +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..40698ccb57 --- /dev/null +++ b/docs/developer/distributed/fsdp_integration.rst @@ -0,0 +1,190 @@ +.. + 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. 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 two challenges for quantized training: + +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. + +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. + +These two challenges are handled by separate mechanisms, described below. + +FSDP2: Quantized Weight All-Gather +------------------------------------ + +**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. + +**Location**: Methods on ``QuantizedTensor`` subclasses in ``transformer_engine/pytorch/tensor/``. + +``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. + * - ``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:: + + 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 + + from torch.distributed._composable.fsdp import fully_shard + import transformer_engine.pytorch as te + + model = te.TransformerLayer(...) + + # FSDP2 automatically handles quantized weight all-gather + fully_shard(model) + + with te.autocast(enabled=True): + 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. + +``_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. + +``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.autocast(enabled=True): + output = model(input) + +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 +----------- + +- 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. +- ``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 +-------- + +- :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..422f74d681 --- /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..674ecbded8 --- /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..760c3ef9d6 --- /dev/null +++ b/docs/developer/distributed/tensor_parallel.rst @@ -0,0 +1,123 @@ +.. + 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 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 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" + +.. figure:: ./img/row_parallel.svg + :align: center + :width: 70% + + 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 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" + +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**: 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). + +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 ``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**: 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 +-------------------- + +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..af08f5e37b --- /dev/null +++ b/docs/developer/index.rst @@ -0,0 +1,47 @@ +.. + 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 PyTorch + 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 + organization. +- **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`. + +.. 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 + testing 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..54ce96fadc --- /dev/null +++ b/docs/developer/jax_frontend/module_system.rst @@ -0,0 +1,133 @@ +.. + 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 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, + kernel_axes=("input", "tp"), + ) + + 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 + + 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 +-------- + +- :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..3922fdb474 --- /dev/null +++ b/docs/developer/jax_frontend/quantization.rst @@ -0,0 +1,116 @@ +.. + 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 (JAX does not yet support generic block scaling 1D/2D) + * - ``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. 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: + +.. 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..84fc000e63 --- /dev/null +++ b/docs/developer/jax_frontend/xla_ffi_primitives.rst @@ -0,0 +1,179 @@ +.. + 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** inner and outer JAX primitives and lower them to XLA FFI calls. +2. **Define abstract evaluation** (shape/dtype inference without running the kernel). +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): + + @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. + +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. + +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/``, organized by +functional area: + +**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 +-------- + +- :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..b164a56e76 --- /dev/null +++ b/docs/developer/linear_walkthrough.rst @@ -0,0 +1,449 @@ +.. + 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. "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_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". + +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 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 +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 +----- + +Here is the basic example of how one could create and call the Linear layer: + +.. code-block:: python + + import transformer_engine.pytorch as te + from transformer_engine.common.recipe import MXFP8BlockScaling + + # Create a Linear module + linear = te.Linear(4096, 16384, bias=True) + + # Enable FP8 with MXFP8 recipe + recipe = MXFP8BlockScaling() + with te.autocast(enabled=True, recipe=recipe): + output = linear(input) # Triggers the flow below + +Phase 1: Module Forward Entry +------------------------------ + +**File**: ``transformer_engine/pytorch/module/linear.py``, ``Linear.forward()`` + +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 ``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. 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`` — 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. + +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()`` + +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 — 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 + # to save GPU memory. + input_quantizer.set_usage(rowwise=True, columnwise=backward_needs_input) + + # 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 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 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): + +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 + + # 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 + +Phase 3: _Linear.forward() — Weight Quantization +-------------------------------------------------- + +**File**: ``transformer_engine/pytorch/module/linear.py``, ``_Linear.forward()`` + +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. + +.. code-block:: python + + # 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, 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=(None if is_first_microbatch is None else "weight"), + update_workspace=is_first_microbatch is None or is_first_microbatch, + ) + +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). + +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`` + +.. code-block:: python + + # Forward GEMM: output = input @ weight^T + # 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) + out_dtype=activation_dtype, + 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, + ) + +The call traverses four layers, matching the :doc:`architecture_overview`: + +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 +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 (see the compute capability checks in +``transformer_engine/common/gemm/`` for the architecture dispatch logic). + +Phase 5: _Linear.forward() — Output Communication +---------------------------------------------------- + +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 (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 +-------------------------------------------------- + +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 + + # 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) + + # 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_func_ctx`` pair handles the fact that +``QuantizedTensorStorage`` is not a ``torch.Tensor`` and therefore cannot be passed +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. + +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. 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()`` + +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 + + # 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) + 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( + 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, + ) + +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) +------------------------------------------------------- + +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 + + # 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 + wgrad, *_ = general_gemm( + inputmat_total, # Columnwise input saved from forward + grad_output, # Columnwise grad_output from Phase 7 + out_dtype=activation_dtype, + grad=True, + ub=ub_obj_wgrad, + ub_type=ub_type_wgrad, + ) + +``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. + +Phase 9: Post-Backward Cleanup +-------------------------------- + +After both GEMMs complete, the backward pass handles recipe-specific bookkeeping. + +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. + +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 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 +-------- + +- :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 +- :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 new file mode 100644 index 0000000000..e186971b15 --- /dev/null +++ b/docs/developer/pytorch_frontend/autograd_integration.rst @@ -0,0 +1,152 @@ +.. + 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 — 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 + +**Backward** (pseudocode): + +.. code-block:: python + + @staticmethod + def backward(ctx, grad_output): + # 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) + + # 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`` 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 +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`` (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 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 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 +------------------------ + +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. + +TE supports activation recomputation (gradient checkpointing) at multiple granularities: + +- **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. Pass + ``checkpoint_core_attention=True`` to ``TransformerLayer.forward()`` to recompute the + core attention operation during backward. + +FP8 Tensors in Autograd +------------------------ + +``QuantizedTensor`` (a ``torch.Tensor`` subclass) can be saved via +``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 +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +``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): ``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 + 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 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/cpp_extensions.rst b/docs/developer/pytorch_frontend/cpp_extensions.rst new file mode 100644 index 0000000000..0ab2027a1c --- /dev/null +++ b/docs/developer/pytorch_frontend/cpp_extensions.rst @@ -0,0 +1,209 @@ +.. + 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 accepts Python objects (``QuantizedTensorStorage``, ``QuantizedTensor``, or +plain ``torch.Tensor``), converts them to ``TensorWrapper`` / ``NVTETensor`` handles, and +calls the C API. + +Architecture +------------ + +.. code-block:: text + + Python (transformer_engine/pytorch/cpp_extensions/) + │ Optional auxiliary processing (e.g. custom recipe handling) + ▼ + pybind11 (transformer_engine/pytorch/csrc/extensions/) + │ py::handle → TensorWrapper conversion, GIL release, C API call + ▼ + C API (transformer_engine/common/include/transformer_engine/) + │ NVTETensor handle → internal Tensor struct + ▼ + CUDA kernels (transformer_engine/common/) + +Python Side +----------- + +**Location**: ``transformer_engine/pytorch/cpp_extensions/`` + +Only a subset of C++ extensions have dedicated Python wrapper modules: + +- ``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 +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 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, quantization_params=None, ...): + # Auxiliary logic (custom tensor dispatch, debug handling, etc.) + if is_custom(A) or is_custom(B): + return custom_gemm(...) + + # 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 out, bias_grad, gelu_input, extra_output + +C++ Side (pybind11) +------------------- + +**Location**: ``transformer_engine/pytorch/csrc/extensions/`` + +The pybind11 module (``transformer_engine/pytorch/csrc/extensions/pybind.cpp``) +registers Python-callable functions. These functions: + +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 + + 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. + +TensorWrapper and NVTETensor +---------------------------- + +``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: + +1. It iterates over a registry of known quantized types (defined in + ``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 + 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. + +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. + +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 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 +---------------------- + +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. (Optional) Add a Python wrapper in ``cpp_extensions/`` if auxiliary processing is + needed before calling the ``tex`` function. 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..db09b50cee --- /dev/null +++ b/docs/developer/pytorch_frontend/index.rst @@ -0,0 +1,19 @@ +.. + 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. + +.. 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..9baf5f3220 --- /dev/null +++ b/docs/developer/pytorch_frontend/module_hierarchy.rst @@ -0,0 +1,96 @@ +.. + 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, two branches: + Branch 1: TransformerEngineBaseModule, with children: + ├── Linear + ├── LayerNormLinear + ├── LayerNormMLP + ├── GroupedLinear + └── DotProductAttention + 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 + └── Fp8Unpadding + +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. +- ``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** + +- 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** + +- ``set_tensor_parallel_group()``: Configure TP process group. +- ``sequence_parallel``: A plain boolean attribute (set in ``__init__``, no dedicated + setter method). +- Manages all-reduce of amax values across distributed ranks. + +**Weight Caching** + +- Provides caching infrastructure for quantized weights that persist across microbatches + during gradient accumulation. + +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 +---------------- + +A typical forward pass through a TE module: + +.. code-block:: text + + 1. autocast() sets global FP8 state + 2. module.forward() called + 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. end_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..248355367f --- /dev/null +++ b/docs/developer/pytorch_frontend/ops_framework.rst @@ -0,0 +1,286 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +.. _ops-framework: + +Op Fusion Framework (``te.ops``) +================================ + +Transformer Engine includes an operation fusion framework (``transformer_engine/pytorch/ops/``) +that enables composing and automatically fusing operations for better performance. + +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 + :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". + +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 + + def save_for_backward(self, *tensors): ... + +Execution Flow +-------------- + +**First call** through ``Sequential.forward()``: + +.. code-block:: text + + 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 + + # Type signature for fusion functions + OperationFusionFunction = Callable[ + [list[FusibleOperation], ...], # ops to scan + list[FusibleOperation], # ops after fusion (same or fewer) + ] + +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() + + op.fuser_forward( + ..., + prev_op_grad_output_quantizer=prev_quantizer, + next_op_input_quantizer=next_quantizer, + ) + +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. + +Available Operations +--------------------- + +Basic operations are located in ``transformer_engine/pytorch/ops/basic/``. Each file +implements a single atomic operation (GEMM, normalization, activation, communication, +etc.). + +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 +------------------------- + +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 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 +-------- + +- :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/developer/pytorch_frontend/transformer_layer.rst b/docs/developer/pytorch_frontend/transformer_layer.rst new file mode 100644 index 0000000000..c2a950b71a --- /dev/null +++ b/docs/developer/pytorch_frontend/transformer_layer.rst @@ -0,0 +1,103 @@ +.. + 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. + +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 ``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 +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/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..ff180903f9 --- /dev/null +++ b/docs/developer/quantization/class_hierarchy.rst @@ -0,0 +1,229 @@ +.. + 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) -> QuantizedTensor: ... + def quantize(self, tensor, *, out=None, dtype=None) -> QuantizedTensor: ... + + # Usage control + def set_usage(self, *, rowwise=False, columnwise=False): ... + +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 +^^^^^^^^^^^^^^^^^^^ + +.. 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 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 + +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: + # 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) -> 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 + 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) +- 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``) +--------------------------------------------------------------------------- + +``QuantizedTensor`` is a ``torch.Tensor`` subclass that wraps a +``QuantizedTensorStorage`` and adds PyTorch integration: + +.. code-block:: python + + class QuantizedTensor(torch.Tensor): + # torch.Tensor subclass machinery + def __torch_dispatch__(cls, func, types, args, kwargs): ... + + # Data access (get_data_tensors() is on concrete Storage subclasses, not here) + 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: + +- **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 + uses ``get_data_tensors()`` to extract raw data directly, avoiding dequantization. +- **Autograd compatible**: Can be saved in autograd's ``ctx.save_for_backward()``. + +Lifecycle +--------- + +A typical quantization lifecycle: + +.. code-block:: python + + 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) + qinput = quantizer(input_tensor) + + # 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() + +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..1480b59a31 --- /dev/null +++ b/docs/developer/quantization/index.rst @@ -0,0 +1,26 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +Quantization +============ + +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 developer documentation here focuses on the internal implementation: + +- 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 + + 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..2dc99f8412 --- /dev/null +++ b/docs/developer/quantization/rowwise_columnwise.rst @@ -0,0 +1,131 @@ +.. + 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 = 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 = 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 +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: rowwise for forward GEMM, columnwise for backward wgrad + fwd_quantizer.set_usage(rowwise=True, columnwise=True) + + # Weight quantizer: rowwise for forward GEMM, columnwise for backward dgrad + weight_quantizer.set_usage(rowwise=True, columnwise=True) + + # Backward grad_output quantizer: rowwise for dgrad, columnwise for wgrad + bwd_quantizer.set_usage(rowwise=True, columnwise=True) + +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 +---------------------------- + +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 + // ... + }; + +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 +------------------------ + +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), + 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, 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 +: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..ab4022046d --- /dev/null +++ b/docs/developer/quantization/scaling_recipes.rst @@ -0,0 +1,144 @@ +.. + 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.autocast(enabled=True, 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 ``autocast`` context manager: + +.. code-block:: python + + # Pseudocode for autocast + @contextmanager + def autocast(enabled, recipe): + FP8GlobalStateManager.set_enabled(enabled) + FP8GlobalStateManager.set_recipe(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 +``prepare_forward()`` / ``end_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. + +.. warning:: + + For the amax all-reduce to work correctly, every rank must participate in the + 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). + +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 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/developer2/design_decisions/index.rst b/docs/developer2/design_decisions/index.rst new file mode 100644 index 0000000000..ff87d1a75d --- /dev/null +++ b/docs/developer2/design_decisions/index.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/index.rst b/docs/developer2/development_workflow/index.rst new file mode 100644 index 0000000000..9b1dad30d6 --- /dev/null +++ b/docs/developer2/development_workflow/index.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/index.rst b/docs/developer2/extending_transformer_engine/index.rst new file mode 100644 index 0000000000..47f01d5a5a --- /dev/null +++ b/docs/developer2/extending_transformer_engine/index.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/implementation_architecture/index.rst b/docs/developer2/implementation_architecture/index.rst new file mode 100644 index 0000000000..507cb4a67c --- /dev/null +++ b/docs/developer2/implementation_architecture/index.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..b59a1503ff --- /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 + + 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/index.rst b/docs/developer2/maintainer_operations/index.rst new file mode 100644 index 0000000000..8dd3fd77a1 --- /dev/null +++ b/docs/developer2/maintainer_operations/index.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/overview/index.rst b/docs/developer2/overview/index.rst new file mode 100644 index 0000000000..6dd55a5747 --- /dev/null +++ b/docs/developer2/overview/index.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 <../setup_build_and_run/index>`, followed by + :doc:`Development Workflow <../development_workflow/index>` + * - Understand the codebase + - :doc:`Project and Architecture <../project_and_architecture/index>` + * - Modify an existing subsystem + - :doc:`Implementation Architecture <../implementation_architecture/index>` + * - Add a new capability + - :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 <../design_decisions/index>` + * - Perform packaging, documentation, or release work + - :doc:`Maintainer Operations <../maintainer_operations/index>` + * - Find a command, term, support boundary, or external reference + - :doc:`Reference <../reference/index>` + +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/project_and_architecture/img/layered_architecture.svg b/docs/developer2/project_and_architecture/img/layered_architecture.svg new file mode 100644 index 0000000000..562458096a --- /dev/null +++ b/docs/developer2/project_and_architecture/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/project_and_architecture/img/pytorch_linear_walkthrough.svg b/docs/developer2/project_and_architecture/img/pytorch_linear_walkthrough.svg new file mode 100644 index 0000000000..ce51bc5bca --- /dev/null +++ b/docs/developer2/project_and_architecture/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/project_and_architecture/index.rst b/docs/developer2/project_and_architecture/index.rst new file mode 100644 index 0000000000..291e0cd1e5 --- /dev/null +++ b/docs/developer2/project_and_architecture/index.rst @@ -0,0 +1,206 @@ +Project and Architecture +======================== + +This section describes where Transformer Engine fits in the software stack, +how its major pieces work together, and how those boundaries shape the code. + +Role in the software stack +-------------------------- + +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. + +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/project_and_architecture/layered_architecture.rst b/docs/developer2/project_and_architecture/layered_architecture.rst new file mode 100644 index 0000000000..e58fe1bc99 --- /dev/null +++ b/docs/developer2/project_and_architecture/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/pytorch_linear_walkthrough.rst b/docs/developer2/project_and_architecture/pytorch_linear_walkthrough.rst new file mode 100644 index 0000000000..90ed879d8f --- /dev/null +++ b/docs/developer2/project_and_architecture/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/reference/index.rst b/docs/developer2/reference/index.rst new file mode 100644 index 0000000000..04b78e4fc0 --- /dev/null +++ b/docs/developer2/reference/index.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/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..5176b7251c --- /dev/null +++ b/docs/developer2/setup_build_and_run/unified_editable_build.rst @@ -0,0 +1,196 @@ +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 +-------------- + +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 + + rm -rf build + pip install -e . -v --no-build-isolation + +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. 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/testing_and_engineering_quality/index.rst b/docs/developer2/testing_and_engineering_quality/index.rst new file mode 100644 index 0000000000..3739c86fad --- /dev/null +++ b/docs/developer2/testing_and_engineering_quality/index.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. diff --git a/docs/index.rst b/docs/index.rst index fcd15a7a11..bbde1608ad 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 @@ -71,3 +71,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 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 diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 2dbf94fc20..f2ab4a8495 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, True]) @@ -130,7 +129,6 @@ def test_dot_product_attention( model_configs, model, ckpt_attn, - workspace_opt, qkv_layout, swa, pad_between_seqs, @@ -209,7 +207,6 @@ def test_dot_product_attention( "UnfusedDotProductAttention", ckpt_attn, qkv_layout, - workspace_opt, pad_between_seqs, is_training, ) @@ -222,7 +219,6 @@ def test_dot_product_attention( "FusedAttention", ckpt_attn, qkv_layout, - workspace_opt, pad_between_seqs, is_training, ) @@ -235,7 +231,6 @@ def test_dot_product_attention( "FlashAttention", ckpt_attn, qkv_layout, - workspace_opt, pad_between_seqs, is_training, ) @@ -267,7 +262,7 @@ def test_dot_product_attention( @pytest.mark.parametrize("model", ["base_1_1", "base_2_1"]) def test_dpa_checkpoint(dtype, model_configs, model): """Test DotProductAttention module with checkpointing""" - test_dot_product_attention(dtype, model_configs, model, True, True, None, False, False) + test_dot_product_attention(dtype, model_configs, model, True, None, False, False) model_configs_max_logit = { @@ -294,7 +289,7 @@ def test_dpa_max_logit(dtype, model_configs, model, qkv_layout): """Test DotProductAttention module with checkpointing""" config = model_configs[model] config.return_max_logit = True - test_dot_product_attention(dtype, model_configs, model, False, True, qkv_layout, False, False) + test_dot_product_attention(dtype, model_configs, model, False, qkv_layout, False, False) model_configs_num_splits = { @@ -315,7 +310,6 @@ def test_dpa_num_splits(dtype, model_configs, model): model_configs, model, False, - True, None, False, False, @@ -349,7 +343,7 @@ def test_dpa_num_splits(dtype, model_configs, model): @pytest.mark.parametrize("model", model_configs_fa4_base.keys()) def test_dpa_fa4_base(dtype, model_configs, model): """Test DotProductAttention with FA4: base configs, GQA, num_splits""" - test_dot_product_attention(dtype, model_configs, model, False, True, None, False, False) + test_dot_product_attention(dtype, model_configs, model, False, None, False, False) # head_dim=256 is supported only on SM100 via FA4's dedicated kernel @@ -374,7 +368,7 @@ def test_dpa_fa4_base(dtype, model_configs, model): @pytest.mark.parametrize("model", model_configs_fa4_hdim256.keys()) def test_dpa_fa4_hdim256(dtype, model_configs, model): """Test DotProductAttention with FA4: head_dim=256 dedicated kernel on SM100""" - test_dot_product_attention(dtype, model_configs, model, False, True, None, False, False) + test_dot_product_attention(dtype, model_configs, model, False, None, False, False) model_configs_fa4_mla = { @@ -398,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, True, "bshd_bshd_bshd", False, False - ) + test_dot_product_attention(dtype, model_configs, model, False, "bshd_bshd_bshd", False, False) model_configs_fa4_swa = { @@ -425,7 +417,7 @@ def test_dpa_fa4_mla(dtype, model_configs, model): @pytest.mark.parametrize("qkv_layout", ["sbhd_sbhd_sbhd", "bshd_bshd_bshd"]) def test_dpa_fa4_sliding_window(dtype, model_configs, model, qkv_layout): """Test DotProductAttention with FA4: sliding window attention""" - test_dot_product_attention(dtype, model_configs, model, False, True, qkv_layout, True, False) + test_dot_product_attention(dtype, model_configs, model, False, qkv_layout, True, False) model_configs_fa4_varlen = { @@ -448,7 +440,7 @@ def test_dpa_fa4_sliding_window(dtype, model_configs, model, qkv_layout): @pytest.mark.parametrize("qkv_layout", ["thd_thd_thd", "bshd_bshd_bshd"]) def test_dpa_fa4_varlen(dtype, model_configs, model, qkv_layout): """Test DotProductAttention with FA4: variable-length sequences (varlen/thd)""" - test_dot_product_attention(dtype, model_configs, model, False, True, qkv_layout, False, False) + test_dot_product_attention(dtype, model_configs, model, False, qkv_layout, False, False) model_configs_fa4_mask = { @@ -472,7 +464,7 @@ def test_dpa_fa4_varlen(dtype, model_configs, model, qkv_layout): @pytest.mark.parametrize("model", model_configs_fa4_mask.keys()) def test_dpa_fa4_mask(dtype, model_configs, model): """Test DotProductAttention with FA4: various attention mask types""" - test_dot_product_attention(dtype, model_configs, model, False, True, None, False, False) + test_dot_product_attention(dtype, model_configs, model, False, None, False, False) model_configs_softmax = { @@ -549,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, 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.") @@ -560,7 +550,7 @@ def test_dpa_softmax(dtype, model_configs, model): @pytest.mark.parametrize("model", model_configs_softmax.keys()) def test_dpa_softmax_thd(dtype, model_configs, model): """Test DotProductAttention module with different softmax types""" - test_dot_product_attention(dtype, model_configs, model, True, True, "thd_thd_thd", False, False) + test_dot_product_attention(dtype, model_configs, model, True, "thd_thd_thd", False, False) model_configs_mla = { @@ -589,7 +579,7 @@ def test_dpa_softmax_thd(dtype, model_configs, model): @pytest.mark.parametrize("model", model_configs_mla.keys()) def test_dpa_mla(dtype, model_configs, model): """Test DotProductAttention module with Multi-Latent Attention (MLA)""" - test_dot_product_attention(dtype, model_configs, model, True, True, None, False, False) + test_dot_product_attention(dtype, model_configs, model, True, None, False, False) model_configs_mask = { @@ -644,7 +634,7 @@ def test_dpa_mla(dtype, model_configs, model): @pytest.mark.parametrize("model", model_configs_mask.keys()) def test_dpa_mask(dtype, model_configs, model): """Test DotProductAttention module with different mask types""" - test_dot_product_attention(dtype, model_configs, model, False, True, None, False, False) + test_dot_product_attention(dtype, model_configs, model, False, None, False, False) model_configs_bias = { @@ -750,7 +740,7 @@ def test_dpa_mask(dtype, model_configs, model): @pytest.mark.parametrize("model", model_configs_bias.keys()) def test_dpa_bias(dtype, model_configs, model): """Test DotProductAttention module with different bias types""" - test_dot_product_attention(dtype, model_configs, model, False, True, None, False, False) + test_dot_product_attention(dtype, model_configs, model, False, None, False, False) model_configs_bias_shapes = { @@ -789,7 +779,7 @@ def test_dpa_bias(dtype, model_configs, model): @pytest.mark.parametrize("model", model_configs_bias_shapes.keys()) def test_dpa_bias_shapes(dtype, model_configs, model): """Test DotProductAttention module with different bias types and shapes""" - test_dot_product_attention(dtype, model_configs, model, False, True, None, False, False) + test_dot_product_attention(dtype, model_configs, model, False, None, False, False) model_configs_swa = { @@ -830,7 +820,7 @@ def test_dpa_bias_shapes(dtype, model_configs, model): @pytest.mark.parametrize("qkv_layout", ["thd_thd_thd", "sbhd_sbhd_sbhd"]) def test_dpa_sliding_window(dtype, model_configs, model, qkv_layout): """Test DotProductAttention module with sliding window attention""" - test_dot_product_attention(dtype, model_configs, model, False, True, qkv_layout, True, False) + test_dot_product_attention(dtype, model_configs, model, False, qkv_layout, True, False) model_configs_alibi_slopes = { @@ -870,7 +860,7 @@ def test_dpa_sliding_window(dtype, model_configs, model, qkv_layout): @pytest.mark.parametrize("model", model_configs_alibi_slopes.keys()) def test_dpa_alibi_slopes(dtype, model_configs, model): """Test DotProductAttention module with ALiBi slopes""" - test_dot_product_attention(dtype, model_configs, model, False, True, None, False, False) + test_dot_product_attention(dtype, model_configs, model, False, None, False, False) qkv_layouts = [ @@ -931,7 +921,7 @@ def test_dpa_alibi_slopes(dtype, model_configs, model): @pytest.mark.parametrize("qkv_layout", qkv_layouts) def test_dpa_qkv_layout(dtype, model_configs, model, qkv_layout): """Test DotProductAttention module with different QKV layouts""" - test_dot_product_attention(dtype, model_configs, model, False, True, qkv_layout, False, False) + test_dot_product_attention(dtype, model_configs, model, False, qkv_layout, False, False) qkv_layouts_thd = ["t3hd", "th3d", "thd_t2hd", "thd_th2d", "thd_thd_thd"] @@ -1006,14 +996,14 @@ def test_dpa_qkv_layout_thd(dtype, model_configs, model, qkv_layout): logging.info("[test_dpa_qkv_layout_thd]: pad_between_seqs = True") pad_between_seqs = True test_dot_product_attention( - dtype, model_configs, model, False, True, qkv_layout, False, pad_between_seqs + dtype, model_configs, model, False, qkv_layout, False, pad_between_seqs ) if get_cudnn_version() >= (9, 3, 0): logging.info("[test_dpa_qkv_layout_thd]: pad_between_seqs = False") # cuDNN 9.3.0+ is required to run pad_between_seqs = False/True in the same run pad_between_seqs = False test_dot_product_attention( - dtype, model_configs, model, False, True, qkv_layout, False, pad_between_seqs + dtype, model_configs, model, False, qkv_layout, False, pad_between_seqs ) @@ -1023,7 +1013,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]]: @@ -1037,7 +1026,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 @@ -1441,7 +1429,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 @@ -1485,7 +1472,6 @@ def test_transformer_layer( "UnfusedDotProductAttention", ckpt_attn, qkv_format, - workspace_opt, fused_qkv_params, RoPE, is_training, @@ -1499,7 +1485,6 @@ def test_transformer_layer( "FusedAttention", ckpt_attn, qkv_format, - workspace_opt, fused_qkv_params, RoPE, is_training, @@ -1513,7 +1498,6 @@ def test_transformer_layer( "FlashAttention", ckpt_attn, qkv_format, - workspace_opt, fused_qkv_params, RoPE, is_training, @@ -1583,7 +1567,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/common/common.h b/transformer_engine/common/common.h index eb4dcc055c..56d994a4a1 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -83,10 +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_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; } 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 a0529c80c0..1033dd93e0 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(); @@ -821,7 +821,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."); } @@ -909,7 +909,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."); } @@ -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_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 0f843019ce..21f034985b 100644 --- a/transformer_engine/common/normalization/layernorm/ln_api.cpp +++ b/transformer_engine/common/normalization/layernorm/ln_api.cpp @@ -27,7 +27,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) + "."); } @@ -85,8 +85,7 @@ void layernorm_fwd(const Tensor& x, // BxSxhidden_size mu->data.dptr, rsigma->data.dptr); } - bool training = - is_delayed_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 07ad3230aa..db799e758a 100644 --- a/transformer_engine/common/normalization/rmsnorm/rmsnorm_api.cpp +++ b/transformer_engine/common/normalization/rmsnorm/rmsnorm_api.cpp @@ -23,7 +23,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) + "."); } @@ -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_delayed_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) { 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 d3ee1a2e2c..dd812f15e9 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 @@ -21,7 +21,6 @@ Float8CurrentScaling, MXFP8BlockScaling, ) -from transformer_engine.pytorch.utils import get_cudnn_version from transformer_engine.pytorch.quantization import ( QuantizerRole, get_fp8_te_dtype, @@ -452,25 +451,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" @@ -1069,11 +1049,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: