wehub-resource-sync
287 lines
11 KiB
Python
287 lines
11 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 copy
|
|
import logging
|
|
from collections import OrderedDict
|
|
from typing import TYPE_CHECKING
|
|
|
|
import paddle
|
|
from paddle.distributed.auto_parallel.static.utils import (
|
|
is_backward_op,
|
|
is_gradient_clip_op,
|
|
is_optimize_op,
|
|
naive_set_dist_op_attr_for_program_by_mesh_and_mapping,
|
|
set_var_dist_attr,
|
|
)
|
|
from paddle.distributed.fleet.meta_optimizers.common import (
|
|
OP_ROLE_KEY,
|
|
OpRole,
|
|
)
|
|
from paddle.framework import core
|
|
from paddle.static import program_guard
|
|
|
|
from ..utils.log_utils import get_logger
|
|
from .pass_base import PassBase, register_pass
|
|
|
|
if TYPE_CHECKING:
|
|
from paddle.base import Variable
|
|
|
|
_supported_optimizer_type = [
|
|
"adam",
|
|
"adamax",
|
|
"adamw",
|
|
"decayed_adagrad",
|
|
"momentum",
|
|
"dgc_momentum",
|
|
"lars_momentum",
|
|
"merged_momentum",
|
|
"lamb",
|
|
"sgd",
|
|
]
|
|
|
|
logger = get_logger(logging.INFO, "MasterGradPass")
|
|
|
|
|
|
def _is_master_grad_cast_op(block, op):
|
|
op_name = op.type
|
|
if op_name != "cast":
|
|
return False
|
|
input_names = op.input_arg_names
|
|
output_names = op.output_arg_names
|
|
|
|
assert len(input_names) == 1
|
|
assert len(output_names) == 1
|
|
|
|
input_var_name = input_names[0]
|
|
|
|
return (
|
|
"@master_grad_fp16" in input_var_name
|
|
or "@master_grad_bf16" in input_var_name
|
|
)
|
|
|
|
|
|
def get_output_in_varlist(op, var_names) -> list[str]:
|
|
grad_names = []
|
|
for output_name in op.output_arg_names:
|
|
if output_name in var_names:
|
|
grad_names.append(output_name)
|
|
return grad_names
|
|
|
|
|
|
@register_pass("auto_parallel_master_grad_pass")
|
|
class MasterGradPass(PassBase):
|
|
"""
|
|
Use the high precision gradient to replace the low precision gradient in optimizer to avoid inf/nan values of low precision.
|
|
The high precision gradient 'master grad' will be used by communication operator, `update_loss_scaling`, `GradClip` and `optimizer`.
|
|
"""
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def _check_self(self):
|
|
return True
|
|
|
|
def _check_conflict(self, other_pass):
|
|
return True
|
|
|
|
def _apply_single_impl(self, main_program, startup_program, context):
|
|
self._completer = self.get_attr("completer")
|
|
dist_context = self.get_attr("dist_context")
|
|
params_grads = self.get_attr("params_grads")
|
|
logger.debug(f"Origin main_program: {main_program}")
|
|
self._add_master_grad(main_program, params_grads, dist_context)
|
|
self._regenerate_optimizer(
|
|
main_program, startup_program, params_grads, dist_context
|
|
)
|
|
logger.debug(f"After main program: {main_program}")
|
|
|
|
def _add_cast_op(self, cur_block, grad_names: list[str], dist_context):
|
|
grad_first_ids = OrderedDict()
|
|
for idx, op in enumerate(cur_block.ops):
|
|
if is_optimize_op(op):
|
|
break
|
|
elif is_backward_op(op):
|
|
var_names = get_output_in_varlist(op, grad_names)
|
|
for var_name in var_names:
|
|
if var_name not in grad_first_ids:
|
|
grad_first_ids[var_name] = idx
|
|
# Communication operators such as 'allreduce_sum' use input var as output.
|
|
else:
|
|
pass
|
|
|
|
# insert cast op
|
|
for grad_name, idx in reversed(grad_first_ids.items()):
|
|
grad_var = cur_block.var(grad_name)
|
|
if (
|
|
grad_var.dtype == paddle.float16
|
|
or grad_var.dtype == paddle.bfloat16
|
|
):
|
|
is_fp16 = grad_var.dtype == paddle.float16
|
|
producer_op = cur_block.ops[idx]
|
|
producer_op_dist_attr = (
|
|
dist_context.get_op_dist_attr_for_program(producer_op)
|
|
)
|
|
assert producer_op_dist_attr is not None, (
|
|
f"The op: '{producer_op}' should be distributed"
|
|
)
|
|
ref_output_dist_attr = (
|
|
producer_op_dist_attr.get_output_dist_attr(grad_name)
|
|
)
|
|
assert ref_output_dist_attr is not None, (
|
|
f"The output: '{grad_name}' should be distributed"
|
|
)
|
|
ref_mesh = ref_output_dist_attr.process_mesh
|
|
ref_dims_mapping = ref_output_dist_attr.dims_mapping
|
|
ref_chunk_id = producer_op_dist_attr.chunk_id
|
|
grad_half_precision_name = (
|
|
grad_name + '@master_grad_fp16'
|
|
if is_fp16
|
|
else grad_name + '@master_grad_bf16'
|
|
)
|
|
grad_half_precision = cur_block.create_var(
|
|
name=grad_half_precision_name,
|
|
dtype=grad_var.dtype,
|
|
shape=grad_var.shape,
|
|
persistable=False,
|
|
stop_gradient=False,
|
|
)
|
|
set_var_dist_attr(
|
|
dist_context,
|
|
grad_half_precision,
|
|
ref_dims_mapping,
|
|
ref_mesh,
|
|
chunk_id=ref_chunk_id,
|
|
)
|
|
|
|
producer_op_dist_attr = (
|
|
dist_context.get_op_dist_attr_for_program(producer_op)
|
|
)
|
|
origin_out_dims_mapping = (
|
|
producer_op_dist_attr.get_output_dims_mapping(grad_name)
|
|
)
|
|
producer_op._rename_output(grad_name, grad_half_precision.name)
|
|
producer_op_dist_attr.set_output_dims_mapping(
|
|
grad_half_precision.name, origin_out_dims_mapping
|
|
)
|
|
grad_var.desc.set_dtype(core.VarDesc.VarType.FP32)
|
|
|
|
cast_op = cur_block._insert_op_without_sync(
|
|
idx + 1,
|
|
type="cast",
|
|
inputs={"X": grad_half_precision},
|
|
outputs={"Out": grad_var},
|
|
attrs={
|
|
"in_dtype": grad_half_precision.dtype,
|
|
"out_dtype": grad_var.dtype,
|
|
},
|
|
)
|
|
cast_op._set_attr(OP_ROLE_KEY, OpRole.Backward)
|
|
naive_set_dist_op_attr_for_program_by_mesh_and_mapping(
|
|
cast_op,
|
|
ref_mesh,
|
|
ref_dims_mapping,
|
|
dist_context,
|
|
chunk_id=ref_chunk_id,
|
|
)
|
|
cur_block._sync_with_cpp()
|
|
|
|
def _regenerate_optimizer(
|
|
self,
|
|
main_program,
|
|
startup_program,
|
|
params_grads: list[tuple[Variable, Variable]],
|
|
dist_context,
|
|
):
|
|
grad_names = [g.name for _, g in params_grads]
|
|
# 1. delete the origin optimizer op
|
|
# 1.1 delete the var and op associated with the optimizer op in main_program
|
|
main_ops = main_program.global_block().ops
|
|
main_ops_len = len(main_ops)
|
|
first_optimize_idx = main_ops_len
|
|
for idx, op in enumerate(main_ops):
|
|
# We don't delete the operators for check_nan_inf
|
|
if is_optimize_op(op) and is_gradient_clip_op(op):
|
|
first_optimize_idx = idx
|
|
break
|
|
assert first_optimize_idx < main_ops_len, (
|
|
"The first optimizer op is not found!"
|
|
)
|
|
deleted_temp_var_names = []
|
|
deleted_persist_var_names = []
|
|
reserved_var_names = []
|
|
for idx in range(main_ops_len - 1, first_optimize_idx - 1, -1):
|
|
op = main_ops[idx]
|
|
inout_arg_names = op.input_arg_names + op.output_arg_names
|
|
if op.type in _supported_optimizer_type:
|
|
param_names = op.input("Param")
|
|
skip_update_names = op.input("SkipUpdate")
|
|
for reserved_name in param_names + skip_update_names:
|
|
if reserved_name not in reserved_var_names:
|
|
reserved_var_names.append(reserved_name)
|
|
for input_name in inout_arg_names:
|
|
if input_name in grad_names:
|
|
continue
|
|
var = main_program.global_block().var(input_name)
|
|
if (
|
|
var.persistable
|
|
and input_name not in deleted_persist_var_names
|
|
):
|
|
deleted_persist_var_names.append(input_name)
|
|
elif (
|
|
not var.persistable
|
|
and input_name not in deleted_temp_var_names
|
|
):
|
|
deleted_temp_var_names.append(input_name)
|
|
main_program.global_block()._remove_op(idx)
|
|
|
|
for var_name in deleted_temp_var_names + deleted_persist_var_names:
|
|
if var_name not in reserved_var_names:
|
|
main_program.global_block()._remove_var(var_name)
|
|
main_program.global_block()._sync_with_cpp()
|
|
|
|
# 1.2 delete the var and op in startup_program
|
|
for reserved_name in reserved_var_names:
|
|
if reserved_name in deleted_persist_var_names:
|
|
deleted_persist_var_names.remove(reserved_name)
|
|
startup_global_block = startup_program.global_block()
|
|
for var_name in deleted_persist_var_names:
|
|
if startup_global_block.has_var(var_name):
|
|
startup_global_block._remove_var(var_name)
|
|
for idx, op in reversed(list(enumerate(startup_global_block.ops))):
|
|
inout_arg_names = op.input_arg_names + op.output_arg_names
|
|
for var_name in inout_arg_names:
|
|
if var_name in deleted_persist_var_names:
|
|
startup_program.global_block()._remove_op(idx)
|
|
break
|
|
|
|
# 2. re-generate new optimizer op
|
|
serial_optimizer = copy.deepcopy(dist_context._serial_optimizer)
|
|
serial_optimizer._learning_rate = (
|
|
dist_context._serial_optimizer._learning_rate
|
|
)
|
|
serial_optimizer._sorted = False
|
|
with (
|
|
program_guard(main_program, startup_program),
|
|
main_program.switch_name_generator_guard("opt_"),
|
|
):
|
|
_ = serial_optimizer.apply_gradients(params_grads)
|
|
self._completer.complete_update_annotation(main_program)
|
|
|
|
def _add_master_grad(self, main_program, params_grads, dist_context):
|
|
grad_names = [g.name for _, g in params_grads]
|
|
for sub_block in main_program.blocks:
|
|
self._add_cast_op(sub_block, grad_names, dist_context)
|