github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-23 10:30:03 +00:00
Code Issues Packages Projects Releases Wiki Activity

[zero] add zero wrappers (#2523)

Browse Source 
* [zero] add zero wrappers

* change names

* add wrapper functions to init
This commit is contained in:
HELSON
2023-01-29 17:52:58 +08:00
committed by GitHub
parent c198c7c0b0
commit b528eea0f0
7 changed files with 128 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
colossalai/nn/parallel/__init__.py
View File

@@ -1,4 +1,5 @@
from .data_parallel import ColoDDP, ZeroDDP
from .gemini_parallel import GeminiDDP
from .zero_wrapper import zero_model_wrapper, zero_optim_wrapper
__all__ = ['ColoDDP', 'ZeroDDP', 'GeminiDDP']
__all__ = ['ColoDDP', 'ZeroDDP', 'GeminiDDP', 'zero_model_wrapper', 'zero_optim_wrapper']