Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
See LICENSE for license information.
Jax
=======
===
Pre-defined Variable of Logical Axes
------------------------------------
...
...
@@ -20,11 +20,11 @@ Variables are available in `transformer_engine.jax.sharding`.
Checkpointing
------------------------------------
-------------
When using checkpointing with Transformer Engine JAX, please be aware of the checkpointing policy being applied to your model. Any JAX checkpointing policy using `dot`, such as `jax.checkpoint_policies.dots_with_no_batch_dims`, may not work with GEMMs provided by Transformer Engine as they do not always use the `jax.lax.dot_general` primitive. Instead, you can use `transformer_engine.jax.checkpoint_policies.dots_and_te_gemms_with_no_batch_dims` or similar policies that are designed to work with Transformer Engine's GEMMs and `jax.lax.dot_general` GEMMs. You may also use any JAX policies that do not filter by primitive, such as `jax.checkpoint_policies.save_only_these_names` or `jax.checkpoint_policies.everything_saveable`.
Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
See LICENSE for license information.
Getting started
==============
===============
.. note::
...
...
@@ -21,7 +21,7 @@ Transformer Engine provides a set of precision debug tools which allow you to ea
There are 4 things one needs to do to use Transformer Engine debug features:
1. Create a configuration YAML file to configure the desired features.
2. Import, initialize, and install the `Nvidia-DL-Framework-Inspect <https://github.com/NVIDIA/nvidia-dlfw-inspect>`_ tool.
2. Import and initialize the `Nvidia-DL-Framework-Inspect <https://github.com/NVIDIA/nvidia-dlfw-inspect>`_ tool, which is installed as a dependency of Transformer Engine.
3. One can pass ``name="..."`` when creating TE layers to easier identify layer names. If this is not provided, names will be inferred automatically.
4. Invoke ``debug_api.step()`` at the end of one forward-backward pass.
...
...
@@ -38,7 +38,7 @@ To start debugging, one needs to create a configuration YAML file. This file lis
one - ``UserProvidedPrecision`` - is a custom feature implemented by the user. Nvidia-DL-Framework-Inspect inserts features into the layers according to the config.
Example training script
----------------------
-----------------------
Let's look at a simple example of training a Transformer layer using Transformer Engine with FP8 precision. This example demonstrates how to set up the layer, define an optimizer, and perform a few training iterations using synthetic data.
...
...
@@ -81,7 +81,7 @@ We will demonstrate two debug features on the code above:
2. Logging statistics for other GEMM operations, such as gradient statistics for data gradient GEMM within the LayerNormLinear sub-layer of the TransformerLayer.
Config file
----------
-----------
We need to prepare the configuration YAML file, as below
...
...
@@ -114,7 +114,8 @@ We need to prepare the configuration YAML file, as below
Further explanation on how to create config files is in the :doc:`next part of the documentation <2_config_file_structure>`.
Adjusting Python file
--------------------
---------------------
.. code-block:: python
...
...
@@ -145,7 +146,8 @@ In the modified code above, the following changes were made:
3. Added ``debug_api.step()`` after each of the forward-backward pass.
Inspecting the logs
------------------
-------------------
Let's look at the files with the logs. Two files will be created:
...
...
@@ -213,7 +215,8 @@ The second log file (``nvdlfw_inspect_statistics_logs/nvdlfw_inspect_globalrank-
INFO - transformer_layer.self_attention.layernorm_qkv_activation_l1_norm iteration=000004 value=130776.7969
Logging using TensorBoard
------------------------
-------------------------
Precision debug tools support logging using `TensorBoard <https://www.tensorflow.org/tensorboard>`_. To enable it, one needs to pass the argument ``tb_writer`` to the ``debug_api.initialize()``. Let's modify ``train.py`` file.
Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
See LICENSE for license information.
Config File Structure
====================
=====================
To enable debug features, create a configuration YAML file to specify the desired behavior, such as determining which GEMMs (General Matrix Multiply operations) should run in higher precision rather than FP8 and defining which statistics to log.
Below, we outline how to structure the configuration YAML file.
General Format
-------------
--------------
A config file can have one or more sections, each containing settings for specific layers and features:
...
...
@@ -55,7 +56,8 @@ Sections may have any name and must contain:
3. Additional fields describing features for those layers.
Layer Specification
------------------
-------------------
Debug layers can be identified by a ``name`` parameter:
...
...
@@ -89,7 +91,8 @@ Examples:
(...)
Names in Transformer Layers
--------------------------
---------------------------
There are three ways to assign a name to a layer in the Transformer Engine:
...
...
@@ -107,6 +110,8 @@ The ``TransformerLayer`` in Transformer Engine is a composition of multiple sub-
depending on the configuration. Some layers, like ``LayerNormLinear``, are fusions of two layers: ``LayerNorm`` and ``Linear``. When referring to such layers in precision debug tools, only the ``Linear`` part is affected.
For `GroupedLinear` layer, the names of underlying GEMMS are of the form `layer_name.gemm_n`, where `n` is the index of the GEMM.
Below is an example ``TransformerLayer`` with four linear layers that can be influenced by the precision debug tools.
.. figure:: ./img/names.svg
...
...
@@ -154,7 +159,7 @@ Below is an example ``TransformerLayer`` with four linear layers that can be inf
Structured Configuration for GEMMs and Tensors
---------------------------------------------
----------------------------------------------
Sometimes a feature is parameterized by a list of tensors or by a list of GEMMs.
There are multiple ways of describing this parameterization.
...
...
@@ -216,7 +221,7 @@ We can use both structs for tensors and GEMMs. The tensors_struct should be nest
gemm_feature_param1: value
Enabling or Disabling Sections and Features
------------------------------------------
-------------------------------------------
Debug features can be enabled or disabled with the ``enabled`` keyword:
Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
See LICENSE for license information.
...
...
@@ -11,7 +11,8 @@ Please refer to the Nvidia-DL-Framework-Inspect `documentation <https://github.c
Below, we outline the steps for debug initialization.
initialize()
-----------
------------
Must be called once on every rank in the global context to initialize Nvidia-DL-Framework-Inspect.
...
...
@@ -34,7 +35,7 @@ Must be called once on every rank in the global context to initialize Nvidia-DL-
log_dir="./log_dir")
set_tensor_reduction_group()
--------------------------
----------------------------
Needed only for logging tensor stats. In multi-GPU training, activation and gradient tensors are distributed across multiple nodes. This method lets you specify the group for the reduction of stats; see the `reduction group section <./4_distributed.rst#reduction-groups>`_ for more details.
...
...
@@ -61,7 +62,7 @@ If the tensor reduction group is not specified, then statistics are reduced acro
# activation/gradient tensor statistics are reduced along pipeline_parallel_group
set_weight_tensor_tp_group_reduce()
---------------------------------
-----------------------------------
By default, weight tensor statistics are reduced within the tensor parallel group. This function allows you to disable that behavior; for more details, see `reduction group section <./4_distributed.rst#reduction-groups>`_.
