# Copyright (c) 2020 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 itertools from contextlib import contextmanager from functools import cached_property from typing import TYPE_CHECKING import numpy as np import paddle import paddle.pir.core as ir_static from paddle import _C_ops from paddle.autograd.backward_utils import ValueDict from paddle.autograd.ir_backward import grad from paddle.base import core, framework from paddle.base.compiler import BuildStrategy from paddle.base.data_feeder import check_type from paddle.base.dygraph.base import switch_to_static_graph from paddle.pir import Value, fake_value, get_fake_value_name, is_fake_value from ..profiler import event_register from .logging_utils import TranslatorLogger from .utils import ( RETURN_NO_VALUE_MAGIC_NUM, Backend, CUDAGraphState, TimeCounter, auto_layout_is_enabled, backend_guard, cse_is_enabled, maybe_dynamic_shape_tensor, use_specialized_device, ) if TYPE_CHECKING: from .program_translator import ConcreteProgram __all__ = [] prog_logger = TranslatorLogger() FAKE_VALUE_NAME = get_fake_value_name() def get_value_name(value): if is_fake_value(value): return FAKE_VALUE_NAME return value.name def apply_general_passes( program, *, enable_cse=True, enable_delete_assert_op=True ): pm = paddle.pir.PassManager(2) if enable_cse: pm.add_pass("common_subexpression_elimination_pass", {}) if enable_delete_assert_op: pm.add_pass("delete_assert_op_pass", {}) pm.run(program) class NestSequence: """ A wrapper class that easily to flatten and restore the nest structure of given sequence. It also remove the duplicate variables in the sequence. For example: >>> t = [v1, v2, v1] >>> m = tolist(t) [v1, v2] >>> m.restore([t1, t2]) [t1, t2, t1] """ def __init__(self, raw_input): self._raw_input = raw_input self._var_map, self._var_list = self._tolist() @property def var_list(self): return self._var_list def _tolist(self): """ Flattens the nested sequences into single list and remove duplicate variables + non-variable elements. """ variable_map = ValueDict() # value -> list idx variable_list = [] for value in paddle.utils.flatten(self._raw_input): if not isinstance(value, Value): continue if value in variable_map: # remove duplicate values. continue variable_map[value] = len(variable_list) variable_list.append(value) return variable_map, variable_list def restore(self, tensor_result_list): """ Restores the nested sequence from tensor list. """ assert len(self._var_list) == len(tensor_result_list) def to_tensor_result(x): if isinstance(x, Value): return tensor_result_list[self._var_map[x]] return x return paddle.utils.pack_sequence_as( self._raw_input, list(map(to_tensor_result, paddle.utils.flatten(self._raw_input))), ) @cached_property def quick_index_map(self): raw_inputs = self._raw_input if len(raw_inputs) == 1: raw_inputs = raw_inputs[0] assert all(isinstance(v, Value) for v in raw_inputs) return [self._var_map[v] for v in raw_inputs] def quick_restore(self, tensor_list): return [tensor_list[idx] for idx in self.quick_index_map] def __getitem__(self, item): return self._var_list[item] class RunnableProgram: """a pir program ready for run_program_op to run. constructed by 3 parts: - pir program (pir::Program) - in_out_values - input_x values ([string | pir::Value]) - input_param values ([string | pir::Value]) - output values ([string | pir::Value]) - forward_backward_ranges - forward_range (tuple(Int, Int)) | None - backward_range (tuple(Int, Int)) | None """ @staticmethod def _get_program_all_values(program): all_values = [] def extend_values(block): all_values.extend(block.kwargs().values()) for op in block.ops: all_values.extend(op.results()) for block in op.blocks(): extend_values(block) extend_values(program.global_block()) return all_values @staticmethod def _get_name_value_map_from_program(program) -> dict[str, Value]: name_to_value_dict: dict[str, Value] = {FAKE_VALUE_NAME: fake_value()} for value in RunnableProgram._get_program_all_values(program): for name in value._names: name_to_value_dict[name] = value return name_to_value_dict @cached_property def name_value_map(self): return RunnableProgram._get_name_value_map_from_program(self.program) def convert_name(self, values): if len(values) == 0: return [] if isinstance(values[0], str): return values return [get_value_name(v) for v in values] @cached_property def x_values(self): return [self.name_value_map[v] for v in self.x_names] @cached_property def param_values(self): return [self.name_value_map[v] for v in self.param_names] @cached_property def out_values(self): return [self.name_value_map[v] for v in self.out_names] @cached_property def x_grad_values(self): return [self.name_value_map[v] for v in self.x_grad_names] @cached_property def param_grad_values(self): return [self.name_value_map[v] for v in self.p_grad_names] @cached_property def out_grad_values(self): return [self.name_value_map[v] for v in self.o_grad_names] def __init__( self, program, in_out_values, out_stop_gradients, grad_in_out_values=None, forward_range=None, backward_range=None, ): assert isinstance(in_out_values, tuple), ( "in_out_values must be tuple with len == 3" ) assert len(in_out_values) == 3, ( "in_out_values must be tuple with len == 3" ) assert isinstance(in_out_values[0], list), ( "in_out_values must be tuple with len == 3" ) self.program = program self.x_names = self.convert_name(in_out_values[0]) self.param_names = self.convert_name(in_out_values[1]) self.out_names = self.convert_name(in_out_values[2]) self.out_stop_gradients = out_stop_gradients self.forward_range = forward_range self.backward_range = backward_range self.has_splited = False self.finish_pass = False if self.forward_range is None: self.forward_range = (0, len(self.program.global_block().ops)) if self.backward_range is None: self.backward_range = ( len(self.program.global_block().ops), len(self.program.global_block().ops), ) if grad_in_out_values is None: grad_in_out_values = [], [], [] self.x_grad_names = self.convert_name(grad_in_out_values[0]) self.p_grad_names = self.convert_name(grad_in_out_values[1]) self.o_grad_names = self.convert_name(grad_in_out_values[2]) # Flag operator, indicating the operator between the forward subgraph and the backward subgraph. After self.program is updated by the pass, it is recommended to use the self.update_op_range interface to update the forward_range and backward_range. self.fwd_end_next_op = ( self.program.global_block().ops[self.forward_range[1]] if self.forward_range[1] < len(self.program.global_block().ops) else None ) self.bwd_start_pre_op = ( self.program.global_block().ops[self.backward_range[0] - 1] if ( self.backward_range[0] > 0 and self.backward_range[0] - 1 < len(self.program.global_block().ops) ) else None ) self.bwd_end_nex_op = ( self.program.global_block().ops[self.backward_range[1]] if self.backward_range[1] < len(self.program.global_block().ops) else None ) def update_op_range(self): if self.fwd_end_next_op is None or self.bwd_start_pre_op is None: self.forward_range = (0, len(self.program.global_block().ops)) self.backward_range = ( len(self.program.global_block().ops), len(self.program.global_block().ops), ) else: fwd_start = self.forward_range[0] fwd_end = self.forward_range[1] bwd_start = self.backward_range[0] bwd_end = self.backward_range[1] for idx, op in enumerate(self.program.global_block().ops): if op == self.fwd_end_next_op: fwd_end = idx if op == self.bwd_start_pre_op: bwd_start = idx + 1 if op == self.bwd_end_nex_op: bwd_end = idx if self.bwd_end_nex_op is None: bwd_end = len(self.program.global_block().ops) self.forward_range = (fwd_start, fwd_end) self.backward_range = (bwd_start, bwd_end) def clone(self): cloned_program, _ = paddle.base.libpaddle.pir.clone_program( self.program ) return RunnableProgram( cloned_program, (self.x_names, self.param_names, self.out_names), self.out_stop_gradients, None, self.forward_range, self.backward_range, ) def split_forward_backward(self): assert self.has_splited is False, ( "Please ensure only split once! don't call split_forward_backward manually." ) self.has_splited = True self.update_op_range() ( [ fwd_prog, bwd_prog, ], prog_attr, ) = paddle.base.libpaddle.pir.split_program( self.program, self.x_values, self.param_values, self.out_values, self.x_grad_values, self.param_grad_values, self.out_grad_values, self.forward_range, self.backward_range, ) return [fwd_prog, bwd_prog], prog_attr def apply_pir_program_pass(self, pass_fn): """ Main entries for pass function, without considering any input/output and forward segmentation. pass_fn' signature is: 1. This function will change forward and backward program. 2. call self.program_attr means start to run. so we can't call this function after program_attr is called. def pass_fn(forward_program, backward_program): return forward_program, backward_program """ origin_fwd = self.forward_program origin_bwd = self.backward_program prog_logger.log( 1, f"******** [JIT] PIR forward program before PIR PASS ********\n{origin_fwd} ", ) prog_logger.log( 1, f"******** [JIT] PIR backward program before PIR PASS ********\n{origin_bwd} ", ) # NOTE(dev): Add this line to trigger program_name_attr logic program_name_attr = self.program_name_attr self.forward_program, self.backward_program = pass_fn( origin_fwd, origin_bwd, program_name_attr ) prog_logger.log( 1, f"******** [JIT] PIR forward program after PIR PASS ********\n{origin_fwd} ", ) prog_logger.log( 1, f"******** [JIT] PIR backward program after PIR PASS ********\n{origin_bwd} ", ) def is_distributed_program(self): for op in self.program.global_block().ops: if op.dist_attr is not None: return True return False def apply_dist_pass_for_origin_program(self): if self.is_distributed_program(): paddle.distributed.auto_parallel.static.mix_to_dist_pass.apply_mix2dist_pass( self.program ) def apply_dist_pass_for_whole_program(self): if self.is_distributed_program(): paddle.distributed.auto_parallel.static.mix_to_dist_pass.apply_mix2dist_pass( self.program ) paddle.distributed.auto_parallel.static.pir_pass.apply_partition_pass( self.program ) paddle.distributed.auto_parallel.static.pir_pass.ReshardPasses.apply_reshard_pass( self.program ) paddle.base.libpaddle.pir.apply_dist2dense_pass(self.program) paddle.distributed.auto_parallel.static.pir_pass.remove_unuseful_comm_op_pass( self.program ) # cached property can ensure program is splited only once. @cached_property def _forward_backward_program(self): return self.split_forward_backward() @cached_property # shouldn't changed when call this once. def program_attr(self): assert self.finish_pass is False, ( "program_attr() is called by PartialProgramLayer, don't call it manually, use program_name_attr instead." ) # can't apply pass after call this function. self.finish_pass = True fwd_map = RunnableProgram._get_name_value_map_from_program( self.forward_program ) program_name_attr = self.program_name_attr no_need_buffer_names = program_name_attr["no_need_buffers"] rename_mapping = {} rename_mapping = RunnableProgram.unify_value_names( self.forward_program, rename_mapping ) rename_mapping = RunnableProgram.unify_value_names( self.backward_program, rename_mapping ) # Update no_need_buffer_names by rename_mapping for original_name, new_name in rename_mapping.items(): if {original_name, new_name} & set(no_need_buffer_names): if original_name in no_need_buffer_names: no_need_buffer_names.remove(original_name) if new_name in no_need_buffer_names: no_need_buffer_names.remove(new_name) RunnableProgram.update_program_name_attr( self.program_name_attr, rename_mapping ) program_attr = {} for k, ns in self.program_name_attr.items(): # Pass output values to create tensors in run program impl if k == "fo": program_attr[f"{k}_values"] = [ fwd_map.get(n, fake_value()) for n in ns ] program_attr[f"{k}_names"] = ns # Restore stop_gradient for output values assert len(program_attr["fo_values"]) == len(self.out_stop_gradients), ( "Output values and stop gradients length mismatch" ) for v, stop_gradient in zip( program_attr["fo_values"], self.out_stop_gradients ): if is_fake_value(v): continue v.stop_gradient = stop_gradient return program_attr @staticmethod def unify_value_names( program, rename_mapping: dict[str, str] ) -> dict[str, str]: """Ensure every value at most has one name in the program.""" rename_mapping = dict(rename_mapping) for value in RunnableProgram._get_program_all_values(program): if not value.has_name: continue new_name = value.name # get first name new_name = rename_mapping.get(new_name, new_name) rename_mapping.update( value._rename(new_name, program.global_block()) ) # Get all values again because some values has been erased. for value in RunnableProgram._get_program_all_values(program): if value.has_name: assert value._has_only_one_name(), ( f"Expected all values in Program have only one name, but {value} has multiple names: {value._names}" ) return rename_mapping @staticmethod def update_program_name_attr( name_attr: dict[str, list[str]], rename_mapping: dict[str, str] ): for k, vs in name_attr.items(): name_attr[k] = [ rename_mapping[v] if v in rename_mapping else v for v in vs ] @cached_property def program_name_attr(self): origin_attr = self._forward_backward_program[1] _program_attr = {} for k, vs in origin_attr.items(): _program_attr[k] = [get_value_name(v) for v in vs] return _program_attr @cached_property def forward_program(self): return self._forward_backward_program[0][0] @cached_property def backward_program(self): return self._forward_backward_program[0][1] class PartialProgramLayerHook: def before_append_backward(self, forward_program, src_vars): return forward_program, src_vars def after_append_backward( self, whole_program, inputs, src_vars, grad_outputs, forward_end_idx, backward_start_idx, ): return whole_program, forward_end_idx, src_vars def after_infer(self, infer_program): return infer_program class OperatorIndexPreservePass: OP_NAME_PREFIX = "preserved_index_" counter = 0 def __init__(self, index, pass_fn): self.name = f"{OperatorIndexPreservePass.OP_NAME_PREFIX}{OperatorIndexPreservePass.counter}" OperatorIndexPreservePass.counter += 1 self.pass_fn = pass_fn self.index = index def __call__(self, program): if len(program.global_block().ops) == 0: assert self.index == 0 return self.pass_fn(program) paddle.base.libpaddle.pir.append_shadow_output( program, program.global_block().ops[0].result(0), self.name, self.index, ) program = self.pass_fn(program) new_index = 0 for op in program.global_block().ops: if ( op.name() == "builtin.shadow_output" and self.name in op.attrs()["output_name"] ): break new_index += 1 # remove forward_backward_separator if new_index >= len(program.global_block().ops): raise RuntimeError( f"Can't find index preserve label {self.name}, don't remove it in pass." ) program.global_block().remove_op(program.global_block().ops[new_index]) self.index = new_index return program class IndicesPreservePass: def __init__(self, indices, pass_fn): self.pass_fn = pass_fn self.indices = indices self.new_indices = None def __call__(self, program): passes = [self.pass_fn] for idx, index in enumerate(self.indices): passes.append(OperatorIndexPreservePass(index, passes[idx])) new_program = passes[-1](program) self.new_indices = [p.index for p in passes[1:]] return new_program class ValuePreservePass: OP_NAME_PREFIX = "preserved_value_" def __init__(self, values, use_cinn_pass): self.values = values self.use_cinn_pass = use_cinn_pass def apply(self, program): raise RuntimeError("Not implemented.") @staticmethod def create_name_generator(prefix): count = 0 def name_gen(): nonlocal count name = f"{prefix}{count}" count += 1 return name return name_gen @staticmethod def attach_preserved_name(value, program, value2name, name_generator): if is_fake_value(value): return None if value in value2name: return value2name[value] name = name_generator() value2name[value] = name paddle.base.libpaddle.pir.append_shadow_output( program, value, name, len(program.global_block().ops), ) return name def __call__(self, program): # create preserved op for args value2name = ValueDict() name_generator = ValuePreservePass.create_name_generator( ValuePreservePass.OP_NAME_PREFIX ) names = paddle.utils.map_structure( lambda value: ValuePreservePass.attach_preserved_name( value, program, value2name, # noqa: F821 name_generator, ), self.values, ) # NOTE(SigureMo): Value maybe removed in pass, don't use value2name after pass del value2name # apply program pass program = self.apply(program) # collect new value name2new_value = {} to_remove_op = [] for op in program.global_block().ops: if op.name() == "builtin.shadow_output": if op.attrs()["output_name"].startswith( ValuePreservePass.OP_NAME_PREFIX ): name2new_value[op.attrs()["output_name"]] = op.operand( 0 ).source() to_remove_op.append(op) # remove old op for op in to_remove_op: program.global_block().remove_op(op) self.values = paddle.utils.map_structure( lambda name: name2new_value.get(name, fake_value()), names ) return program class FullGraphPreProcessPass(ValuePreservePass): def apply(self, program): program = paddle.base.libpaddle.pir.apply_bn_add_act_pass(program) if self.use_cinn_pass: # NOTE(gongshaotian): execute infer_symbolic_shape_pass before reduce_as_sum_pass pm = paddle.base.libpaddle.pir.PassManager() pm.add_pass("delete_assert_op_pass", {}) paddle.base.libpaddle.pir.infer_symbolic_shape_pass(pm, program) paddle.base.libpaddle.pir.reduce_as_sum_pass(pm, program) pm.run(program) return program class PartialProgramLayer: """ PartialProgramLayer wraps all the ops from layers decorated by `@to_static` and execute them as a static subgraph. .. note:: **1. This is a very low level API. Users should not use this API directly. Please use `partial_program_from(concrete_program)` to create it. **2. TensorArray is not currently supported in the output. Args: main_program(Program): The main program that contains ops need to be executed. inputs(list[Variable]): The input list of the decorated function by `@to_static`. outputs(list[Variable]): The output list of the decorated function by `@to_static`. parameters(list[Tensor]|None): All trainable parameters included in the program. Default None. constraints(list[tuple[str, int|None, int|None]]): A list to specify the constraints of the program. Default None. Returns: Layer: A Layer object that run all ops internally in static graph mode. """ HOOKED_RUN_IMPL = None def __init__( self, main_program, inputs, outputs, parameters=None, *, constraints=None, **kwargs, ): super().__init__() self._inputs = NestSequence(inputs) self._outputs = NestSequence(outputs) # Avoid mutable default argument pitfall (new list per instance) self._constraints = constraints if constraints is not None else [] self._params, self._param_values = ( parameters if parameters is not None else ([], []) ) self._build_strategy = kwargs.get('build_strategy', BuildStrategy()) assert isinstance(self._build_strategy, BuildStrategy) self._origin_main_program = self._verify_program( main_program, self._outputs ) if parameters is not None: parameters[0][:] = self._params parameters[1][:] = self._param_values # Set default mode to train self.training = True self._program_extra_info = {} amp_dtype, custom_white_list, custom_black_list = None, None, None tracer = framework._dygraph_tracer() if tracer: custom_white_list, custom_black_list = tracer._get_amp_op_list() amp_dtype = tracer._amp_dtype if amp_dtype is not None and amp_dtype in ['float16', 'bfloat16']: # For AMP training self._amp_list = ( paddle.static.amp.fp16_lists.AutoMixedPrecisionLists( custom_white_list=custom_white_list, custom_black_list=custom_black_list, dtype=amp_dtype, ) ) # program_id -> list(scope) self._scope_cache = {} self._hookers = [] self._backend = kwargs.get('backend', Backend.PHI) self._grad_var_names = {} self._compile_time_counter = TimeCounter() self._prog_attrs_map_cache = {} @staticmethod def run_impl(partial_program_layer, inputs, parameters, attrs): prog_attrs, cuda_graph_attrs = attrs scope_cache_key = paddle.base.core.calc_scope_cache_key( paddle.base.core.get_program_id_from_attrs(prog_attrs), inputs, cuda_graph_attrs["cuda_graph_state"] != CUDAGraphState.DISABLE, cuda_graph_attrs["cuda_graph_dispatch_key"], ) return _C_ops.run_program( PartialProgramLayer._valid_vars(inputs), PartialProgramLayer._valid_vars(parameters), partial_program_layer._create_scope_vec( cache_key=scope_cache_key, use_scope_cache=True, ), prog_attrs, cuda_graph_attrs, ) def __call__(self, inputs): """ Execute static graph by Interpreter and Return dynamic Tensors. """ attrs = self._prepare_attributes(in_sot_mode=False) inputs = self._prepare_inputs(inputs) out = self.call_run_impl_with_hook( inputs, self._params, attrs, ) restored_nest_out = self._restore_out(out) return self._remove_no_value(restored_nest_out) @event_register("sot call partial_program") def sot_call(self, inputs): """ In sot, inputs and outputs of partial program only contain tensors, so we can skip some step to speed up """ attrs = self._prepare_attributes(in_sot_mode=True) out = self.call_run_impl_with_hook( inputs, self._params, attrs, ) return self._outputs.quick_restore(out) def call_run_impl_with_hook( self, inputs, parameters, attrs, ): if PartialProgramLayer.HOOKED_RUN_IMPL is None: return PartialProgramLayer.run_impl.__get__(self)( inputs, parameters, attrs, ) else: return PartialProgramLayer.HOOKED_RUN_IMPL( PartialProgramLayer.run_impl.__get__(self), inputs, parameters, attrs, ) @cached_property def origin_runnable_program(self) -> RunnableProgram: inputs = list(self._inputs.var_list) outputs = list(self._outputs.var_list) # NOTE(SigureMo): Record original stop gradient for output values to avoid # losing during optimization passes. out_stop_gradients = [v.stop_gradient for v in outputs] params = self._param_values paddle.base.libpaddle.pir.append_shadow_outputs( self._origin_main_program, outputs, len(self._origin_main_program.global_block().ops), "output_", ) return RunnableProgram( self._origin_main_program, (inputs, params, outputs), out_stop_gradients, ) def add_hooker(self, hooker): self._hookers.append(hooker) def _get_scope(self, cache_key=None, use_scope_cache=False): if not use_scope_cache: return core.Scope() if cache_key not in self._scope_cache: self._scope_cache[cache_key] = [] cached_scopes = self._scope_cache[cache_key] for scope in cached_scopes: if scope._can_reused: return scope scope = core.Scope() cached_scopes.append(scope) return scope # whole @switch_to_static_graph def _create_program(self, is_infer_mode=False) -> RunnableProgram: if is_infer_mode: def pass_fn(forward_program, backward_program, program_name_attr): # if-else pass if self._backend.is_cinn(): apply_general_passes( forward_program, enable_cse=cse_is_enabled(), enable_delete_assert_op=self._backend.is_cinn(), ) paddle.base.libpaddle.pir.bind_symbolic_constraints( forward_program, self._constraints ) paddle.base.libpaddle.pir.apply_cinn_pass(forward_program) elif self._backend.is_pcc(): paddle.base.libpaddle.pir.bind_symbolic_constraints( forward_program, self._constraints ) paddle.base.libpaddle.pir.apply_pcc_pass(forward_program) else: apply_general_passes( forward_program, enable_cse=cse_is_enabled(), enable_delete_assert_op=self._backend.is_cinn(), ) paddle.base.libpaddle.pir.check_infer_symbolic_if_need( forward_program ) return forward_program, backward_program # TODO(xiongkun) who to transfer the pruning program? infer_program = self.origin_runnable_program.clone() if auto_layout_is_enabled() and self._backend.is_cinn(): pm = paddle.pir.PassManager(2) pm.add_pass("auto_layout_pass", {}) pm.run(infer_program.program) for hooker in self._hookers: hooker.after_infer(infer_program) infer_program.apply_pir_program_pass(pass_fn) return infer_program else: train_program: RunnableProgram = ( self.origin_runnable_program.clone() ) train_program.apply_dist_pass_for_origin_program() # Author(liujinnan): auto_layout_pass should be applied to the original_program, before append backward. So we put it here. if auto_layout_is_enabled() and self._backend.is_cinn(): pm = paddle.pir.PassManager(2) pm.add_pass("auto_layout_pass", {}) pm.run(train_program.program) train_program = self._append_backward(train_program) # Note: Only set grad type once after initializing train program. So we put it here. self._set_grad_type(self._params, train_program) def pass_fn(forward_program, backward_program, program_name_attr): def init_backward_program_shape_analysis( forward_program, backward_program ): forward_shape_analysis = paddle.base.libpaddle.pir.get_shape_constraint_ir_analysis( forward_program ) backward_shape_analysis = paddle.base.libpaddle.pir.get_shape_constraint_ir_analysis( backward_program ) backward_shape_analysis.register_symbol_cstr_from_shape_analysis( forward_shape_analysis ) forward_name_value_map = { name: item for item in forward_program.list_vars() for name in item._names } def share_symbol_shape_from_forward_to_backward( forward_value, backward_value ): backward_shape_analysis.set_shape_or_data_for_var( backward_value, forward_shape_analysis.get_shape_or_data_for_var( forward_value ), ) def get_kwargs_forward_matched_value(kw_name, kw_value): if kw_name in program_name_attr['bo_g']: idx = program_name_attr['bo_g'].index(kw_name) return forward_name_value_map[ program_name_attr['fo'][idx] ] elif kw_name in forward_name_value_map: return forward_name_value_map[kw_name] else: raise Exception(f"kw_args: {kw_name} not found") for [kw_name, kw_value] in ( backward_program.global_block().kwargs().items() ): forward_matched_value = ( get_kwargs_forward_matched_value(kw_name, kw_value) ) share_symbol_shape_from_forward_to_backward( forward_matched_value, kw_value ) apply_general_passes( forward_program, enable_cse=cse_is_enabled(), enable_delete_assert_op=self._backend.is_cinn(), ) apply_general_passes( backward_program, enable_cse=cse_is_enabled(), enable_delete_assert_op=self._backend.is_cinn(), ) if self._backend.is_cinn(): paddle.base.libpaddle.pir.bind_symbolic_constraints( forward_program, self._constraints ) paddle.base.libpaddle.pir.apply_cinn_pass(forward_program) init_backward_program_shape_analysis( forward_program, backward_program ) paddle.base.libpaddle.pir.apply_cinn_pass(backward_program) elif self._backend.is_pcc(): paddle.base.libpaddle.pir.bind_symbolic_constraints( forward_program, self._constraints ) paddle.base.libpaddle.pir.apply_pcc_pass(forward_program) else: paddle.base.libpaddle.pir.check_infer_symbolic_if_need( forward_program ) return forward_program, backward_program train_program.apply_pir_program_pass(pass_fn) return train_program @cached_property def _train_program_id(self): program_id = paddle.utils._hash_with_id(self.train_program, self) return program_id @cached_property def _infer_program_id(self): return paddle.utils._hash_with_id(self.infer_program, self) @property def program(self) -> RunnableProgram: """ Return current train or eval program. """ if self.training: return self.train_program else: return self.infer_program @property def program_id(self): """ Return current train or eval program hash id. """ if self.training: return self._train_program_id else: return self._infer_program_id @cached_property def train_program(self) -> RunnableProgram: with backend_guard(self._backend), self._compile_time_counter.record(): return self._create_program() @cached_property def infer_program(self) -> RunnableProgram: with backend_guard(self._backend), self._compile_time_counter.record(): return self._create_program(is_infer_mode=True) def _verify_program(self, main_program, outputs): """ Verify that the program parameter is initialized, prune some unused params, and remove redundant op callstack. """ # Check all params from main program can be found in self._params self._check_params_all_inited(main_program) return main_program @switch_to_static_graph def _append_backward( self, train_runnable_program: RunnableProgram ) -> RunnableProgram: program = train_runnable_program.program targets = train_runnable_program.out_values # TODO(@zhuoge): refine the interface, use runnable_program to apply passes. for hooker in self._hookers: program, targets = hooker.before_append_backward(program, targets) inputs = train_runnable_program.x_values params = train_runnable_program.param_values combined_inputs = list(itertools.chain(inputs, params)) forward_end_idx = len(program.global_block().ops) forward_end_op = None if forward_end_idx > 0: forward_end_op = program.global_block().ops[-1] grad_info_map = [None] * len(combined_inputs) with backend_guard(self._backend): check_type( targets, 'targets', (Value, list, tuple), 'paddle.static.gradients', ) with ir_static.program_guard(program, None): # create outputs_grad for backward to avoid full and full_like op. forward_outputs_grads = [] for out_value in targets: if out_value.stop_gradient is True: forward_outputs_grads.append(fake_value()) else: value = paddle.full_like( out_value, fill_value=1.0, dtype=out_value.dtype, ) forward_outputs_grads.append(value) paddle.base.libpaddle.pir.append_shadow_outputs( program, forward_outputs_grads, len(program.global_block().ops), "grad_input_", ) op_between_forward_and_backward = ( len(program.global_block().ops) - forward_end_idx ) # call grad to get backward ops. if ( len( list( filter(lambda x: x.stop_gradient is False, targets) ) ) > 0 ): grad_info_map = grad( inputs=combined_inputs, outputs=list( filter(lambda x: x.stop_gradient is False, targets) ), grad_outputs=list( filter( lambda x: not is_fake_value(x), forward_outputs_grads, ) ), ) if forward_end_op is not None: for idx, op in enumerate(program.global_block().ops): if op == forward_end_op: forward_end_idx = idx + 1 break for hooker in self._hookers: ( program, forward_end_idx, targets, ) = hooker.after_append_backward( program, combined_inputs, targets, forward_outputs_grads, forward_end_idx, forward_end_idx + op_between_forward_and_backward, ) mapping_value = lambda x: x if isinstance(x, Value) else fake_value() inputs_size = len(inputs) x_grad_value = list(map(mapping_value, grad_info_map[0:inputs_size])) p_grad_value = list(map(mapping_value, grad_info_map[inputs_size:])) o_grad_value = list(map(mapping_value, forward_outputs_grads)) # insert grads name for RunnableProgram (we need name for grad_inputs and grad_outputs) input_grads_to_append = list( filter(lambda x: not is_fake_value(x), o_grad_value) ) output_grads_to_append = list( filter(lambda x: not is_fake_value(x), x_grad_value + p_grad_value) ) backward_end_op_index = len(program.global_block().ops) paddle.base.libpaddle.pir.append_shadow_outputs( program, output_grads_to_append, backward_end_op_index, "grad_output_", ) backward_start_op_index = ( forward_end_idx + op_between_forward_and_backward ) # construct a runnable program. full_graph_pre_process_pass = FullGraphPreProcessPass( [inputs, params, targets, x_grad_value, p_grad_value, o_grad_value], self._backend.is_cinn(), ) forward_index_pass = IndicesPreservePass( [forward_end_idx, backward_start_op_index, backward_end_op_index], full_graph_pre_process_pass, ) program = forward_index_pass(program) ( inputs, params, targets, x_grad_value, p_grad_value, o_grad_value, ) = full_graph_pre_process_pass.values ( forward_end_idx, backward_start_op_index, backward_end_op_index, ) = forward_index_pass.new_indices whole_program = RunnableProgram( program, (inputs, params, targets), train_runnable_program.out_stop_gradients, (x_grad_value, p_grad_value, o_grad_value), (0, forward_end_idx), (backward_start_op_index, backward_end_op_index), ) whole_program.apply_dist_pass_for_whole_program() return whole_program def _prepare_attributes(self, in_sot_mode=False): prog_attr_key = (self.program_id, self.training, in_sot_mode) if prog_attr_key not in self._prog_attrs_map_cache: prog_attrs = { 'forward_program': self.program.forward_program, 'backward_program': self.program.backward_program, 'is_test': not self.training, 'program_id': self.program_id, 'in_sot_mode': in_sot_mode, } | self.program.program_attr self._prog_attrs_map_cache[prog_attr_key] = ( paddle.base.core.construct_program_attribute_map(prog_attrs) ) cuda_graph_attrs = { 'cuda_graph_state': CUDAGraphState.DISABLE, # default value for not use cuda graph 'cuda_graph_dispatch_key': 0, # default value for not use cuda graph } return self._prog_attrs_map_cache[prog_attr_key], cuda_graph_attrs def _prepare_inputs(self, inputs): """ Prepare inputs, outputs, attrs. """ assert isinstance(inputs, (tuple, list)) # Flatten inputs with nested structure into single list. flatten_inputs = paddle.utils.flatten(inputs) # Convert variable into Tensor and feed in training data. input_vars = [] expected_place = framework._current_expected_place() for i, value in enumerate(flatten_inputs): if isinstance(value, np.ndarray): var = None var = core.eager.Tensor( value=value, persistable=False, place=expected_place, zero_copy=True, ) elif isinstance(value, core.eager.Tensor): # NOTE(Aurelius84): If var is on CPUPlace, it will be transformed multi times # into CUDAPlace when it's as input of multi Ops. so we move it in advance to avoid this problem. if ( value.stop_gradient and not value.place._equals(expected_place) and not use_specialized_device() and not maybe_dynamic_shape_tensor(value) ): var = value._copy_to(expected_place, False) var.stop_gradient = True else: var = value else: continue input_vars.append(var) return input_vars def _create_scope_vec(self, cache_key=None, use_scope_cache=False): inner_scope = self._get_scope( cache_key=cache_key, use_scope_cache=use_scope_cache ) return [inner_scope] def _restore_out(self, out_vars): """ Restores same nested outputs by only replacing the Variable with Tensor. """ outs = self._outputs.restore(out_vars) if outs is not None and len(outs) == 1: outs = outs[0] return outs @switch_to_static_graph def _clone_for_test(self, main_program): return main_program.clone(for_test=True) def _is_no_value(self, var): if isinstance(var, core.eager.Tensor) and var.shape == [1]: # NOTE: .numpy() will insert MemcpySync operation, it hits performance. if var.numpy()[0] == RETURN_NO_VALUE_MAGIC_NUM: return True return False def _remove_no_value(self, out_vars): """ Removes invalid value for various-length return statement """ if isinstance(out_vars, core.eager.Tensor): if self._is_no_value(out_vars): return None return out_vars elif isinstance(out_vars, (tuple, list)): if isinstance(out_vars, tuple): res = tuple( var for var in out_vars if not self._is_no_value(var) ) else: res = [var for var in out_vars if not self._is_no_value(var)] has_removed = len(out_vars) > len(res) # len(out_vars) > len(res) means we have removed var. This is # preventing out_vars is empty or just one element at the beginning if len(res) == 0 and has_removed: return None elif len(res) == 1 and has_removed: return res[0] return res return out_vars def _set_grad_type(self, params, train_program: RunnableProgram): # NOTE: if user set sparse gradient mode, the param's gradient # will be SelectedRows, not DenseTensor. But tracer will just # set param grad Tensor by forward Tensor(DenseTensor) # If we don't change grad_var type here, RunProgramOp need # transform SelectedRows to DenseTensor forcibly, it may not # be user wanted result. forward_params_grads = train_program.param_grad_values train_program = train_program.program for param, value in zip(params, forward_params_grads): if is_fake_value(value): continue if value.is_selected_row_type(): param._set_grad_type( paddle.base.core.VarDesc.VarType.SELECTED_ROWS ) elif value.is_dense_tensor_type(): param._set_grad_type( paddle.base.core.VarDesc.VarType.DENSE_TENSOR ) else: raise NotImplementedError( "only support selected_row and dense_tensor grad type." ) def _check_params_all_inited(self, main_program): """ Check all params from main program are already initialized, see details as follows: 1. all parameters in self._params should be type `framework.EagerParamBase` which are created in dygraph. 2. all parameters from transformed program can be found in self._params. Because they share same data with EagerParamBase of original dygraph. """ if not isinstance(self._params, (list, tuple)): raise TypeError( f"Type of self._params in PartialProgramLayer should be list or tuple, but received {type(self._params)}." ) param_and_buffer_names_set = set() for i, var in enumerate(self._params): # self._params contains parameters and buffers with persistable=True. if not isinstance(var, core.eager.Tensor): raise TypeError( f'Type of self._params[{i}] in PartialProgramLayer should be Parameter or Variable, but received {type(var)}.' ) param_and_buffer_names_set.add(var.name) @staticmethod def _valid_vars(vars): return vars if vars else None @contextmanager def replace_run_impl_guard(new_run_impl): """ A context manager to temporarily replace the run_impl of PartialProgramLayer. This is used for testing purposes. """ old_run_impl = PartialProgramLayer.HOOKED_RUN_IMPL PartialProgramLayer.HOOKED_RUN_IMPL = new_run_impl try: yield finally: PartialProgramLayer.HOOKED_RUN_IMPL = old_run_impl def partial_program_from( concrete_program: ConcreteProgram, from_method: bool = False ) -> PartialProgramLayer: inputs = concrete_program.inputs # NOTE(SigureMo): Remove the first arg `self` from method args. if inputs and from_method: inputs = inputs[1:] return PartialProgramLayer( concrete_program.main_program, inputs, concrete_program.outputs, concrete_program.parameters, constraints=concrete_program.constraints, **concrete_program.kwargs, )