Implementation:NVIDIA TransformerEngine ONNX Export

Field	Value
Sources	TransformerEngine
Domains	Deep_Learning, PyTorch
Last Updated	2026-02-07 14:00 GMT

Overview

Provides ONNX export utilities for Transformer Engine models, including a context manager for export mode and ONNX translation table registration.

Description

This module provides the infrastructure for exporting Transformer Engine models to ONNX format:

• onnx_export() -- Context manager that enables ONNX export mode. When active, TE operations use ONNX-compatible implementations instead of custom CUDA kernels. Requires PyTorch >= 2.4.

• is_in_onnx_export_mode() -- Global state check used by operations (LayerNorm, RMSNorm, softmax) to select ONNX-compatible code paths.

• assert_warmed_up() -- Validates that the model has been run at least once before export (required for FP8 state initialization).

For PyTorch >= 2.4, the module also imports ONNX extensions including custom ops for FP8/MXFP8 quantize/dequantize, GEMM, LayerNorm, attention mask, and a te_translation_table.

Usage

Wrap torch.onnx.export calls inside the onnx_export context manager. Use te_translation_table as the custom translation table.

Code Reference

Source Location

Repository

NVIDIA/TransformerEngine

File

transformer_engine/pytorch/export.py

Lines

1--71

Signature

@contextmanager
def onnx_export(enabled: bool = False) -> Generator[None, None, None]: ...
def is_in_onnx_export_mode() -> bool: ...
def assert_warmed_up(module: torch.nn.Module) -> None: ...

Import

from transformer_engine.pytorch.export import onnx_export, is_in_onnx_export_mode, te_translation_table

I/O Contract

Inputs

Name	Type	Required	Description
enabled	bool	No	Whether to enable ONNX export mode (default False)
module	torch.nn.Module	Yes	Model to validate (for assert_warmed_up)

Outputs

Name	Type	Description
(context)	None	Context manager yields nothing; sets global state
is_export	bool	is_in_onnx_export_mode() returns current state

Usage Examples

from transformer_engine.pytorch.export import onnx_export, te_translation_table

with onnx_export(enabled=True):
    torch.onnx.export(
        model,
        dynamo=True,
        custom_translation_table=te_translation_table,
    )

Related Pages

• Environment:NVIDIA_TransformerEngine_CUDA_Toolkit_Requirements