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

add doc

Browse Source
This commit is contained in:
oahzxl
2023-01-10 09:59:47 +08:00
parent 1be0ac3cbf
commit 7d4abaa525
3 changed files with 113 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
colossalai/autochunk/reorder_graph.py
View File

@@ -3,6 +3,10 @@ from .utils import find_idx_by_name
class ReorderGraph(object):
"""
Reorder node list and indice trace list
"""
def __init__(self, trace_indice: TraceIndice) -> None:
self.trace_indice = trace_indice
self.all_reorder_map = {
@@ -60,7 +64,9 @@ class ReorderGraph(object):
def _reorder_idx_trace(self, reorder_map):
# reorder list
new_idx_trace_list = [None for _ in range(len(self.trace_indice.indice_trace_list))]
new_idx_trace_list = [
None for _ in range(len(self.trace_indice.indice_trace_list))
]
for old_idx, new_idx in reorder_map.items():
new_idx_trace_list[new_idx] = self.trace_indice.indice_trace_list[old_idx]
self.trace_indice.indice_trace_list = new_idx_trace_list