from copy import deepcopy from typing import Any, List, Tuple from torch import Tensor from torch.fx import Graph, Node from colossalai.auto_parallel.passes.runtime_apply_pass import runtime_apply, runtime_comm_spec_apply from colossalai.fx.codegen.activation_checkpoint_codegen import _find_nested_ckpt_regions from colossalai.fx.profiler import ( activation_size, calculate_bwd_time, calculate_fwd_out, calculate_fwd_time, calculate_fwd_tmp, ) from colossalai.logging import get_dist_logger from .ckpt_solver_base import CheckpointSolverBase from .operation import Backward, Chain, ForwardCheck, ForwardEnable, ForwardNograd, Loss, Sequence __all__ = ["CheckpointSolverRotor"] class CheckpointSolverRotor(CheckpointSolverBase): def __init__( self, graph: Graph, free_memory: float = -1, cnode: List[str] = None, memory_slots: int = 500, optim_multiplier: float = 1.0, ): """This is the simple implementation of dynamic programming algorithm rotor in https://hal.inria.fr/hal-02352969. Some code are adapted from https://gitlab.inria.fr/hiepacs/rotor. Usage: Assume that we have a ``GraphModule``, and we have already done the extractions to the graph to retrieve all information needed, then we could use the following code to find a solution using ``CheckpointSolverRotor``: >>> solver = CheckpointSolverRotor(gm.graph, free_memory=torch.cuda.mem_get_info(device=0)[0]) >>> rotor_graph = solver.solve(force_python=True) # otherwise use C solver >>> gm.graph = rotor_graph # set the graph to a new graph Args: graph (Graph): The computing graph to be optimized. free_memory (float, optional): Memory constraint for the solution, unit is byte. Use ``torch.cuda.mem_get_info(device=0)[0]`` to estimate the free_memory. Defaults to -1. cnode (List[str], optional): Common node List, should be the subset of input. Defaults to None. memory_slots (int, optional): Number of slots for discretizing memory budget. Defaults to 500. optim_multiplier (float, optional): The multiplier of extra weight storage for the ``torch.optim.Optimizer``. Default to 1.0. """ super().__init__(graph, free_memory, True, cnode, optim_multiplier) self.memory_slots = memory_slots # construct chain unit = self.free_memory // self.memory_slots self.chain = self._construct_chain(self.graph, self.node_list) self.chain.discretize_all(unit) self.cost_table = None self.back_ptr = None self.sequence = None def solve(self, force_python: bool = False, verbose: bool = False) -> Graph: """Solve the checkpointing problem using rotor algorithm. Args: force_python (bool, optional): Use Python version of solver, else use C version. Defaults to False. verbose (bool, optional): Print verbose information. Defaults to False. Returns: graph (Graph): The optimized graph, should be a copy of the original graph. """ chain = self.chain # compute cost table if force_python: self.cost_table, self.back_ptr = self._compute_table(chain, self.memory_slots) else: self.cost_table, self.back_ptr = self._compute_table_c(chain, self.memory_slots) if verbose: self.print_chain() # backtrack try: self.sequence = self._backtrack( chain, 0, len(chain), self.memory_slots - chain.x[0], self.cost_table, self.back_ptr ) self._annotate_from_sequence(self.sequence, self.node_list) except ValueError as e: # using logger to annonce that the solver is failed logger = get_dist_logger() logger.warning(f"Checkpoint solver failed: {e}") raise ValueError if verbose: self.print_sequence() return deepcopy(self.graph) def print_chain(self): print("[input]", self.chain.x[0], self.chain.xbar[0], self.chain.ftmp[0], self.chain.btmp[0]) for idx in range(len(self.node_list) - 1): print( self.node_list[idx], self.chain.x[idx + 1], self.chain.xbar[idx + 1], self.chain.ftmp[idx], self.chain.btmp[idx], ) print(f"Chain = {self.chain}") def print_sequence(self): print(f"Sequence = {self.sequence}") @classmethod def _construct_chain(cls, graph: Graph, node_list: List[List[Node]]) -> Chain: input_tensors = cls._extract_input(graph) ftime, btime, ftmp, btmp = list(), list(), list(), list() xbar, x = [activation_size(input_tensors)], [activation_size(input_tensors)] for node in node_list: node_info = cls._extract_node_info(node) ftime.append(node_info[0]) btime.append(node_info[1]) x.append(node_info[2]) xbar.append(node_info[3]) ftmp.append(node_info[4]) btmp.append(node_info[5]) # currently we view loss backward temp as zero btime.append(0) btmp.append(0) return Chain(ftime, btime, x, xbar, ftmp, btmp) @classmethod def _extract_node_info(cls, node: List[Node]) -> Tuple[int, ...]: """Extract node info from a list of nodes""" xbar = 0 ftime = 0 btime = 0 fwd_mem_peak = 0 for n in node: assert isinstance(n, Node), f"{n} is not a Node" if n.target == runtime_apply or n.target == runtime_comm_spec_apply: # in this case we need to calculate memory usage directly based on the statics that hooked in node.meta xbar += n.meta["fwd_mem_out"] fwd_mem_peak = max(fwd_mem_peak, xbar + n.meta["fwd_mem_tmp"]) else: xbar += calculate_fwd_tmp(n) + calculate_fwd_out(n) fwd_mem_peak = max(fwd_mem_peak, xbar + n.meta["fwd_mem_tmp"] + cls._extract_unused_output(n)) # minimum flop count is required ftime += max(calculate_fwd_time(n), 1.0) btime += max(calculate_bwd_time(n), 1.0) x = calculate_fwd_out(node[-1]) xbar = max(x, xbar) ftmp = fwd_mem_peak - xbar btmp = cls._extract_btmp(node) return ftime, btime, x, xbar, ftmp, btmp @staticmethod def _extract_input(graph: Graph) -> Tuple[Tensor, ...]: """Extract input tensors from a Graph""" input_tensors = [] for node in graph.nodes: if node.op == "placeholder": input_tensors.append(node.meta["fwd_out"]) return input_tensors @staticmethod def _extract_unused_output(node: Node) -> int: """Extract unused output from `torch.fx.Node`""" return activation_size(node.meta["fwd_out"]) - calculate_fwd_out(node) @staticmethod def _extract_btmp(node: List[Node]) -> int: """Extract btmp from a list of nodes""" def _extract_deps_size(): deps_size = 0 for k, v in deps.items(): k: Node if v > 0: deps_size += k.meta["bwd_mem_out"] if v == float("-inf"): deps_size -= calculate_fwd_tmp(k) + calculate_fwd_out(k) return deps_size btmp = 0 deps = {} for n in reversed(node): deps[n] = len(n.all_input_nodes) btmp = max(btmp, _extract_deps_size() + n.meta["bwd_mem_tmp"]) for child in n.users: if child in deps: deps[child] -= 1 if deps[child] <= 0: deps[child] = float("-inf") # free return btmp @staticmethod def _compute_table(chain: Chain, mmax: int) -> Tuple: """Compute the table using dynamic programming. Returns the cost table and the backtracking pointer. Args: chain (Chain): A basic linearized structure for solving the dynamic programming problem. mmax (int): Maximum number of memory slots. Returns: cost_table (List): cost_table[m][lhs][rhs] indicates the optimal cost of the subproblem from lhs to rhs with m memory slots. back_ptr (List): back_ptr[m][lhs][rhs] indicates the best operation at this point. It is (True,) if the optimal choice is a chain checkpoint, it is (False, j) if the optimal choice is a leaf checkpoint of length j """ ftime = chain.ftime + [0.0] btime = chain.btime x = chain.x + [0] xbar = chain.xbar + [0] ftmp = chain.ftmp + [0] btmp = chain.btmp + [0] # Build table cost_table = [[{} for _ in range(len(chain) + 1)] for _ in range(mmax + 1)] back_ptr = [[{} for _ in range(len(chain) + 1)] for _ in range(mmax + 1)] # Initialize corner cases where length of sequence equals to 1, i.e. lhs == rhs for m in range(mmax + 1): for i in range(len(chain) + 1): limit = max(x[i + 1] + xbar[i + 1] + ftmp[i], x[i + 1] + xbar[i + 1] + btmp[i]) if m >= limit: cost_table[m][i][i] = ftime[i] + btime[i] else: cost_table[m][i][i] = float("inf") # Compute tables for m in range(mmax + 1): for d in range(1, len(chain) + 1): for i in range(len(chain) + 1 - d): idx = i + d mmin = x[idx + 1] + x[i + 1] + ftmp[i] if idx > i + 1: mmin = max(mmin, x[idx + 1] + max(x[j] + x[j + 1] + ftmp[j] for j in range(i + 1, idx))) if m < mmin: cost_table[m][i][idx] = float("inf") else: leaf_checkpoints = [ (j, sum(ftime[i:j]) + cost_table[m - x[j]][j][idx] + cost_table[m][i][j - 1]) for j in range(i + 1, idx + 1) if m >= x[j] ] if leaf_checkpoints: best_leaf = min(leaf_checkpoints, key=lambda t: t[1]) else: best_leaf = None if m >= xbar[i + 1]: chain_checkpoint = cost_table[m][i][i] + cost_table[m - xbar[i + 1]][i + 1][idx] else: chain_checkpoint = float("inf") if best_leaf and best_leaf[1] <= chain_checkpoint: cost_table[m][i][idx] = best_leaf[1] back_ptr[m][i][idx] = (False, best_leaf[0]) else: cost_table[m][i][idx] = chain_checkpoint back_ptr[m][i][idx] = (True,) return cost_table, back_ptr @staticmethod def _compute_table_c(chain: Chain, mmax: int) -> Tuple: try: from .rotorc import compute_table # build module if module not found except ModuleNotFoundError: import os import subprocess import sys logger = get_dist_logger() logger.info("rotorc hasn't been built! Building library...", ranks=[0]) this_dir = os.path.dirname(os.path.abspath(__file__)) result = subprocess.Popen( [ f"{sys.executable}", f"{os.path.join(this_dir, 'build_c_ext.py')}", "build_ext", f"--build-lib={this_dir}", ], stdout=subprocess.PIPE, stderr=subprocess.PIPE, ) if result.wait() == 0: logger.info("rotorc has been built!", ranks=[0]) from .rotorc import compute_table else: logger.warning("rotorc built failed! Using python version!", ranks=[0]) return CheckpointSolverRotor._compute_table(chain, mmax) return compute_table(chain, mmax) @staticmethod def _backtrack( chain: Chain, lhs: int, rhs: int, budget: int, cost_table: List[Any], back_ptr: List[Any] ) -> "Sequence": """Backtrack the cost table and retrieve the optimal checkpointing strategy. Args: chain (Chain): A basic linearized structure for solving the dynamic programming problem. lhs (int): The left index of the interval to backtrack. rhs (int): The right index of the interval to backtrack. budget (int): The memory budget for processing this interval. cost_table (List[Any]): See ``._compute_table()`` for definitions back_ptr (List[Any]): See ``._compute_table()`` for definitions Raises: ValueError: Can not process the chain. Returns: sequence (Sequence): The sequence of executing nodes with checkpoints. """ if budget <= 0: raise ValueError(f"Can not process a chain with negative memory {budget}") elif cost_table[budget][lhs][rhs] == float("inf"): raise ValueError(f"Can not process this chain from index {lhs} to {rhs} with memory {budget}") sequence = Sequence() if rhs == lhs: if lhs == len(chain): sequence += [Loss()] else: sequence += [ForwardEnable(lhs), Backward(lhs)] return sequence if back_ptr[budget][lhs][rhs][0]: sequence += [ ForwardEnable(lhs), CheckpointSolverRotor._backtrack( chain, lhs + 1, rhs, budget - chain.xbar[lhs + 1], cost_table, back_ptr ), Backward(lhs), ] else: best_leaf = back_ptr[budget][lhs][rhs][1] sequence += [ForwardCheck(lhs)] sequence += [ForwardNograd(k) for k in range(lhs + 1, best_leaf)] sequence += [ CheckpointSolverRotor._backtrack( chain, best_leaf, rhs, budget - chain.x[best_leaf], cost_table, back_ptr ), CheckpointSolverRotor._backtrack(chain, lhs, best_leaf - 1, budget, cost_table, back_ptr), ] return sequence @staticmethod def _annotate_from_sequence(sequence: Sequence, node_list: List[List[Node]]): """Annotate the nodes in the ``node_list`` with activation checkpoint from the sequence. Args: sequence (Sequence): The sequence of executing nodes with activation checkpoint annotations. node_list (List[List[Node]]): The list of nodes to annotate. """ op_list = sequence.list_operations() loss_op = next(op for op in op_list if isinstance(op, Loss)) fwd_list = op_list[: op_list.index(loss_op)] bwd_list = op_list[op_list.index(loss_op) + 1 :] ckpt_idx = 0 in_ckpt = False ckpt_region = [] # forward annotation for idx, op in enumerate(fwd_list, 0): if in_ckpt: if isinstance(op, ForwardNograd): ckpt_region.append(idx) elif isinstance(op, ForwardEnable): in_ckpt = False for node_idx in ckpt_region: for n in node_list[node_idx]: n.meta["activation_checkpoint"] = [ckpt_idx] ckpt_idx += 1 ckpt_region = [] elif isinstance(op, ForwardCheck): for node_idx in ckpt_region: for n in node_list[node_idx]: n.meta["activation_checkpoint"] = [ckpt_idx] ckpt_idx += 1 ckpt_region = [idx] else: if isinstance(op, ForwardCheck): in_ckpt = True ckpt_region.append(idx) # annotate the backward if there is any nested activation checkpoint in_recompute = False for op in bwd_list: if in_recompute: if isinstance(op, ForwardNograd): ckpt_region.append(op.index) elif isinstance(op, ForwardEnable): for node_idx in ckpt_region: for n in node_list[node_idx]: n.meta["activation_checkpoint"].append(ckpt_idx) ckpt_idx += 1 ckpt_region = [] elif isinstance(op, ForwardCheck): for node_idx in ckpt_region: for n in node_list[node_idx]: n.meta["activation_checkpoint"].append(ckpt_idx) ckpt_idx += 1 ckpt_region = [op.index] elif isinstance(op, Backward): for node_idx in ckpt_region: for n in node_list[node_idx]: n.meta["activation_checkpoint"].append(ckpt_idx) in_recompute = False else: if not isinstance(op, Backward): in_recompute = True ckpt_idx = 0 ckpt_region = [] if isinstance(op, ForwardCheck): ckpt_region.append(op.index) # postprocess, make sure every activation checkpoint label in the # same activation checkpoint region (level = 0) has the same length op_list = [] for node in node_list: op_list += node ckpt_regions = _find_nested_ckpt_regions(op_list) for start_idx, end_idx in ckpt_regions: nested_length = max( len(op_list[idx].meta["activation_checkpoint"]) for idx in range(start_idx, end_idx + 1) ) for idx in range(start_idx, end_idx + 1): op_list[idx].meta["activation_checkpoint"] += [None] * ( nested_length - len(op_list[idx].meta["activation_checkpoint"]) )