Lightning-AI/pytorch-lightning · error · TypeError

`activation_checkpointing_policy` must be a set, found {poli

Error message

`activation_checkpointing_policy` must be a set, found {policy}. You can try defining and passing `auto_wrapper_callable` instead.

What it means

The activation_checkpointing_policy argument of XLAFSDPStrategy must be a set of module types (used to build the internal auto_wrapper). Passing any other type (a list, tuple, string, or a single class) triggers this TypeError, with a hint to fall back to auto_wrapper_callable for more flexible behavior.

Source

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

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:
            return nullcontext()
        from torch_xla.distributed.fsdp import XlaFullyShardedDataParallel as XLAFSDP

        if not isinstance(module, XLAFSDP):

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Wrap the policy in a set: activation_checkpointing_policy={MyBlock, OtherBlock}
2. If you need a callable predicate (e.g. name-based matching), define auto_wrapper_callable instead and omit the policy

Example fix

# before
strategy = XLAFSDPStrategy(activation_checkpointing_policy=[TransformerBlock])

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

assert isinstance(activation_checkpointing_policy, set), 'policy must be a set of module types'

def is_valid_policy(policy) -> bool:
    return policy is None or isinstance(policy, set)

When it happens

Trigger: XLAFSDPStrategy(activation_checkpointing_policy=[MyBlock]) (list), =MyBlock (bare class), or ='MyBlock' (string) — anything that is not a Python set.

Common situations: Users naturally writing a list of layer classes; copying policies from tutorials that use different container types; passing a policy helper's return value that isn't a set.

Understand the failure class

Background: Invalid argument type errors: "must be of type string", "expected X, got Y", and ERR_INVALID_ARG_TYPE explained — this error's family across 15 libraries.

Related errors

• You cannot set both `auto_wrapper_callable` and `activation_
• Blocking backward sync is only possible if the module passed
• 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

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