Lightning-AI/pytorch-lightning · error · TypeError

Gradient clipping with FSDP is only possible if the module p

Error message

Gradient clipping with FSDP is only possible if the module passed to `{type(self).__name__}.clip_gradients_norm` is wrapped in `FullyShardedDataParallel`. Got: {module.__class__.__name__}.

What it means

FSDP gradient clipping by norm must run through FullyShardedDataParallel.clip_grad_norm_ (it needs the sharded/unsharded coordination across ranks). If the module passed to clip_gradients_norm is not an FSDP-wrapped module (typically the unwrapped root), TypeError is raised.

Source

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

        obj = [obj]
        torch.distributed.broadcast_object_list(obj, src, group=_group.WORLD)
        return obj[0]

    @override
    def clip_gradients_norm(
        self,
        module: Module,
        optimizer: Optimizer,
        max_norm: Union[float, int],
        norm_type: Union[float, int] = 2.0,
        error_if_nonfinite: bool = True,
    ) -> Tensor:
        """Clip gradients by norm."""
        from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel

        if not isinstance(module, FullyShardedDataParallel):
            # the root must be wrapped
            raise TypeError(
                "Gradient clipping with FSDP is only possible if the module passed to"
                f" `{type(self).__name__}.clip_gradients_norm` is wrapped in `FullyShardedDataParallel`."
                f" Got: {module.__class__.__name__}."
            )
        self.precision.unscale_gradients(optimizer)
        return module.clip_grad_norm_(max_norm=max_norm, norm_type=norm_type)

    @override
    def save_checkpoint(
        self,
        path: _PATH,
        state: dict[str, Union[Module, Optimizer, Any]],
        storage_options: Optional[Any] = None,
        filter: Optional[dict[str, Callable[[str, Any], bool]]] = None,
    ) -> None:
        """Save model, optimizer, and other state to a checkpoint on disk.

        If the state-dict-type is ``'full'``, the checkpoint will be written to a single file containing the weights,

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Use the wrapped module returned by fabric.setup/setup_module for clipping
2. Ensure clipping happens after setup, and you keep the returned (wrapped) reference
3. Clip manually with the FSDP API on the wrapped root if you need custom behavior

Example fix

# before
raw_model = MyModel()
optimizer = ...
strategy.clip_gradients_norm(raw_model, optimizer, max_norm=1.0)

# after
model, optimizer = fabric.setup(raw_model, optimizer)
strategy.clip_gradients_norm(model, optimizer, max_norm=1.0)

from torch.distributed.fsdp import FullyShardedDataParallel
assert isinstance(module, FullyShardedDataParallel), "clip after fabric.setup, on the wrapped module"

from torch.distributed.fsdp import FullyShardedDataParallel
def is_fsdp_wrapped(module) -> bool:
    return isinstance(module, FullyShardedDataParallel)

When it happens

Trigger: strategy.clip_gradients_norm(module, optimizer, ...) where module is the original unwrapped model (e.g. before setup_module/setup) or a plain nn.Module.

Common situations: Calling fabric.clip_gradients on the raw model reference before setup; storing the pre-wrap model and clipping after; partially refactored custom loops.

Related errors

• The FSDP strategy can only work with the `FSDPPrecision` plu
• `activation_checkpointing_policy` must be a set, found {poli
• Blocking backward sync is only possible if the module passed
• `gradient_clip_val` should be an int or a float. Got {gradie
• `gradient_clip_algorithm='norm'` is currently not supported

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