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:
@@ -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",
|
||||
]
|
||||
|
||||
Reference in New Issue
Block a user