---
title: TransformerEnginePrecision
description: |-
  Plugin for training with fp8 precision via nvidia’s
  Transformer Engine.
canonical: >-
  https://lightning.ai/docs/fabric/stable/full-api-reference/api/generated/lightning.fabric.plugins.precision.TransformerEnginePrecision
---
# TransformerEnginePrecision {#transformerengineprecision}

### `class lightning.fabric.plugins.precision.TransformerEnginePrecision(*, weights_dtype, recipe=None, replace_layers=None, fallback_compute_dtype=None)` {#lightning.fabric.plugins.precision.TransformerEnginePrecision}

Bases: [`Precision`](/docs/fabric/2.6.6/full-api-reference/api/generated/lightning.fabric.plugins.precision.Precision.md#lightning.fabric.plugins.precision.Precision "lightning.fabric.plugins.precision.precision.Precision")

Plugin for training with fp8 precision via nvidia’s [Transformer Engine](https://docs.nvidia.com/deeplearning/transformer-engine).

Warning

This is an [experimental](/docs/pytorch/stable/versioning#experimental-api "(in PyTorch Lightning v2.7.0.dev)") feature.

Parameters:

*   **weights\_dtype** ([`dtype`](https://docs.pytorch.org/docs/stable/tensor_attributes.html#torch.dtype "(in PyTorch v2.15)")) – The weights dtype to use.
    
*   **recipe** (`Union`\[[`Mapping`](https://docs.python.org/3/library/collections.abc.html#collections.abc.Mapping "(in Python v3.15)")\[[`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)"), [`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")\], `DelayedScaling`, [`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")\]) – Recipe for the DelayedScaling [configuration](https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/api/common.html#transformer_engine.common.recipe.DelayedScaling). In dict format or the dataclass format.
    
*   **replace\_layers** ([`Optional`](https://docs.python.org/3/library/typing.html#typing.Optional "(in Python v3.15)")\[[`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")\]) – Whether to replace `Linear` and `LayerNorm` layers automatically with their Transformer Engine alternatives. Note that they don’t subclass the torch equivalents so checks like `isinstance(l, torch.nn.Linear)` will not pass.
    
*   **fallback\_compute\_dtype** ([`Optional`](https://docs.python.org/3/library/typing.html#typing.Optional "(in Python v3.15)")\[[`dtype`](https://docs.pytorch.org/docs/stable/tensor_attributes.html#torch.dtype "(in PyTorch v2.15)")\]) – The compute dtype to use for operations that don’t support fp8 autocast. Defaults to the same as `weights_dtype`.
    

Note

Support for FP8 in the linear layers with this plugin is currently limited to tensors with shapes where the dimensions are divisible by 8 and 16 respectively. You might want to add padding to your inputs to conform to this restriction.

#### `convert_input(data)` {#lightning.fabric.plugins.precision.TransformerEnginePrecision.convert_input}

Convert model inputs (forward) to the floating point precision type of this plugin.

This is a no-op in the base precision plugin, since we assume the data already has the desired type (default is torch.float32).

Return type:

[`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")

#### `convert_module(module)` {#lightning.fabric.plugins.precision.TransformerEnginePrecision.convert_module}

Convert the module parameters to the precision type this plugin handles.

This is optional and depends on the precision limitations during optimization.

Return type:

[`Module`](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "(in PyTorch v2.15)")

#### `convert_output(data)` {#lightning.fabric.plugins.precision.TransformerEnginePrecision.convert_output}

Convert outputs to the floating point precision type expected after model’s forward.

This is a no-op in the base precision plugin, since we assume the data already has the desired type (default is torch.float32).

Return type:

[`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")

#### `forward_context()` {#lightning.fabric.plugins.precision.TransformerEnginePrecision.forward_context}

A contextmanager for managing model forward/training\_step/evaluation\_step/predict\_step.

Return type:

[`AbstractContextManager`](https://docs.python.org/3/library/contextlib.html#contextlib.AbstractContextManager "(in Python v3.15)")

#### `module_init_context()` {#lightning.fabric.plugins.precision.TransformerEnginePrecision.module_init_context}

Instantiate module parameters or tensors in the precision type this plugin handles.

This is optional and depends on the precision limitations during optimization.

Return type:

[`AbstractContextManager`](https://docs.python.org/3/library/contextlib.html#contextlib.AbstractContextManager "(in Python v3.15)")

#### `tensor_init_context()` {#lightning.fabric.plugins.precision.TransformerEnginePrecision.tensor_init_context}

Controls how tensors get created (device, dtype).

Return type:

[`AbstractContextManager`](https://docs.python.org/3/library/contextlib.html#contextlib.AbstractContextManager "(in Python v3.15)")

---

**Navigation**

- Parent: [Generated](/docs/fabric/2.6.6/full-api-reference/api/generated.md)