def _get_pserver_grad_param_var(var, var_dict):
"""
Return pserver side grad/param variable, return None
if the variable is not grad/param, e.g.
a@GRAD -> [email protected]
a@GRAD -> a@GRAD (a is not split)
fc_0.w_0 -> fc_0.w_0.block_0
fc_0.w_0 -> fc_0.w_0 (weight is not split)
_generated_var_123 -> None
"""
grad_block = None
for _, g in six.iteritems(var_dict):
if _orig_varname(g.name) == _orig_varname(var.name):
# skip per trainer vars
if g.name.find(".trainer_") == -1:
# only param or grads have split blocks
ovar_name = _orig_varname(g.name)
if ovar_name in config.param_grad_ep_mapping:
grad_block = g
break
elif ovar_name in config.grad_param_mapping:
grad_block = g
break
return grad_block
def _append_pserver_grad_merge_ops(optimize_block, grad_varname_for_block,
endpoint, grad_to_block_id):
trainers = config.get_trainers()
program = optimize_block.program
pserver_block = program.global_block()
grad_block = None
for g in config.param_grad_ep_mapping[endpoint]["grads"]:
if _orig_varname(g.name) == \
_orig_varname(grad_varname_for_block):
grad_block = g
break
if not grad_block:
# do not append this op if current endpoint
# is not dealing with this grad block
return None
orig_varname, block_name, trainer_name = _get_varname_parts(
grad_block.name)
if block_name:
merged_var_name = '.'.join([orig_varname, block_name])
else:
merged_var_name = orig_varname
merged_var = pserver_block.create_var(
name=grad_block.name,
persistable=True,
type=grad_block.type,
dtype=grad_block.dtype,
shape=grad_block.shape)
grad_to_block_id.append(merged_var.name + ":" + str(optimize_block.idx))
if config.is_sync_mode() and trainers > 1:
vars2merge = []
for i in range(trainers):
per_trainer_name = "%s.trainer_%d" % \
(merged_var_name, i)
per_trainer_var = pserver_block.create_var(
name=per_trainer_name,
persistable=False,
type=grad_block.type,
dtype=grad_block.dtype,
shape=grad_block.shape)
vars2merge.append(per_trainer_var)
optimize_block.append_op(
type="sum",
inputs={"X": vars2merge},
outputs={"Out": merged_var},
attrs={"use_mkldnn": False})
optimize_block.append_op(
type="scale",
inputs={"X": merged_var},
outputs={"Out": merged_var},
attrs={"scale": 1.0 / float(trainers)})
return merged_var
def _is_opt_op_on_pserver(endpoint, op):
param_names = [
p.name for p in config.param_grad_ep_mapping[endpoint]["params"]
]
unmerged_varnames = []
merged_varnames = []
merged_ordernames = []
for name in param_names:
orig_varname = _orig_varname(name)
for pairs in config.merged_variables_pairs:
merged_p = pairs[0]
if merged_p.merged_var.name == orig_varname:
if merged_p.merged_var.name == merged_p.ordered_vars[
0].name:
unmerged_varnames.append(merged_p.ordered_vars[0].name)
else:
merged_varnames.append(merged_p.merged_var.name)
merged_ordernames.append(merged_p.ordered_vars[0].name)
break
param = op.input("Param")[0]
if param in unmerged_varnames:
return True
for i in range(len(merged_ordernames)):
if param == merged_ordernames[i]:
merged_p = merged_varnames[i]
merged_g = "{}@GRAD".format(merged_varnames[i])
op._set_attr(OP_ROLE_VAR_ATTR_NAME, [merged_p, merged_g])
return True
return False
def get_entry_attr(param_name):
origin_name = _orig_varname(param_name)
o_main_program = config.get_origin_main_program()
for op in o_main_program.global_block().ops:
if is_distributed_sparse_op(op) and get_sparse_tablename(
op) == origin_name:
entry = op.attr("entry")
return entry
def get_distributed_from_listen_and_serv(program, origin_program):
op = get_op_by_type(program.global_block(), "listen_and_serv")
sparse_varnames = get_sparse_tablenames(origin_program, True)
sparse_params = []
grad_to_params = op.attr('sparse_grad_to_param')
for grad_to_param in grad_to_params:
_, param = grad_to_param.split(":")
if _orig_varname(param) in sparse_varnames:
sparse_params.append(param)
return sparse_params
def add_geo_optimizer_pass(program, config):
endpoint = config.get_ps_endpoint()
params = [p for p in config.param_grad_ep_mapping[endpoint]["params"]]
sparse_tablenames = get_sparse_tablenames(config.get_origin_main_program(),
False)
for param in params:
_clone_var(program.global_block(), param)
optimize_block = []
sparse_grad_to_param = []
param_to_block_id = []
pre_block_idx = program.num_blocks - 1
for param in params:
per_opt_block = program._create_block(pre_block_idx)
optimize_block.append(per_opt_block)
var_name = param.name
pserver_block = per_opt_block.program.global_block()
param = pserver_block.vars[var_name]
delta_var_name = "%s.delta" % (param.name)
origin_varname = _orig_varname(param.name)
if origin_varname in sparse_tablenames:
sparse_grad_to_param.append(":".join([delta_var_name, param.name]))
delta_var = pserver_block.create_var(
name=delta_var_name,
persistable=False,
type=param.type,
dtype=param.dtype,
shape=param.shape)
per_opt_block.append_op(
type="sum",
inputs={"X": [param, delta_var]},
outputs={"Out": param})
param_to_block_id.append(delta_var_name + ":" + str(per_opt_block.idx))
op = get_op_by_type(program.global_block(), "listen_and_serv")
op._set_attr("optimize_blocks", optimize_block)
op._set_attr("grad_to_block_id", param_to_block_id)
op._set_attr("sparse_grad_to_param", sparse_grad_to_param)
return program
def get_initializer_attrs(acture_value_names):
l_sep = ","
l_in = "&"
init_attrs = []
o_startup_program = config.get_origin_startup_program()
for value_name in acture_value_names:
origin_var_name = _orig_varname(value_name)
for op in o_startup_program.global_block().ops:
if op.type in opt_init_map.keys(
) and origin_var_name == op.output("Out")[0]:
init_attr = [op.type]
for attr in opt_init_map[op.type]:
init_attr.append(str(op.attr(attr)))
init_attrs.append(l_in.join(init_attr))
break
return l_sep.join(init_attrs)
def _get_param_block(opt_op):
# param is already created on global program
unmerged_vars = []
merged_vars = []
merged_ordervars = []
param_vars = [
p for p in config.param_grad_ep_mapping[endpoint]["params"]
]
for var in param_vars:
name = var.name
orig_varname = _orig_varname(name)
for pairs in config.merged_variables_pairs:
merged_p = pairs[0]
if merged_p.merged_var.name == orig_varname:
if merged_p.merged_var.name == merged_p.ordered_vars[
0].name:
unmerged_vars.append(merged_p.ordered_vars[0])
else:
merged_vars.append(merged_p.merged_var)
merged_ordervars.append(merged_p.ordered_vars[0])
break
param_name = opt_op.input("Param")[0]
for i in range(len(unmerged_vars)):
if _same_or_split_var(param_name, unmerged_vars[i].name):
for var in param_vars:
if _same_or_split_var(var.name, unmerged_vars[i].name):
return var
for i in range(len(merged_ordervars)):
if _same_or_split_var(param_name, merged_ordervars[i].name):
for var in param_vars:
if _same_or_split_var(var.name, merged_vars[i].name):
return var
return None
def build_pserver_startup_program_pass(program, p_main_program, config):
ps_endpoint = config.get_ps_endpoint()
o_startup_program = config.get_origin_startup_program()
program.random_seed = o_startup_program.random_seed
params = config.param_grad_ep_mapping[ps_endpoint]["params"]
merged_ordervars = []
for var in params:
name = var.name
orig_varname = _orig_varname(name)
for pairs in config.merged_variables_pairs:
merged_p = pairs[0]
if merged_p.merged_var.name == orig_varname:
if merged_p.merged_var.name != merged_p.ordered_vars[0].name:
merged_ordervars.append(merged_p.ordered_vars[0])
break
def _get_splited_name_and_shape(varname):
for splited_param in params:
pname = splited_param.name
if _same_or_split_var(pname, varname) and varname != pname:
return pname, splited_param.shape
for idx, ordered in enumerate(merged_ordervars):
if _same_or_split_var(varname, ordered.name):
return pname, splited_param.shape
return "", []
# 1. create vars in pserver program to startup program
pserver_vars = p_main_program.global_block().vars
created_var_map = collections.OrderedDict()
for _, var in six.iteritems(pserver_vars):
tmpvar = program.global_block()._clone_variable(var)
created_var_map[var.name] = tmpvar
# 2. rename op outputs
for op in o_startup_program.global_block().ops:
new_outputs = collections.OrderedDict()
# do not append startup op if var is not on this pserver
op_on_pserver = False
# TODO(gongwb) : remove this line.
if op.type not in ["recv", "fetch_barrier", "concat"]:
for key in op.output_names:
newname, _ = _get_splited_name_and_shape(op.output(key)[0])
if newname:
op_on_pserver = True
new_outputs[key] = created_var_map[newname]
elif op.output(key)[0] in pserver_vars:
op_on_pserver = True
new_outputs[key] = pserver_vars[op.output(key)[0]]
if op_on_pserver:
# most startup program ops have no inputs
new_inputs = _get_input_map_from_op(pserver_vars, op)
if op.type in [
"gaussian_random", "fill_constant", "uniform_random",
"truncated_gaussian_random"
]:
op._set_attr("shape", list(new_outputs["Out"].shape))
program.global_block().append_op(
type=op.type,
inputs=new_inputs,
outputs=new_outputs,
attrs=op.all_attrs())
return program
def large_scale_sparse_pass(program, main_program, config, is_startup=False):
opt_value_map = {}
opt_value_map["sgd"] = ["Param"]
opt_value_map["adam"] = ["Param", "Moment1", "Moment2"]
opt_value_map["adagrad"] = ["Param", "Moment"]
opt_value_map["adamax"] = ["Param", "Moment", "InfNorm"]
opt_value_map["momentum"] = ["Param", "Velocity"]
opt_value_map["lars_momentum"] = ["Param", "Velocity"]
opt_value_map["rmsprop"] = ["Param", "Moment", "MeanSquare"]
opt_value_map["decayed_adagrad"] = ["Param", "Moment"]
opt_value_map["ftrl"] = ["Param", "SquaredAccumulator", "LinearAccumulator"]
geo_value_map = {}
geo_value_map["sum"] = "Param"
opt_init_map = {}
opt_init_map["gaussian_random"] = ["seed", "mean", "std"]
opt_init_map["fill_constant"] = ["value"]
opt_init_map["uniform_random"] = ["seed", "min", "max"]
opt_init_map["truncated_gaussian_random"] = ["seed", "mean", "std"]
def get_entry_attr(param_name):
origin_name = _orig_varname(param_name)
o_main_program = config.get_origin_main_program()
for op in o_main_program.global_block().ops:
if is_distributed_sparse_op(op) and get_sparse_tablename(
op) == origin_name:
entry = op.attr("entry")
return entry
def get_initializer_attrs(acture_value_names):
l_sep = ","
l_in = "&"
init_attrs = []
o_startup_program = config.get_origin_startup_program()
for value_name in acture_value_names:
origin_var_name = _orig_varname(value_name)
for op in o_startup_program.global_block().ops:
if op.type in opt_init_map.keys(
) and origin_var_name == op.output("Out")[0]:
init_attr = [op.type]
for attr in opt_init_map[op.type]:
init_attr.append(str(op.attr(attr)))
init_attrs.append(l_in.join(init_attr))
break
return l_sep.join(init_attrs)
def get_optimizer_values(block):
value_names = []
acture_names = []
value_dims = []
grad = None
opt_idx = -1
fuse = False
for op in block.ops:
opt_idx += 1
if op.type not in opt_value_map.keys():
continue
if op.type in ["sgd", "adam"]:
fuse = True
grad = main_program.global_block().vars[op.input("Grad")[0]]
for value in opt_value_map[op.type]:
var = main_program.global_block().vars[op.input(value)[0]]
if len(var.shape) != 2:
raise ValueError("sparse param's dimension must be 2")
value_names.append(value)
value_dims.append(var.shape[1])
acture_names.append(var.name)
if value_names:
break
return grad, opt_idx, value_names, value_dims, acture_names, fuse
def add_fuse_large_scale_op(block, global_block, table_name, value_names,
acture_names, grad, is_entry, opt_idx):
op = block.ops[opt_idx]
if op.type == "sgd":
grad = main_program.global_block().vars[op.input("Grad")[0]]
lr = main_program.global_block().vars[op.input("LearningRate")[0]]
block._insert_op(
opt_idx,
type="lookup_sparse_table_fuse_sgd",
inputs={"Grad": grad,
"LearningRate": lr},
attrs={
"is_entry": is_entry,
"tablename": table_name,
"value_names": value_names
})
elif op.type == "adam":
grad = main_program.global_block().vars[op.input("Grad")[0]]
lr = main_program.global_block().vars[op.input("LearningRate")[0]]
beta1_pow = main_program.global_block().vars[op.input("Beta1Pow")[
0]]
beta2_pow = main_program.global_block().vars[op.input("Beta2Pow")[
0]]
beta1_pow_o = main_program.global_block().vars[op.output(
"Beta1PowOut")[0]]
beta2_pow_o = main_program.global_block().vars[op.output(
"Beta2PowOut")[0]]
beta1 = op.attr('beta1')
beta2 = op.attr('beta2')
epsilon = op.attr('epsilon')
block._insert_op(
opt_idx,
type="lookup_sparse_table_fuse_adam",
inputs={
"Grad": grad,
"LearningRate": lr,
"Beta1Pow": beta1_pow,
"Beta2Pow": beta2_pow
},
outputs={
"Beta1PowOut": beta1_pow_o,
"Beta2PowOut": beta2_pow_o
},
attrs={
"beta1": beta1,
"beta2": beta2,
"epsilon": epsilon,
"is_entry": is_entry,
"tablename": table_name,
"value_names": value_names
})
else:
raise ValueError("only support sgd/adam optimizer now")
def add_large_scale_op(block, global_block, table_name, value_names,
acture_names, grad, is_entry, opt_idx):
ids = global_block.create_var(
name="kSparseIDs@{}".format(table_name),
persistable=False,
dtype="int64",
shape=[1, 1],
lod_level=0)
# insert grad split to ids and tensor op
block._insert_op(
opt_idx,
type="lookup_sparse_table_grad_split",
inputs={"Grad": grad},
outputs={"Row": ids,
"Value": grad},
attrs={"tablename": table_name,
"is_entry": is_entry})
# insert read at first
vars = [global_block.vars[acture_name] for acture_name in acture_names]
block._insert_op(
opt_idx + 1,
type="lookup_sparse_table_read",
inputs={"Ids": ids},
outputs={"Out": vars},
attrs={"tablename": table_name,
"value_names": value_names})
# append write at last
inputs = {"Ids": ids, "In": vars}
block.append_op(
type="lookup_sparse_table_write",
inputs=inputs,
outputs={},
attrs={"tablename": table_name,
"value_names": value_names})
op = get_op_by_type(main_program.global_block(), "listen_and_serv")
param_blockid_map = {}
grad_blockid_map = {}
grad_to_params = op.attr('sparse_grad_to_param')
grad_to_block_ids = op.attr('grad_to_block_id')
origin_program = config.get_origin_main_program()
sparse_varnames = get_sparse_tablenames(origin_program, False)
for grad_to_block_id in grad_to_block_ids:
grad, blockid = grad_to_block_id.split(":")
grad_blockid_map[grad] = int(blockid)
for grad_to_param in grad_to_params:
grad, param = grad_to_param.split(":")
if _orig_varname(param) in sparse_varnames:
continue
param_blockid_map[param] = grad_blockid_map[grad]
if not is_startup:
for param, blockid in param_blockid_map.items():
opt_block = program.block(blockid)
grad, opt_idx, value_names, value_dims, acture_names, fuse = \
get_optimizer_values(opt_block)
entry_attr = get_entry_attr(param)
is_entry = False if entry_attr == "none" else True
if fuse:
add_fuse_large_scale_op(opt_block,
program.global_block(), param,
value_names, acture_names, grad,
is_entry, opt_idx)
else:
add_large_scale_op(opt_block,
program.global_block(), param, value_names,
acture_names, grad, is_entry, opt_idx)
else:
large_scale_kv_metas = []
for param, blockid in param_blockid_map.items():
opt_block = main_program.block(blockid)
grad, opt_idx, value_names, value_dims, acture_names, fuse = \
get_optimizer_values(opt_block)
entry_attr = get_entry_attr(param)
if fuse:
# remove origin optimzier op
opt_block._remove_op(opt_idx)
# training/infer
mode = "0"
names_str = ",".join(value_names)
dims_str = ",".join([str(dim) for dim in value_dims])
ids_name = "kSparseIDs@{}".format(param)
cached_str = ",".join(acture_names + [ids_name])
init_attr_str = get_initializer_attrs(acture_names)
meta_str = ":".join([
param, names_str, dims_str, mode, grad.name, cached_str,
init_attr_str, entry_attr
])
print("large_scale_metas: {}".format(meta_str))
large_scale_kv_metas.append(meta_str)
program.global_block().append_op(
type="lookup_sparse_table_init",
inputs=None,
outputs=None,
attrs={"large_scale_metas": large_scale_kv_metas})
# todo: need delete unused var.
return program