Lightning-AI/pytorch-lightning · error · ValueError

You cannot set both `auto_wrapper_callable` and `activation_

Error message

You cannot set both `auto_wrapper_callable` and `activation_checkpointing_policy`. Choose one

What it means

Raised by _activation_checkpointing_kwargs (called from XLAFSDPStrategy kwargs parsing) when the user sets both activation_checkpointing_policy and an auto_wrapper_callable in the FSDP kwargs. The policy is implemented by building an auto_wrapper callable internally, so an explicit auto_wrapper_callable conflicts with it and only one may be provided.

Source

Thrown at src/lightning/fabric/strategies/xla_fsdp.py:659

        # this is not transformer specific despite the name
        policy = partial(transformer_auto_wrap_policy, transformer_layer_cls=policy)
    kwargs["auto_wrap_policy"] = policy
    return kwargs


def _activation_checkpointing_auto_wrapper(policy: _POLICY_SET, module: Module, *args: Any, **kwargs: Any) -> Module:
    from torch_xla.distributed.fsdp import XlaFullyShardedDataParallel as XLAFSDP
    from torch_xla.distributed.fsdp import checkpoint_module

    module = checkpoint_module(module) if isinstance(module, tuple(policy)) else module
    return XLAFSDP(module, *args, **kwargs)


def _activation_checkpointing_kwargs(policy: Optional[_POLICY_SET], kwargs: dict) -> dict:
    if not policy:
        return kwargs
    if "auto_wrapper_callable" in kwargs:
        raise ValueError(
            "You cannot set both `auto_wrapper_callable` and `activation_checkpointing_policy`. Choose one"
        )
    if not isinstance(policy, set):
        raise TypeError(
            f"`activation_checkpointing_policy` must be a set, found {policy}. You can try defining and"
            " passing `auto_wrapper_callable` instead."
        )
    auto_wrapper_callable = partial(_activation_checkpointing_auto_wrapper, policy)
    kwargs["auto_wrapper_callable"] = auto_wrapper_callable
    return kwargs


class _XLAFSDPBackwardSyncControl(_BackwardSyncControl):
    @override
    def no_backward_sync(self, module: Module, enabled: bool) -> AbstractContextManager:
        """Blocks gradient synchronization inside the :class:`~torch_xla.distributed.fsdp.XlaFullyShardedDataParallel`
        wrapper."""
        if not enabled:

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Remove auto_wrapper_callable and keep activation_checkpointing_policy if you just want selective activation checkpointing
2. Alternatively drop activation_checkpointing_policy and implement the behavior entirely inside your own auto_wrapper_callable

Example fix

# before
strategy = XLAFSDPStrategy(
    activation_checkpointing_policy={MyBlock},
    auto_wrapper_callable=partial(apply_activation_checkpointing, check_fn=lambda m: isinstance(m, MyBlock)),
)

# after
strategy = XLAFSDPStrategy(activation_checkpointing_policy={MyBlock})

kwargs = dict(fsdp_kwargs)
if 'auto_wrapper_callable' in kwargs:
    assert not activation_checkpointing_policy, 'choose one of auto_wrapper_callable or activation_checkpointing_policy'

When it happens

Trigger: XLAFSDPStrategy(activation_checkpointing_policy={...}, auto_wrapper_callable=my_wrapper) or passing auto_wrapper_callable through XLAFSDPStrategy(..., **fsdp_kwargs) alongside the policy.

Common situations: Copy-pasting FSDP configs from other examples that include auto_wrapper_callable and then adding the Lightning convenience argument activation_checkpointing_policy; migrating from DDP/FSDP strategies where both were combined.

Related errors

• `activation_checkpointing_policy` must be a set, found {poli
• Could not find a XLAFSDP model in the provided checkpoint st
• Found multiple XLAFSDP modules in the given state. Saving ch
• Multihost setups do not have a shared filesystem, so the che
• Got `XLAFSDPStrategy.load_checkpoint(..., state={state!r})`

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