2023-09-19 06:20:26 +00:00
|
|
|
from typing import Any, Dict, List, Union
|
2023-01-06 06:14:45 +00:00
|
|
|
|
|
|
|
from torch.fx.node import Node
|
|
|
|
|
2023-01-20 03:18:17 +00:00
|
|
|
from colossalai.logging import get_dist_logger
|
|
|
|
|
2023-01-31 08:00:06 +00:00
|
|
|
NON_COMPUTE_OP = ["placeholder", "get_attr", "output"]
|
|
|
|
NON_COMPUTE_NAME = ["getattr", "eq", "_assert_is_none", "_assert", "finfo", "size"]
|
2023-01-20 03:18:17 +00:00
|
|
|
logger = get_dist_logger()
|
|
|
|
|
|
|
|
|
2023-02-01 05:18:51 +00:00
|
|
|
class NodeMgr(object):
|
2023-03-08 08:22:30 +00:00
|
|
|
def __init__(self, nodes_list: List[Node]) -> None:
|
|
|
|
self._node_list = nodes_list
|
2023-02-01 05:18:51 +00:00
|
|
|
self._node_dict = {}
|
|
|
|
self._set_node_dict()
|
|
|
|
|
|
|
|
def _set_node_dict(self) -> None:
|
|
|
|
"""
|
|
|
|
create a dict {node_name: node_idx}
|
|
|
|
"""
|
|
|
|
self._node_dict.clear()
|
|
|
|
for idx, node in enumerate(self._node_list):
|
|
|
|
self._node_dict[node.name] = idx
|
|
|
|
|
|
|
|
def find_node_idx(self, node: Node) -> int:
|
|
|
|
"""
|
|
|
|
find node's index
|
|
|
|
"""
|
|
|
|
return self._node_dict[node.name]
|
|
|
|
|
|
|
|
def find_node_idx_by_name(self, node_name: str) -> int:
|
|
|
|
"""
|
|
|
|
find node's index
|
|
|
|
"""
|
|
|
|
return self._node_dict[node_name]
|
|
|
|
|
|
|
|
def get_node_by_idx(self, idx: int) -> Node:
|
|
|
|
"""
|
|
|
|
get a node by index
|
|
|
|
"""
|
|
|
|
return self._node_list[idx]
|
|
|
|
|
|
|
|
def get_node_slice_by_idx(self, start: int, end: int) -> List[Node]:
|
|
|
|
"""
|
|
|
|
get a slice of node by index
|
|
|
|
"""
|
|
|
|
return self._node_list[start:end]
|
|
|
|
|
|
|
|
def get_node_list(self) -> List:
|
|
|
|
"""
|
|
|
|
get full node list
|
|
|
|
"""
|
|
|
|
return self._node_list
|
|
|
|
|
|
|
|
def update_node_list(self, node_list: List) -> None:
|
|
|
|
"""
|
|
|
|
update node list, reset node dict
|
|
|
|
"""
|
|
|
|
self._node_list = node_list
|
|
|
|
self._set_node_dict()
|
|
|
|
|
|
|
|
|
2023-01-31 08:00:06 +00:00
|
|
|
def get_logger() -> Any:
|
2023-01-20 03:18:17 +00:00
|
|
|
return logger
|
|
|
|
|
2023-01-06 06:14:45 +00:00
|
|
|
|
2023-01-20 02:13:03 +00:00
|
|
|
def flat_list(inputs: Any) -> List:
|
2023-01-16 11:25:05 +00:00
|
|
|
"""
|
2023-01-19 03:41:00 +00:00
|
|
|
flat a list by recursion
|
2023-01-16 11:25:05 +00:00
|
|
|
"""
|
2023-01-20 02:13:03 +00:00
|
|
|
if not (isinstance(inputs, list) or isinstance(inputs, set) or isinstance(inputs, tuple)):
|
|
|
|
return [inputs]
|
2023-01-16 11:25:05 +00:00
|
|
|
res = []
|
|
|
|
for i in inputs:
|
|
|
|
if isinstance(i, list) or isinstance(i, set) or isinstance(i, tuple):
|
2023-01-19 03:41:00 +00:00
|
|
|
res.extend(flat_list(i))
|
2023-03-08 08:22:30 +00:00
|
|
|
elif isinstance(i, dict):
|
|
|
|
res.extend(flat_list(list(i.keys())))
|
2023-01-16 11:25:05 +00:00
|
|
|
else:
|
|
|
|
res.append(i)
|
|
|
|
return res
|
|
|
|
|
|
|
|
|
2023-01-20 02:13:03 +00:00
|
|
|
def find_first_tensor_arg(node: Node) -> Node:
|
2023-01-16 11:25:05 +00:00
|
|
|
"""
|
|
|
|
Find the first input tensor arg for a node
|
|
|
|
"""
|
|
|
|
for arg in node.args:
|
|
|
|
if type(arg) == type(node):
|
|
|
|
return arg
|
|
|
|
raise RuntimeError()
|
|
|
|
|
|
|
|
|
2023-01-20 02:13:03 +00:00
|
|
|
def is_non_compute_node(node: Node) -> bool:
|
2023-01-31 08:00:06 +00:00
|
|
|
if any(i == node.op for i in NON_COMPUTE_OP) or any(i == get_node_name(node) for i in NON_COMPUTE_NAME):
|
2023-01-19 03:41:00 +00:00
|
|
|
return True
|
|
|
|
if "getitem" in node.name:
|
2023-02-01 05:18:51 +00:00
|
|
|
if get_node_shape(node) is not None:
|
|
|
|
return False
|
2023-01-19 03:41:00 +00:00
|
|
|
node_args = flat_list(node.args[1:])
|
|
|
|
for node_arg in node_args:
|
|
|
|
if any(i == str(node_arg) for i in ["None", "Ellipsis"]):
|
|
|
|
return False
|
2023-01-20 02:13:03 +00:00
|
|
|
if "slice" in str(node_arg):
|
|
|
|
return False
|
2023-01-06 06:14:45 +00:00
|
|
|
return True
|
|
|
|
return False
|
|
|
|
|
|
|
|
|
2023-03-10 02:23:26 +00:00
|
|
|
def get_node_shape(node: Node) -> Any:
|
|
|
|
"""
|
|
|
|
return node data shape
|
|
|
|
"""
|
|
|
|
if get_node_name(node) in ["split", "unbind"]:
|
2023-02-01 05:18:51 +00:00
|
|
|
return node.meta["tensor_meta"][0].shape
|
2023-01-06 06:14:45 +00:00
|
|
|
if hasattr(node.meta["tensor_meta"], "shape"):
|
|
|
|
return node.meta["tensor_meta"].shape
|
|
|
|
return None
|
|
|
|
|
|
|
|
|
2023-01-20 02:13:03 +00:00
|
|
|
def is_non_memory_node(node: Node) -> bool:
|
|
|
|
if "getitem" in node.name:
|
|
|
|
return True
|
|
|
|
if "output" in node.op:
|
|
|
|
return True
|
|
|
|
return is_non_compute_node(node)
|
|
|
|
|
|
|
|
|
2023-01-31 08:00:06 +00:00
|
|
|
def is_non_compute_node_except_placeholder(node: Node) -> bool:
|
2023-01-19 03:41:00 +00:00
|
|
|
if "placeholder" in node.op:
|
|
|
|
return False
|
|
|
|
return is_non_compute_node(node)
|
2023-01-06 06:14:45 +00:00
|
|
|
|
|
|
|
|
2023-01-31 08:00:06 +00:00
|
|
|
def is_non_compute_node_except_placeholder_output(node: Node) -> bool:
|
2023-01-19 03:41:00 +00:00
|
|
|
if "output" in node.op:
|
|
|
|
return False
|
|
|
|
return is_non_compute_node_except_placeholder(node)
|
2023-01-06 06:14:45 +00:00
|
|
|
|
|
|
|
|
2023-01-31 08:00:06 +00:00
|
|
|
def delete_free_var_from_last_use(user_to_last_uses: Dict) -> None:
|
2023-01-06 06:14:45 +00:00
|
|
|
for key, value in user_to_last_uses.items():
|
|
|
|
for n in value:
|
|
|
|
if n.op == "placeholder":
|
|
|
|
user_to_last_uses[key].remove(n)
|
|
|
|
|
|
|
|
|
2023-01-31 08:00:06 +00:00
|
|
|
def find_chunk_all_input_nodes(nodes: List[Node]) -> List:
|
2023-01-06 06:14:45 +00:00
|
|
|
"""
|
|
|
|
Find non-compute input and output node names.
|
|
|
|
input nodes are nodes used in the list
|
|
|
|
output nodes are nodes will use nodes in the list
|
|
|
|
"""
|
|
|
|
input_nodes = []
|
|
|
|
for node in nodes:
|
|
|
|
for input_node in node._input_nodes.keys():
|
|
|
|
if input_node not in nodes and input_node not in input_nodes:
|
|
|
|
input_nodes.append(input_node)
|
|
|
|
return input_nodes
|
|
|
|
|
|
|
|
|
2023-01-31 08:00:06 +00:00
|
|
|
def find_chunk_compute_input_and_output_nodes(nodes: List[Node]) -> Union[List, List]:
|
2023-01-06 06:14:45 +00:00
|
|
|
"""
|
|
|
|
Find non-compute input and output node names.
|
|
|
|
input nodes are nodes used in the list
|
|
|
|
output nodes are nodes will use nodes in the list
|
|
|
|
"""
|
|
|
|
input_nodes = []
|
|
|
|
output_nodes = []
|
|
|
|
|
|
|
|
# if a node has an input node which is not in the node list
|
|
|
|
# we treat that input node as the input of the checkpoint function
|
|
|
|
for node in nodes:
|
|
|
|
for input_node in node._input_nodes.keys():
|
2023-09-19 06:20:26 +00:00
|
|
|
if (
|
|
|
|
input_node not in nodes
|
|
|
|
and input_node not in input_nodes
|
|
|
|
and not is_non_compute_node_except_placeholder(input_node)
|
|
|
|
):
|
2023-01-06 06:14:45 +00:00
|
|
|
input_nodes.append(input_node)
|
|
|
|
|
|
|
|
# if a node has a user node which is not in the node list
|
|
|
|
# we treat that user node as the node receiving the current node output
|
|
|
|
for node in nodes:
|
|
|
|
for output_node in node.users.keys():
|
2023-09-19 06:20:26 +00:00
|
|
|
if (
|
|
|
|
output_node not in nodes
|
|
|
|
and node not in output_nodes
|
|
|
|
and not is_non_compute_node_except_placeholder_output(output_node)
|
|
|
|
):
|
2023-01-06 06:14:45 +00:00
|
|
|
output_nodes.append(node)
|
|
|
|
|
|
|
|
return input_nodes, output_nodes
|
2023-01-31 08:00:06 +00:00
|
|
|
|
|
|
|
|
|
|
|
def get_module_node_name(node: Node) -> str:
|
|
|
|
"""
|
|
|
|
get module class name
|
|
|
|
"""
|
|
|
|
node_targets = node.target.split(".")
|
|
|
|
module = node.graph.owning_module
|
|
|
|
for i in node_targets:
|
|
|
|
module = getattr(module, i)
|
|
|
|
module_name = str(module.__class__).split(".")[-1][:-2]
|
|
|
|
module_name = module_name.lower()
|
|
|
|
return module_name
|
|
|
|
|
|
|
|
|
|
|
|
def get_node_name(node: Node) -> str:
|
|
|
|
"""
|
|
|
|
get node name
|
|
|
|
"""
|
|
|
|
node_name = node.name
|
|
|
|
if "_" in node_name:
|
|
|
|
for i in range(len(node_name) - 1, -1, -1):
|
|
|
|
if node_name[i] == "_":
|
|
|
|
node_name = node_name[:i]
|
|
|
|
break
|
|
|
|
elif node_name[i] in ["1", "2", "3", "4", "5", "6", "7", "8", "9", "0"]:
|
|
|
|
continue
|
|
|
|
else:
|
|
|
|
break
|
|
|
|
return node_name
|
2023-02-01 05:18:51 +00:00
|
|
|
|
|
|
|
|
|
|
|
def find_tensor_node(node_list: List[Node]) -> List[Node]:
|
|
|
|
"""
|
|
|
|
find tensor nodes from a node list
|
|
|
|
"""
|
|
|
|
out = []
|
|
|
|
for node in node_list:
|
|
|
|
if get_node_shape(node) is not None:
|
|
|
|
out.append(node)
|
|
|
|
return out
|
|
|
|
|
|
|
|
|
|
|
|
def find_tensor_shape_node(node_list: List[Node]) -> List[Node]:
|
|
|
|
"""
|
|
|
|
find tensor and shape nodes from a node list
|
|
|
|
"""
|
|
|
|
out = []
|
|
|
|
for node in node_list:
|
|
|
|
if get_node_shape(node) is not None:
|
|
|
|
out.append(node)
|
2023-09-19 06:20:26 +00:00
|
|
|
elif (
|
|
|
|
len(node.meta["fwd_out"]) > 0
|
|
|
|
and isinstance(node.meta["fwd_out"], list)
|
|
|
|
and isinstance(node.meta["fwd_out"][0], int)
|
|
|
|
):
|
2023-02-01 05:18:51 +00:00
|
|
|
out.append(node)
|
|
|
|
return out
|