{"record":{"id":"7c98b3e7e2bf13da","repo":"Lightning-AI/pytorch-lightning","slug":"activation-checkpointing-policy-must-be-a-set-f","errorCode":null,"errorMessage":"`activation_checkpointing_policy` must be a set, found {policy}. You can try defining and passing `auto_wrapper_callable` instead.","messagePattern":"`activation_checkpointing_policy` must be a set, found (.+?)\\. You can try defining and passing `auto_wrapper_callable` instead\\.","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"src/lightning/fabric/strategies/xla_fsdp.py","lineNumber":663,"sourceCode":"\n\ndef _activation_checkpointing_auto_wrapper(policy: _POLICY_SET, module: Module, *args: Any, **kwargs: Any) -> Module:\n    from torch_xla.distributed.fsdp import XlaFullyShardedDataParallel as XLAFSDP\n    from torch_xla.distributed.fsdp import checkpoint_module\n\n    module = checkpoint_module(module) if isinstance(module, tuple(policy)) else module\n    return XLAFSDP(module, *args, **kwargs)\n\n\ndef _activation_checkpointing_kwargs(policy: Optional[_POLICY_SET], kwargs: dict) -> dict:\n    if not policy:\n        return kwargs\n    if \"auto_wrapper_callable\" in kwargs:\n        raise ValueError(\n            \"You cannot set both `auto_wrapper_callable` and `activation_checkpointing_policy`. Choose one\"\n        )\n    if not isinstance(policy, set):\n        raise TypeError(\n            f\"`activation_checkpointing_policy` must be a set, found {policy}. You can try defining and\"\n            \" passing `auto_wrapper_callable` instead.\"\n        )\n    auto_wrapper_callable = partial(_activation_checkpointing_auto_wrapper, policy)\n    kwargs[\"auto_wrapper_callable\"] = auto_wrapper_callable\n    return kwargs\n\n\nclass _XLAFSDPBackwardSyncControl(_BackwardSyncControl):\n    @override\n    def no_backward_sync(self, module: Module, enabled: bool) -> AbstractContextManager:\n        \"\"\"Blocks gradient synchronization inside the :class:`~torch_xla.distributed.fsdp.XlaFullyShardedDataParallel`\n        wrapper.\"\"\"\n        if not enabled:\n            return nullcontext()\n        from torch_xla.distributed.fsdp import XlaFullyShardedDataParallel as XLAFSDP\n\n        if not isinstance(module, XLAFSDP):","sourceCodeStart":645,"sourceCodeEnd":681,"githubUrl":"https://github.com/Lightning-AI/pytorch-lightning/blob/9fed5c27d2a62ff0efd6c3573599921d6ff67c14/src/lightning/fabric/strategies/xla_fsdp.py#L645-L681","documentation":"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.","triggerScenarios":"XLAFSDPStrategy(activation_checkpointing_policy=[MyBlock]) (list), =MyBlock (bare class), or ='MyBlock' (string) — anything that is not a Python set.","commonSituations":"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.","solutions":["Wrap the policy in a set: activation_checkpointing_policy={MyBlock, OtherBlock}","If you need a callable predicate (e.g. name-based matching), define auto_wrapper_callable instead and omit the policy"],"exampleFix":"# before\nstrategy = XLAFSDPStrategy(activation_checkpointing_policy=[TransformerBlock])\n\n# after\nstrategy = XLAFSDPStrategy(activation_checkpointing_policy={TransformerBlock})","handlingStrategy":"type-guard","validationCode":"assert isinstance(activation_checkpointing_policy, set), 'policy must be a set of module types'","typeGuard":"def is_valid_policy(policy) -> bool:\n    return policy is None or isinstance(policy, set)","tryCatchPattern":null,"preventionTips":["Write policies with set literals {BlockA, BlockB}","For predicate-based selection use auto_wrapper_callable instead"],"tags":["xla","fsdp","activation-checkpointing","type-error","lightning-fabric"],"backgroundTag":"invalid-argument-type","analyzedSha":"9fed5c27d2a62ff0efd6c3573599921d6ff67c14","analyzedAt":"2026-08-28T11:52:41.083Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}