ColossalAI/colossalai/autochunk/estimate_memory.py

from typing import Dict, List

import torch
from torch.fx.node import Node

from .utils import NodeMgr, get_node_shape, is_non_memory_node


class EstimateMemory(object):
    """
    Estimate memory with chunk
    """

    def __init__(self) -> None:
        pass

    def _get_node_size(self, x: Node) -> float:
        """
        return node size in MB
        """
        x = x.meta["tensor_meta"]
        if not hasattr(x, "numel"):
            out = sum([i.numel * torch.tensor([], dtype=i.dtype).element_size() for i in x])
        else:
            out = x.numel * torch.tensor([], dtype=x.dtype).element_size()
        out = float(out) / 1024**2
        return out

    def _add_active_node(self, n: Node, active_nodes: Dict, chunk_ratio: float) -> None:
        """
        add an active node and its shape to active node dict
        """
        if get_node_shape(n) is None:
            return
        if n.op == "placeholder":
            return
        if n not in active_nodes:
            node_size = self._get_node_size(n) * chunk_ratio
            active_nodes[n] = node_size

    def _build_delete_node_dict(self, node_mgr: NodeMgr) -> Dict:
        """
        build delete node dict, means node should be deleted at what time
        """
        delete_node_dict = {}
        for idx, node in enumerate(node_mgr.get_node_list()):
            # skip non shape node
            if get_node_shape(node) is None:
                continue
            # dont remove free nodes
            elif node.op == "placeholder":
                delete_node_dict[node] = len(node_mgr.get_node_list())
            # node no user
            elif len(node.users) == 0:
                delete_node_dict[node] = idx
            # log max use
            else:
                node_user_idx = [node_mgr.find_node_idx(i) for i in node.users.keys()]
                delete_node_dict[node] = max(node_user_idx)
        return delete_node_dict

    def _remove_deactive_node(
        self, user_idx: int, user: Node, active_nodes: List, delete_node_dict: List, kept_nodes: List = None
    ) -> None:
        """
        remove deactivate nodes from active nodes
        """
        if kept_nodes is None:
            kept_nodes = []
        if user.op in ("output",):
            return

        for node in list(active_nodes.keys()):
            # dont delete kept nodes
            if node in kept_nodes:
                continue
            # should be deleted
            if delete_node_dict[node] <= user_idx:
                active_nodes.pop(node)

    def _get_tmp_memory(self, node, not_contiguous_list, delete=False):
        mem = 0
        not_contiguous_ops = ["permute"]

        if node.op == "call_function" and any(n in node.name for n in ["matmul", "reshape"]):
            for n in node.args:
                if n in not_contiguous_list:
                    # matmul won't change origin tensor, but create a tmp copy
                    mem += self._get_node_size(n)
        elif node.op == "call_module":
            for n in node.args:
                if n in not_contiguous_list:
                    # module will just make origin tensor to contiguous
                    if delete:
                        not_contiguous_list.remove(n)
        elif node.op == "call_method" and any(i in node.name for i in not_contiguous_ops):
            if node not in not_contiguous_list:
                not_contiguous_list.append(node)
        return mem

    def _get_chunk_ratio(self, node, chunk_node_dim, chunk_size):
        if node not in chunk_node_dim:
            return 1.0
        node_shape = get_node_shape(node)
        chunk_dim = chunk_node_dim[node]["chunk_dim"]
        if chunk_dim is None:
            return 1.0
        else:
            return chunk_size / float(node_shape[chunk_dim])

    def _print_compute_op_mem_log(self, log, nodes, title=None):
        if title:
            print(title)
        for idx, (l, n) in enumerate(zip(log, nodes)):
            if n.op in ["placeholder", "get_attr", "output"]:
                continue
            if any(i in n.name for i in ["getitem", "getattr"]):
                continue
            print("%s:%.2f \t" % (n.name, l), end="")
            if (idx + 1) % 3 == 0:
                print("")
        print("\n")

    def _add_active_nodes_from_list(self, active_nodes: List, nodes: List) -> List:
        """
        add active nodes from nodes
        """
        for n in nodes:
            self._add_active_node(n, active_nodes, 1)

    def _get_memory_from_active_nodes(self, active_nodes: Dict) -> float:
        """
        sum all memory of active nodes
        """
        out = [i for i in active_nodes.values()]
        out = sum(out)
        return out

    def estimate_chunk_inference_mem(self, node_list: List, chunk_infos: Dict = None, print_mem: bool = False):
        """
        Estimate inference memory with chunk

        Args:
            node_list (List): _description_
            chunk_infos (Dict): Chunk information. Defaults to None.
            print_mem (bool): Wether to print peak memory of every node. Defaults to False.

        Returns:
            act_memory_peak_log (List): peak memory of every node
            act_memory_after_node_log (List): memory after executing every node
            active_node_list_log (List): active nodes of every node. active nodes refer to
                nodes generated but not deleted.
        """
        act_memory = 0.0
        act_memory_peak_log = []
        act_memory_after_node_log = []
        active_nodes = {}
        active_nodes_log = []
        not_contiguous_list = []
        node_mgr = NodeMgr(node_list)
        delete_node_dict = self._build_delete_node_dict(node_mgr)

        use_chunk = True if chunk_infos is not None else False
        chunk_within = False
        chunk_region_idx = None
        chunk_ratio = 1  # use it to estimate chunk mem
        chunk_inputs_all = []

        if use_chunk:
            chunk_regions = [i["region"] for i in chunk_infos]
            chunk_starts = [i[0] for i in chunk_regions]
            chunk_ends = [i[1] for i in chunk_regions]
            chunk_inputs = [i["inputs"] for i in chunk_infos]
            chunk_inputs_non_chunk = [i["inputs_non_chunk"] for i in chunk_infos]
            chunk_inputs_all = [j for i in chunk_inputs for j in i] + [j for i in chunk_inputs_non_chunk for j in i]
            chunk_outputs = [i["outputs"] for i in chunk_infos]
            chunk_node_dim = [i["node_chunk_dim"] for i in chunk_infos]
            chunk_sizes = [i["chunk_size"] if "chunk_size" in i else 1 for i in chunk_infos]

        for idx, node in enumerate(node_mgr.get_node_list()):
            # if node in chunk start nodes, change chunk ratio and add chunk_tensor
            if use_chunk and idx in chunk_starts:
                chunk_within = True
                chunk_region_idx = chunk_starts.index(idx)
                self._add_active_nodes_from_list(active_nodes, chunk_outputs[chunk_region_idx])

            # determine chunk ratio for current node
            if chunk_within:
                chunk_ratio = self._get_chunk_ratio(
                    node, chunk_node_dim[chunk_region_idx], chunk_sizes[chunk_region_idx]
                )

            # add current node as active node
            self._add_active_node(node, active_nodes, chunk_ratio)
            act_memory = self._get_memory_from_active_nodes(active_nodes)

            # if node is placeholder, just add the size of the node
            if node.op == "placeholder":
                act_memory_peak_log.append(act_memory)
            # skip output
            elif node.op == "output":
                continue
            # no change for non compute node
            elif is_non_memory_node(node):
                act_memory_peak_log.append(act_memory)
            # node is a compute op, calculate tmp
            else:
                # forward memory
                # TODO: contiguous_memory still not accurate for matmul, view, reshape and transpose
                tmp_memory = self._get_tmp_memory(node, not_contiguous_list, delete=True) * chunk_ratio
                # record max act memory
                act_memory_peak_log.append(act_memory + tmp_memory)

            # remove_deactive_node
            self._remove_deactive_node(idx, node, active_nodes, delete_node_dict, kept_nodes=chunk_inputs_all)

            # if node in chunk end nodes, restore chunk settings
            if use_chunk and idx in chunk_ends:
                self._remove_deactive_node(idx, node, active_nodes, delete_node_dict)  # dont provide kept nodes now
                chunk_within = False
                chunk_ratio = 1
                chunk_region_idx = None

            act_memory = self._get_memory_from_active_nodes(active_nodes)
            act_memory_after_node_log.append(act_memory)
            active_nodes_log.append(active_nodes.copy())

        if print_mem:
            print("with chunk" if use_chunk else "without chunk")
            self._print_compute_op_mem_log(act_memory_peak_log, node_mgr.get_node_list(), "peak")

        # param_memory = parameter_size(gm)
        # all_memory = act_memory + param_memory
        return act_memory_peak_log, act_memory_after_node_log, active_nodes_log