Lightning-AI/pytorch-lightning · error · ModuleNotFoundError

str(_TRANSFORMER_ENGINE_AVAILABLE)

Error message

str(_TRANSFORMER_ENGINE_AVAILABLE)

What it means

TransformerEnginePrecision requires the NVIDIA Transformer Engine package. The module-level import check failed and stored the exception message in _TRANSFORMER_ENGINE_AVAILABLE; __init__ re-raises it as ModuleNotFoundError when the plugin is instantiated on a machine without the package (or without a compatible GPU/build).

Source

Thrown at src/lightning/fabric/plugins/precision/transformer_engine.py:76

        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.

    """

    precision: Literal["transformer-engine", "transformer-engine-float16"] = "transformer-engine"

    def __init__(
        self,
        *,
        weights_dtype: torch.dtype,
        recipe: Optional[Union[Mapping[str, Any], "DelayedScaling"]] = None,
        replace_layers: Optional[bool] = None,
        fallback_compute_dtype: Optional[torch.dtype] = None,
    ) -> None:
        if not _TRANSFORMER_ENGINE_AVAILABLE:
            raise ModuleNotFoundError(str(_TRANSFORMER_ENGINE_AVAILABLE))
        from transformer_engine.common.recipe import DelayedScaling

        if recipe is None:
            recipe = DelayedScaling()
        elif isinstance(recipe, Mapping):
            recipe = dict(recipe)  # copy
            if "fp8_format" in recipe:
                from transformer_engine.common.recipe import Format

                recipe["fp8_format"] = getattr(Format, recipe["fp8_format"])
            recipe = DelayedScaling(**recipe)

        self.weights_dtype = weights_dtype
        self.recipe = recipe
        self.replace_layers = replace_layers
        self.fallback_compute_dtype = fallback_compute_dtype or weights_dtype

    @override

View on GitHub (pinned to 9fed5c27d2)

Solutions

  1. pip install transformer_engine (preferably a build matching your torch and CUDA version, e.g. the +cu12xyz variant)
  2. Verify with 'python -c "import transformer_engine"' that the import actually succeeds in the target environment
  3. If FP8 is not needed, switch to a different precision plugin such as MixedPrecision('bf16-mixed')

Example fix

# before
precision = TransformerEnginePrecision(weights_dtype=torch.float8_e4m3fn, ...)
# after (shell)
pip install transformer_engine
# or in code, fall back:
try:
    precision = TransformerEnginePrecision(...)
except ModuleNotFoundError:
    precision = MixedPrecision("bf16-mixed")
Defensive patterns

Strategy: fallback

Validate before calling

try:
    import transformer_engine  # noqa
    available = True
except ImportError:
    available = False
if not available:
    # choose another precision plugin instead of constructing TransformerEnginePrecision

Type guard

import importlib.util

def transformer_engine_available() -> bool:
    return importlib.util.find_spec("transformer_engine") is not None

Try / catch

try:
    precision = TransformerEnginePrecision(weights_dtype=torch.float8_e4m3fn)
except ModuleNotFoundError:
    precision = MixedPrecision("bf16-mixed")

Prevention

When it happens

Trigger: Instantiating TransformerEnginePrecision(...) when 'import transformer_engine' fails — package not installed, CPU-only machine, or CUDA version unsupported by the wheels you installed.

Common situations: Running a training script written for H100/FP8 on a dev box or CI runner without transformer_engine installed; installing a transformer_engine build mismatched with the installed torch/CUDA version so the import fails.

Understand the failure class

Background: "X is not installed. Please install it with pip install Y": missing optional dependency errors — ImportError/ValueError raised when a library's optional extra was never installed — this error's family across 22 libraries.

Related errors


AI-assisted analysis of Lightning-AI/pytorch-lightning@9fed5c27d2 (2026-08-28). Data as JSON: /api/errors/5d61ebcebdc06982. Report an issue: GitHub.