Lightning-AI/pytorch-lightning · error · ValueError

The optimizer does not seem to reference any FSDP parameters

Error message

The optimizer does not seem to reference any FSDP parameters. HINT: Make sure to create the optimizer after setting up the model.

What it means

With use_orig_params unset/False, FSDP setup_optimizer requires the optimizer to reference the flat (flattened) FSDP parameter groups created by setup_module. _optimizer_has_flat_params detects no FlatParameter in the optimizer, meaning the optimizer was created before the model was wrapped, and raises ValueError.

Source

Thrown at src/lightning/fabric/strategies/fsdp.py:335

        # activation checkpointing needs to be set up after wrapping the model
        _setup_activation_checkpointing(module, self._activation_checkpointing_kwargs)

        return module

    @override
    def setup_optimizer(self, optimizer: Optimizer) -> Optimizer:
        """Set up an optimizer for a model wrapped with FSDP.

        This setup method doesn't modify the optimizer or wrap the optimizer. The only thing it currently does is verify
        that the optimizer was created after the model was wrapped with :meth:`setup_module` with a reference to the
        flattened parameters.

        """
        if self._fsdp_kwargs.get("use_orig_params"):
            return super().setup_optimizer(optimizer)
        if not _optimizer_has_flat_params(optimizer):
            # We avoid this limitation by setting `use_orig_params=True`
            raise ValueError(
                "The optimizer does not seem to reference any FSDP parameters. HINT: Make sure to create the optimizer"
                " after setting up the model."
            )
        return optimizer

    @override
    def module_to_device(self, module: Module) -> None:
        pass

    @override
    def module_init_context(self, empty_init: Optional[bool] = None) -> AbstractContextManager:
        precision_init_ctx = self.precision.module_init_context()
        module_sharded_ctx = self.module_sharded_context()
        stack = ExitStack()
        if empty_init:
            # Materialization happens in `setup`. When modules get wrapped by FSDP, the sequence of operations is:
            # 1) materialize module 2) call `reset_parameters()` 3) shard the module.
            # These operations are applied to each submodule 'bottom up' in the module hierarchy.

View on GitHub (pinned to 9fed5c27d2)

Solutions

  1. Recreate the optimizer after setup_module so it sees the FSDP parameters, or just use fabric.setup(model, optimizer)
  2. Set FSDPStrategy(use_orig_params=True) so original parameter references stay valid
  3. Pass module.parameters() to a new optimizer created after the module is set up

Example fix

# before
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
model = fabric.setup_module(model)
optimizer = fabric.setup_optimizer(optimizer)  # ValueError

# after
model, optimizer = fabric.setup(model, torch.optim.Adam(model.parameters(), lr=1e-3))
Defensive patterns

Strategy: fallback

Validate before calling

# create optimizer only after module setup
model = fabric.setup_module(model)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)  # now references FSDP params
optimizer = fabric.setup_optimizer(optimizer)

Try / catch

try:
    optimizer = fabric.setup_optimizer(optimizer)
except ValueError as e:
    if "does not seem to reference any FSDP parameters" in str(e):
        optimizer = type(optimizer)(module.parameters(), lr=optimizer.defaults["lr"])
        optimizer = fabric.setup_optimizer(optimizer)

Prevention

When it happens

Trigger: Creating torch.optim.Adam(model.parameters(), ...) on the raw model and calling fabric.setup_optimizer(optimizer) after FSDP wrapped the model with use_orig_params not True — the optimizer's params are stale references.

Common situations: Standard PyTorch ordering (optimizer before setup) copied into Fabric FSDP code; refactoring from DDP where param identity is preserved; partial refactor where setup_module was called but the optimizer object predates it.

Related errors


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