ColossalAI/colossalai/autochunk/estimate_memory.py

import copy
from typing import Any, Callable, Dict, Iterable, List, Tuple

import torch
from torch.fx.node import Node, map_arg

from colossalai.fx.profiler import activation_size, parameter_size

from .utils import (
    delete_free_var_from_last_use,
    find_idx_by_name,
    get_node_shape,
    is_non_compute_node_except_placeholder,
)


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

    def __init__(self) -> None:
        pass

    def _get_meta_node_size(self, x):
        x = x.meta["tensor_meta"]
        x = x.numel * torch.tensor([], dtype=x.dtype).element_size()
        return x

    def _get_output_node(self, n):
        out_size = activation_size(n.meta["fwd_out"])
        out_node = [n.name] if out_size > 0 else []
        return out_size, out_node

    def _get_output_node_size(self, n):
        return self._get_output_node(n)[0]

    def _add_active_node(self, n, active_list):
        new_active = self._get_output_node(n)[1]
        if n.op == "placeholder":
            new_active.append(n.name)
        for i in new_active:
            if i not in active_list:
                active_list.append(i)

    def _get_delete_node(self, user, user_to_last_uses, to_keep=None):
        delete_size = 0
        delete_node = []
        if user.op not in ("output",):
            nodes_to_delete = user_to_last_uses.get(user, [])
            if to_keep is not None:
                keep_list = []
                for n in nodes_to_delete:
                    if n.name in to_keep:
                        keep_list.append(n)
                for n in keep_list:
                    if n in nodes_to_delete:
                        nodes_to_delete.remove(n)
            if len(nodes_to_delete):
                out_node = [self._get_output_node(i) for i in nodes_to_delete]
                delete_size = sum([i[0] for i in out_node])
                for i in range(len(out_node)):
                    if out_node[i][0] > 0:
                        delete_node.append(out_node[i][1][0])
                    elif nodes_to_delete[i].op == "placeholder":
                        delete_node.append(nodes_to_delete[i].name)
                    # elif any(j in nodes_to_delete[i].name for j in ['transpose', 'permute', 'view']):
                    #     delete_node.append(nodes_to_delete[i].name)
        return delete_size, delete_node

    def _get_delete_node_size(self, user, user_to_last_uses, to_keep):
        return self._get_delete_node(user, user_to_last_uses, to_keep)[0]

    def _remove_deactive_node(self, user, user_to_last_uses, active_list):
        delete_node = self._get_delete_node(user, user_to_last_uses)[1]
        for i in delete_node:
            if i in active_list:
                active_list.remove(i)

    def _get_chunk_inputs_size(
        self, chunk_inputs, chunk_inputs_non_chunk, node_list, chunk_end_idx
    ):
        nodes_to_delete = []
        for chunk_input in chunk_inputs + chunk_inputs_non_chunk:
            chunk_input_users = chunk_input.users.keys()
            chunk_input_users_idx = [
                find_idx_by_name(i.name, node_list) for i in chunk_input_users
            ]
            if all(i <= chunk_end_idx for i in chunk_input_users_idx):
                if chunk_input not in nodes_to_delete:
                    nodes_to_delete.append(chunk_input)
        out_node = [self._get_output_node(i) for i in nodes_to_delete]
        delete_size = sum([i[0] for i in out_node])
        return delete_size

    def _get_last_usr(self, nodes):
        node_to_last_use: Dict[Node, Node] = {}
        user_to_last_uses: Dict[Node, List[Node]] = {}

        def register_last_uses(n: Node, user: Node):
            if n not in node_to_last_use:
                node_to_last_use[n] = user
                user_to_last_uses.setdefault(user, []).append(n)

        for node in reversed(nodes):
            map_arg(node.args, lambda n: register_last_uses(n, node))
            map_arg(node.kwargs, lambda n: register_last_uses(n, node))
        return user_to_last_uses

    def _get_contiguous_memory(self, node, not_contiguous_list, delete=False):
        mem = 0
        not_contiguous_ops = ["permute"]
        inherit_contiguous_ops = ["transpose", "view"]

        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_output_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 float(chunk_size) / node_shape[chunk_dim]

    def _get_chunk_delete_node_size(
        self, user, user_to_last_uses, chunk_ratio, chunk_inputs_names
    ):
        # if any(j in user.name for j in ['transpose', 'permute', 'view']):
        #     return 0
        if user.op in ("placeholder", "output"):
            return 0
        nodes_to_delete = user_to_last_uses.get(user, [])
        delete_size = 0
        for n in nodes_to_delete:
            if n.name in chunk_inputs_names:
                continue
            delete_size += self._get_output_node_size(n) * chunk_ratio
        return delete_size

    def _print_mem_log(self, log, nodes, title=None):
        if title:
            print(title)
        for idx, (l, n) in enumerate(zip(log, nodes)):
            print("%s:%.2f \t" % (n.name, l), end="")
            if (idx + 1) % 3 == 0:
                print("")
        print("\n")

    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 estimate_chunk_inference_mem(
        self,
        node_list: List,
        chunk_infos=None,
        print_mem=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 excuting 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_node_list = []
        active_node_list_log = []
        not_contiguous_list = []
        user_to_last_uses = self._get_last_usr(node_list)
        user_to_last_uses_no_free_var = self._get_last_usr(node_list)
        delete_free_var_from_last_use(user_to_last_uses_no_free_var)

        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_names = []

        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_names = [j.name for i in chunk_inputs for j in i] + [
                j.name for i in chunk_inputs_non_chunk for j in i
            ]
            chunk_outputs = [i["outputs"][0] 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_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)
                act_memory += self._get_output_node_size(
                    chunk_outputs[chunk_region_idx]
                ) / (1024**2)

            # 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],
                )

            # if node is placeholder, just add the size of the node
            if node.op == "placeholder":
                act_memory += self._get_meta_node_size(node) * chunk_ratio / (1024**2)
                act_memory_peak_log.append(act_memory)
            # skip output
            elif node.op == "output":
                continue
            # no change for non compute node
            elif is_non_compute_node_except_placeholder(node):
                act_memory_peak_log.append(act_memory)
            # node is a compute op
            # calculate tmp, output node and delete node memory
            else:
                # forward memory
                # TODO: contiguous_memory still not accurate for matmul, view, reshape and transpose
                act_memory += (
                    self._get_contiguous_memory(node, not_contiguous_list)
                    * chunk_ratio
                    / (1024**2)
                )
                act_memory += (
                    self._get_output_node_size(node) * chunk_ratio / (1024**2)
                )
                # record max act memory
                act_memory_peak_log.append(act_memory)
                # delete useless memory
                act_memory -= (
                    self._get_contiguous_memory(node, not_contiguous_list, delete=True)
                    * chunk_ratio
                    / (1024**2)
                )
                # delete unused vars not in chunk_input_list
                # we can't delete input nodes until chunk ends
                if chunk_within:
                    act_memory -= self._get_chunk_delete_node_size(
                        node,
                        user_to_last_uses_no_free_var,
                        chunk_ratio,
                        chunk_inputs_names,
                    ) / (1024**2)
                else:
                    act_memory -= self._get_delete_node_size(
                        node, user_to_last_uses_no_free_var, chunk_inputs_names
                    ) / (1024**2)

            # log active node, only effective without chunk
            self._add_active_node(node, active_node_list)
            self._remove_deactive_node(node, user_to_last_uses, active_node_list)

            # if node in chunk end nodes, restore chunk settings
            if use_chunk and idx in chunk_ends:
                act_memory -= (
                    self._get_output_node_size(node) * chunk_ratio / (1024**2)
                )
                act_memory -= self._get_chunk_inputs_size(
                    chunk_inputs[chunk_region_idx],
                    chunk_inputs_non_chunk[chunk_region_idx],
                    node_list,
                    chunk_regions[chunk_region_idx][1],
                ) / (1024**2)
                chunk_within = False
                chunk_ratio = 1
                chunk_region_idx = None

            act_memory_after_node_log.append(act_memory)
            active_node_list_log.append(copy.deepcopy(active_node_list))

        if print_mem:
            print("with chunk" if use_chunk else "without chunk")
            # self._print_mem_log(act_memory_peak_log, node_list, "peak")
            # self._print_mem_log(act_memory_after_node_log, node_list, "after")
            self._print_compute_op_mem_log(act_memory_peak_log, node_list, "peak")
            # self._print_compute_op_mem_log(
            #     act_memory_after_node_log, node_list, "after"
            # )

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