chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,13 @@
|
||||
# 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.
|
||||
@@ -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()
|
||||
+259
@@ -0,0 +1,259 @@
|
||||
# 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
|
||||
|
||||
__all__ = []
|
||||
|
||||
|
||||
class GradientClipHelper:
|
||||
def __init__(self, mp_ring_id):
|
||||
self.mp_ring_id = mp_ring_id
|
||||
|
||||
def _is_gradient_clip_op(self, op):
|
||||
return op.desc.has_attr("op_namescope") and op.desc.attr(
|
||||
"op_namescope"
|
||||
).startswith("/gradient_clip")
|
||||
|
||||
def prune_gradient_clip(self, block, shard, ring_ids):
|
||||
"""
|
||||
prune gradient_clip related ops for params that not belong to cur shard
|
||||
prune: square, reduce_sum, elementwise_mul
|
||||
keep: sum, sqrt, elementwise_max, elementwise_div
|
||||
"""
|
||||
deprecated_vars = set()
|
||||
deprecate_op_idx = set()
|
||||
reversed_x_paramname = []
|
||||
global_norm_sum_op_idx = -1
|
||||
for idx, op in enumerate(block.ops):
|
||||
if not self._is_gradient_clip_op(op):
|
||||
continue
|
||||
if op.type == "sum":
|
||||
global_norm_sum_op_idx = idx
|
||||
continue
|
||||
deprecate_op = False
|
||||
for input_name in op.desc.input_arg_names():
|
||||
if input_name in deprecated_vars:
|
||||
deprecate_op = True
|
||||
# TODO (JZ-LIANG) revise this for uniform mixed parallelism
|
||||
param_name = input_name.removesuffix("@MERGED").removesuffix(
|
||||
"@GRAD"
|
||||
)
|
||||
if shard.is_param(param_name) and not shard.has_param(
|
||||
param_name
|
||||
):
|
||||
deprecate_op = True
|
||||
elif shard.is_param(param_name):
|
||||
reversed_x_paramname.append(param_name)
|
||||
|
||||
if deprecate_op:
|
||||
deprecate_op_idx.add(idx)
|
||||
for output_name in op.desc.output_arg_names():
|
||||
if output_name not in op.desc.input_arg_names():
|
||||
deprecated_vars.add(output_name)
|
||||
|
||||
# NOTE(wangxi): If only have 2 sharding, and 1 param.
|
||||
# sharding 0 will not deprecated_vars, will return, only
|
||||
# sharding 1 will insert allreduce, then hang.
|
||||
if not deprecated_vars and global_norm_sum_op_idx == -1:
|
||||
# got no gradient_clip op
|
||||
return
|
||||
|
||||
for idx, op in reversed(list(enumerate(block.ops))):
|
||||
if not self._is_gradient_clip_op(op):
|
||||
continue
|
||||
if idx in deprecate_op_idx:
|
||||
block._remove_op(idx, sync=False)
|
||||
continue
|
||||
if op.type == "sum":
|
||||
reversed_inputs = []
|
||||
global_norm_sum_op_idx = idx
|
||||
for input_name in op.desc.input_arg_names():
|
||||
if input_name not in deprecated_vars:
|
||||
reversed_inputs.append(input_name)
|
||||
|
||||
op.desc.set_input("X", reversed_inputs)
|
||||
assert len(op.desc.output_arg_names()) == 1
|
||||
sum_res = op.desc.output_arg_names()[0]
|
||||
|
||||
# NOTE(wangxi): If we have 2 param, but sharding is 4,
|
||||
# then the sum op in some cards will not have input.
|
||||
# So we use fill_constant_op to set `sum_var` to zero,
|
||||
# which does not affect correctness.
|
||||
if len(reversed_inputs) == 0:
|
||||
sum_var = block.var(sum_res)
|
||||
namescope = op.attr("op_namescope")
|
||||
|
||||
block._remove_op(idx, sync=False)
|
||||
op = block._insert_op_without_sync(
|
||||
idx,
|
||||
type='fill_constant',
|
||||
inputs={},
|
||||
outputs={'Out': sum_res},
|
||||
attrs={
|
||||
'shape': sum_var.shape,
|
||||
'dtype': sum_var.dtype,
|
||||
'value': 0.0,
|
||||
OP_ROLE_KEY: OpRole.Optimize,
|
||||
},
|
||||
)
|
||||
op._set_attr('op_namescope', namescope)
|
||||
|
||||
# allreduce(mp)->allreduce(sharding)->allreduce(pp)
|
||||
idx_offset = 1
|
||||
for ring_id in ring_ids:
|
||||
if ring_id == -1:
|
||||
continue
|
||||
# this allreduce should not overlap with calc and should be scheduled in calc stream
|
||||
block._insert_op_without_sync(
|
||||
idx + idx_offset,
|
||||
type='all_reduce',
|
||||
inputs={'x': sum_res},
|
||||
outputs={'out': sum_res},
|
||||
attrs={
|
||||
'ring_id': ring_id,
|
||||
'op_namescope': "/gradient_clip_model_parallelism",
|
||||
'reduce_type': paddle.distributed.ReduceOp.Sum,
|
||||
OP_ROLE_KEY: OpRole.Optimize,
|
||||
},
|
||||
)
|
||||
idx_offset += 1
|
||||
|
||||
# the grad sum here 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}]"
|
||||
)
|
||||
|
||||
for var_name in deprecated_vars:
|
||||
block._remove_var(var_name, sync=False)
|
||||
block._sync_with_cpp()
|
||||
return
|
||||
|
||||
# TODO (JZ-LIANG) revise this for uniform mixed parallelism
|
||||
def sync_global_norm(self, block, ring_ids, mp_rank):
|
||||
"""
|
||||
prune gradient_clip related ops for params that not belong to cur shard
|
||||
prune: square, reduce_sum, elementwise_mul
|
||||
keep: sum, sqrt, elementwise_max, elementwise_div
|
||||
"""
|
||||
is_clip_grad_by_global_norm = False
|
||||
for idx, op in list(enumerate(block.ops)):
|
||||
if not self._is_gradient_clip_op(op):
|
||||
continue
|
||||
if op.type == 'sum':
|
||||
is_clip_grad_by_global_norm = True
|
||||
break
|
||||
if not is_clip_grad_by_global_norm:
|
||||
# TODO(Yuang Liu): need some extra handles when clip_grad_norm for mp
|
||||
return
|
||||
|
||||
removed_op_idx = set()
|
||||
removed_tmp_var = set()
|
||||
for idx, op in list(enumerate(block.ops)):
|
||||
if not self._is_gradient_clip_op(op):
|
||||
continue
|
||||
if op.type == 'sum':
|
||||
break
|
||||
for input_name in op.input_arg_names:
|
||||
input_var = block.var(input_name)
|
||||
# NOTE: when mp_degree > 1, some vars will be split into each mp rank.
|
||||
# However, there still some vars such as Scale, Bias are not split.
|
||||
# Those not be split vars should only be counted once during grad clip
|
||||
# by global norm. Those vars either doesn't have is_distributed attr
|
||||
# or the is_distributed attr has been set as False.
|
||||
# Therefore, we prune those duplicated vars for grad clip.
|
||||
if mp_rank >= 1 and (
|
||||
not (
|
||||
hasattr(input_var, 'is_distributed')
|
||||
and input_var.is_distributed
|
||||
)
|
||||
):
|
||||
removed_op_idx.add(idx)
|
||||
for output_name in op.output_arg_names:
|
||||
removed_tmp_var.add(output_name)
|
||||
|
||||
for idx, op in reversed(list(enumerate(block.ops))):
|
||||
if not self._is_gradient_clip_op(op):
|
||||
continue
|
||||
if idx in removed_op_idx:
|
||||
block._remove_op(idx, sync=False)
|
||||
|
||||
for var_name in removed_tmp_var:
|
||||
block._remove_var(var_name, sync=False)
|
||||
|
||||
for idx, op in list(enumerate(block.ops)):
|
||||
if not self._is_gradient_clip_op(op):
|
||||
continue
|
||||
if op.type == 'sum':
|
||||
# If mp_rank == 0, no extra handles, just allreduce
|
||||
# If mp_rank >= 1, some extra handles is needed
|
||||
sum_rst_var = block.var(op.output_arg_names[0])
|
||||
if mp_rank >= 1:
|
||||
reserved_vars = []
|
||||
for input_name in op.input_arg_names:
|
||||
if input_name not in removed_tmp_var:
|
||||
reserved_vars.append(input_name)
|
||||
|
||||
if len(reserved_vars) > 0:
|
||||
op.desc.set_input("X", reserved_vars)
|
||||
else:
|
||||
# If all input of sum op should be removed, then remove the sum op.
|
||||
# And set the output's value of sum to 0.
|
||||
namescope = op.attr("op_namescope")
|
||||
block._remove_op(idx, sync=False)
|
||||
fill_constant_op = block._insert_op_without_sync(
|
||||
idx,
|
||||
type='fill_constant',
|
||||
inputs={},
|
||||
outputs={'Out': sum_rst_var},
|
||||
attrs={
|
||||
'shape': sum_rst_var.shape,
|
||||
'dtype': sum_rst_var.dtype,
|
||||
'value': 0.0,
|
||||
OP_ROLE_KEY: OpRole.Optimize,
|
||||
},
|
||||
)
|
||||
fill_constant_op._set_attr('op_namescope', namescope)
|
||||
self._insert_allreduce(block, ring_ids, idx, sum_rst_var)
|
||||
break
|
||||
|
||||
@staticmethod
|
||||
def _insert_allreduce(block, ring_ids, idx, var):
|
||||
for ring_id in ring_ids:
|
||||
if ring_id == -1:
|
||||
continue
|
||||
|
||||
idx = idx + 1
|
||||
block._insert_op_without_sync(
|
||||
idx,
|
||||
type='all_reduce',
|
||||
inputs={'x': var},
|
||||
outputs={'out': var},
|
||||
attrs={
|
||||
'ring_id': ring_id,
|
||||
'op_namescope': "/gradient_clip_model_parallelism",
|
||||
'reduce_type': paddle.distributed.ReduceOp.SUM,
|
||||
OP_ROLE_KEY: OpRole.Optimize,
|
||||
},
|
||||
)
|
||||
+575
@@ -0,0 +1,575 @@
|
||||
# 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.framework import core
|
||||
from paddle.utils import unique_name
|
||||
|
||||
from ..common import OP_ROLE_KEY, OpRole, is_optimizer_op, is_update_op
|
||||
|
||||
__all__ = []
|
||||
|
||||
|
||||
class PlaceType:
|
||||
# sync with memcpy op, maybe not a good design
|
||||
CPU = 0
|
||||
CUDA = 1
|
||||
CUDA_PINNED = 2
|
||||
XPU = 3 # unsupported for now
|
||||
|
||||
@staticmethod
|
||||
def default_device():
|
||||
if core.is_compiled_with_cuda():
|
||||
return PlaceType.CUDA
|
||||
return PlaceType.CPU
|
||||
|
||||
@staticmethod
|
||||
def default_pinned():
|
||||
if core.is_compiled_with_cuda():
|
||||
return PlaceType.CUDA_PINNED
|
||||
return PlaceType.CPU
|
||||
|
||||
|
||||
class OffloadHelper:
|
||||
cpu_place_type = 0
|
||||
cuda_place_type = PlaceType.default_device()
|
||||
cuda_pinned_place_type = PlaceType.default_pinned()
|
||||
|
||||
def __init__(self, mp_ring_id=None, dp_ring_id=None):
|
||||
self.mp_ring_id = mp_ring_id
|
||||
self.dp_ring_id = dp_ring_id
|
||||
|
||||
def _insert_cast_op(self, block, idx, src_name, dst_name):
|
||||
src_var = block.var(src_name)
|
||||
if not block.has_var(dst_name):
|
||||
block.create_var(
|
||||
name=dst_name,
|
||||
shape=src_var.shape,
|
||||
dtype=core.VarDesc.VarType.FP16,
|
||||
persistable=True,
|
||||
)
|
||||
dst_var = block.var(dst_name)
|
||||
assert dst_var.dtype == paddle.float16
|
||||
block._insert_op_without_sync(
|
||||
idx,
|
||||
type='cast',
|
||||
inputs={'X': src_var},
|
||||
outputs={'Out': dst_var},
|
||||
attrs={
|
||||
'in_dtype': src_var.dtype,
|
||||
'out_dtype': dst_var.dtype,
|
||||
OP_ROLE_KEY: OpRole.Optimize,
|
||||
},
|
||||
)
|
||||
|
||||
def _insert_broadcast_op(self, block, idx, param_name):
|
||||
rings = []
|
||||
|
||||
if self.dp_ring_id is not None:
|
||||
rings.append(self.dp_ring_id)
|
||||
|
||||
# need sync non distributed param in mp group
|
||||
if self.mp_ring_id is not None:
|
||||
param = block.var(param_name)
|
||||
if not hasattr(param, 'is_distributed') or not param.is_distributed:
|
||||
rings.append(self.mp_ring_id)
|
||||
|
||||
# the insert op order is: mp, dp
|
||||
for ring in rings:
|
||||
block._insert_op_without_sync(
|
||||
idx,
|
||||
type="broadcast",
|
||||
inputs={'x': param_name},
|
||||
outputs={'out': param_name},
|
||||
attrs={
|
||||
'ring_id': ring,
|
||||
'root': 0,
|
||||
OP_ROLE_KEY: OpRole.Forward,
|
||||
},
|
||||
)
|
||||
|
||||
def _insert_memcpy_op(self, block, idx, src_name, dst_name, dst_place_type):
|
||||
src_var = block.var(src_name)
|
||||
dst_var = block.var(dst_name)
|
||||
block._insert_op_without_sync(
|
||||
idx,
|
||||
type='memcpy',
|
||||
inputs={'X': src_var},
|
||||
outputs={'Out': dst_var},
|
||||
attrs={
|
||||
'dst_place_type': dst_place_type,
|
||||
OP_ROLE_KEY: OpRole.Optimize,
|
||||
},
|
||||
)
|
||||
|
||||
def _insert_fetch_op(self, block, idx, src_name, dst_name):
|
||||
self._insert_memcpy_op(
|
||||
block, idx, src_name, dst_name, OffloadHelper.cuda_place_type
|
||||
)
|
||||
|
||||
def _insert_offload_op(self, block, idx, src_name, dst_name):
|
||||
self._insert_memcpy_op(
|
||||
block, idx, src_name, dst_name, OffloadHelper.cuda_pinned_place_type
|
||||
)
|
||||
|
||||
def _get_offload_var_name(self, name):
|
||||
return unique_name.generate(name + '@offload')
|
||||
|
||||
def _create_offload_var(self, var_name, offload_var_name, blocks):
|
||||
for block in blocks:
|
||||
var = block.var(var_name)
|
||||
var.persistable = False
|
||||
offload_var = block.create_var(
|
||||
name=offload_var_name,
|
||||
shape=var.shape,
|
||||
dtype=var.dtype,
|
||||
persistable=True,
|
||||
)
|
||||
|
||||
def offload_fp32param(self, block, startup_block, offload=True):
|
||||
"""
|
||||
(p_fp16) = cast(p)
|
||||
(p_fp16_recompute) = cast(p)
|
||||
(pout,) = adam(p)
|
||||
===========================>
|
||||
rename(p_fp16_recompute, p_fp16)
|
||||
|
||||
(p,) = prefetch(p@offload)
|
||||
(pout,) = adam(p)
|
||||
(p_fp16) = cast(p)
|
||||
(p@offload) = memcpy(p)
|
||||
"""
|
||||
param_to_idx = {}
|
||||
param_to_fp16 = {}
|
||||
# recompute_var which need rename to fp16_param
|
||||
fp16_param_to_recompute = {}
|
||||
recompute_to_fp16 = {}
|
||||
|
||||
def remove_param(input_name):
|
||||
param_to_idx.pop(input_name)
|
||||
if input_name in param_to_fp16:
|
||||
fp16_param = param_to_fp16.pop(input_name)
|
||||
if fp16_param in fp16_param_to_recompute:
|
||||
recompute = fp16_param_to_recompute.pop(fp16_param)
|
||||
recompute_to_fp16.pop(recompute)
|
||||
|
||||
# step1: record param
|
||||
for idx, op in reversed(list(enumerate(block.ops))):
|
||||
if is_update_op(op):
|
||||
param = op.desc.input("Param")[0]
|
||||
param_to_idx[param] = idx
|
||||
|
||||
# step2: remove param which can't offload and
|
||||
# record param->fp16param, fp16param->recompute_var
|
||||
for idx, op in enumerate(block.ops):
|
||||
if is_optimizer_op(op):
|
||||
break
|
||||
# TODO (Yuang Liu): tmp solution for fuse_grad_merge + optimize_cast
|
||||
if not offload and op.type == 'coalesce_tensor':
|
||||
continue
|
||||
for input_name in op.desc.input_arg_names():
|
||||
if input_name not in param_to_idx:
|
||||
continue
|
||||
|
||||
# param which will be used by fp32 op
|
||||
if op.type != 'cast':
|
||||
remove_param(input_name)
|
||||
continue
|
||||
|
||||
# param is only used by cast op,
|
||||
# which to cast fp32_param to fp16_param
|
||||
output_name = op.output_arg_names[0]
|
||||
if 'cast_fp16' not in output_name:
|
||||
remove_param(input_name)
|
||||
continue
|
||||
|
||||
if 'subprog' not in output_name:
|
||||
assert output_name == input_name + '.cast_fp16'
|
||||
assert input_name not in param_to_fp16, (
|
||||
"There must be only one cast op from fp32 param to fp16 param."
|
||||
)
|
||||
param_to_fp16[input_name] = output_name
|
||||
else:
|
||||
# fp16-->recompute_var
|
||||
assert input_name in param_to_fp16, (
|
||||
"param must first be cast to fp16"
|
||||
)
|
||||
fp16_param = param_to_fp16[input_name]
|
||||
fp16_param_to_recompute[fp16_param] = output_name
|
||||
recompute_to_fp16[output_name] = fp16_param
|
||||
|
||||
param_name_to_offload_name = {}
|
||||
# step3: main_block add offload, cast op
|
||||
# change recompute to fp16, remove cast(param) to fp16
|
||||
for idx, op in reversed(list(enumerate(block.ops))):
|
||||
if is_update_op(op):
|
||||
param = op.desc.input("Param")[0]
|
||||
if param not in param_to_idx:
|
||||
continue
|
||||
# step3.1: create offload_var
|
||||
offload_var_name = self._get_offload_var_name(param)
|
||||
param_name_to_offload_name[param] = offload_var_name
|
||||
if offload:
|
||||
self._create_offload_var(
|
||||
param, offload_var_name, [block, startup_block]
|
||||
)
|
||||
|
||||
# step3.2: insert cast op and offload op
|
||||
self._insert_offload_op(
|
||||
block, idx + 1, param, offload_var_name
|
||||
)
|
||||
|
||||
assert param in param_to_fp16
|
||||
fp16_param_name = param_to_fp16[param]
|
||||
fp16_param_var = block.var(fp16_param_name)
|
||||
fp16_param_var.persistable = True
|
||||
self._insert_cast_op(
|
||||
block, idx + 1, param, param_to_fp16[param]
|
||||
)
|
||||
|
||||
if offload:
|
||||
# step3.3: insert fetch op
|
||||
self._insert_fetch_op(block, idx, offload_var_name, param)
|
||||
continue
|
||||
|
||||
# step3.4: remove cast op
|
||||
if op.type == 'cast':
|
||||
input_name = op.desc.input_arg_names()[0]
|
||||
if input_name in param_to_idx:
|
||||
block._remove_op(idx, sync=False)
|
||||
continue
|
||||
|
||||
# step3.5: change recompute_param to fp16_param
|
||||
for input_name in op.desc.input_arg_names():
|
||||
if input_name in recompute_to_fp16:
|
||||
op._rename_input(input_name, recompute_to_fp16[input_name])
|
||||
for output_name in op.desc.output_arg_names():
|
||||
if output_name in recompute_to_fp16:
|
||||
op._rename_output(
|
||||
output_name, recompute_to_fp16[output_name]
|
||||
)
|
||||
|
||||
# step4: remove recompute_param
|
||||
for name in recompute_to_fp16.keys():
|
||||
block._remove_var(name, sync=False)
|
||||
|
||||
# step5: startup_block add offload
|
||||
visited_vars = set()
|
||||
# FIXME(wangxi): should insert in idx, need move comm init to the head.
|
||||
insert_idx = len(startup_block.ops)
|
||||
for idx, op in reversed(list(enumerate(startup_block.ops))):
|
||||
for out_name in op.output_arg_names:
|
||||
if out_name in visited_vars:
|
||||
continue
|
||||
|
||||
if out_name in param_name_to_offload_name:
|
||||
var_name = out_name
|
||||
if offload:
|
||||
offload_var_name = param_name_to_offload_name[var_name]
|
||||
self._insert_offload_op(
|
||||
startup_block,
|
||||
insert_idx,
|
||||
var_name,
|
||||
offload_var_name,
|
||||
)
|
||||
self._insert_cast_op(
|
||||
startup_block,
|
||||
insert_idx,
|
||||
var_name,
|
||||
param_to_fp16[var_name],
|
||||
)
|
||||
# NOTE(wangxi): cast and offload should insert after broadcast param.
|
||||
# the insert op order is: {mp, dp}broadcast, cast, offload
|
||||
self._insert_broadcast_op(
|
||||
startup_block, insert_idx, var_name
|
||||
)
|
||||
|
||||
visited_vars.add(out_name)
|
||||
|
||||
block._sync_with_cpp()
|
||||
startup_block._sync_with_cpp()
|
||||
|
||||
def cast_fp32param_in_optimize(self, block, startup_block):
|
||||
"""
|
||||
(p_fp16) = cast(p)
|
||||
(p_fp16_recompute) = cast(p)
|
||||
(pout,) = adam(p)
|
||||
===========================>
|
||||
rename(p_fp16_recompute, p_fp16)
|
||||
|
||||
(pout,) = adam(p)
|
||||
(p_fp16) = cast(p)
|
||||
"""
|
||||
self.offload_fp32param(block, startup_block, offload=False)
|
||||
|
||||
def offload(self, block, startup_block):
|
||||
"""
|
||||
(m1, m2) = prefetch(m1@offload, m2@offload)
|
||||
(m1out, m2out, pout) = adam(m1, m2, p)
|
||||
(m1@offload, m2@offload) = memcpy(m1, m2)
|
||||
"""
|
||||
vars_name_to_offload_name = {}
|
||||
|
||||
# main_block add offload
|
||||
for idx, op in reversed(list(enumerate(block.ops))):
|
||||
if not is_optimizer_op(op):
|
||||
break
|
||||
|
||||
vars_name = []
|
||||
if op.type == "adam" or op.type == "adamw":
|
||||
# {Moment1Out = [''], Moment2Out = [''], ParamOut = ['']} =
|
||||
# adam(inputs={Moment1 = [''], Moment2 = [''], Param = ['']})
|
||||
vars_name.append(op.desc.input("Moment1")[0])
|
||||
vars_name.append(op.desc.input("Moment2")[0])
|
||||
elif op.type == 'momentum':
|
||||
pass
|
||||
elif op.type == 'lars':
|
||||
pass
|
||||
elif op.type == 'lamb':
|
||||
pass
|
||||
|
||||
# step1: create and init offload_var
|
||||
for var_name in vars_name:
|
||||
assert var_name not in vars_name_to_offload_name
|
||||
|
||||
offload_var_name = self._get_offload_var_name(var_name)
|
||||
vars_name_to_offload_name[var_name] = offload_var_name
|
||||
|
||||
self._create_offload_var(
|
||||
var_name, offload_var_name, [block, startup_block]
|
||||
)
|
||||
|
||||
# step2: insert offload op
|
||||
for var_name in vars_name:
|
||||
offload_var_name = vars_name_to_offload_name[var_name]
|
||||
self._insert_offload_op(
|
||||
block, idx + 1, var_name, offload_var_name
|
||||
)
|
||||
|
||||
# step3: insert fetch op
|
||||
for var_name in vars_name:
|
||||
offload_var_name = vars_name_to_offload_name[var_name]
|
||||
self._insert_fetch_op(block, idx, offload_var_name, var_name)
|
||||
|
||||
# startup_block add offload
|
||||
visited_vars = set()
|
||||
for idx, op in reversed(list(enumerate(startup_block.ops))):
|
||||
for out_name in op.output_arg_names:
|
||||
if out_name in visited_vars:
|
||||
continue
|
||||
|
||||
if out_name in vars_name_to_offload_name:
|
||||
var_name = out_name
|
||||
offload_var_name = vars_name_to_offload_name[var_name]
|
||||
# insert offload op after var is generated
|
||||
self._insert_offload_op(
|
||||
startup_block, idx + 1, var_name, offload_var_name
|
||||
)
|
||||
visited_vars.add(out_name)
|
||||
|
||||
block._sync_with_cpp()
|
||||
startup_block._sync_with_cpp()
|
||||
|
||||
def opt_sharding_cast_fp32param(
|
||||
self, block, startup_block, params, offload=False
|
||||
):
|
||||
"""
|
||||
(p_fp16) = cast(p)
|
||||
(p_fp16_recompute) = cast(p)
|
||||
(pout,) = adam(p)
|
||||
===========================>
|
||||
rename(p_fp16_recompute, p_fp16)
|
||||
|
||||
(pout,) = adam(p)
|
||||
(p_fp16) = cast(p)
|
||||
broadcast(p_fp16)
|
||||
"""
|
||||
global_params = set()
|
||||
local_params = set()
|
||||
param_to_fp16 = {}
|
||||
# recompute_var which need rename to fp16_param
|
||||
fp16_param_to_recompute = {}
|
||||
recompute_to_fp16 = {}
|
||||
|
||||
def remove_param(input_name):
|
||||
global_params.remove(input_name)
|
||||
if input_name in local_params: # noqa: FURB132
|
||||
local_params.remove(input_name)
|
||||
if input_name in param_to_fp16:
|
||||
fp16_param = param_to_fp16.pop(input_name)
|
||||
if fp16_param in fp16_param_to_recompute:
|
||||
recompute = fp16_param_to_recompute.pop(fp16_param)
|
||||
recompute_to_fp16.pop(recompute)
|
||||
|
||||
# step1: record param
|
||||
global_params = set(params)
|
||||
for idx, op in reversed(list(enumerate(block.ops))):
|
||||
if is_update_op(op):
|
||||
param = op.desc.input("Param")[0]
|
||||
local_params.add(param)
|
||||
|
||||
# step2: remove param which can't offload and
|
||||
# record param->fp16param, fp16param->recompute_var
|
||||
for idx, op in enumerate(block.ops):
|
||||
if is_optimizer_op(op):
|
||||
break
|
||||
# TODO (Yuang Liu): tmp solution for fuse_grad_merge + optimize_cast
|
||||
if op.type == 'coalesce_tensor':
|
||||
continue
|
||||
for input_name in op.desc.input_arg_names():
|
||||
if input_name not in global_params:
|
||||
continue
|
||||
|
||||
# param which will be used by fp32 op
|
||||
if op.type != 'cast':
|
||||
remove_param(input_name)
|
||||
continue
|
||||
|
||||
# param is only used by cast op,
|
||||
# which to cast fp32_param to fp16_param
|
||||
output_name = op.output_arg_names[0]
|
||||
if 'cast_fp16' not in output_name:
|
||||
remove_param(input_name)
|
||||
continue
|
||||
|
||||
if 'subprog' not in output_name:
|
||||
assert output_name == input_name + '.cast_fp16'
|
||||
assert input_name not in param_to_fp16, (
|
||||
"There must be only one cast op from fp32 param to fp16 param."
|
||||
)
|
||||
param_to_fp16[input_name] = output_name
|
||||
else:
|
||||
# fp16-->recompute_var
|
||||
assert input_name in param_to_fp16, (
|
||||
"param must first be cast to fp16"
|
||||
)
|
||||
fp16_param = param_to_fp16[input_name]
|
||||
fp16_param_to_recompute[fp16_param] = output_name
|
||||
recompute_to_fp16[output_name] = fp16_param
|
||||
|
||||
param_name_to_offload_name = {}
|
||||
# step3: main_block add offload, cast op
|
||||
# change recompute to fp16, remove cast(param) to fp16
|
||||
for idx, op in reversed(list(enumerate(block.ops))):
|
||||
if is_update_op(op):
|
||||
param = op.desc.input("Param")[0]
|
||||
if param not in global_params:
|
||||
continue
|
||||
# step3.1: create offload_var
|
||||
offload_var_name = self._get_offload_var_name(param)
|
||||
param_name_to_offload_name[param] = offload_var_name
|
||||
if offload:
|
||||
self._create_offload_var(
|
||||
param, offload_var_name, [block, startup_block]
|
||||
)
|
||||
|
||||
# step3.2: insert cast op and offload op
|
||||
self._insert_offload_op(
|
||||
block, idx + 1, param, offload_var_name
|
||||
)
|
||||
|
||||
assert param in param_to_fp16
|
||||
fp16_param_name = param_to_fp16[param]
|
||||
fp16_param_var = block.var(fp16_param_name)
|
||||
fp16_param_var.persistable = True
|
||||
self._insert_cast_op(
|
||||
block, idx + 1, param, param_to_fp16[param]
|
||||
)
|
||||
|
||||
if offload:
|
||||
# step3.3: insert fetch op
|
||||
self._insert_fetch_op(block, idx, offload_var_name, param)
|
||||
|
||||
continue
|
||||
|
||||
# step3.4: remove cast op
|
||||
if op.type == 'cast':
|
||||
input_name = op.desc.input_arg_names()[0]
|
||||
if input_name in global_params:
|
||||
block._remove_op(idx, sync=False)
|
||||
continue
|
||||
|
||||
# step3.5: change recompute_param to fp16_param
|
||||
for input_name in op.desc.input_arg_names():
|
||||
if input_name in recompute_to_fp16:
|
||||
op._rename_input(input_name, recompute_to_fp16[input_name])
|
||||
for output_name in op.desc.output_arg_names():
|
||||
if output_name in recompute_to_fp16:
|
||||
op._rename_output(
|
||||
output_name, recompute_to_fp16[output_name]
|
||||
)
|
||||
|
||||
# step4: remove recompute_param
|
||||
for name in recompute_to_fp16.keys():
|
||||
block._remove_var(name, sync=False)
|
||||
|
||||
# step5: remove fp32 param which not need
|
||||
for idx, op in enumerate(block.ops):
|
||||
if op.type not in ['coalesce_tensor', 'c_broadcast', 'broadcast']:
|
||||
continue
|
||||
for input_name in op.desc.input_arg_names():
|
||||
if input_name in param_to_fp16:
|
||||
op._rename_input(input_name, param_to_fp16[input_name])
|
||||
for output_name in op.desc.output_arg_names():
|
||||
if output_name in param_to_fp16:
|
||||
op._rename_output(output_name, param_to_fp16[output_name])
|
||||
|
||||
for param in global_params:
|
||||
assert param in param_to_fp16
|
||||
fp16_param_name = param_to_fp16[param]
|
||||
fp16_param_var = block.var(fp16_param_name)
|
||||
fp16_param_var.persistable = True
|
||||
|
||||
if param not in local_params:
|
||||
block._remove_var(param, sync=False)
|
||||
|
||||
# step6: startup_block add offload
|
||||
visited_vars = set()
|
||||
insert_idx = len(startup_block.ops)
|
||||
for idx, op in reversed(list(enumerate(startup_block.ops))):
|
||||
for out_name in op.output_arg_names:
|
||||
if out_name in visited_vars:
|
||||
continue
|
||||
|
||||
if out_name in param_to_fp16:
|
||||
var_name = out_name
|
||||
if offload:
|
||||
self._insert_offload_op(
|
||||
startup_block,
|
||||
idx + 1,
|
||||
var_name,
|
||||
param_name_to_offload_name[var_name],
|
||||
)
|
||||
|
||||
self._insert_cast_op(
|
||||
startup_block,
|
||||
insert_idx,
|
||||
var_name,
|
||||
param_to_fp16[var_name],
|
||||
)
|
||||
|
||||
# NOTE(wangxi): cast and offload should insert after broadcast param.
|
||||
# the insert op order is: {mp, dp}broadcast, cast, offload
|
||||
self._insert_broadcast_op(
|
||||
startup_block, insert_idx, var_name
|
||||
)
|
||||
|
||||
if var_name not in local_params:
|
||||
param = startup_block.var(out_name)
|
||||
param.persistable = False
|
||||
|
||||
visited_vars.add(out_name)
|
||||
|
||||
block._sync_with_cpp()
|
||||
startup_block._sync_with_cpp()
|
||||
@@ -0,0 +1,153 @@
|
||||
# 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.
|
||||
|
||||
__all__ = []
|
||||
|
||||
|
||||
class ProgramDeps:
|
||||
def __init__(self, block, start_vars, end_vars):
|
||||
self._block = block
|
||||
# vars where to start to build the deps
|
||||
self._start_vars = start_vars
|
||||
# vars where to stop to build the deps
|
||||
self._end_vars = end_vars
|
||||
# var name -> op idxs which depends on this var
|
||||
self._var_to_use_op = {}
|
||||
# sub block deps which is a subset of this topo
|
||||
self._sub_block_deps = {}
|
||||
# var name -> op idxs which generate var
|
||||
self._var_to_generate_op = {}
|
||||
self._should_removed_var = set()
|
||||
self._father_block_deps = None
|
||||
self._build_deps()
|
||||
|
||||
def get_sub_block_deps(self, idx):
|
||||
if idx in self._sub_block_deps:
|
||||
return self._sub_block_deps[idx]
|
||||
else:
|
||||
return None
|
||||
|
||||
def get_var_deps(self, var_name):
|
||||
if var_name in self._var_to_use_op:
|
||||
return self._var_to_use_op[var_name]
|
||||
else:
|
||||
return None
|
||||
|
||||
def _build_deps(
|
||||
self,
|
||||
):
|
||||
for var_name in self._start_vars:
|
||||
self._var_to_use_op[var_name] = []
|
||||
self._var_to_generate_op[var_name] = []
|
||||
|
||||
for idx, op in enumerate(self._block.ops):
|
||||
if op.type in [
|
||||
"c_sync_comm_stream",
|
||||
"c_calc_comm_stream",
|
||||
'all_reduce',
|
||||
]:
|
||||
continue
|
||||
input_vars = op.desc.input_arg_names()
|
||||
output_vars = op.desc.output_arg_names()
|
||||
deps_reduce = False
|
||||
for input_name in input_vars:
|
||||
if input_name in self._var_to_use_op:
|
||||
deps_reduce = True
|
||||
if not deps_reduce:
|
||||
continue
|
||||
for input_name in input_vars:
|
||||
if input_name in self._var_to_use_op:
|
||||
self._var_to_use_op[input_name].append(idx)
|
||||
for output_name in output_vars:
|
||||
if output_name not in self._var_to_use_op:
|
||||
self._var_to_use_op[output_name] = []
|
||||
if output_name not in self._var_to_generate_op:
|
||||
self._var_to_generate_op[output_name] = [idx]
|
||||
else:
|
||||
self._var_to_generate_op[output_name].append(idx)
|
||||
if op.type == "conditional_block":
|
||||
# subblock
|
||||
assert op.desc.has_attr("sub_block")
|
||||
subblock_idx = op.desc.attr("sub_block").id
|
||||
subblock_deps = ProgramDeps(
|
||||
self._block.program.block(subblock_idx),
|
||||
op.desc.input_arg_names(),
|
||||
op.desc.output_arg_names(),
|
||||
)
|
||||
self._sub_block_deps[subblock_idx] = subblock_deps
|
||||
subblock_deps._father_block_deps = self
|
||||
|
||||
def crop_input_var_from_op(self, op_idx, var_name):
|
||||
if var_name in self._var_to_use_op:
|
||||
# update var -> dep_var_op
|
||||
if self._var_to_use_op[var_name] != []:
|
||||
if op_idx not in self._var_to_use_op[var_name]:
|
||||
raise ValueError(
|
||||
f"op_idx: {op_idx} is not in self._var_to_use_op[{var_name}], "
|
||||
f"self._var_to_use_op[{var_name}] is {self._var_to_use_op[var_name]}"
|
||||
)
|
||||
self._var_to_use_op[var_name].remove(op_idx)
|
||||
# update _should_removed_var
|
||||
if var_name in self._start_vars:
|
||||
self._should_removed_var.discard(var_name)
|
||||
elif (
|
||||
self._var_to_use_op[var_name] == []
|
||||
): # no more deps of this var
|
||||
self._should_removed_var.add(var_name)
|
||||
elif (
|
||||
self._var_to_generate_op[var_name][-1]
|
||||
>= self._var_to_use_op[var_name][-1]
|
||||
):
|
||||
# there are circle in the graph
|
||||
self._should_removed_var.add(var_name)
|
||||
else: # input_name should not be deleted
|
||||
self._should_removed_var.discard(var_name)
|
||||
|
||||
def crop_output_var_from_op(self, op_idx, var_name):
|
||||
if var_name in self._var_to_generate_op:
|
||||
assert op_idx in self._var_to_generate_op[var_name]
|
||||
self._var_to_generate_op[var_name].remove(op_idx)
|
||||
if self._block.has_var(var_name):
|
||||
if (
|
||||
var_name not in self._var_to_generate_op
|
||||
or self._var_to_generate_op[var_name] == []
|
||||
):
|
||||
self._block._remove_var(var_name, sync=False)
|
||||
|
||||
def remove_op(self, op_idx, reserved_vars=None):
|
||||
# update deps
|
||||
op = self._block.ops[op_idx]
|
||||
for input_name in op.desc.input_arg_names():
|
||||
if reserved_vars is not None and input_name in reserved_vars:
|
||||
continue
|
||||
self.crop_input_var_from_op(op_idx, input_name)
|
||||
for output_name in op.desc.output_arg_names():
|
||||
if reserved_vars is not None and output_name in reserved_vars:
|
||||
continue
|
||||
self.crop_output_var_from_op(op_idx, output_name)
|
||||
self._block._remove_op(op_idx, sync=False)
|
||||
|
||||
def should_remove_op(self, op_idx):
|
||||
op = self._block.ops[op_idx]
|
||||
|
||||
# NOTE: At present, it is found that the OP without output is
|
||||
# only send_v2 and partial_send op, which will be used in
|
||||
# all device
|
||||
if len(op.desc.output_arg_names()) == 0:
|
||||
return False
|
||||
|
||||
for output_name in op.desc.output_arg_names():
|
||||
if output_name not in self._should_removed_var:
|
||||
return False
|
||||
return True
|
||||
@@ -0,0 +1,175 @@
|
||||
# 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 re
|
||||
|
||||
from paddle.distributed.fleet.meta_optimizers.common import is_optimizer_op
|
||||
from paddle.distributed.fleet.meta_optimizers.sharding.fp16_helper import (
|
||||
FP16Utils,
|
||||
)
|
||||
from paddle.distributed.fleet.meta_optimizers.sharding.utils import get_var_size
|
||||
|
||||
__all__ = []
|
||||
|
||||
|
||||
class Shard:
|
||||
def __init__(
|
||||
self,
|
||||
):
|
||||
self.global_params = set()
|
||||
self.worker_idx = -1
|
||||
self.worker_num = -1
|
||||
self.global_param2device = {}
|
||||
self.device2global_params = {}
|
||||
|
||||
def setup(self, params_grads, worker_idx, worker_num):
|
||||
# param names of all devices
|
||||
self.global_params = {x[0].name for x in params_grads}
|
||||
# _param(str) -> device_id(int)
|
||||
self.worker_idx = worker_idx
|
||||
self.worker_num = worker_num
|
||||
# global_param2device contains fp32 params and fp16 params
|
||||
# device2global_params only contains fp32 params
|
||||
(
|
||||
self.global_param2device,
|
||||
self.device2global_params,
|
||||
) = self._split_params(params_grads, worker_idx, worker_num)
|
||||
|
||||
def has_param(self, var_name):
|
||||
return (
|
||||
var_name in self.global_param2device
|
||||
and self._var_device_id(var_name) == self.worker_idx
|
||||
)
|
||||
|
||||
def has_opt_var(self, var_name):
|
||||
return self._var_device_id(var_name) == self.worker_idx
|
||||
|
||||
def has_var(self, var_name):
|
||||
return (
|
||||
self._var_device_id(var_name) == -1
|
||||
or self._var_device_id(var_name) == self.worker_idx
|
||||
)
|
||||
|
||||
def _split_params(self, params_grads, worker_idx, worker_num):
|
||||
param2device = {}
|
||||
total_param_mem = 0.0
|
||||
param2mem = []
|
||||
for param in [x[0] for x in params_grads]:
|
||||
mem = get_var_size(param)
|
||||
total_param_mem += mem
|
||||
param2mem.append((param.name, mem))
|
||||
device2params = {x: [] for x in range(worker_num)}
|
||||
device_idx = 0
|
||||
mem_accu = 0.0
|
||||
for param_name, mem in param2mem:
|
||||
if mem_accu > total_param_mem * 1.0 * (device_idx + 1) / worker_num:
|
||||
device_idx += 1
|
||||
device2params[device_idx].append(param_name)
|
||||
param2device[param_name] = device_idx
|
||||
mem_accu += mem
|
||||
return param2device, device2params
|
||||
|
||||
def _var_device_id(self, var_name):
|
||||
if var_name in self.global_param2device:
|
||||
return self.global_param2device[var_name]
|
||||
for suffix in [
|
||||
"_moment1_0",
|
||||
"_moment2_0",
|
||||
"_beta1_pow_acc_0",
|
||||
"_beta2_pow_acc_0",
|
||||
"_velocity_0",
|
||||
]:
|
||||
base_name = re.sub(suffix, '', var_name)
|
||||
if base_name in self.global_param2device:
|
||||
return self.global_param2device[base_name]
|
||||
return -1
|
||||
|
||||
def find_broadcast_params(self, block):
|
||||
broadcast_vars = set()
|
||||
fp16_params = set()
|
||||
fp16_to_fp32 = {}
|
||||
|
||||
param_usage = dict.fromkeys(self.global_params, 0)
|
||||
for op in block.ops:
|
||||
if is_optimizer_op(op):
|
||||
continue
|
||||
for input_name in op.desc.input_arg_names():
|
||||
if input_name in self.global_params:
|
||||
param_usage[input_name] += 1
|
||||
|
||||
for op in block.ops:
|
||||
if not FP16Utils.is_fp16_cast_op(block, op, self.global_params):
|
||||
continue
|
||||
input_name = op.input_arg_names[0]
|
||||
output_name = op.output_arg_names[0]
|
||||
broadcast_vars.add(output_name)
|
||||
fp16_params.add(output_name)
|
||||
fp16_to_fp32[output_name] = input_name
|
||||
param_usage[input_name] -= 1
|
||||
self.global_param2device[output_name] = self.global_param2device[
|
||||
input_name
|
||||
]
|
||||
|
||||
for param, usage in param_usage.items():
|
||||
if usage > 0:
|
||||
broadcast_vars.add(param)
|
||||
return broadcast_vars
|
||||
|
||||
def device(self, var_name):
|
||||
return self._var_device_id(var_name)
|
||||
|
||||
def is_param(self, var_name):
|
||||
return var_name in self.global_params
|
||||
|
||||
def is_opti_var(self, var_name):
|
||||
if var_name in self.global_params:
|
||||
return True
|
||||
for suffix in [
|
||||
"_moment1_0",
|
||||
"_moment2_0",
|
||||
"_beta1_pow_acc_0",
|
||||
"_beta2_pow_acc_0",
|
||||
"_velocity_0",
|
||||
]:
|
||||
base_name = re.sub(suffix, '', var_name)
|
||||
if base_name in self.global_params:
|
||||
return True
|
||||
return False
|
||||
|
||||
def filter_grads(self, grads):
|
||||
grads_in_shard = []
|
||||
for grad in grads:
|
||||
param = grad.split("@")[0]
|
||||
if self.has_param(param):
|
||||
grads_in_shard.append(grad)
|
||||
return grads_in_shard
|
||||
|
||||
|
||||
class ProgramSegment:
|
||||
def __init__(self, block):
|
||||
self._block = block
|
||||
self._allreduce_vars = []
|
||||
# sub program start idx
|
||||
self._start_idx = -1
|
||||
# sub program end idx
|
||||
self._end_idx = -1
|
||||
# param name to broadcast name
|
||||
self._param2broadcast = {}
|
||||
self._broadcast_vars = []
|
||||
# cast op pairs, fp16 name (str) -> fp32 name (str)
|
||||
self._cast_ops = {}
|
||||
# fill constant vars
|
||||
self._fill_constant_vars = []
|
||||
# parameter mems
|
||||
self._param_mem = 0.0
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,41 @@
|
||||
# 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 paddle.distributed.fleet.meta_optimizers.common import OP_ROLE_VAR_KEY
|
||||
|
||||
__all__ = []
|
||||
|
||||
|
||||
class WeightDecayHelper:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def _is_weight_decay_op(self, op):
|
||||
return op.desc.has_attr("op_namescope") and op.desc.attr(
|
||||
"op_namescope"
|
||||
).startswith("/regularization")
|
||||
|
||||
def prune_weight_decay(self, block, shard):
|
||||
for idx, op in reversed(list(enumerate(block.ops))):
|
||||
if not self._is_weight_decay_op(op):
|
||||
continue
|
||||
if OP_ROLE_VAR_KEY not in op.attr_names:
|
||||
raise ValueError(
|
||||
"The Weight Decay op should hold op_role_var attribute"
|
||||
f"but the {op.type} op does not hold op_role_var"
|
||||
)
|
||||
op_role_var = op.all_attrs()[OP_ROLE_VAR_KEY]
|
||||
if not shard.has_param(op_role_var[0]):
|
||||
block._remove_op(idx, sync=False)
|
||||
block._sync_with_cpp()
|
||||
Reference in New Issue
Block a user