Lightning-AI/pytorch-lightning · error · RuntimeError

Failed to determine the arguments that were used to compile

Error message

Failed to determine the arguments that were used to compile the module. Make sure to import lightning before `torch.compile` is used.

What it means

_unwrap_compiled() detects a torch.compile'd module (OptimizedModule) and reads its _compile_kwargs to preserve the compile settings when re-wrapping. Lightning patches torch.compile to record these kwargs; if _compile_kwargs is missing it means torch.compile ran before lightning was imported, so the settings cannot be recovered and a RuntimeError is raised.

Source

Thrown at src/lightning/fabric/wrappers.py:356

        if isinstance(obj, _FabricDataLoader):
            return obj._dataloader
        return obj

    types = [_FabricModule, _FabricOptimizer, _FabricDataLoader]
    types.append(OptimizedModule)

    return apply_to_collection(collection, dtype=tuple(types), function=_unwrap)


def _unwrap_compiled(obj: Union[Any, OptimizedModule]) -> tuple[Union[Any, nn.Module], Optional[dict[str, Any]]]:
    """Removes the :class:`torch._dynamo.OptimizedModule` around the object if it is wrapped.

    Use this function before instance checks against e.g. :class:`_FabricModule`.

    """
    if isinstance(obj, OptimizedModule):
        if (compile_kwargs := getattr(obj, "_compile_kwargs", None)) is None:
            raise RuntimeError(
                "Failed to determine the arguments that were used to compile the module. Make sure to import"
                " lightning before `torch.compile` is used."
            )
        return obj._orig_mod, compile_kwargs
    return obj, None


def _to_compiled(module: nn.Module, compile_kwargs: dict[str, Any]) -> OptimizedModule:
    return torch.compile(module, **compile_kwargs)  # type: ignore[return-value]


def _backward_hook(requires_backward: bool, *_: Any) -> None:
    if requires_backward and not _in_fabric_backward:
        raise RuntimeError(
            "The current strategy and precision selection requires you to call `fabric.backward(loss)`"
            " instead of `loss.backward()`."
        )

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Ensure `import lightning` (or `import lightning.fabric`) appears before any torch.compile call, including in transitively imported modules that compile at import time
2. For HF transformers, avoid torch_compile=True in from_pretrained unless lightning is imported first, and compile manually after setup instead
3. Compile the model after fabric.setup() rather than at import time

Example fix

# before
# utils.py (imported first)
model = torch.compile(model)  # lightning not imported yet
import lightning as L

# after
import lightning as L  # first
model = L.Fabric().setup(model)
model = torch.compile(model)

import lightning  # must run before any torch.compile
import torch
model = torch.compile(model)
assert not isinstance(model, torch._dynamo.eval_frame.OptimizedModule) or getattr(model, '_compile_kwargs', None) is not None, 'compile ran before lightning import'

When it happens

Trigger: Calling torch.compile(model) at module import time in a script that imports lightning later (or not at all before compile), then fabric.setup()/load/backward() which internally unwraps the OptimizedModule.

Common situations: A third-party library (e.g. transformers with torch_compile=True, or a utils module) compiles models at import time before `import lightning` executes; reordered imports after refactoring; worker processes that compile models without importing lightning first.

Related errors

• Using a compiled model is incompatible with the current stra
• Device should be CPU, got {device} instead.
• `devices` selected with `CPUAccelerator` should be an int >
• Device should be CUDA, got {device} instead.
• You requested to find {num_devices} devices but there are no

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