ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

style: apply ruff formatting to FSDP2 changes

Browse Source
This commit is contained in:
ViperEkura
2026-07-23 16:30:10 +08:00
parent 4e38183e86
commit 021e6f3788
2 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/parallel/executor.py
+1 -3
View File
@@ -363,9 +363,7 @@ class FSDP2Executor(BaseExecutor):
@contextmanager @contextmanager
def _no_sync(self, model: nn.Module): def _no_sync(self, model: nn.Module):
fsdp_modules = [ fsdp_modules = [m for m in model.modules() if isinstance(m, FSDPModule)]
m for m in model.modules() if isinstance(m, FSDPModule)
]
if fsdp_modules: if fsdp_modules:
for m in fsdp_modules: for m in fsdp_modules:
m.set_requires_gradient_sync(False, recurse=True) m.set_requires_gradient_sync(False, recurse=True)
scripts/tools/train.py
+3 -1
View File
@@ -472,7 +472,9 @@ def train(
] ]
assert os.path.exists(param_path) assert os.path.exists(param_path)
if nprocs > 1 and parallel_mode == "none": if nprocs > 1 and parallel_mode == "none":
raise ValueError("--nprocs > 1 requires --parallel_mode to be 'ddp', 'fsdp', or 'fsdp2'") raise ValueError(
"--nprocs > 1 requires --parallel_mode to be 'ddp', 'fsdp', or 'fsdp2'"
)
# Load config # Load config
config_path = os.path.join(param_path, "config.json") config_path = os.path.join(param_path, "config.json")