chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,273 @@
|
||||
# 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()
|
||||
Reference in New Issue
Block a user