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

[shardformer] support vision transformer (#4096)

Browse Source 
* first v of vit shardformer

* keep vit

* update

* vit shard add vitattention vitlayer

* update num head shard para

* finish test for vit

* add new_model_class & postprocess

* add vit readme

* delete old files & fix the conflict

* fix sth
This commit is contained in:
Kun Lin
2023-06-28 13:28:18 +08:00
committed by Frank Lee
parent ac80937138
commit 8af29ee47a
10 changed files with 159 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tests/test_device/test_device_mesh.py
View File

@@ -86,4 +86,4 @@ def test_device_mesh_from_process_group():
if __name__ == '__main__':
test_device_mesh()
test_device_mesh_from_process_group()
test_device_mesh_from_process_group()