feat: 新增FSDP并行后端
- FSDPExecutor通过**fsdp_kwargs直传FSDP参数 - unwrap_model同时支持DDP和FSDP - parallel_mode新增fsdp选项
This commit is contained in:
@@ -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
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user