wehub-resource-sync
1153 lines
41 KiB
Python
1153 lines
41 KiB
Python
# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
|
||
#
|
||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||
# you may not use this file except in compliance with the License.
|
||
# You may obtain a copy of the License at
|
||
#
|
||
# http://www.apache.org/licenses/LICENSE-2.0
|
||
#
|
||
# Unless required by applicable law or agreed to in writing, software
|
||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
# See the License for the specific language governing permissions and
|
||
# limitations under the License.
|
||
from __future__ import annotations
|
||
|
||
import functools
|
||
import logging
|
||
import math
|
||
import os
|
||
import time
|
||
from typing import TYPE_CHECKING
|
||
|
||
import paddle
|
||
from paddle import pir
|
||
from paddle.autograd import backward_utils
|
||
from paddle.base import core
|
||
from paddle.base.framework import in_cinn_debug_mode
|
||
|
||
if TYPE_CHECKING:
|
||
from collections.abc import Sequence
|
||
|
||
_PADDLE_DTYPE_2_NBYTES = {
|
||
core.DataType.BOOL: 1,
|
||
core.DataType.FLOAT16: 2,
|
||
core.DataType.BFLOAT16: 2,
|
||
core.DataType.FLOAT32: 4,
|
||
core.DataType.FLOAT64: 8,
|
||
core.DataType.FLOAT8_E4M3FN: 1,
|
||
core.DataType.FLOAT8_E5M2: 1,
|
||
core.DataType.INT8: 1,
|
||
core.DataType.INT16: 2,
|
||
core.DataType.INT32: 4,
|
||
core.DataType.INT64: 8,
|
||
core.DataType.UINT8: 1,
|
||
core.DataType.COMPLEX64: 8,
|
||
core.DataType.COMPLEX128: 16,
|
||
}
|
||
|
||
# define the default recompute ops that can be fused between pairs
|
||
DEFAULT_RECOMPUTABLE_OPS: list[str] = [
|
||
"pd_op.full_int_array",
|
||
"pd_op.full",
|
||
# "pd_op.sum",
|
||
"pd_op.divide",
|
||
"pd_op.subtract",
|
||
"pd_op.add",
|
||
"pd_op.multiply",
|
||
"pd_op.elementwise_pow",
|
||
"pd_op.rsqrt",
|
||
"pd_op.reshape",
|
||
"pd_op.full_like",
|
||
"pd_op.assign",
|
||
"pd_op.expand",
|
||
"pd_op.scale",
|
||
"pd_op.exp",
|
||
"pd_op.sin",
|
||
"pd_op.cos",
|
||
"pd_op.add_n",
|
||
# "pd_op.any",
|
||
"pd_op.cast",
|
||
"pd_op.concat",
|
||
"pd_op.full_with_tensor",
|
||
"pd_op.gather_nd",
|
||
"pd_op.logical_and",
|
||
"pd_op.logical_not",
|
||
"pd_op.where",
|
||
"pd_op.pow",
|
||
"pd_op.shape",
|
||
"pd_op.shape64",
|
||
"pd_op.slice",
|
||
"pd_op.squeeze",
|
||
"pd_op.unsqueeze",
|
||
"pd_op.transpose",
|
||
# "pd_op.prod",
|
||
"pd_op.log",
|
||
"pd_op.log1p",
|
||
"pd_op.logit",
|
||
# "pd_op.max",
|
||
# "pd_op.min",
|
||
"pd_op.expand_as",
|
||
"pd_op.split",
|
||
"pd_op.arange",
|
||
"pd_op.put_along_axis",
|
||
"pd_op.tanh",
|
||
"pd_op.atan",
|
||
"pd_op.atanh",
|
||
"pd_op.sinh",
|
||
"pd_op.asin",
|
||
"pd_op.asinh",
|
||
"pd_op.cosh",
|
||
"pd_op.acos",
|
||
"pd_op.acosh",
|
||
"pd_op.abs",
|
||
"pd_op.sign",
|
||
"pd_op.expm1",
|
||
"pd_op.erf",
|
||
"pd_op.erfinv",
|
||
"pd_op.ceil",
|
||
"pd_op.floor",
|
||
"pd_op.frac",
|
||
"pd_op.round",
|
||
"pd_op.trunc",
|
||
"pd_op.angle",
|
||
"pd_op.as_complex",
|
||
"pd_op.as_real",
|
||
"pd_op.complex",
|
||
"pd_op.real",
|
||
"pd_op.imag",
|
||
"pd_op.conj",
|
||
"pd_op.greater_equal",
|
||
"pd_op.greater_than",
|
||
"pd_op.not_equal",
|
||
"pd_op.equal",
|
||
"pd_op.less_equal",
|
||
"pd_op.less_than",
|
||
"pd_op.bitwise_and",
|
||
"pd_op.bitwise_or",
|
||
"pd_op.bitwise_xor",
|
||
"pd_op.bitwise_not",
|
||
"pd_op.isinf",
|
||
"pd_op.isnan",
|
||
# "pd_op.gather",
|
||
"pd_op.sigmoid",
|
||
]
|
||
|
||
# define the ops that are tending to recompute.These ops are more likely to save memory and get fused.
|
||
TENDING_TO_RECOMPUTE_OPS: list[str] = [
|
||
"pd_op.full_int_array",
|
||
"pd_op.full",
|
||
]
|
||
|
||
VIEW_OPS: list[str] = []
|
||
|
||
RANDOM_OPS: list[str] = ["pd_op.randint", "pd_op.uniform", "pd_op.dropout"]
|
||
|
||
COMPUTE_INTENSIVE_OPS: list[str] = [
|
||
"pd_op.matmul",
|
||
"pd_op.conv2d",
|
||
"pd_op.layer_norm",
|
||
"pd_op.batchnorm",
|
||
"pd_op.softmax",
|
||
"pd_op.all_reduce_",
|
||
"pd_op.c_broadcast_",
|
||
"pd_op.reduce_",
|
||
]
|
||
|
||
IGNORE_OPS: list[str] = [
|
||
"cf.stack_create",
|
||
]
|
||
|
||
AGGRESSIVE_RECOMPUTATION = False
|
||
# Restricts the amount of computation recompute can do.
|
||
MAX_DIST_FROM_BW = 3
|
||
|
||
MINIMUM_WEIGHT = 0.1
|
||
|
||
|
||
def DebugPrint(*args):
|
||
flag = os.getenv("FLAGS_print_auto_recompute_debug")
|
||
if flag and str(flag).lower() in ("1", "true"):
|
||
print(*args, flush=True)
|
||
|
||
|
||
class JudgeFusionLoop:
|
||
def __init__(self, program, unrecomputable_ops):
|
||
self.ops = program.global_block().ops
|
||
self.unrecomputable_ops = unrecomputable_ops
|
||
self.downstream_unrecomputable_ops_map = {op: set() for op in self.ops}
|
||
self.upstream_unrecomputable_ops_map = {op: set() for op in self.ops}
|
||
self._set_has_unfusible_on_path_map()
|
||
|
||
def _set_has_unfusible_on_path_map(self):
|
||
def _get_used_external_value(op):
|
||
defined_values = set()
|
||
used_values = []
|
||
_get_used_external_value_impl(defined_values, used_values, op)
|
||
return used_values
|
||
|
||
def _get_used_external_value_impl(defined_values, used_values, op):
|
||
for operand in op.operands_source():
|
||
if operand not in defined_values:
|
||
used_values.append(operand)
|
||
defined_values.add(operand)
|
||
for block in op.blocks():
|
||
for value in block.args():
|
||
defined_values.add(value)
|
||
for _, value in block.kwargs():
|
||
defined_values.add(value)
|
||
for block in op.blocks():
|
||
for inner_op in block.ops:
|
||
_get_used_external_value_impl(
|
||
defined_values, used_values, inner_op
|
||
)
|
||
for result_value in op.results():
|
||
defined_values.add(result_value)
|
||
|
||
def _get_producer_ops(op):
|
||
producers = set()
|
||
for operand in _get_used_external_value(op):
|
||
if operand.get_defining_op() is None:
|
||
continue
|
||
source_op = operand.get_defining_op()
|
||
if source_op.get_parent_block() == op.get_parent_block():
|
||
producers.add(source_op)
|
||
return producers
|
||
|
||
def _get_consumer_ops(op):
|
||
consumers = set()
|
||
for result in op.results():
|
||
for parent_op in result.all_used_ops_in_same_block():
|
||
if parent_op is not None:
|
||
consumers.add(parent_op)
|
||
return consumers
|
||
|
||
def _get_upstream_ops_recursively(cur):
|
||
upstream_unrecomputable_ops = set()
|
||
for new_op in _get_producer_ops(cur):
|
||
upstream_unrecomputable_ops |= (
|
||
self.upstream_unrecomputable_ops_map[new_op]
|
||
)
|
||
if cur.name() in self.unrecomputable_ops:
|
||
upstream_unrecomputable_ops.add(cur)
|
||
return upstream_unrecomputable_ops
|
||
|
||
def _get_downstream_ops_recursively(cur):
|
||
downstream_unrecomputable_ops = set()
|
||
for new_op in _get_consumer_ops(cur):
|
||
downstream_unrecomputable_ops |= (
|
||
self.downstream_unrecomputable_ops_map[new_op]
|
||
)
|
||
if cur.name() in self.unrecomputable_ops:
|
||
downstream_unrecomputable_ops.add(cur)
|
||
return downstream_unrecomputable_ops
|
||
|
||
for op in self.ops:
|
||
self.upstream_unrecomputable_ops_map[op] |= (
|
||
_get_upstream_ops_recursively(op)
|
||
)
|
||
for op in reversed(self.ops):
|
||
self.downstream_unrecomputable_ops_map[op] |= (
|
||
_get_downstream_ops_recursively(op)
|
||
)
|
||
|
||
def _has_unfusible_op_on_any_path(self, op1, op2):
|
||
no_unfusible_op_on_path = (
|
||
len(
|
||
self.downstream_unrecomputable_ops_map[op1]
|
||
& self.upstream_unrecomputable_ops_map[op2]
|
||
)
|
||
== 0
|
||
and len(
|
||
self.downstream_unrecomputable_ops_map[op2]
|
||
& self.upstream_unrecomputable_ops_map[op1]
|
||
)
|
||
== 0
|
||
)
|
||
return (
|
||
not no_unfusible_op_on_path
|
||
if op1 is not None and op2 is not None
|
||
else False
|
||
)
|
||
|
||
|
||
class Op2IdxMap:
|
||
def __init__(self, program):
|
||
self.op_to_idx_map = {}
|
||
for idx, op_iter in enumerate(program.global_block().ops):
|
||
self.op_to_idx_map[op_iter] = idx
|
||
|
||
def get_idx(self, op):
|
||
if self.op_to_idx_map.get(op, None):
|
||
return self.op_to_idx_map[op]
|
||
raise RuntimeError("op not found in program")
|
||
|
||
|
||
def auto_recompute(
|
||
program: paddle.static.Program,
|
||
inputs: Sequence[pir.Value],
|
||
outputs: Sequence[pir.Value],
|
||
grad_outputs: Sequence[pir.Value],
|
||
fwd_op_end_idx: int,
|
||
backward_op_start_idx: int,
|
||
recomputable_ops: Sequence[str] | None = None,
|
||
) -> tuple[paddle.static.Program, int]:
|
||
'''
|
||
Considering the compiler fuse strategy, we model the pir graph.
|
||
Convert the pir calculation graph into a networkx calculation
|
||
graph. Find the cut point through the min-cut algorithm,
|
||
which is the value to be saved in pir forward calculation graph.
|
||
|
||
Recompute the forward computation graph to replace intermediate
|
||
variables in the forward graph held by the backward graph.
|
||
|
||
.. warning::
|
||
This API is experimental and likely to change.
|
||
|
||
Args:
|
||
program (Program): The program to be recomputed.
|
||
inputs:(list[Value]|tuple(Value)): The input Values
|
||
of the forward graph.
|
||
outputs:(list[Value]|tuple(Value)): The out Values
|
||
of the forward graph.
|
||
grad_outputs:(list[Value]|tuple(Value)): initial gradient values
|
||
of `outputs` .
|
||
forward_op_end_idx(int): The index of the last forward op.
|
||
backward_op_start_idx(int): The index of the start backward op.
|
||
recomputable_ops(list[str]|tuple(str)|None): The op names that can
|
||
be recomputed. If 'recompute_ops' is None, we will use the
|
||
default recomputable_ops. Default None.
|
||
Returns:
|
||
recomputed_program(Program): The recomputed program.
|
||
fwd_op_end_idx(int): The index of the last forward op in recomputed program.
|
||
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> import numpy as np
|
||
>>> import paddle
|
||
>>> from paddle.autograd.ir_backward import grad as ir_grad
|
||
>>> from paddle.base import core
|
||
>>> from paddle.decomposition import decompose
|
||
>>> def forward(x):
|
||
... y = paddle.sin(x)
|
||
... z = paddle.cos(y)
|
||
... return z
|
||
|
||
>>> np_x = np.random.random(size=[4096, 4096]).astype("float32")
|
||
>>> paddle.enable_static()
|
||
>>> core._set_prim_all_enabled(True)
|
||
>>> main_program = paddle.static.Program()
|
||
>>> with paddle.static.program_guard(main_program):
|
||
>>> x = paddle.static.data(
|
||
>>> name="x", shape=[4096, 4096], dtype="float32"
|
||
>>> )
|
||
>>> x.stop_gradient = False
|
||
>>> out = forward(x)
|
||
>>> out_grad = paddle.full(
|
||
>>> shape=out.shape, fill_value=3, dtype="float32"
|
||
>>> )
|
||
>>> [out] = decompose(main_program, [out])
|
||
>>> [dx] = ir_grad(out, [x], out_grad)
|
||
>>> main_program, _ = paddle.decomposition.auto_recompute(
|
||
>>> main_program,
|
||
>>> [x],
|
||
>>> [out],
|
||
>>> grad_outputs=[out_grad],
|
||
>>> fwd_op_end_idx=2,
|
||
>>> backward_op_start_idx=4
|
||
>>> )
|
||
>>> exe = paddle.static.Executor(paddle.CUDAPlace(0))
|
||
>>> res = exe.run(
|
||
>>> feed={'x': np_x},
|
||
>>> fetch_list=[dx],
|
||
>>> )
|
||
>>> print(main_program)
|
||
{
|
||
(%0) = "pd_op.data" () {dtype:(pd_op.DataType)float32,name:"x",place:(pd_op.Place)Place(undefined:0),shape:(pd_op.IntArray)[4096,4096],stop_gradient:[false]} : () -> pd_op.tensor<4096x4096xf32>
|
||
(%1) = "pd_op.sin" (%0) {stop_gradient:[false]} : (pd_op.tensor<4096x4096xf32>) -> pd_op.tensor<4096x4096xf32>
|
||
(%2) = "pd_op.cos" (%1) {stop_gradient:[false]} : (pd_op.tensor<4096x4096xf32>) -> pd_op.tensor<4096x4096xf32>
|
||
(%3) = "pd_op.full" () {dtype:(pd_op.DataType)float32,place:(pd_op.Place)Place(undefined:0),shape:(pd_op.IntArray)[4096,4096],stop_gradient:[true],value:(Float)3} : () -> pd_op.tensor<4096x4096xf32>
|
||
(%4) = "pd_op.sin" (%0) {stop_gradient:[false]} : (pd_op.tensor<4096x4096xf32>) -> pd_op.tensor<4096x4096xf32>
|
||
(%5) = "pd_op.sin" (%4) {stop_gradient:[false]} : (pd_op.tensor<4096x4096xf32>) -> pd_op.tensor<4096x4096xf32>
|
||
(%6) = "pd_op.full" () {dtype:(pd_op.DataType)float32,place:(pd_op.Place)Place(cpu),shape:(pd_op.IntArray)[1],stop_gradient:[true],value:(Float)-1} : () -> pd_op.tensor<1xf32>
|
||
(%7) = "pd_op.scale" (%5, %6) {bias:(Float)0,bias_after_scale:true,stop_gradient:[false]} : (pd_op.tensor<4096x4096xf32>, pd_op.tensor<1xf32>) -> pd_op.tensor<4096x4096xf32>
|
||
(%8) = "pd_op.multiply" (%7, %3) {stop_gradient:[false]} : (pd_op.tensor<4096x4096xf32>, pd_op.tensor<4096x4096xf32>) -> pd_op.tensor<4096x4096xf32>
|
||
(%9) = "pd_op.cos" (%0) {stop_gradient:[false]} : (pd_op.tensor<4096x4096xf32>) -> pd_op.tensor<4096x4096xf32>
|
||
(%10) = "pd_op.multiply" (%9, %8) {stop_gradient:[false]} : (pd_op.tensor<4096x4096xf32>, pd_op.tensor<4096x4096xf32>) -> pd_op.tensor<4096x4096xf32>
|
||
(%11) = "pd_op.fetch" (%10) {col:(Int32)0,is_persistable:[true],name:"fetch0",stop_gradient:[false]} : (pd_op.tensor<4096x4096xf32>) -> pd_op.tensor<4096x4096xf32>
|
||
}
|
||
'''
|
||
DebugPrint("program before recompute:", program)
|
||
# 1. find smart recompute needed saved values by min-cut algorithm
|
||
# 1.1 classify value nodes
|
||
import networkx as nx
|
||
|
||
start_time = time.time()
|
||
|
||
# model value as graph's node, op as graph's edge
|
||
(
|
||
required_fw_value_nodes,
|
||
required_bw_value_nodes,
|
||
unclaimed_value_nodes,
|
||
) = classify_value_node(program, grad_outputs, fwd_op_end_idx)
|
||
|
||
if len(required_bw_value_nodes) == 0 or backward_op_start_idx >= len(
|
||
program.global_block().ops
|
||
):
|
||
return program, fwd_op_end_idx
|
||
|
||
all_ops = program.global_block().ops
|
||
# 1.2 cal value nodes dist to backward
|
||
dist_from_bw = cal_value_nodes_dist_to_backward(
|
||
all_ops, required_fw_value_nodes
|
||
)
|
||
|
||
# 1.3 classify ops
|
||
default_recomputable_ops = DEFAULT_RECOMPUTABLE_OPS
|
||
view_ops = VIEW_OPS
|
||
|
||
default_recomputable_ops += view_ops
|
||
|
||
recomputable_ops = (
|
||
set(recomputable_ops)
|
||
if recomputable_ops is not None
|
||
else set(default_recomputable_ops)
|
||
)
|
||
|
||
random_ops = RANDOM_OPS
|
||
compute_intensive_ops = COMPUTE_INTENSIVE_OPS
|
||
tending_to_recompute_ops = TENDING_TO_RECOMPUTE_OPS
|
||
|
||
unrecomputable_ops = random_ops + compute_intensive_ops
|
||
|
||
fusible_ops = recomputable_ops | set(random_ops)
|
||
|
||
# 1.4 Model pir graph. Convert the pir calculation graph into a networkx calculation graph.
|
||
outputs = backward_utils.ValueSet(outputs)
|
||
inputs = backward_utils.ValueSet(inputs)
|
||
placeholder_value_nodes = inputs | outputs
|
||
|
||
value_id_dict = {}
|
||
nx_graph = nx.DiGraph()
|
||
|
||
judge_fusion_loop = JudgeFusionLoop(program, unrecomputable_ops)
|
||
forward_ops = set(program.global_block().ops[: fwd_op_end_idx + 1])
|
||
|
||
def _get_bw_no_need_buffer_values(program, backward_op_start_idx):
|
||
need_buffer_values = backward_utils.ValueSet()
|
||
all_values = backward_utils.ValueSet()
|
||
for op in program.global_block().ops[backward_op_start_idx:]:
|
||
for op_operand_source in op.operands_source():
|
||
all_values.add(op_operand_source)
|
||
if op.is_no_need_buffer(op_operand_source):
|
||
continue
|
||
need_buffer_values.add(op_operand_source)
|
||
bw_no_need_buffer_values = all_values - need_buffer_values
|
||
return bw_no_need_buffer_values
|
||
|
||
bw_no_need_buffer_values = _get_bw_no_need_buffer_values(
|
||
program, backward_op_start_idx
|
||
)
|
||
|
||
def _is_fusible(value_node1, value_node2):
|
||
return (
|
||
value_node1.get_defining_op().name() in fusible_ops
|
||
and value_node2.get_defining_op().name() in fusible_ops
|
||
)
|
||
|
||
def _is_materialized_backwards(value_node):
|
||
cur_value_nodes = backward_utils.ValueSet()
|
||
cur_value_nodes.add(value_node)
|
||
while len(cur_value_nodes) > 0:
|
||
cur_value_node = cur_value_nodes.pop()
|
||
users = find_value_node_users(
|
||
cur_value_node, bw_no_need_buffer_values, True, forward_ops
|
||
)
|
||
for user in users:
|
||
if user not in required_fw_value_nodes and not _is_fusible(
|
||
cur_value_node, user
|
||
):
|
||
return True
|
||
if (
|
||
user not in required_fw_value_nodes
|
||
and get_real_define_op_name(user) in view_ops
|
||
):
|
||
cur_value_nodes.add(user)
|
||
return False
|
||
|
||
def _is_materialized(value_node, placeholder_value_nodes):
|
||
if value_node in placeholder_value_nodes:
|
||
return True
|
||
users = find_value_node_users(
|
||
value_node, bw_no_need_buffer_values, True, forward_ops
|
||
)
|
||
return not all(_is_fusible(value_node, user) for user in users)
|
||
|
||
def _get_node_weight(value_node, placeholder_value_nodes):
|
||
mem_sz = cal_value_node_size(value_node)
|
||
|
||
if (
|
||
value_node.get_defining_op().name() in tending_to_recompute_ops
|
||
and mem_sz == 0
|
||
):
|
||
return MINIMUM_WEIGHT
|
||
|
||
# Heuristic to bias towards nodes closer to the backwards pass
|
||
mem_sz = int(
|
||
mem_sz * (1.1 ** max(min(dist_from_bw[value_node], 100), 1))
|
||
)
|
||
|
||
if _is_materialized(value_node, placeholder_value_nodes):
|
||
return mem_sz
|
||
else:
|
||
return mem_sz * 2
|
||
|
||
def _ban_recomputation(value_node):
|
||
if AGGRESSIVE_RECOMPUTATION:
|
||
return value_node.get_defining_op().name() in unrecomputable_ops
|
||
else:
|
||
if value_node.get_defining_op().name() in tending_to_recompute_ops:
|
||
return False
|
||
|
||
if value_node.get_defining_op().name() not in recomputable_ops:
|
||
return True
|
||
|
||
# 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(value_node):
|
||
return True
|
||
|
||
if dist_from_bw[value_node] > MAX_DIST_FROM_BW:
|
||
return True
|
||
# If the output of an op is 4x smaller (arbitrary choice),
|
||
# then we don't allow recomputation.
|
||
output_size = cal_value_node_size(value_node)
|
||
inputs = get_real_input_nodes(value_node)
|
||
inputs_size = sum(cal_value_node_size(i) for i in inputs)
|
||
return output_size * 4 < inputs_size
|
||
|
||
for value_node in (
|
||
required_fw_value_nodes
|
||
| required_bw_value_nodes
|
||
| unclaimed_value_nodes
|
||
):
|
||
if not value_node.initialized():
|
||
continue
|
||
|
||
if value_node.get_defining_op().name() == "builtin.combine":
|
||
continue
|
||
|
||
if value_node.get_defining_op().name() in IGNORE_OPS:
|
||
continue
|
||
|
||
if len(
|
||
value_node.all_used_ops_in_same_block()
|
||
) == 1 and value_node.all_used_ops_in_same_block()[0].name() in [
|
||
"builtin.split",
|
||
"builtin.slice",
|
||
]:
|
||
continue
|
||
|
||
if value_node in required_bw_value_nodes:
|
||
DebugPrint(
|
||
"add edge link from: ", value_node.id, " -> ", "sink", " (inf) "
|
||
)
|
||
nx_graph.add_edge(value_node.id + "_in", "sink", capacity=math.inf)
|
||
value_id_dict[value_node.id] = value_node
|
||
continue
|
||
|
||
if value_node in inputs:
|
||
DebugPrint(
|
||
"add edge link from: ",
|
||
" source ",
|
||
" -> ",
|
||
value_node.id,
|
||
" (inf)",
|
||
)
|
||
nx_graph.add_edge(
|
||
"source", value_node.id + "_in", capacity=math.inf
|
||
)
|
||
value_id_dict[value_node.id] = value_node
|
||
|
||
# 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(value_node)
|
||
and value_node in required_fw_value_nodes
|
||
):
|
||
DebugPrint(
|
||
"add edge link from: ",
|
||
" source ",
|
||
" -> ",
|
||
value_node.id,
|
||
"(inf)",
|
||
)
|
||
nx_graph.add_edge(
|
||
"source", value_node.id + "_in", capacity=math.inf
|
||
)
|
||
value_id_dict[value_node.id] = value_node
|
||
|
||
weight = _get_node_weight(
|
||
value_node,
|
||
placeholder_value_nodes,
|
||
)
|
||
|
||
# Creates the weights on the "node" edge
|
||
nx_graph.add_edge(
|
||
value_node.id + "_in", value_node.id + "_out", capacity=weight
|
||
)
|
||
value_id_dict[value_node.id] = value_node
|
||
|
||
users = find_value_node_users(
|
||
value_node, bw_no_need_buffer_values, True, forward_ops
|
||
)
|
||
for user in users:
|
||
DebugPrint(
|
||
"add edge link from: ",
|
||
value_node.id,
|
||
" -> ",
|
||
user.id,
|
||
" (inf) ",
|
||
)
|
||
nx_graph.add_edge(
|
||
value_node.id + "_out", user.id + "_in", capacity=math.inf
|
||
)
|
||
for user in value_node.all_used_ops_in_same_block():
|
||
if user in forward_ops:
|
||
if judge_fusion_loop._has_unfusible_op_on_any_path(
|
||
value_node.get_defining_op(), user
|
||
):
|
||
DebugPrint(
|
||
"add edge link from: ",
|
||
" source ",
|
||
" -> ",
|
||
value_node.id,
|
||
"(inf)",
|
||
)
|
||
nx_graph.add_edge(
|
||
"source", value_node.id + "_in", capacity=math.inf
|
||
)
|
||
|
||
DebugPrint(
|
||
"add edge link from: ",
|
||
value_node.id,
|
||
" -> ",
|
||
" sink ",
|
||
"(inf)",
|
||
)
|
||
|
||
nx_graph.add_edge(
|
||
value_node.id + "_out", "sink", capacity=math.inf
|
||
)
|
||
|
||
# 1.5 find saved values by minimum cut.
|
||
cut_value, partition = nx.minimum_cut(nx_graph, "source", "sink")
|
||
DebugPrint("Cut Value:", cut_value)
|
||
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_value_nodes = backward_utils.ValueSet()
|
||
for value_node_in, value_node_out in cutset:
|
||
assert value_node_in[:-3] == value_node_out[:-4]
|
||
value_node = value_id_dict[value_node_in[:-3]]
|
||
cut_value_nodes.add(value_node)
|
||
|
||
saved_values = cut_value_nodes
|
||
# (TODO: wanghao107): remove it and fix model
|
||
# saved_values = cut_value_nodes | inputs
|
||
saved_values = cut_value_nodes
|
||
# 2.partition the joint graph by saved values.
|
||
(
|
||
program_after_recompute,
|
||
fwd_op_end_idx_after_recompute,
|
||
) = partition_joint_graph(
|
||
program,
|
||
saved_values,
|
||
inputs,
|
||
outputs,
|
||
bw_no_need_buffer_values,
|
||
fwd_op_end_idx,
|
||
backward_op_start_idx,
|
||
)
|
||
DebugPrint("program after recompute:", program_after_recompute)
|
||
end_time = time.time()
|
||
if in_cinn_debug_mode():
|
||
logger = logging.getLogger("auto-recompute")
|
||
logger.setLevel(logging.INFO)
|
||
logger.info(
|
||
f"Time of auto recompute program: ***** [ {end_time - start_time} ] ***** seconds."
|
||
)
|
||
return program_after_recompute, fwd_op_end_idx_after_recompute
|
||
|
||
|
||
def partition_joint_graph(
|
||
program: paddle.static.Program,
|
||
saved_values: list[pir.Value],
|
||
inputs: list[pir.Value],
|
||
outputs: list[pir.Value],
|
||
bw_no_need_buffer_values: list[pir.Value],
|
||
fwd_op_end_idx: int,
|
||
backward_op_start_idx: int,
|
||
) -> tuple[paddle.static.Program, int]:
|
||
"""
|
||
Partition the joint graph, recompute the intermediate values
|
||
by saved values to save memory.
|
||
Args:
|
||
program(Program): The program to be recomputed.
|
||
saved_values(list[valueiable]): The saved values
|
||
of forward graph which used by backward graph.
|
||
inputs:(list[Value]|tuple(Value)): The input Values
|
||
of the forward graph.
|
||
outputs(list[valueiable]): The out values
|
||
of the forward graph.
|
||
forward_op_end_idx(int): The index of the last forward op.
|
||
backward_op_start_idx(int): The index of the start backward op.
|
||
Returns:
|
||
recomputed_program(Program): The recomputed program.
|
||
fwd_op_end_idx(int): The index of the last forward op in
|
||
recomputed program.
|
||
"""
|
||
saved_values = backward_utils.ValueSet(saved_values)
|
||
outputs = backward_utils.ValueSet(outputs)
|
||
|
||
# 1. Analyze the program, get all forward program mid hold values
|
||
mid_hold_values = analyze_mid_hold_values(
|
||
program,
|
||
saved_values,
|
||
inputs,
|
||
outputs,
|
||
bw_no_need_buffer_values,
|
||
fwd_op_end_idx,
|
||
backward_op_start_idx,
|
||
)
|
||
DebugPrint("saved values: ")
|
||
DebugPrint([f"({v}, {v.get_defining_op().id()})" for v in saved_values])
|
||
DebugPrint("mid values: ")
|
||
DebugPrint([f"({v}, {v.get_defining_op().id()})" for v in mid_hold_values])
|
||
|
||
mem = 0
|
||
for mid in mid_hold_values:
|
||
mem += cal_value_node_size(mid)
|
||
DebugPrint("Saved Memory is: ", mem / 1024 / 1024 / 1024, "GB")
|
||
|
||
# 2. Extract the recompute subgraph and replace forward mid hold values with recompute subgraph's outputs
|
||
program, fwd_op_end_idx = replace_mid_values_with_forward_subgraph(
|
||
program,
|
||
saved_values,
|
||
mid_hold_values,
|
||
fwd_op_end_idx,
|
||
backward_op_start_idx,
|
||
)
|
||
|
||
return program, fwd_op_end_idx
|
||
|
||
|
||
def replace_mid_values_with_forward_subgraph(
|
||
program, saved_values, mid_values, fwd_op_end_idx, backward_op_start_idx
|
||
):
|
||
def _extract_forward_recompute_subgraph_for_backward(
|
||
saved_values, mid_values
|
||
):
|
||
def _find_recompute_ops(
|
||
recompute_value,
|
||
saved_values,
|
||
marked_recompute_ops,
|
||
needed_saved_values,
|
||
chain,
|
||
):
|
||
new_chain = list(chain)
|
||
new_chain.append(recompute_value)
|
||
define_op = recompute_value.get_defining_op()
|
||
if define_op in marked_recompute_ops or define_op is None:
|
||
return
|
||
if define_op.name() in [
|
||
"builtin.parameter",
|
||
"pd_op.data",
|
||
]:
|
||
if recompute_value not in needed_saved_values:
|
||
needed_saved_values.add(recompute_value)
|
||
return
|
||
op_inputs = define_op.operands_source()
|
||
if len(op_inputs) == 0 and define_op.name() not in [
|
||
"pd_op.full",
|
||
"pd_op.full_int_array",
|
||
]:
|
||
raise Exception(
|
||
f"Every path to recompute value {recompute_value} must have saved value or starting point of the path is one of op in [pd_op.full, pd_op.full_int_array], but find {define_op.name()} op, op ir is {define_op}"
|
||
)
|
||
for op_input in op_inputs:
|
||
if op_input in saved_values:
|
||
if op_input not in needed_saved_values:
|
||
needed_saved_values.add(op_input)
|
||
continue
|
||
_find_recompute_ops(
|
||
op_input,
|
||
saved_values,
|
||
marked_recompute_ops,
|
||
needed_saved_values,
|
||
new_chain,
|
||
)
|
||
marked_recompute_ops.add(define_op)
|
||
|
||
return
|
||
|
||
recompute_subgraph_ops = set()
|
||
recompute_subgraph_inputs = backward_utils.ValueSet()
|
||
recompute_subgraph_outputs_backward_needed = mid_values
|
||
|
||
for recompute_value in mid_values:
|
||
_find_recompute_ops(
|
||
recompute_value,
|
||
saved_values,
|
||
recompute_subgraph_ops,
|
||
recompute_subgraph_inputs,
|
||
[],
|
||
)
|
||
|
||
DebugPrint("Recompute Ops: ", len(recompute_subgraph_ops))
|
||
DebugPrint("Recompute Ops: ", recompute_subgraph_ops)
|
||
recompute_subgraph = {
|
||
"inputs": recompute_subgraph_inputs,
|
||
"recompute_ops": recompute_subgraph_ops,
|
||
"outputs": recompute_subgraph_outputs_backward_needed,
|
||
}
|
||
return recompute_subgraph
|
||
|
||
op_2_id_map = Op2IdxMap(program)
|
||
|
||
forward_ops = set(program.global_block().ops[: fwd_op_end_idx + 1])
|
||
backward_ops = set(program.global_block().ops[backward_op_start_idx:])
|
||
first_backward_op = program.global_block().ops[backward_op_start_idx]
|
||
|
||
# 1. find forward subgraph to recompute mid values that backward need to hold.
|
||
recompute_forward_subgraph = (
|
||
_extract_forward_recompute_subgraph_for_backward(
|
||
saved_values, mid_values
|
||
)
|
||
)
|
||
|
||
# 2. clone subgraph which need to be recomputed
|
||
origin_ops = recompute_forward_subgraph["recompute_ops"]
|
||
origin_subgraph_inputs = recompute_forward_subgraph["inputs"]
|
||
origin_subgraph_outputs = recompute_forward_subgraph["outputs"]
|
||
cloned_ops, value_map, cloned_op_first_grad_user_map = clone_graph(
|
||
program,
|
||
origin_ops,
|
||
origin_subgraph_inputs,
|
||
first_backward_op,
|
||
backward_ops,
|
||
op_2_id_map,
|
||
)
|
||
|
||
for origin_op in origin_ops:
|
||
origin_op.set_bool_attr("is_recompute_op", True)
|
||
for cloned_op in cloned_ops:
|
||
cloned_op.set_bool_attr("is_recompute_bw_op", True)
|
||
|
||
# 3. replace mid values that backward need to hold with recompute subgraph's outputs
|
||
cloned_subgraph_outputs = backward_utils.ValueSet()
|
||
for origin_value in origin_subgraph_outputs:
|
||
cloned_value = value_map.look_up(origin_value)
|
||
origin_value.replace_grad_users_with(cloned_value, backward_ops)
|
||
cloned_subgraph_outputs.add(cloned_value)
|
||
|
||
# 4. reset recomputed ops location in program
|
||
for op in reversed(cloned_ops):
|
||
first_subgraph_grad_user = cloned_op_first_grad_user_map.get(op, None)
|
||
for op_outputs in op.results():
|
||
for child in op_outputs.all_used_ops_in_same_block():
|
||
if cloned_op_first_grad_user_map.get(child, 0):
|
||
if first_subgraph_grad_user is None or op_2_id_map.get_idx(
|
||
cloned_op_first_grad_user_map[child]
|
||
) < op_2_id_map.get_idx(first_subgraph_grad_user):
|
||
first_subgraph_grad_user = (
|
||
cloned_op_first_grad_user_map[child]
|
||
)
|
||
assert first_subgraph_grad_user is not None
|
||
cloned_op_first_grad_user_map[op] = first_subgraph_grad_user
|
||
|
||
for cloned_op in cloned_ops:
|
||
cloned_op.move_before(cloned_op_first_grad_user_map[cloned_op])
|
||
return program, fwd_op_end_idx
|
||
|
||
|
||
def classify_value_node(program, grad_outputs, fwd_op_end_idx):
|
||
all_ops = program.global_block().ops
|
||
required_fw_ops = set(all_ops[: fwd_op_end_idx + 1])
|
||
|
||
required_fw_op_idxs = list(range(0, fwd_op_end_idx + 1))
|
||
required_fw_value_nodes = backward_utils.ValueSet(
|
||
program.global_block().get_values_by_op_idx(required_fw_op_idxs)
|
||
)
|
||
|
||
required_bw_op_idxs = list(range(fwd_op_end_idx + 1, len(all_ops)))
|
||
required_bw_value_nodes = backward_utils.ValueSet(
|
||
program.global_block().get_values_by_op_idx(required_bw_op_idxs)
|
||
)
|
||
|
||
# TODO(chenxi67) optimize classify algorithm by using unclaimed_ops. Remove them to fasten bw_ops detecting time.
|
||
# unclaimed_ops = {
|
||
# op
|
||
# for op in all_ops
|
||
# if op not in required_fw_ops and op not in required_bw_ops
|
||
# }
|
||
|
||
# unclaimed_op_idxs = []
|
||
# for idx, op in enumerate(all_ops):
|
||
# if op in unclaimed_ops:
|
||
# unclaimed_op_idxs.append(idx)
|
||
# unclaimed_value_nodes = backward_utils.ValueSet(
|
||
# program.global_block().get_values_by_op_idx(unclaimed_op_idxs)
|
||
# )
|
||
|
||
return (
|
||
required_fw_value_nodes,
|
||
required_bw_value_nodes,
|
||
backward_utils.ValueSet(),
|
||
)
|
||
|
||
|
||
# Sometimes we need to discard no_need_buffer values because they‘re not REAL tensor users.
|
||
def find_value_node_users(
|
||
value_node,
|
||
bw_no_need_buffer_values={},
|
||
without_no_need_buffer=False,
|
||
forward_ops={},
|
||
):
|
||
'''
|
||
Find all the value nodes which use the same value node to be computed.
|
||
'''
|
||
users = backward_utils.ValueSet()
|
||
ops = value_node.all_used_ops_in_same_block()
|
||
if without_no_need_buffer:
|
||
if value_node in bw_no_need_buffer_values:
|
||
ops = [op for op in ops if op in forward_ops]
|
||
for op in ops:
|
||
if op.name() == "builtin.combine":
|
||
combine_result = op.results()[0]
|
||
for (
|
||
combine_res_used_op
|
||
) in combine_result.all_used_ops_in_same_block():
|
||
results = combine_res_used_op.results()
|
||
for result in results:
|
||
if len(
|
||
result.all_used_ops_in_same_block()
|
||
) == 1 and result.all_used_ops_in_same_block()[
|
||
0
|
||
].name() in [
|
||
"builtin.split",
|
||
"builtin.slice",
|
||
]:
|
||
split_results = result.all_used_ops_in_same_block()[
|
||
0
|
||
].results()
|
||
users |= backward_utils.ValueSet(split_results)
|
||
else:
|
||
users.add(result)
|
||
else:
|
||
results = op.results()
|
||
for result in results:
|
||
if len(
|
||
result.all_used_ops_in_same_block()
|
||
) == 1 and result.all_used_ops_in_same_block()[0].name() in [
|
||
"builtin.split",
|
||
"builtin.slice",
|
||
]:
|
||
split_results = result.all_used_ops_in_same_block()[
|
||
0
|
||
].results()
|
||
users |= backward_utils.ValueSet(split_results)
|
||
else:
|
||
users.add(result)
|
||
return users
|
||
|
||
|
||
def get_real_input_nodes(output_value_node):
|
||
real_input_nodes = backward_utils.ValueSet()
|
||
define_op = output_value_node.get_defining_op()
|
||
if define_op.name() in ["builtin.split", "builtin.slice"]:
|
||
op_input = define_op.operands_source()[0]
|
||
real_define_op = op_input.get_defining_op()
|
||
input_value_nodes = real_define_op.operands_source()
|
||
else:
|
||
input_value_nodes = define_op.operands_source()
|
||
for input_value_node in input_value_nodes:
|
||
if (
|
||
input_value_node.get_defining_op()
|
||
and input_value_node.get_defining_op().name() == "builtin.combine"
|
||
):
|
||
real_input_nodes |= backward_utils.ValueSet(
|
||
input_value_node.get_defining_op().operands_source()
|
||
)
|
||
else:
|
||
real_input_nodes.add(input_value_node)
|
||
return real_input_nodes
|
||
|
||
|
||
def get_real_define_op_name(value_node):
|
||
define_op = value_node.get_defining_op()
|
||
if define_op.name() in ["builtin.split", "builtin.slice"]:
|
||
op_input = define_op.operands_source()[0]
|
||
return op_input.get_defining_op().name()
|
||
else:
|
||
return define_op.name()
|
||
|
||
|
||
def is_dynamic_value_node(value_node):
|
||
try:
|
||
return -1 in value_node.shape
|
||
except:
|
||
raise ValueError(f"value node not found in program: {value_node} ")
|
||
|
||
|
||
def is_vector_value_node(value_node):
|
||
try:
|
||
return value_node.type().as_vec_type() is not None
|
||
except:
|
||
raise ValueError(f"value node illegal: {value_node} ")
|
||
|
||
|
||
def cal_value_node_size_impl(value_node):
|
||
if is_dynamic_value_node(value_node):
|
||
value_node_shape = [i for i in value_node.shape if i != -1]
|
||
else:
|
||
value_node_shape = value_node.shape
|
||
return (
|
||
functools.reduce(lambda x, y: x * y, value_node_shape, 1)
|
||
* _PADDLE_DTYPE_2_NBYTES[value_node.dtype]
|
||
)
|
||
|
||
|
||
def cal_value_node_size(value_node):
|
||
if is_vector_value_node(value_node):
|
||
value_vec = value_node.type().as_vec_type().as_list()
|
||
sum_res = 0
|
||
for child_node in value_vec:
|
||
sum_res += cal_value_node_size_impl(child_node)
|
||
return sum_res
|
||
return cal_value_node_size_impl(value_node)
|
||
|
||
|
||
def cal_value_nodes_dist_to_backward(all_ops, required_fw_value_nodes):
|
||
dist_from_bw = backward_utils.ValueDict()
|
||
# calculate value node the shortest dist to backward graph
|
||
for op in reversed(all_ops):
|
||
if op.name() == "builtin.combine":
|
||
continue
|
||
op_results = op.results()
|
||
for op_result in op_results:
|
||
used_ops = op_result.all_used_ops_in_same_block()
|
||
if len(used_ops) == 1 and used_ops[0].name() in [
|
||
"builtin.split",
|
||
"builtin.slice",
|
||
]:
|
||
continue
|
||
real_users = find_value_node_users(op_result)
|
||
if op_result not in required_fw_value_nodes:
|
||
dist_from_bw[op_result] = 0
|
||
else:
|
||
dist_from_bw[op_result] = int(1e9)
|
||
for user in real_users:
|
||
dist_from_bw[op_result] = min(
|
||
dist_from_bw[op_result], dist_from_bw[user] + 1
|
||
)
|
||
return dist_from_bw
|
||
|
||
|
||
def all_used_op_consider_combine(program, value):
|
||
def filter_unused_combine(op):
|
||
if (
|
||
op.name() == "builtin.combine"
|
||
and len(op.result(0).all_used_ops_in_same_block()) == 0
|
||
):
|
||
return False
|
||
return True
|
||
|
||
return list(
|
||
filter(filter_unused_combine, value.all_used_ops_in_same_block())
|
||
)
|
||
|
||
|
||
def analyze_mid_hold_values(
|
||
program,
|
||
saved_values,
|
||
inputs,
|
||
outputs,
|
||
no_need_buffer_values,
|
||
fwd_op_end_idx,
|
||
backward_op_start_idx,
|
||
):
|
||
forward_ops = set(program.global_block().ops[: fwd_op_end_idx + 1])
|
||
backward_ops = set(program.global_block().ops[backward_op_start_idx:])
|
||
mid_hold_values = backward_utils.ValueSet()
|
||
for op in forward_ops:
|
||
for result in op.results():
|
||
all_used_ops = all_used_op_consider_combine(program, result)
|
||
if (
|
||
any(used_op in backward_ops for used_op in all_used_ops)
|
||
and result not in saved_values
|
||
and result not in outputs
|
||
and result not in inputs
|
||
and result not in no_need_buffer_values
|
||
and op.name() not in IGNORE_OPS
|
||
):
|
||
mid_hold_values.add(result)
|
||
return mid_hold_values
|
||
|
||
|
||
def get_first_backward_use_op(fwd_op, backward_ops, op_2_id_map):
|
||
first_backward_use_op = None
|
||
for user_op in fwd_op.results()[0].all_used_ops_in_same_block():
|
||
if user_op in backward_ops and (
|
||
first_backward_use_op is None
|
||
or op_2_id_map.get_idx(user_op)
|
||
< op_2_id_map.get_idx(first_backward_use_op)
|
||
):
|
||
first_backward_use_op = user_op
|
||
return first_backward_use_op
|
||
|
||
|
||
def clone_graph(
|
||
program,
|
||
origin_ops,
|
||
graph_inputs,
|
||
clone_insertion_op,
|
||
backward_ops,
|
||
op_2_id_map,
|
||
):
|
||
pir.set_insertion_point(clone_insertion_op)
|
||
all_ops = program.global_block().ops
|
||
value_map = paddle.pir.IrMapping()
|
||
origin_ops = set(origin_ops)
|
||
cloned_ops = []
|
||
cloned_op_first_grad_user_map = {}
|
||
for input_value in graph_inputs:
|
||
value_map.add(input_value, input_value)
|
||
for op in all_ops:
|
||
if op in origin_ops:
|
||
new_op = op.clone(
|
||
value_map, paddle.pir.CloneOptions(False, True, True)
|
||
)
|
||
first_backward_use_op = get_first_backward_use_op(
|
||
op, backward_ops, op_2_id_map
|
||
)
|
||
if (
|
||
first_backward_use_op is not None
|
||
and first_backward_use_op.has_attr('op_role')
|
||
and first_backward_use_op.has_attr('chunk_id')
|
||
):
|
||
new_op.set_int_attr("op_role", first_backward_use_op.op_role)
|
||
new_op.set_int_attr("chunk_id", first_backward_use_op.chunk_id)
|
||
cloned_ops.append(new_op)
|
||
if first_backward_use_op is not None:
|
||
cloned_op_first_grad_user_map[new_op] = first_backward_use_op
|
||
pir.set_insertion_point_to_block_end(program.global_block())
|
||
return cloned_ops, value_map, cloned_op_first_grad_user_map
|