feat: 新增FSDP并行后端

- FSDPExecutor通过**fsdp_kwargs直传FSDP参数
- unwrap_model同时支持DDP和FSDP
- parallel_mode新增fsdp选项
This commit is contained in:
2026-05-25 19:43:14 +08:00
parent 82a3f2626f
commit 7df6eb9211
4 changed files with 39 additions and 2 deletions
+3 -1
View File
@@ -8,15 +8,17 @@ import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.nn.parallel import DistributedDataParallel as DDP
from astrai.factory import BaseFactory
def unwrap_model(model: nn.Module) -> nn.Module:
"""Unwrap DDP wrapper if present to get the original model."""
if isinstance(model, DDP):
return model.module
if isinstance(model, FSDP):
return model._fsdp_wrapped_module
return model