Lightning-AI/pytorch-lightning · error · TypeError

Blocking backward sync is only possible if the module passed

Error message

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__}.

What it means

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.

Source

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

            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):
            raise TypeError(
                "Blocking backward sync is only possible if the module passed to"
                f" `{self.__class__.__name__}.no_backward_sync` is wrapped in `XlaFullyShardedDataParallel`."
                f" Got: {module.__class__.__name__}."
            )
        return module.no_sync()

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Ensure the module passed is the fabric.setup() output (the XLAFSDP-wrapped model)
2. Check with isinstance(module, XlaFullyShardedDataParallel) before entering the context
3. If using Fabric's built-in gradient accumulation, prefer fabric.no_backward_sync instead, which resolves the right module internally

Example fix

# before
model = MyModel()
with fabric.strategy.no_backward_sync(model, enabled=True):
    ...

# after
model = fabric.setup(MyModel())
with fabric.no_backward_sync(model, enabled=True):
    ...

from torch_xla.distributed.fsdp import XlaFullyShardedDataParallel as XLAFSDP
assert isinstance(model, XLAFSDP), 'model must be the fabric.setup() output'

from torch_xla.distributed.fsdp import XlaFullyShardedDataParallel as XLAFSDP

def is_xlafsdp_module(obj) -> bool:
    return isinstance(obj, XLAFSDP)

When it happens

Trigger: 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.

Common situations: 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.

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

• `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/5bae6b842267fa0d. Report an issue: GitHub.