github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-02 17:46:42 +00:00
Code Issues Packages Projects Releases Wiki Activity

[booster] add low level zero plugin (#3594)

Browse Source 
* [booster] add low level zero plugin

* [booster] fix gemini plugin test

* [booster] fix precision

* [booster] add low level zero plugin test

* [test] fix booster plugin test oom

* [test] fix booster plugin test oom

* [test] fix googlenet and inception output trans

* [test] fix diffuser clip vision model

* [test] fix torchaudio_wav2vec2_base

* [test] fix low level zero plugin test
This commit is contained in:
Hongxin Liu
2023-04-26 14:37:25 +08:00
committed by GitHub
parent b9a8dff7e5
commit 4b3240cb59
9 changed files with 476 additions and 81 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
tests/kit/model_zoo/diffusers/diffusers.py
View File

@@ -18,6 +18,7 @@ data_vae_fn = lambda: dict(sample=torch.randn(2, 3, 32, 32))
data_unet_fn = lambda: dict(sample=torch.randn(2, 3, 32, 32), timestep=3)
identity_output = lambda x: x
clip_vision_model_output = lambda x: dict(pooler_output=x[1])
def data_clip_model():
@@ -65,7 +66,7 @@ model_zoo.register(name='diffusers_clip_text_model',
model_zoo.register(name='diffusers_clip_vision_model',
model_fn=partial(transformers.CLIPVisionModel, config=transformers.CLIPVisionConfig()),
data_gen_fn=data_clip_vision,
output_transform_fn=identity_output)
output_transform_fn=clip_vision_model_output)
model_zoo.register(name='diffusers_unet2d_model',
model_fn=diffusers.UNet2DModel,