frozenleaves/accelerate
1
0
Fork 0
mirror of https://github.com/huggingface/accelerate.git synced 2025-10-20 18:13:46 +08:00
Code Issues Packages Projects Releases Wiki Activity

Fix: prepare works even if nothing except tp specified (rare) (#3707)

Browse Source
This commit is contained in:
Matej Sirovatka
2025-08-05 13:07:37 +02:00
committed by GitHub
parent cb343c63d7
commit c173b4fdd6
1 changed files with 9 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
src/accelerate/accelerator.py
View File

@ -1584,9 +1584,17 @@ class Accelerator:
return result if len(result) > 1 else result[0] return result if len(result) > 1 else result[0]
def _prepare_tp(self, *args): 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 device_mesh = self.torch_device_mesh
for arg in args: for arg in result:
if not isinstance(arg, torch.nn.Module): if not isinstance(arg, torch.nn.Module):
continue continue