feat: add FSDP2Executor using torch.distributed.fsdp.fully_shard API

- New FSDP2Executor registers as 'fsdp2' in ExecutorFactory, using per-module fully_shard() instead of FSDP1 FlatParameter wrapper
- FSDP2 preserves original Parameter objects as DTensors, eliminating use_orig_params=True hack
- FSDP2Executor implements _no_sync via set_requires_gradient_sync, clip_grad_norm via unshard, unwrap_model via DTensor.full_tensor
- Drop **_extra/**_ddp_only_kwargs fallbacks in BaseExecutor/FSDPExecutor/FSDP2Executor, replaced by parallel_mode-aware executor_kwargs dispatch in train.py (ddp-only kwargs only passed when parallel_mode=ddp)
- Export FSDP2Executor in astrai.parallel.__init__
This commit is contained in:
2026-07-20 01:46:25 +08:00
parent 74b9308883
commit 0b6a17330f
3 changed files with 94 additions and 8 deletions
+2
View File
@@ -4,6 +4,7 @@ from astrai.parallel.executor import (
BaseExecutor,
DDPExecutor,
ExecutorFactory,
FSDP2Executor,
FSDPExecutor,
GradientState,
NoneExecutor,
@@ -35,4 +36,5 @@ __all__ = [
"NoneExecutor",
"DDPExecutor",
"FSDPExecutor",
"FSDP2Executor",
]