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

[autochunk] support parsing blocks (#2506)

Browse Source
This commit is contained in:
oahzxl
2023-01-20 11:18:17 +08:00
committed by GitHub
parent 35c0c0006e
commit c04f183237
7 changed files with 314 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
colossalai/autochunk/trace_flow.py
View File

@@ -281,7 +281,10 @@ class TraceFlow(object):
if chunk_dim is not None:
user_source = self.trace_indice._find_source_trace_from_node(user)[chunk_dim]
if input_node_idx in user_source:
input_dict[user_idx] = user_source[input_node_idx]
if get_node_shape(input_node)[user_source[input_node_idx][0]] == 1:
input_dict[user_idx] = [None]
else:
input_dict[user_idx] = user_source[input_node_idx]
else:
return None, None
if len(input_dict) == 0: