# mypy: ignore-errors from torch.fx.experimental.proxy_tensor import is_sym_node, py_sym_types from torch.fx.experimental.sym_node import magic_methods, method_to_operator from torch.fx.experimental.symbolic_shapes import ( hint_int, free_symbols, is_symbol_binding_fx_node, find_symbol_binding_fx_nodes ) from torch.fx.experimental._backward_state import BackwardState import torch import torch.fx as fx import operator import math import torch.utils._pytree as pytree import copy import os import itertools import sympy from collections import defaultdict from torch.fx.passes import graph_drawer from typing import List, Optional, Set, Tuple, Union from .compile_utils import fx_graph_cse, get_aten_target from . import config import functools AOT_PARTITIONER_DEBUG = config.debug_partitioner def must_recompute(node): return node.meta.get("recompute", False) def has_recomputable_ops(fx_g): found = False for node in fx_g.graph.nodes: if must_recompute(node): return True return False def has_recomputable_rng_ops(fx_g): for node in fx_g.graph.nodes: if must_recompute(node) and hasattr(node.target, "tags") and torch.Tag.nondeterministic_seeded in node.target.tags: return True return False def sym_node_size(node): if isinstance(node.meta["val"], (torch.SymInt, torch.SymBool)): return 1 assert isinstance(node.meta["val"], torch.SymFloat) return 4 class InvalidNodeBase: def __repr__(self): return "Invalid Node" InvalidNode = InvalidNodeBase() def _extract_graph_with_inputs_outputs(joint_graph, inputs, outputs): """ Given a graph, extracts out a subgraph that takes the specified nodes as inputs and returns the specified outputs. This includes specifying non-placeholder nodes as inputs. The general strategy is to initialize all inputs with proxies as we encounter them, and trace through the graph, only keeping values which take in valid proxies. Then, all dead code is eliminated. """ new_graph = fx.Graph() env = {} # Add new placeholder nodes in the order specified by the inputs for node in inputs: new_node = new_graph.placeholder(node.name) # Can't use node_copy here as we may be turning previous call_function into placeholders new_node.meta = node.meta env[node] = new_node for node in joint_graph.nodes: if node in inputs: continue elif node.op == 'placeholder': env[node] = InvalidNode elif node.op == 'call_function': all_args = pytree.arg_tree_leaves(*node.args, **node.kwargs) all_args = [isinstance(env[x], InvalidNodeBase) for x in all_args if isinstance(x, fx.Node)] if any(all_args): env[node] = InvalidNode continue env[node] = new_graph.node_copy(node, lambda x: env[x]) elif node.op == 'get_attr': env[node] = new_graph.node_copy(node, lambda x: env[x]) elif node.op == 'output': pass output_values = [] for x in outputs: if isinstance(x, fx.Node): if x not in env: raise RuntimeError(f"Node {x} couldn't be found in env") assert not isinstance(env[x], InvalidNodeBase), f"Node {x} was invalid, but is output" output_values.append(env[x]) else: output_values.append(x) new_graph.output(output_values) new_graph.eliminate_dead_code() new_graph.lint() return new_graph def _is_primal(node): return ( node.op == "placeholder" and "tangents" not in node.target and not _is_bwd_seed_offset(node) and not _is_fwd_seed_offset(node) ) def _is_tangent(node): return node.op == "placeholder" and "tangents" in node.target def _is_bwd_seed_offset(node): return node.op == "placeholder" and ("bwd_seed" in node.target or "bwd_base_offset" in node.target) def _is_fwd_seed_offset(node): return node.op == "placeholder" and ("fwd_seed" in node.target or "fwd_base_offset" in node.target) def _is_backward_state(node): return node.op == "placeholder" and isinstance(node.meta.get("val"), BackwardState) def _extract_fwd_bwd_outputs(joint_module: fx.GraphModule, *, num_fwd_outputs): outputs = pytree.arg_tree_leaves(*(node.args for node in joint_module.graph.nodes if node.op == 'output')) fwd_outputs = outputs[:num_fwd_outputs] bwd_outputs = outputs[num_fwd_outputs:] return fwd_outputs, bwd_outputs def _remove_by_name(saved_values, name): for saved_value in saved_values: if saved_value.name == name: saved_values.remove(saved_value) break def _placeholders(nodes): # Avoid making an entire pass over the graph if we only care about the input placeholders result = [] for node in nodes: if node.op == 'placeholder': result.append(node) else: break # placeholders are all at the start of graph return result def _extract_fwd_bwd_modules(joint_module: fx.GraphModule, saved_values, saved_sym_nodes, *, num_fwd_outputs): fwd_outputs, bwd_outputs = _extract_fwd_bwd_outputs(joint_module, num_fwd_outputs=num_fwd_outputs) placeholders = _placeholders(joint_module.graph.nodes) primal_inputs = [*filter(_is_primal, placeholders)] tangent_inputs = [*filter(_is_tangent, placeholders)] fwd_seed_offset_inputs = [*filter(_is_fwd_seed_offset, placeholders)] bwd_seed_offset_inputs = [*filter(_is_bwd_seed_offset, placeholders)] backward_state_inputs = [*filter(_is_backward_state, placeholders)] bwd_graph = _extract_graph_with_inputs_outputs( joint_module.graph, saved_sym_nodes + saved_values + tangent_inputs + bwd_seed_offset_inputs, bwd_outputs ) for node in _placeholders(bwd_graph.nodes): assert node.op == 'placeholder' # This is to filter out saved values that don't actually end up being used by the backwards pass if not node.users: _remove_by_name(saved_values, node.name) _remove_by_name(saved_sym_nodes, node.name) elif _is_backward_state(node): # BackwardState is saved directly _remove_by_name(saved_values, node.name) assert backward_state_inputs # Now that we have the finalized list of saved values, we need to ensure # we propagate all symbols which are referenced by backwards inputs. # These are not directly used in the graph but are required for downstream # sizevar assignment saved_symbols: Set[sympy.Symbol] = set() saved_sym_nodes_binding = [] saved_sym_nodes_derived = [] # Some symbols may already be bound in the directly saved_sym_nodes, # keep track of them so we don't re-bind them for node in saved_sym_nodes: symbol = is_symbol_binding_fx_node(node) if symbol: saved_symbols.add(symbol) saved_sym_nodes_binding.append(node) else: saved_sym_nodes_derived.append(node) # Now go through all of the prospective backward inputs and track any # other symbols we need to bind symbol_bindings = find_symbol_binding_fx_nodes(joint_module.graph) for node in itertools.chain(saved_sym_nodes_derived, saved_values, tangent_inputs): if "val" not in node.meta: continue new_symbols = free_symbols(node.meta["val"]) - saved_symbols # NB: Deterministic order please! for s in sorted(new_symbols, key=lambda s: s.name): # NB: For well formed graphs, the symbol should always be present, # but we also have ways to produce ill-formed graphs, e.g., direct # make_fx usages, so don't choke in this case if s not in symbol_bindings: continue saved_sym_nodes_binding.append(symbol_bindings[s]) saved_symbols |= new_symbols # Update saved_sym_nodes that are now reordered to have all bindings at # front. This can also be used later on to figure out the position of saved # sym nodes in the output of fwd graph. saved_sym_nodes.clear() saved_sym_nodes.extend(saved_sym_nodes_binding + saved_sym_nodes_derived) # Now, we re-generate the fwd/bwd graphs. # NB: This might increase compilation time, but I doubt it matters fwd_graph = _extract_graph_with_inputs_outputs( joint_module.graph, primal_inputs + fwd_seed_offset_inputs, fwd_outputs + saved_values + saved_sym_nodes ) bwd_graph = _extract_graph_with_inputs_outputs( joint_module.graph, saved_sym_nodes + saved_values + tangent_inputs + bwd_seed_offset_inputs + backward_state_inputs, bwd_outputs ) fwd_module = fx._lazy_graph_module._make_graph_module(joint_module, fwd_graph) bwd_module = fx._lazy_graph_module._make_graph_module(joint_module, bwd_graph) return fwd_module, bwd_module def default_partition( joint_module: fx.GraphModule, _joint_inputs, *, num_fwd_outputs ) -> Tuple[fx.GraphModule, fx.GraphModule]: """ Partitions the :attr:`joint_module` in a manner that closely resembles the behavior observed in the original ``.forward()`` and ``.backward()`` of the callable, i.e., the resulting forward graph contains those operators that are executed in the original ``.forward()`` callable passed to :func:`aot_function`. The default partitioner collects the operators that are between the forward inputs and the forward outputs. This helps in finding the tensors which have to be stashed for the backward pass. These stashed tensors become the output of the generated forward graph. The remaining operators are then placed in the backward graph. .. warning:: This API is experimental and likely to change. Args: joint_module(fx.GraphModule): The joint forward and backward graph. This is the result of AOT Autograd tracing. Returns: Returns the generated forward and backward Fx graph modules. """ if has_recomputable_ops(joint_module): return min_cut_rematerialization_partition(joint_module, _joint_inputs, num_fwd_outputs=num_fwd_outputs) primal_inputs = list(filter(_is_primal, joint_module.graph.nodes)) fwd_seed_offset_inputs = list(filter(_is_fwd_seed_offset, joint_module.graph.nodes)) inputs = primal_inputs + fwd_seed_offset_inputs fwd_outputs, bwd_outputs = _extract_fwd_bwd_outputs(joint_module, num_fwd_outputs=num_fwd_outputs) forward_only_graph = _extract_graph_with_inputs_outputs(joint_module.graph, inputs, fwd_outputs) forward_node_names = {node.name for node in forward_only_graph.nodes if node.op != 'output'} saved_values = [] saved_sym_nodes = [] for node in joint_module.graph.nodes: if node.name not in forward_node_names: continue if is_sym_node(node): # Symints must be kept separate from tensors so that PythonFunction only calls # save_for_backward on tensors and stashes symints in autograd .ctx saved_sym_nodes.append(node) elif ( 'tensor_meta' not in node.meta and node.op == 'call_function' ): # Since we can't save tuple of tensor values, we need to flatten out what we're saving users = node.users assert all(user.target == operator.getitem for user in users) saved_values.extend(users) else: backward_usages = [n for n in node.users if n.name not in forward_node_names] if 'tensor_meta' in node.meta and all(is_sym_node(n) for n in backward_usages): # If we have a tensor in the forward, where only its sizes/strides are needed in the backward, # and not the actual tensor data, # then it will be a lot cheaper to save only the sizes/strides, and not the actual tensor. # # Note that saving the tensor could also cause compilation problems: # If the user mutated an input in the forward and uses its sizes/strides in the backward, # then we would be obligated to clone the input before saving it to appease autograd. # (This is how we originally found this bug). saved_sym_nodes.extend(backward_usages) else: saved_values.append(node) saved_values = list(dict.fromkeys(saved_values).keys()) saved_sym_nodes = list(dict.fromkeys(saved_sym_nodes).keys()) return _extract_fwd_bwd_modules(joint_module, saved_values, saved_sym_nodes=saved_sym_nodes, num_fwd_outputs=num_fwd_outputs) def _prod(x): s = 1 for i in x: s *= i return s def _tensor_nbytes(numel, dtype): return numel * dtype.itemsize def _size_of(node: fx.Node) -> int: if 'val' in node.meta: val = node.meta['val'] if isinstance(val, py_sym_types): if isinstance(val, torch.SymInt): return 1 else: return 999999 # NB: The fallback values here are meaningless, maybe we should respect # torch._inductor.config.unbacked_symint_fallback (but this is a # layering violation) elif isinstance(val, (list, tuple)): return sum(_tensor_nbytes(hint_int(n.numel(), fallback=4098), n.dtype) for n in val if isinstance(n, torch.Tensor)) elif isinstance(val, torch.Tensor): return _tensor_nbytes(hint_int(val.numel(), fallback=4098), val.dtype) raise RuntimeError(f"Unknown metadata type {type(val)}") # Only needed since we don't always trace with fake tensors. if 'tensor_meta' in node.meta: metadata = node.meta['tensor_meta'] # TODO: What is to_size_hint suppose to be? numel = _prod(map(to_size_hint, metadata.shape)) # noqa: F821 dtype = metadata.dtype else: return 0 return _tensor_nbytes(numel, dtype) # Used for some investigative purposes def _count_ops(graph): from collections import defaultdict cnt = defaultdict(int) for node in graph.nodes: if node.op == 'call_function': cnt[node.target.__name__] += 1 print(sorted(cnt.items(), key=lambda x: x[1], reverse=True)) @functools.lru_cache(None) def pointwise_ops(): ops = [] for attr_name in dir(torch.ops.aten): opoverloadpacket = getattr(torch.ops.aten, attr_name) if not isinstance(opoverloadpacket, torch._ops.OpOverloadPacket): continue for overload in opoverloadpacket.overloads(): op_overload = getattr(opoverloadpacket, overload) if torch.Tag.pointwise in op_overload.tags: # currently aot autograd uses packet not overload ops.append(opoverloadpacket) break return ops def get_depth(node, depth_map): if node in depth_map: return depth_map[node] # Base case if node.op == "placeholder": depth_map[node] = 0 return depth_map[node] # Handle output node if node.op == "output": args = node.args[0] for arg in args: if isinstance(arg, torch.fx.node.Node): get_depth(arg, depth_map) return # Get the depth of args and set the depth of this node arg_depths = [get_depth(arg, depth_map) for arg in node.all_input_nodes if isinstance(arg, torch.fx.node.Node)] # factory ops like full, rand might not have any input args if len(arg_depths) == 0: arg_depths = [0] depth_map[node] = max(arg_depths) + 1 return depth_map[node] def sort_depths(args, depth_map): arg_depths = {arg: depth_map[arg] for arg in args if isinstance(arg, torch.fx.node.Node)} return sorted(arg_depths.items(), key=lambda x: x[1], reverse=True) def reordering_to_mimic_autograd_engine(gm): """ This pass finds the first bwd node in the graph (by looking at users of tangents) and then reorders the graph by walking from this node to all the way to the end of the graph. At each op in this traveral, we insert this op in a new graph and try to bring only the relevant subgraph from the other non-bwd edges relevant for this op. This closely mimics the behavior of autograd engine. Why is this pass required in the first place? This is an artifact of how partitioners work today. The starting point of partitioner is a joint graph, which is fwd and then bwd graph. In the case of checkpointing, we keep portions of fwd graph in their original place in the joint graph, while obtaining a bwd graph. As a result, the resulting bwd graph has copies of recomputed fwd subgraphs followed by the original bwd graph. If we run this naively, this leads to bad memory footprint, because the fwd subgraphs are live for way longer duration than necessary. This pass reorders the operations such that we prioritize the ops for the original bwd graph while only realizing those ops from the fwd graph that are necessary at any given point in the graph. """ new_graph = fx.Graph() env = {} # Add new placeholder nodes in the order specified by the inputs for node in gm.graph.nodes: if node.op == "placeholder": new_node = new_graph.placeholder(node.name) # Can't use node_copy here as we may be turning previous call_function into placeholders new_node.meta = node.meta env[node] = new_node order = {} for idx, node in enumerate(gm.graph.nodes): order[node] = idx # Populate depth for the nodes. Depth is the distance from the inputs. depths = {} output_node = next(node for node in gm.graph.nodes if node.op == "output") get_depth(output_node, depths) def insert_node_in_graph(node): if node in env: return env[node] # Bias traversal towards the nodes that have higher depth - prioritizes # critical path first. for arg, _ in sort_depths(node.all_input_nodes, depths): env[arg] = insert_node_in_graph(arg) env[node] = new_graph.node_copy(node, lambda x: env[x]) return env[node] # Find first bwd node in the graph tangent_inputs = list(filter(_is_tangent, gm.graph.nodes)) first_node_in_bwd = None minimum_order = math.inf for tangent in tangent_inputs: for user in tangent.users: if order[user] < minimum_order: minimum_order = order[user] first_node_in_bwd = user assert first_node_in_bwd is not None # Build the graph op-by-op by starting from the node all the way to the end for node in list(gm.graph.nodes)[order[first_node_in_bwd]:]: insert_node_in_graph(node) # The output node is already built by the traversal. new_gm = torch.fx.GraphModule(gm, new_graph) return new_gm def functionalize_rng_ops(joint_module, fw_module, bw_module, num_sym_nodes): # During user-driven activation checkpointing, we have to ensure that a rng # op in fwd yields the same output as the recomputed rng op in the bwd. To # do this, we use functionalize wrappers to wrap the random ops and share # rng state between the fwd and bwd graphs. # There are 3 main steps to do this # Step 1 - Construct a mapping of rng node between the fwd and its counterpart in bwd. # Step 2 - Modify the fwd pass such that # 1) Replace rand with run_and_save_rng_state wrapper # 2) Replace the users of the original op with the output[1] of this op. # 3) Collect all the rng_state - output[0] of each op, and make them # output nodes. Special care needs to be taken here because fwd outputs # has symints at the very end. # Step 3 - Modify the bwd pass such that # 1) Add the input nodes just before the tangents for the stashed rng states # 2) Replace rand with run_with_save_rng_state wrappers # 3) Use the stashed states as inputs to these ops # Unique id to generate name uid = itertools.count() def get_rng_ops(gmod): random_nodes = {} for node in gmod.graph.nodes: if ( node.op == "call_function" and hasattr(node.target, "tags") and torch.Tag.nondeterministic_seeded in node.target.tags ): random_nodes[node.name] = node return random_nodes def get_device(node): """ Check the example value of the node outputs to find the device type. """ if "val" not in node.meta: return None candidates = node.meta["val"] if not isinstance(candidates, tuple): candidates = (candidates,) for candidate in candidates: if isinstance(candidate, torch.Tensor): if candidate.device.type == "cuda": return "cuda" return "cpu" def get_sample_rng_state(device): if device == "cuda": return torch.cuda.get_rng_state() return torch.get_rng_state() # Step 1 - Construct a mapping of rng node between the fwd and its counterpart in bwd. joint_graph_rng_ops = get_rng_ops(joint_module) fw_graph_rng_ops = get_rng_ops(fw_module) bw_graph_rng_ops = get_rng_ops(bw_module) recomputable_rng_ops_map = dict() for node in joint_module.graph.nodes: if ( must_recompute(node) and hasattr(node.target, "tags") and torch.Tag.nondeterministic_seeded in node.target.tags ): base_node = joint_graph_rng_ops[node.name] fw_node = fw_graph_rng_ops[node.name] bw_node = bw_graph_rng_ops[node.name] recomputable_rng_ops_map[base_node] = {"fwd": fw_node, "bwd": bw_node} run_and_save_rng = torch._prims.rng_prims.run_and_save_rng_state run_with_rng_state = torch._prims.rng_prims.run_with_rng_state for node in bw_module.graph.nodes: if node.op == "placeholder" and "tangent" in node.name: bw_tangent_start_node = node break fw_rng_state_outputs = [] for base_node, node_pair in recomputable_rng_ops_map.items(): # Step 2 - Modify the fwd pass such that fw_node = node_pair["fwd"] bw_node = node_pair["bwd"] fw_graph = fw_module.graph with fw_graph.inserting_before(fw_node): functional_fw_node = fw_graph.create_node( "call_function", run_and_save_rng, args=(fw_node.target, *fw_node.args), kwargs=fw_node.kwargs ) state = fw_graph.create_node("call_function", operator.getitem, args=(functional_fw_node, 0), kwargs={}) rng_output = fw_graph.create_node("call_function", operator.getitem, args=(functional_fw_node, 1,), kwargs={}) fw_node.replace_all_uses_with(rng_output) fw_graph.erase_node(fw_node) fw_rng_state_outputs.append(state) # Step 3 - Modify the bwd pass such that bw_graph = bw_module.graph with bw_graph.inserting_before(bw_tangent_start_node): state_name = f"rng_state_output_{next(uid)}" bw_rng_state_node = bw_graph.placeholder(state_name) bw_rng_state_node.meta["val"] = get_sample_rng_state(get_device(fw_node)) with bw_graph.inserting_before(bw_node): rng_output = bw_graph.create_node( "call_function", run_with_rng_state, args=(bw_rng_state_node, bw_node.target, *bw_node.args), kwargs=bw_node.kwargs ) bw_node.replace_all_uses_with(rng_output) bw_graph.erase_node(bw_node) # Add the rng states in the output of the fwd graph. AOT Autograd assumes # that symints are at the end of forward graph outputs. So, insert the new # rng states accordingly. fw_output_node = next(node for node in fw_module.graph.nodes if node.op == "output") fw_outputs = fw_output_node.args[0] sym_node_start_idx = len(fw_outputs) - num_sym_nodes outputs = fw_outputs[:sym_node_start_idx] + fw_rng_state_outputs + fw_outputs[sym_node_start_idx:] fw_module.graph.output(outputs) fw_module.graph.erase_node(fw_output_node) fw_module.recompile() bw_module.recompile() return fw_module, bw_module def cleanup_recompute_tags(joint_module): """ If there are two consecutive checkpointed blocks with no operator in between, we would still want to stash the tensor at the boundary of checkpointed blocks. The following pass makes the last output node non-recomputable to allow for that. """ for node in joint_module.graph.nodes: if must_recompute(node): for user in node.users: if must_recompute(user) and user.meta["recompute"] > node.meta["recompute"]: node.meta["recompute"] = 0 return joint_module def min_cut_rematerialization_partition( joint_module: fx.GraphModule, _joint_inputs, compiler="inductor", recomputable_ops=None, *, num_fwd_outputs ) -> Tuple[fx.GraphModule, fx.GraphModule]: """ Partitions the joint graph such that the backward recomputes the forward. Recomputing helps in trading off memory bandwidth with computation. To create the fwd and bwd graph, we copy the joint graph, manually set the outputs to just original forward or backward outputs. And then we run the resulting graphs through dead code elimination. .. warning:: This API is experimental and likely to change. Args: joint_module(fx.GraphModule): The joint forward and backward graph. This is the result of AOT Autograd tracing. _joint_inputs: The inputs to the joint graph. This is unused. compiler: This option determines the default set of recomputable ops. Currently, there are two options: ``nvfuser`` and ``inductor``. recomputable_ops: This is an optional set of recomputable ops. If this is not None, then this set of ops will be used instead of the default set of ops. num_fwd_outputs: The number of outputs from the forward graph. Returns: Returns the generated forward and backward Fx graph modules. """ try: import networkx as nx except ImportError as e: raise RuntimeError("Need networkx installed to perform smart recomputation " "heuristics") from e joint_module.graph.eliminate_dead_code() joint_module.recompile() fx_g = joint_module.graph # add the CSE pass if config.cse: cse_graph = fx_graph_cse(fx_g) joint_module.graph = cse_graph full_bw_graph = joint_module.graph graph_has_recomputable_ops = has_recomputable_ops(joint_module) graph_has_recomputable_rng_ops = has_recomputable_rng_ops(joint_module) if graph_has_recomputable_ops: joint_module = cleanup_recompute_tags(joint_module) name_to_node = {} for node in joint_module.graph.nodes: name_to_node[node.name] = node def classify_nodes(joint_module): required_bw_nodes = set() for node in joint_module.graph.nodes: if node.op == 'placeholder' and "tangents" in node.target: required_bw_nodes.add(node) if node in required_bw_nodes: for user in node.users: required_bw_nodes.add(user) primal_inputs = list(filter(_is_primal, joint_module.graph.nodes)) fwd_seed_offset_inputs = list(filter(_is_fwd_seed_offset, joint_module.graph.nodes)) inputs = primal_inputs + fwd_seed_offset_inputs fwd_outputs, bwd_outputs = _extract_fwd_bwd_outputs(joint_module, num_fwd_outputs=num_fwd_outputs) required_bw_nodes.update(o for o in bwd_outputs if o is not None) forward_only_graph = _extract_graph_with_inputs_outputs(joint_module.graph, inputs, fwd_outputs) required_fw_nodes = {name_to_node[node.name] for node in forward_only_graph.nodes if node.op != 'output'} unclaimed_nodes = {node for node in joint_module.graph.nodes if node not in required_fw_nodes and node not in required_bw_nodes} return fwd_outputs, required_fw_nodes, required_bw_nodes, unclaimed_nodes, inputs orig_fw_outputs, required_fw_nodes, required_bw_nodes, unclaimed_nodes, inputs = classify_nodes(joint_module) # networkx blows up on graphs with no required backward nodes # Since there's nothing to partition anyway, and the default partitioner can "handle" # this case, send our graph over to the default partitioner. if len(required_bw_nodes) == 0: return default_partition(joint_module, _joint_inputs, num_fwd_outputs=num_fwd_outputs) for node in reversed(joint_module.graph.nodes): if node not in required_fw_nodes: node.dist_from_bw = 0 else: node.dist_from_bw = int(1e9) for user in node.users: node.dist_from_bw = min(node.dist_from_bw, user.dist_from_bw + 1) aten = torch.ops.aten prims = torch.ops.prims # compiler == "nvfuser" is the default set of recomputable ops default_recomputable_ops = [aten.add, aten.sub, aten.div, aten.atan2, aten.mul, aten.max, aten.min, aten.pow, aten.remainder, aten.fmod, aten.__and__, aten.__or__, aten.__xor__, aten.__lshift__, aten.__rshift__, aten.eq, aten.ne, aten.ge, aten.gt, aten.le, aten.lt, aten.abs, aten.bitwise_not, aten.ceil, aten.floor, aten.frac, aten.neg, aten.relu, aten.round, aten.silu, aten.trunc, aten.log, aten.log10, aten.log1p, aten.log2, aten.lgamma, aten.exp, aten.expm1, aten.erf, aten.erfc, aten.cos, aten.acos, aten.cosh, aten.sin, aten.asin, aten.sinh, aten.tan, aten.atan, aten.tanh, aten.atanh, aten.sqrt, aten.rsqrt, aten.reciprocal, aten.sigmoid, aten.softplus, aten.threshold, aten.threshold_backward, aten.clamp, aten.where, aten.lerp, aten.addcmul, aten.gelu, aten.gelu_backward, aten.sum, aten.mean, aten._grad_sum_to_size, aten.sum_to_size, aten.amax, aten.to, aten.type_as, operator.getitem, aten.squeeze, aten.unsqueeze, aten.rsub, aten._to_copy] # noqa: E501,B950 view_ops = [aten.squeeze, aten.unsqueeze, aten.alias] if compiler == "inductor": default_recomputable_ops += [prims.div, prims.convert_element_type, aten.clone, aten._to_copy, aten.full_like, prims.var, prims.sum, aten.var, aten.std, prims.broadcast_in_dim, aten.select, aten.permute, aten._unsafe_view, aten.view, aten.expand, aten.slice, aten.reshape, aten.broadcast_tensors, aten.scalar_tensor, aten.ones, aten.new_zeros, aten.lift_fresh_copy, aten.arange, aten.triu, aten.var_mean, aten.isinf, aten.any, aten.full, aten.as_strided, aten.zeros, aten.argmax, aten.maximum] # noqa: E501,B950 view_ops += [aten.view, aten.slice, aten.permute, aten.t, prims.broadcast_in_dim, aten.expand, aten.as_strided] # Natalia said that we should allow recomputing indexing :) default_recomputable_ops += [aten.index] default_recomputable_ops += view_ops default_recomputable_ops += pointwise_ops() default_recomputable_ops += [ aten.zeros_like, ] default_recomputable_ops += [ method_to_operator(m) for m in magic_methods ] recomputable_ops = set(recomputable_ops) if recomputable_ops is not None else set(default_recomputable_ops) random_ops = [aten.native_dropout, aten.rand_like, aten.randn_like] compute_intensive_ops = [aten.mm, aten.convolution, aten.convolution_backward, aten.bmm, aten.addmm, aten.upsample_bilinear2d, aten._softmax, aten._softmax_backward_data, aten.native_layer_norm, aten.native_layer_norm_backward, aten.native_batch_norm, aten.native_batch_norm_backward, aten._native_batch_norm_legit] # noqa: E501,B950 fusible_ops = recomputable_ops | set(random_ops) if AOT_PARTITIONER_DEBUG: joint_module_ops = { str(node.target._overloadpacket) for node in joint_module.graph.nodes if node.op == "call_function" and hasattr(node.target, "_overloadpacket") } ops_ignored = joint_module_ops - {str(i) for i in recomputable_ops} print("Ops banned from rematerialization: ", ops_ignored) print() def is_materialized_backwards(node): cur_nodes = {node} while len(cur_nodes) > 0: cur = cur_nodes.pop() for user in cur.users: if user not in required_fw_nodes and not is_fusible(cur, user): return True if user not in required_fw_nodes and get_aten_target(user) in view_ops: cur_nodes.add(user) return False def ban_recomputation(node): if "recompute" in node.meta: return node.meta["recompute"] == 0 elif config.aggressive_recomputation: ignored_ops = random_ops + compute_intensive_ops return (node.op == 'call_function' and get_aten_target(node) in ignored_ops) else: if node.op != 'call_function': return False if get_aten_target(node) not in recomputable_ops: return True if node.target == operator.getitem: return False if node.target in [aten.lift_fresh_copy.default, aten.lift_fresh.default]: return False # If a node *must* be materialized in the backwards pass, then we # should never recompute it. This is a pretty subtle point. In # general, the assumption we make is that recomputing a node in the # backwards pass is "free". However, if a node must be materialized # in the backwards pass, then recomputing it is never free. if is_materialized_backwards(node): return True # Arbitrary hack that sometimes seems to help things. The above # modification appears to have made this heuristic a lot less critical # for performance. # TODO: Investigate why this hack helps. # TODO: Investigate the interaction with compiler assisted # activation checkpointing. Removing the heuristic improves both # memory footprint and speedup. if not graph_has_recomputable_ops: if compiler == "inductor" and node.dist_from_bw > config.max_dist_from_bw: return True # If the output of an op is 4x smaller (arbitrary choice), # then we don't allow recomputation. input_tensors_size = sum(_size_of(i) for i in node.args if isinstance(i, fx.Node)) output_size = _size_of(node) return (output_size * 4 < input_tensors_size) def is_fusible(a, b): # We can perform "memory fusion" into a cat, but cat cannot be a # producer to a fusion if get_aten_target(b) == aten.cat: return True return get_aten_target(a) in fusible_ops and get_aten_target(b) in fusible_ops def is_materialized(node): if node.op == 'placeholder': return True return not all(is_fusible(node, user) for user in node.users) def get_node_weight(node) -> int: mem_sz = _size_of(node) # Heuristic to bias towards nodes closer to the backwards pass # Complete guess about current value mem_sz = int(mem_sz * (1.1 ** max(min(node.dist_from_bw, 100), 1))) # mem_sz = int(mem_sz + node.dist_from_bw) if is_materialized(node): return mem_sz else: return mem_sz * 2 nx_graph = nx.DiGraph() for node in full_bw_graph.nodes: if node.op == 'output': continue if node in required_bw_nodes: if node not in inputs: nx_graph.add_edge(node.name + "_in", "sink", capacity=math.inf) continue # If someone saves a input for backward as-is and backward # returns that tensor as-is as a grad input, then the node x would # be both a required_bw_node and an input. In this case we # (1) connect x_in to to the source, (2) x_out to the sink, and # (3) assign the proper weight to the x_in-x_out edge, so that # x would be part of cut nodes. A case where this happens is if # NestedTensor saves a offset tensor as part of the singleton int # in sizes. nx_graph.add_edge(node.name + "_out", "sink", capacity=math.inf) if _is_primal(node) or _is_fwd_seed_offset(node): nx_graph.add_edge("source", node.name + "_in", capacity=math.inf) # If a node can't be recomputed (too expensive or involves randomness), # we prevent it from being recomputed by adding an inf edge to the source # We only need to ban nodes in the fw pass, as those are the only ones that would be recomputed. if ban_recomputation(node) and node in required_fw_nodes: nx_graph.add_edge("source", node.name + "_in", capacity=math.inf) # Checks if a node is actually a tuple. Can be simplified to just an isinstance check if we always use faketensors. is_non_tensor_node = (('val' not in node.meta and 'tensor_meta' not in node.meta) or ('val' in node.meta and not isinstance(node.meta['val'], torch.Tensor))) if is_sym_node(node): weight = sym_node_size(node) elif is_non_tensor_node: weight = 0 if isinstance(node.meta.get("val"), BackwardState) else math.inf else: weight = get_node_weight(node) # Creates the weights on the "node" edge nx_graph.add_edge(node.name + "_in", node.name + "_out", capacity=weight) for user in node.users: nx_graph.add_edge(node.name + "_out", user.name + "_in", capacity=math.inf) try: cut_value, partition = nx.minimum_cut(nx_graph, "source", "sink") except Exception: print('Failed to compute min-cut on following graph:') print('\n'.join(nx.readwrite.edgelist.generate_edgelist(nx_graph))) raise reachable, non_reachable = partition cutset = set() for u, nbrs in ((n, nx_graph[n]) for n in reachable): cutset.update((u, v) for v in nbrs if v in non_reachable) cut_nodes = set() for node_in, node_out in cutset: assert node_in[:-3] == node_out[:-4] node_name = node_in[:-3] cut_nodes.add(node_name) # To make this stuff deterministic node_idx = {node: idx for idx, node in enumerate(joint_module.graph.nodes)} saved_values = sorted((name_to_node[node] for node in cut_nodes), key=lambda x: node_idx[x]) # save_for_backward on tensors and stashes symints in autograd .ctx saved_sym_nodes = list(filter(is_sym_node, saved_values)) saved_values = list(filter(lambda n: not is_sym_node(n), saved_values)) # NB: saved_sym_nodes will be mutated to reflect the actual saved symbols fw_module, bw_module = _extract_fwd_bwd_modules( joint_module, saved_values, saved_sym_nodes=saved_sym_nodes, num_fwd_outputs=num_fwd_outputs) if graph_has_recomputable_ops: if graph_has_recomputable_rng_ops: fw_module, bw_module = functionalize_rng_ops( joint_module, fw_module, bw_module, len(saved_sym_nodes) ) bw_module = reordering_to_mimic_autograd_engine(bw_module) if AOT_PARTITIONER_DEBUG: print("Theoretical Activations Stored: ", sum([_size_of(i) for i in saved_values]) / 1e9) fw_module_nodes = {node.name for node in fw_module.graph.nodes if node.op == 'call_function'} bw_module_nodes = {node.name for node in bw_module.graph.nodes if node.op == 'call_function'} remat_nodes = fw_module_nodes & bw_module_nodes counts = defaultdict(int) for node in fw_module.graph.nodes: if node.name in remat_nodes and hasattr(node.target, '_overloadpacket'): counts[str(node.target._overloadpacket)] += 1 print(f"# remat/fw/bw: {len(remat_nodes)}/{len(fw_module_nodes)}/{len(bw_module_nodes)}") print("Count of Ops Rematerialized: ", sorted(counts.items(), key=lambda x: x[1], reverse=True)) return fw_module, bw_module def draw_graph( traced: torch.fx.GraphModule, fname: str, figname: str = "fx_graph", clear_meta: bool = True, prog: Union[str, List[str]] = None, parse_stack_trace: bool = False, dot_graph_shape: Optional[str] = None, ) -> None: if clear_meta: new_graph = copy.deepcopy(traced.graph) traced = fx.GraphModule(traced, new_graph) for node in traced.graph.nodes: node.meta = {} base, ext = os.path.splitext(fname) if not ext: ext = ".svg" print(f"Writing FX graph to file: {base}{ext}") g = graph_drawer.FxGraphDrawer( traced, figname, parse_stack_trace=parse_stack_trace, dot_graph_shape=dot_graph_shape, ) x = g.get_main_dot_graph() write_method = getattr(x, "write_" + ext.lstrip(".")) fname = f"{base}{ext}" if prog is None: write_method(fname) else: write_method(fname, prog=prog) def draw_joint_graph( graph: torch.fx.GraphModule, joint_inputs, file_name: str = "full_graph.png", dot_graph_shape: Optional[str] = None, ): draw_graph(graph, file_name, dot_graph_shape=dot_graph_shape) return default_partition(graph, joint_inputs)