{"record":{"id":"5bae6b842267fa0d","repo":"Lightning-AI/pytorch-lightning","slug":"blocking-backward-sync-is-only-possible-if-the-mod-5bae6b","errorCode":null,"errorMessage":"Blocking backward sync is only possible if the module passed to `{self.__class__.__name__}.no_backward_sync` is wrapped in `XlaFullyShardedDataParallel`. Got: {module.__class__.__name__}.","messagePattern":"Blocking backward sync is only possible if the module passed to `(.+?)\\.no_backward_sync` is wrapped in `XlaFullyShardedDataParallel`\\. Got: (.+?)\\.","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"src/lightning/fabric/strategies/xla_fsdp.py","lineNumber":682,"sourceCode":"            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):\n            raise TypeError(\n                \"Blocking backward sync is only possible if the module passed to\"\n                f\" `{self.__class__.__name__}.no_backward_sync` is wrapped in `XlaFullyShardedDataParallel`.\"\n                f\" Got: {module.__class__.__name__}.\"\n            )\n        return module.no_sync()\n","sourceCodeStart":664,"sourceCodeEnd":688,"githubUrl":"https://github.com/Lightning-AI/pytorch-lightning/blob/9fed5c27d2a62ff0efd6c3573599921d6ff67c14/src/lightning/fabric/strategies/xla_fsdp.py#L664-L688","documentation":"XLAFSDPStrategy.no_backward_sync(module, enabled=True) returns a context manager that defers gradient synchronization by calling module.no_sync(), which only exists on XlaFullyShardedDataParallel wrappers. If the passed module is not XLAFSDP-wrapped (e.g. a raw nn.Module, or wrapped by a different strategy), the TypeError is raised.","triggerScenarios":"with fabric.strategy.no_backward_sync(raw_model, enabled=True): loss.backward() where raw_model never went through fabric.setup under the XLAFSDP strategy; passing the wrong model in a multi-model setup.","commonSituations":"Calling no_backward_sync before fabric.setup(model); porting gradient-accumulation code from the DDP strategy (which accepts any module) to XLAFSDP; referencing the unwrapped module variable after setup.","solutions":["Ensure the module passed is the fabric.setup() output (the XLAFSDP-wrapped model)","Check with isinstance(module, XlaFullyShardedDataParallel) before entering the context","If using Fabric's built-in gradient accumulation, prefer fabric.no_backward_sync instead, which resolves the right module internally"],"exampleFix":"# before\nmodel = MyModel()\nwith fabric.strategy.no_backward_sync(model, enabled=True):\n    ...\n\n# after\nmodel = fabric.setup(MyModel())\nwith fabric.no_backward_sync(model, enabled=True):\n    ...","handlingStrategy":"type-guard","validationCode":"from torch_xla.distributed.fsdp import XlaFullyShardedDataParallel as XLAFSDP\nassert isinstance(model, XLAFSDP), 'model must be the fabric.setup() output'","typeGuard":"from torch_xla.distributed.fsdp import XlaFullyShardedDataParallel as XLAFSDP\n\ndef is_xlafsdp_module(obj) -> bool:\n    return isinstance(obj, XLAFSDP)","tryCatchPattern":null,"preventionTips":["Prefer fabric.no_backward_sync over strategy.no_backward_sync with a raw module","Only reference the wrapped model variable after setup in accumulation loops"],"tags":["xla","fsdp","gradient-accumulation","no-sync","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"}