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

[fx] methods to get fx graph property. (#1246)

Browse Source 
* [CLI] add CLI launcher

* Revert "[CLI] add CLI launcher"

This reverts commit df7e6506d4.

* manipulation

* [fx]add graph manipulation methods.

* [fx]methods to get fx graph property.

* add unit test

* add docstring to explain top node and leaf node in this context
This commit is contained in:
YuliangLiu0306
2022-07-12 14:10:37 +08:00
committed by GitHub
parent 30b4fc0eb0
commit 97d713855a
3 changed files with 190 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
colossalai/fx/passes/adding_split_node_pass.py
View File

@@ -10,6 +10,7 @@ def pipe_split():
def balanced_split_pass(gm: torch.fx.GraphModule, pp_size: int):
# TODO(lyl): balanced policy V2, split module by node size(weight+bias+output)
mod_graph = gm.graph
total_param_amount = 0
for param in mod_graph.owning_module.parameters():
@@ -68,6 +69,9 @@ def uniform_split_pass(gm: torch.fx.GraphModule, pp_size: int):
def split_with_split_nodes_pass(annotated_gm: torch.fx.GraphModule):
# TODO(lyl): use partition IR to assign partition ID to each node.
# Currently: analyzing graph -> annotate graph by inserting split node -> use split module pass to split graph
# In future: graph to partitions -> analyzing partition IR -> recombining partitions to get best performance -> assign partition ID to each node
part_idx = 0
def split_callback(n: torch.fx.Node):