diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py index 0bc7650f682..a1cbfaa8d68 100755 --- a/src/accelerate/accelerator.py +++ b/src/accelerate/accelerator.py @@ -1792,6 +1792,12 @@ def prepare_model( if device_placement is None: device_placement = self.device_placement and self.distributed_type != DistributedType.FSDP + # Ensure we can't double wrap a model + if getattr(model, "_is_accelerate_prepared", False): + if model not in self._models: + self._models.append(model) + return model + self._models.append(model) # TODO: Look at enabling native TP training directly with a proper config @@ -2054,6 +2060,7 @@ def prepare_model( model = compile_regions(model, **self.state.dynamo_plugin.to_kwargs()) else: model = torch.compile(model, **self.state.dynamo_plugin.to_kwargs()) + model._is_accelerate_prepared = True return model def _prepare_ao(self, *args): diff --git a/tests/test_accelerator.py b/tests/test_accelerator.py index c56b00ec3f5..ebab8a8c057 100644 --- a/tests/test_accelerator.py +++ b/tests/test_accelerator.py @@ -464,6 +464,17 @@ def test_is_accelerator_prepared(self): "Valid Dataloader is missing `_is_accelerator_prepared` or is set to `False`" ) + def test_prepare_model_twice_does_not_double_wrap(self): + accelerator = Accelerator() + model = torch.nn.Linear(10, 2) + prepared_model = accelerator.prepare_model(model) + num_models_before = len(accelerator._models) + reprepared_model = accelerator.prepare_model(prepared_model) + assert len(accelerator._models) == num_models_before, ( + "prepare_model should not add duplicate entries to _models" + ) + assert reprepared_model is prepared_model, "prepare_model should return the same object when called twice" + @require_cuda_or_xpu @slow @require_bnb