From 021e6f37880b4befbe197f8728e273818de41453 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Thu, 23 Jul 2026 16:30:10 +0800 Subject: [PATCH] style: apply ruff formatting to FSDP2 changes --- astrai/parallel/executor.py | 4 +--- scripts/tools/train.py | 4 +++- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/astrai/parallel/executor.py b/astrai/parallel/executor.py index 209615e..5c006db 100644 --- a/astrai/parallel/executor.py +++ b/astrai/parallel/executor.py @@ -363,9 +363,7 @@ class FSDP2Executor(BaseExecutor): @contextmanager def _no_sync(self, model: nn.Module): - fsdp_modules = [ - m for m in model.modules() if isinstance(m, FSDPModule) - ] + fsdp_modules = [m for m in model.modules() if isinstance(m, FSDPModule)] if fsdp_modules: for m in fsdp_modules: m.set_requires_gradient_sync(False, recurse=True) diff --git a/scripts/tools/train.py b/scripts/tools/train.py index 68fcaea..7488718 100644 --- a/scripts/tools/train.py +++ b/scripts/tools/train.py @@ -472,7 +472,9 @@ def train( ] assert os.path.exists(param_path) 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 config_path = os.path.join(param_path, "config.json")