# Copyright (c) Microsoft Corporation. # SPDX-License-Identifier: Apache-2.0 # DeepSpeed Team from typing import Callable, Any, List, Dict, Optional from collections import defaultdict import torch from torch.fx import Node, Graph, GraphModule from torch.fx.node import map_aggregate from .util import get_last_uses def get_output_node(graph: Graph): for v in graph.nodes: if v.target == "output": return v raise ValueError("No output node found") def should_release_reduce_buckets(graph_order, graph_id: int) -> bool: backward_graph_ids = [g_id for g_id, needs_backward in graph_order if needs_backward] return not backward_graph_ids or graph_id == backward_graph_ids[0] def add_end_backward(graph: Graph, graph_id: int, release_reduce_buckets: bool = True): reduce_nodes = [n for n in graph.nodes if n.target == torch.ops.dc.reduce_grad.default] if len(reduce_nodes) == 0: return with graph.inserting_before(get_output_node(graph)): graph.create_node("call_function", torch.ops.dc.end_backward.default, (reduce_nodes, graph_id, release_reduce_buckets)) def replace_reduce_outputs_with_none(graph: Graph): output_node = get_output_node(graph) new_outputs = map_aggregate( output_node.args[0], lambda n: None if isinstance(n, Node) and n.target == torch.ops.dc.reduce_grad.default else n) output_node.args = (new_outputs, ) def move_primals_to_head(graph: Graph): # Move primals to the head of the graph primals = [n for n in graph.nodes if n.op == "placeholder"] non_primals = [n for n in graph.nodes if n.op != "placeholder"] all_nodes = primals + non_primals new_graph = Graph() env = {} for node in all_nodes: new_node = new_graph.node_copy(node, lambda n: env[n.name]) env[node.name] = new_node new_graph.lint() return new_graph def add_args_process(graph: Graph, node: Node, fn: Callable[..., Any], extra_args: List[int] = [], name=None, meta={}) -> List[Node]: # Apply fn to all args of node new_nodes = [] with graph.inserting_before(node): target_args = [arg for arg in node.args if isinstance(arg, Node)] for arg in target_args: new_node = graph.create_node('call_function', fn, (arg, ) + tuple(extra_args), name=name) for k, v in meta.items(): new_node.meta[k] = v node.replace_input_with(arg, new_node) new_nodes.append(new_node) return new_nodes def add_postprocess(graph: Graph, node: Node, fn: Callable[..., Any], extra_args: List[Any] = [], extra_kwargs: Dict[str, Any] = {}, name=None, meta={}) -> Node: # https://github.com/pytorch/examples/blob/main/fx/wrap_output_dynamically.py with graph.inserting_after(node): args = (node, ) for a in extra_args: # To add ds_id args += (a, ) node_users = node.users.keys() new_node = graph.create_node('call_function', fn, args, extra_kwargs, name=name) users = {} for u in node_users: if u != new_node: users[u] = (node, new_node) for u, (old_in, new_in) in users.items(): u.replace_input_with(old_in, new_in) for k, v in meta.items(): new_node.meta[k] = v return new_node def _make_node_meta(node: Node, ds_id: int, comm: bool): meta = {"param_name": node.name, "ds_id": ds_id, "comm": comm} if "tensor_meta" in node.meta: meta["tensor_meta"] = node.meta["tensor_meta"] return meta def add_free_activations(graph_id: int, graph: Graph, activation_node_names: List[str]): node_to_last_use, _ = get_last_uses(graph) activation_nodes_set = set([n for n in graph.nodes if n.op == "placeholder" and n.name in activation_node_names]) offload_id_to_node = {} node_to_wait_reload = {} for node in graph.nodes: if node.target == torch.ops.dc.reload_tensor.default: offload_act = node.args[0] # node_to_offload_id[offload_act] = node.args[2] offload_id_to_node[node.args[2]] = offload_act elif node.target == torch.ops.dc.wait_reload.default: offload_id = node.args[2] node_to_wait_reload[offload_id_to_node[offload_id]] = node activation_nodes_set = set(node_to_wait_reload[n] if n in node_to_wait_reload else n for n in activation_nodes_set) last_user_to_uses = defaultdict(list) for node, last_user in node_to_last_use.items(): last_user_to_uses[last_user].append(node) def _should_free(node: Node) -> bool: if not hasattr(node, "meta"): return False if "tensor_meta" not in node.meta: return False return True def free_tensors(tensors: List[torch.Tensor]): for a in tensors: if a.numel() > 10_000_000: a.data = torch.empty([0], device=a.device, dtype=a.dtype) for last_user, used_nodes in last_user_to_uses.items(): activation_args = [an for an in used_nodes if an in activation_nodes_set and _should_free(an)] if len(activation_args) == 0: continue node_name = f"free_activations_{[n.name for n in used_nodes]}" with graph.inserting_after(last_user): args = (activation_args, ) graph.create_node('call_function', torch.ops.dc.free_tensors.default, args, {}, name=node_name) # Python version for debugging # graph.create_node('call_function', free_tensors, args, {}, name=node_name) def find_node_by_name(gm: GraphModule, name: str) -> Optional[Node]: for node in gm.graph.nodes: if node.name == name: return node return None def get_node_shape_meta(node: Node) -> Optional[torch.Tensor]: return node.meta.get("val") or node.meta.get("example_value") def find_node_by_tag(gm: GraphModule, tag: str) -> Optional[Node]: input_id_node = None for node in gm.graph.nodes: # https://github.com/pytorch/pytorch/blob/085b71eab05cbc7d474a173884269c62d2778f77/torch/_dynamo/utils.py#L5048 tensor_dict = node.meta.get('tensor_dict') if tensor_dict and tensor_dict.get('tag') == tag: input_id_node = node break return input_id_node def replace_node_users(node: Node, replacement: Node, exclude: Optional[List[Node]] = None): exclude = exclude or [] to_replace = [u for u in node.users if u not in exclude] for user in to_replace: user.replace_input_with(node, replacement)