wehub-resource-sync
274 lines
9.8 KiB
Python
Executable File
274 lines
9.8 KiB
Python
Executable File
# 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.
|
|
|
|
import paddle
|
|
from paddle.distributed.fleet.meta_optimizers.common import (
|
|
OP_ROLE_KEY,
|
|
OpRole,
|
|
is_optimizer_op,
|
|
)
|
|
from paddle.framework import core
|
|
|
|
__all__ = []
|
|
|
|
|
|
class FP16Utils:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@staticmethod
|
|
def is_fp16_cast_op(block, op, params):
|
|
if op.type != "cast":
|
|
return False
|
|
if is_optimizer_op(op):
|
|
return False
|
|
assert len(op.desc.input_arg_names()) == 1
|
|
assert len(op.desc.output_arg_names()) == 1
|
|
input_name, output_name = (
|
|
op.desc.input_arg_names()[0],
|
|
op.desc.output_arg_names()[0],
|
|
)
|
|
if input_name not in params:
|
|
return False
|
|
input_var = block.var(input_name)
|
|
output_var = block.var(output_name)
|
|
if (
|
|
input_var.dtype != core.VarDesc.VarType.FP32
|
|
or output_var.dtype != core.VarDesc.VarType.FP16
|
|
):
|
|
return False
|
|
return True
|
|
|
|
@staticmethod
|
|
def is_fp32_cast_op(block, op):
|
|
if op.type != "cast":
|
|
return False
|
|
if not is_optimizer_op(op):
|
|
return False
|
|
assert len(op.desc.input_arg_names()) == 1
|
|
assert len(op.desc.output_arg_names()) == 1
|
|
input_name, output_name = (
|
|
op.desc.input_arg_names()[0],
|
|
op.desc.output_arg_names()[0],
|
|
)
|
|
input_var = block.var(input_name)
|
|
output_var = block.var(output_name)
|
|
if (
|
|
input_var.dtype != core.VarDesc.VarType.FP16
|
|
or output_var.dtype != core.VarDesc.VarType.FP32
|
|
):
|
|
return False
|
|
return True
|
|
|
|
@staticmethod
|
|
def remove_cast_op(block, params, segment, offset):
|
|
inserted_op_num = 0
|
|
for op_idx in reversed(
|
|
range(offset + segment._start_idx, offset + segment._end_idx)
|
|
):
|
|
op = block.ops[op_idx]
|
|
if FP16Utils.is_fp16_cast_op(block, op, params):
|
|
block._remove_op(op_idx, sync=False)
|
|
inserted_op_num -= 1
|
|
block._sync_with_cpp()
|
|
return inserted_op_num
|
|
|
|
@staticmethod
|
|
def prune_fp16(block, shard, reduced_grads_to_param, ring_ids):
|
|
"""
|
|
1. prune all cast_fp16_to_fp32 ops if the param not belongs to this shard
|
|
2. revise amp inifine grad checking for sharding
|
|
"""
|
|
# remove cast
|
|
for idx, op in reversed(list(enumerate(block.ops))):
|
|
if not FP16Utils.is_fp32_cast_op(block, op):
|
|
continue
|
|
output_name = op.desc.output_arg_names()[0]
|
|
# TODO (JZ-LIANG) revise this for uniform mixed parallelism
|
|
param_name = output_name.removesuffix("@MERGED").removesuffix(
|
|
"@GRAD"
|
|
)
|
|
if param_name not in shard.global_params:
|
|
raise ValueError(
|
|
"Output 'X' of cast_op must be a grad of"
|
|
f"model param, but {output_name} is not a grad"
|
|
)
|
|
if output_name in reduced_grads_to_param:
|
|
continue
|
|
if shard.has_param(param_name):
|
|
continue
|
|
block._remove_op(idx, sync=False)
|
|
block._remove_var(output_name, sync=False)
|
|
|
|
block._sync_with_cpp()
|
|
update_loss_scaling_op_idx = -1
|
|
inf_var_name = ''
|
|
for idx, op in reversed(list(enumerate(block.ops))):
|
|
if op.type == "update_loss_scaling":
|
|
update_loss_scaling_op_idx = idx
|
|
inf_var_name = op.desc.input('FoundInfinite')[0]
|
|
if op.type in ["check_finite_and_unscale", "update_loss_scaling"]:
|
|
reversed_x = []
|
|
reversed_x_paramname = []
|
|
for input_name in op.desc.input('X'):
|
|
# TODO (JZ-LIANG) revise this for uniform mixed parallelism
|
|
param_name = input_name.removesuffix(
|
|
"@MERGED"
|
|
).removesuffix("@GRAD")
|
|
if param_name not in shard.global_params:
|
|
raise ValueError(
|
|
"Input 'X' of check_finite_and_unscale must"
|
|
f"be grads, but {input_name} is not a grad"
|
|
)
|
|
if shard.has_param(param_name):
|
|
reversed_x.append(input_name)
|
|
reversed_x_paramname.append(param_name)
|
|
op.desc.set_input('X', reversed_x)
|
|
op.desc.set_output('Out', reversed_x)
|
|
|
|
# the grad checking should take the all and only param in the current shard
|
|
to_check_param = set(reversed_x_paramname)
|
|
should_check_param = set(shard.global_params).intersection(
|
|
{
|
|
param
|
|
for param, worker_idx in shard.global_param2device.items()
|
|
if worker_idx == shard.worker_idx
|
|
}
|
|
)
|
|
assert to_check_param == should_check_param, (
|
|
f"amp \
|
|
check_finite_and_unscale checking miss [{should_check_param - to_check_param}] and got unexpected [{to_check_param - should_check_param}]"
|
|
)
|
|
|
|
if update_loss_scaling_op_idx == -1:
|
|
return
|
|
inf_var = block.var(inf_var_name)
|
|
inf_var_int32 = block.create_var(
|
|
name=inf_var_name + "@cast_int32",
|
|
shape=inf_var.shape,
|
|
dtype=core.VarDesc.VarType.INT32,
|
|
)
|
|
|
|
block._insert_op_without_sync(
|
|
update_loss_scaling_op_idx,
|
|
type='cast',
|
|
inputs={'X': inf_var},
|
|
outputs={'Out': inf_var_int32},
|
|
attrs={
|
|
"in_dtype": inf_var.dtype,
|
|
"out_dtype": inf_var_int32.dtype,
|
|
OP_ROLE_KEY: OpRole.Optimize,
|
|
},
|
|
)
|
|
update_loss_scaling_op_idx += 1
|
|
|
|
# allreduce(mp)->allreduce(sharding)->allreduce(pp)
|
|
for ring_id in ring_ids:
|
|
if ring_id == -1:
|
|
continue
|
|
# this allreduce communication should not overlap with calc
|
|
block._insert_op_without_sync(
|
|
update_loss_scaling_op_idx,
|
|
type='all_reduce',
|
|
inputs={'x': inf_var_int32},
|
|
outputs={'out': inf_var_int32},
|
|
attrs={
|
|
'ring_id': ring_id,
|
|
'op_type': paddle.distributed.ReduceOp.MAX,
|
|
OP_ROLE_KEY: OpRole.Optimize,
|
|
},
|
|
)
|
|
update_loss_scaling_op_idx += 1
|
|
|
|
block._insert_op_without_sync(
|
|
update_loss_scaling_op_idx,
|
|
type='cast',
|
|
inputs={'X': inf_var_int32},
|
|
outputs={'Out': inf_var},
|
|
attrs={
|
|
"in_dtype": inf_var_int32.dtype,
|
|
"out_dtype": inf_var.dtype,
|
|
OP_ROLE_KEY: OpRole.Optimize,
|
|
},
|
|
)
|
|
update_loss_scaling_op_idx += 1
|
|
block._sync_with_cpp()
|
|
|
|
# TODO (JZ-LIANG) revise this for uniform mixed parallelism
|
|
@staticmethod
|
|
def sync_amp_check_nan_inf(block, ring_ids):
|
|
update_loss_scaling_op_idx = -1
|
|
|
|
for idx, op in reversed(list(enumerate(block.ops))):
|
|
if op.type == "update_loss_scaling":
|
|
update_loss_scaling_op_idx = idx
|
|
inf_var_name = op.desc.input('FoundInfinite')[0]
|
|
break
|
|
|
|
# not use amp
|
|
if update_loss_scaling_op_idx == -1:
|
|
return
|
|
# 0. inf_var_int32 = cast(inf_var)
|
|
# 1. inf_var_int32 = allreduce_max(inf_var_int32)
|
|
# 3. inf_var = cast(inf_var_int32)
|
|
inf_var = block.var(inf_var_name)
|
|
inf_var_int32 = block.create_var(
|
|
name=inf_var_name + "@cast_int32",
|
|
shape=inf_var.shape,
|
|
dtype=core.VarDesc.VarType.INT32,
|
|
)
|
|
block._insert_op_without_sync(
|
|
update_loss_scaling_op_idx,
|
|
type='cast',
|
|
inputs={'X': inf_var},
|
|
outputs={'Out': inf_var_int32},
|
|
attrs={
|
|
"in_dtype": inf_var.dtype,
|
|
"out_dtype": inf_var_int32.dtype,
|
|
OP_ROLE_KEY: OpRole.Optimize,
|
|
},
|
|
)
|
|
update_loss_scaling_op_idx += 1
|
|
|
|
# allreduce(mp)->allreduce(pp)
|
|
for ring_id in ring_ids:
|
|
if ring_id == -1:
|
|
continue
|
|
block._insert_op_without_sync(
|
|
update_loss_scaling_op_idx,
|
|
type='all_reduce',
|
|
inputs={'x': inf_var_int32},
|
|
outputs={'out': inf_var_int32},
|
|
attrs={
|
|
'ring_id': ring_id,
|
|
'op_type': paddle.distributed.ReduceOp.MAX,
|
|
OP_ROLE_KEY: OpRole.Optimize,
|
|
},
|
|
)
|
|
update_loss_scaling_op_idx += 1
|
|
|
|
block._insert_op_without_sync(
|
|
update_loss_scaling_op_idx,
|
|
type='cast',
|
|
inputs={'X': inf_var_int32},
|
|
outputs={'Out': inf_var},
|
|
attrs={
|
|
"in_dtype": inf_var_int32.dtype,
|
|
"out_dtype": inf_var.dtype,
|
|
OP_ROLE_KEY: OpRole.Optimize,
|
|
},
|
|
)
|
|
update_loss_scaling_op_idx += 1
|
|
block._sync_with_cpp()
|