github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

try import deepspeed when using zero (#130)

pull/133/head
ver217 2022-01-07 17:24:57 +08:00 committed by GitHub
parent b7975d2bcd
commit 9ef05ed1fc
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
colossalai/zero/__init__.py
View File

@ -29,6 +29,7 @@ def convert_to_zero(model: nn.Module,
:return: (model, optimizer)
:rtype: Tuple
"""
import deepspeed
assert level == 2 or level == 3, 'Only ZERO Optimizer Level 2 and 3 are provided'
model = NaiveAMPModel(model, output_to_fp32=False)