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

[autoparallel] added liveness analysis (#1516)

Browse Source 
* [autoparallel] added liveness analysis

* remove memory cost
This commit is contained in:
Frank Lee
2022-08-30 15:54:37 +08:00
committed by GitHub
parent 9a9ef65313
commit a0436a62ee
4 changed files with 247 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
colossalai/fx/passes/utils.py
View File

@@ -160,3 +160,20 @@ def assign_bfs_level_to_nodes(graph: Graph):
new_process_list.extend(get_all_consumers(graph, node))
nodes_to_process = new_process_list
current_level += 1
def get_node_module(node) -> torch.nn.Module:
"""
Find the module associated with the given node.
Args:
node (torch.fx.Node): a torch.fx.Node object in the fx computation graph
Returns:
torch.nn.Module: the module associated with the given node
"""
assert node.graph.owning_module is not None, 'Cannot find the owning_module for node.graph, please make sure the graph is associated with a GraphModule object'
assert node.op == 'call_module', f'Expected node.op to be call_module, but found {node.op}'
module = node.graph.owning_module.get_submodule(node.target)
return module