def test_no_backward_sync():
    """Test that the backward sync control disables gradient sync on modules that benefit from it."""
    from torch.distributed._composable.fsdp import FSDPModule

    strategy = ModelParallelStrategy(parallelize_fn=(lambda m, _: m))
    assert isinstance(strategy._backward_sync_control, _ParallelBackwardSyncControl)

    fsdp_layer = Mock(spec=FSDPModule)
    other_layer = nn.Linear(2, 2)
    module = Mock()
    module.modules = Mock(return_value=[fsdp_layer, other_layer])

    with strategy._backward_sync_control.no_backward_sync(module, True):
        fsdp_layer.set_requires_gradient_sync.assert_called_with(False, recurse=False)
    fsdp_layer.set_requires_gradient_sync.assert_called_with(True, recurse=False)

    with strategy._backward_sync_control.no_backward_sync(module, False):
        fsdp_layer.set_requires_gradient_sync.assert_called_with(True, recurse=False)
    fsdp_layer.set_requires_gradient_sync.assert_called_with(False, recurse=False)
Test that the backward sync control disables gradient sync on modules that benefit from it.
tests/tests_fabric/strategies/test_model_parallel.py