Allow preparing a model without an optimizer under FSDP2

src/accelerate/accelerator.py

@@ -1474,9 +1474,9 @@ def prepare(self, *args, device_placement=None):
                     optimizer_count += 1
 
             # This needs to be written as such, so that passing other objects other than models/optimizers doesn't raise an error
-            if (model_count < 1 and optimizer_count > 0) or (model_count > 0 and optimizer_count < 1):
+            if model_count < 1 and optimizer_count > 0:
                 raise ValueError(
-                    "When using FSDP2, a model and optimizer must be passed together to `Accelerator.prepare()`"
+                    "When using FSDP2, an optimizer must be passed together with its model to `Accelerator.prepare()`"
                     " as the optimizer needs to have its parameters modified after the model is converted."
                 )
             if model_count > 1:
@@ -1644,15 +1644,17 @@ def _prepare_fsdp2(self, *args):
             if isinstance(obj, torch.nn.Module):
                 model_index, model = i, obj
 
-        # Invariant: if we have a model, we also have an optimizer (checked in `prepare`)
+        # Invariant: if we have an optimizer, we also have a model (checked in `prepare`)
         if model_index is None:
             return tuple(result)
 
         # Needs to be done first, to make sure AC + fully_shard will work as expected
         self.state.fsdp_plugin.set_auto_wrap_policy(model)
 
-        # Apply AC if needed
-        if self.state.fsdp_plugin.activation_checkpointing:
+        # Apply AC if needed. A model already prepared, e.g. without its optimizer for evaluation, has it already
+        from torch.distributed.fsdp import FSDPModule
+
+        if self.state.fsdp_plugin.activation_checkpointing and not isinstance(model, FSDPModule):
             model = fsdp2_apply_ac(self, model)
 
         # Apply compile if needed, has to be *after* applying AC

tests/fsdp/test_fsdp.py

@@ -561,6 +561,32 @@ def test_clip_grad_value(self):
             accelerator.clip_grad_value_(model.parameters(), 1.0)
             assert model[0].weight.grad.full_tensor().tolist() == [[1.0, 0.5, -1.0, -0.5]] * 4
 
+    def test_prepare_model_without_optimizer(self):
+        """A model can be prepared without an optimizer, for evaluation, and is sharded."""
+        from torch.distributed.fsdp import FSDPModule
+
+        with patch_environment(**self.fsdp_envs[2]):
+            accelerator = Accelerator()
+            model = accelerator.prepare(torch.nn.Sequential(torch.nn.Linear(4, 4)))
+            assert isinstance(model, FSDPModule)
+            model(torch.ones(1, 4, device=accelerator.device))
+
+    def test_prepare_model_again_with_optimizer(self):
+        """A model prepared without an optimizer, for evaluation, can then be prepared with its optimizer for training,
+        without applying activation checkpointing a second time."""
+        env = self.fsdp_envs[2].copy()
+        env["FSDP_ACTIVATION_CHECKPOINTING"] = "true"
+        env["FSDP_AUTO_WRAP_POLICY"] = "TRANSFORMER_BASED_WRAP"
+        env["FSDP_TRANSFORMER_CLS_TO_WRAP"] = "Linear"
+        with patch_environment(**env):
+            accelerator = Accelerator()
+            model = accelerator.prepare(torch.nn.Sequential(torch.nn.Linear(4, 4), torch.nn.Linear(4, 4)))
+            names = [name for name, _ in model.named_parameters()]
+            model, optimizer = accelerator.prepare(model, torch.optim.SGD(model.parameters(), lr=0.1))
+            assert [name for name, _ in model.named_parameters()] == names
+            accelerator.backward(model(torch.ones(1, 4, device=accelerator.device)).sum())
+            optimizer.step()
+
     def test_param_mapping_error_handling(self):
         """Test FSDP2's defensive error handling for parameter mapping failures in tied/non-tied cases."""
         from unittest.mock import Mock, patch