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

[elixir] add elixir plugin and its unit test (#3865)

Browse Source
This commit is contained in:
Haichen Huang
2023-05-31 12:10:44 +08:00
committed by GitHub
parent 206280408a
commit dbb9659099
10 changed files with 386 additions and 96 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
tests/kit/model_zoo/transformers/bert.py
View File

@@ -19,13 +19,18 @@ def data_gen_fn():
output_transform_fn = lambda x: x
def output_bert(x):
return dict(pooler_output=x.get('pooler_output'))
config = transformers.BertConfig(hidden_size=128, num_hidden_layers=2, num_attention_heads=4, intermediate_size=256)
# register the BERT variants
model_zoo.register(name='transformers_bert',
model_fn=lambda: transformers.BertModel(config),
data_gen_fn=data_gen_fn,
output_transform_fn=output_transform_fn,
output_transform_fn=output_bert,
model_attribute=ModelAttribute(has_control_flow=True))
model_zoo.register(name='transformers_bert_for_pretraining',
model_fn=lambda: transformers.BertForPreTraining(config),