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

seperate reorder

Browse Source
This commit is contained in:
oahzxl
2023-01-06 16:53:01 +08:00
parent 6685a9d022
commit c3d72f7db9
4 changed files with 29 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
colossalai/autochunk/autochunk_codegen.py
View File

@@ -103,7 +103,7 @@ def emit_code_with_chunk(
nodes,
emit_node_func,
delete_unused_value_func,
chunk_region_search,
chunk_region_search: ChunkRegionSearch,
chunk_infos,
):
"""Emit code with nested activation checkpoint
@@ -133,7 +133,7 @@ def emit_code_with_chunk(
chunk_outputs = [i["outputs"][0] for i in chunk_infos]
chunk_outputs_dim = [i["outputs_dim"] for i in chunk_infos]
node_list = chunk_region_search.index_tracer.reorder_node_list(node_list)
node_list = chunk_region_search.reorder_graph.reorder_node_list(node_list)
node_idx = 0
region_idx = 0
within_chunk_region = False