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

[autochunk] support vit (#3084)

Browse Source 
support vit for autochunk
* support some new ops for vit
* fix some bugs
* add test for vit
This commit is contained in:
Xuanlei Zhao
2023-03-10 10:23:26 +08:00
committed by GitHub
parent e58a3c804c
commit 10c61de2f7
8 changed files with 445 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
colossalai/autochunk/utils.py
View File

@@ -109,8 +109,11 @@ def is_non_compute_node(node: Node) -> bool:
return False
def get_node_shape(node: Node) -> List:
if get_node_name(node) == "split":
def get_node_shape(node: Node) -> Any:
"""
return node data shape
"""
if get_node_name(node) in ["split", "unbind"]:
return node.meta["tensor_meta"][0].shape
if hasattr(node.meta["tensor_meta"], "shape"):
return node.meta["tensor_meta"].shape