From 6c398ec21b019a344f96bc3eb60e7a73391f3715 Mon Sep 17 00:00:00 2001 From: S1ro1 Date: Mon, 4 Aug 2025 00:42:31 +0000 Subject: [PATCH] Fix: prepare works even if nothing except tp specified (rare) --- src/accelerate/accelerator.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py index 3949e0afe56..1814d2dde3c 100755 --- a/src/accelerate/accelerator.py +++ b/src/accelerate/accelerator.py @@ -1584,9 +1584,17 @@ def prepare(self, *args, device_placement=None): return result if len(result) > 1 else result[0] def _prepare_tp(self, *args): + # First pass: prepare everything except schedulers (and model, which is prepared separately below) + result = [ + self._prepare_one(obj, first_pass=True) if not isinstance(obj, torch.nn.Module) else obj for obj in args + ] + + # Second pass: prepare schedulers + result = [self._prepare_one(obj) if not isinstance(obj, torch.nn.Module) else obj for obj in result] + device_mesh = self.torch_device_mesh - for arg in args: + for arg in result: if not isinstance(arg, torch.nn.Module): continue