mirror of
https://github.com/deepspeedai/DeepSpeed.git
synced 2025-10-20 15:33:51 +08:00
25 lines
886 B
Python
Executable File
25 lines
886 B
Python
Executable File
'''
|
|
Copyright 2020 The Microsoft DeepSpeed Team
|
|
'''
|
|
from .containers import HFGPT2LayerPolicy
|
|
from .containers import HFBertLayerPolicy
|
|
from .containers import BLOOMLayerPolicy
|
|
from .containers import HFGPTJLayerPolicy
|
|
from .containers import HFGPTNEOLayerPolicy
|
|
from .containers import GPTNEOXLayerPolicy
|
|
from .containers import HFOPTLayerPolicy
|
|
from .containers import MegatronLayerPolicy
|
|
from .containers import HFDistilBertLayerPolicy
|
|
from .containers import HFCLIPLayerPolicy
|
|
from .containers import UNetPolicy
|
|
from .containers import VAEPolicy
|
|
|
|
# transformer-based policies
|
|
replace_policies = [
|
|
HFBertLayerPolicy, HFGPTNEOLayerPolicy, GPTNEOXLayerPolicy, HFGPTJLayerPolicy, MegatronLayerPolicy,
|
|
HFGPT2LayerPolicy, BLOOMLayerPolicy, HFOPTLayerPolicy, HFCLIPLayerPolicy, HFDistilBertLayerPolicy
|
|
]
|
|
|
|
# non-transformer-based policies
|
|
generic_policies = [UNetPolicy, VAEPolicy]
|