github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-01 09:07:51 +00:00
Code Issues Packages Projects Releases Wiki Activity

[zero] add strict ddp mode (#2508)

Browse Source 
* [zero] add strict ddp mode

* [polish] add comments for strict ddp mode

* [zero] fix test error
This commit is contained in:
HELSON
2023-01-20 14:04:38 +08:00
committed by GitHub
parent c04f183237
commit 2d1a7dfe5f
5 changed files with 31 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
examples/language/gpt/gemini/commons/model_zoo.py
View File

@@ -53,6 +53,14 @@ def gpt2_24b(checkpoint=True):
return GPTLMModel(hidden_size=8192, num_layers=30, num_attention_heads=16, checkpoint=checkpoint)
def gpt2_30b(checkpoint=True):
return GPTLMModel(hidden_size=8192, num_layers=37, num_attention_heads=16, checkpoint=checkpoint)
def gpt2_40b(checkpoint=True):
return GPTLMModel(hidden_size=8192, num_layers=50, num_attention_heads=16, checkpoint=checkpoint)
def model_builder(model_size: str) -> callable:
if model_size == "gpt2_medium":
return gpt2_medium
@@ -66,6 +74,10 @@ def model_builder(model_size: str) -> callable:
return gpt2_20b
elif model_size == "gpt2_24b":
return gpt2_24b
elif model_size == "gpt2_30b":
return gpt2_30b
elif model_size == "gpt2_40b":
return gpt2_40b
else:
raise TypeError(f"model_builder {model_size}")