wehub-resource-sync
355 lines
15 KiB
Python
355 lines
15 KiB
Python
# 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.
|
|
"""Parameter Server utils"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import warnings
|
|
from typing import TYPE_CHECKING
|
|
|
|
import paddle
|
|
|
|
if TYPE_CHECKING:
|
|
from paddle import Tensor
|
|
from paddle.distributed.fleet.base.role_maker import RoleMakerBase
|
|
from paddle.static import Executor, Program
|
|
|
|
__all__ = []
|
|
|
|
|
|
class DistributedInfer:
|
|
"""
|
|
Utility class for distributed infer of PaddlePaddle.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
main_program: Program | None = None,
|
|
startup_program: Program | None = None,
|
|
) -> None:
|
|
if main_program:
|
|
self.origin_main_program = main_program.clone()
|
|
else:
|
|
self.origin_main_program = (
|
|
paddle.static.default_main_program().clone()
|
|
)
|
|
|
|
if startup_program:
|
|
self.origin_startup_program = startup_program
|
|
else:
|
|
self.origin_startup_program = (
|
|
paddle.static.default_startup_program()
|
|
)
|
|
self.sparse_table_maps = None
|
|
|
|
def init_distributed_infer_env(
|
|
self,
|
|
exe: Executor,
|
|
loss: Tensor,
|
|
role_maker: RoleMakerBase | None = None,
|
|
dirname: str | None = None,
|
|
) -> None:
|
|
from paddle.distributed import fleet
|
|
|
|
if fleet.fleet._runtime_handle is None:
|
|
fleet.init(role_maker=role_maker)
|
|
|
|
fake_optimizer = paddle.optimizer.SGD()
|
|
strategy = fleet.DistributedStrategy()
|
|
strategy.a_sync = True
|
|
optimizer = fleet.distributed_optimizer(
|
|
fake_optimizer, strategy=strategy
|
|
)
|
|
optimizer.minimize(
|
|
loss, startup_program=self.origin_startup_program
|
|
)
|
|
|
|
if fleet.is_server():
|
|
fleet.init_server(dirname=dirname)
|
|
fleet.run_server()
|
|
else:
|
|
exe.run(paddle.static.default_startup_program())
|
|
fleet.init_worker()
|
|
self._init_dense_params(exe, dirname)
|
|
global_startup_program = paddle.static.default_startup_program()
|
|
global_startup_program = self.origin_startup_program
|
|
global_main_program = paddle.static.default_main_program()
|
|
global_main_program = self.origin_main_program
|
|
|
|
def _get_sparse_table_map(self):
|
|
from paddle.distributed import fleet
|
|
|
|
if self.sparse_table_maps is None:
|
|
self.sparse_table_maps = {}
|
|
send_ctx = fleet.fleet._runtime_handle._send_ctx
|
|
for gradname, ctx in send_ctx.items():
|
|
if ctx.is_sparse:
|
|
param = gradname.strip("@GRAD")
|
|
self.sparse_table_maps[param] = ctx.table_id()
|
|
else:
|
|
continue
|
|
return self.sparse_table_maps
|
|
|
|
def _init_dense_params(self, exe=None, dirname=None):
|
|
sparse_table_maps = self._get_sparse_table_map()
|
|
|
|
if dirname is not None and exe is not None:
|
|
all_persist_vars = [
|
|
v
|
|
for v in self.origin_main_program.list_vars()
|
|
if paddle.static.io.is_persistable(v)
|
|
]
|
|
dense_persist_vars = [
|
|
(v.name, v)
|
|
for v in all_persist_vars
|
|
if v.name not in sparse_table_maps
|
|
]
|
|
need_load_vars = [
|
|
v[1]
|
|
for v in dense_persist_vars
|
|
if os.path.isfile(os.path.join(dirname, v[0]))
|
|
]
|
|
paddle.static.load_vars(
|
|
exe,
|
|
dirname,
|
|
main_program=self.origin_main_program,
|
|
vars=need_load_vars,
|
|
)
|
|
|
|
def get_dist_infer_program(self) -> Program:
|
|
varname2tables = self._get_sparse_table_map()
|
|
convert_program = self._convert_program(
|
|
self.origin_main_program, varname2tables
|
|
)
|
|
return convert_program
|
|
|
|
def _convert_program(self, main_program, varname2tables):
|
|
def distributed_ops_pass(program):
|
|
SPARSE_OP_TYPE_DICT = {"lookup_table": "W", "lookup_table_v2": "W"}
|
|
|
|
def _get_pull_sparse_ops(_program):
|
|
pull_sparse_ops = {}
|
|
for op in _program.global_block().ops:
|
|
if (
|
|
op.type in SPARSE_OP_TYPE_DICT.keys()
|
|
and op.attr('remote_prefetch') is True
|
|
):
|
|
param_name = op.input(SPARSE_OP_TYPE_DICT[op.type])[0]
|
|
ops = pull_sparse_ops.get(param_name, [])
|
|
ops.append(op)
|
|
pull_sparse_ops[param_name] = ops
|
|
return pull_sparse_ops
|
|
|
|
def _pull_sparse_fuse(_program, pull_sparse_ops):
|
|
def dag_check_up_and_reorder(program, inputs, outputs):
|
|
global_block = program.global_block()
|
|
min_output_index = len(global_block.ops)
|
|
max_input_index = -1
|
|
input_indexes = [0] * len(global_block.ops)
|
|
output_indexes = [0] * len(global_block.ops)
|
|
for idx, op in enumerate(global_block.ops):
|
|
for i in range(0, len(op.output_names)):
|
|
if input_indexes[idx] == 1:
|
|
break
|
|
outs = op.output(op.output_names[i])
|
|
for in_id, in_var in enumerate(inputs):
|
|
if in_var.name in outs:
|
|
input_indexes[idx] = 1
|
|
max_input_index = max(max_input_index, idx)
|
|
break
|
|
|
|
for i in range(0, len(op.input_names)):
|
|
if output_indexes[idx] == 1:
|
|
break
|
|
ins = op.input(op.input_names[i])
|
|
for out_id, out_var in enumerate(outputs):
|
|
if out_var.name in ins:
|
|
output_indexes[idx] = 1
|
|
min_output_index = min(
|
|
min_output_index, idx
|
|
)
|
|
|
|
for i in range(len(global_block.ops)):
|
|
if input_indexes[i] == 1 and output_indexes[i] == 1:
|
|
warnings.warn(
|
|
"unable to re-arrange dags order to combine distributed embedding ops because a op both needs embedding table's output as input and produces ids as the same embedding table's input"
|
|
)
|
|
return
|
|
|
|
if min_output_index < max_input_index:
|
|
move_ops = []
|
|
for i in range(
|
|
min_output_index + 1, len(input_indexes)
|
|
):
|
|
if input_indexes[i] == 1:
|
|
move_ops.append((global_block.ops[i], i))
|
|
for i, op in enumerate(move_ops):
|
|
queue = []
|
|
visited = set()
|
|
queue.append(op[1])
|
|
visited.add(op[0])
|
|
start = 0
|
|
while start < len(queue):
|
|
pos = queue[start]
|
|
op = global_block.ops[pos]
|
|
op_inputs = []
|
|
for k in range(0, len(op.input_names)):
|
|
ins = op.input(op.input_names[k])
|
|
op_inputs.append(ins)
|
|
for j in range(
|
|
pos - 1, min_output_index - 1, -1
|
|
):
|
|
op1 = global_block.ops[j]
|
|
if op1 in visited:
|
|
continue
|
|
found = False
|
|
for k in range(0, len(op1.output_names)):
|
|
outs = op1.output(op1.output_names[k])
|
|
for t in range(len(op_inputs)):
|
|
for y in op_inputs[t]:
|
|
if y in outs:
|
|
found = True
|
|
break
|
|
if found:
|
|
break
|
|
if found:
|
|
break
|
|
if found:
|
|
if output_indexes[j]:
|
|
warnings.warn(
|
|
"unable to re-arrange dags order to combine distributed embedding ops"
|
|
)
|
|
return
|
|
queue.append(j)
|
|
visited.add(global_block.ops[j])
|
|
start = start + 1
|
|
|
|
queue.sort()
|
|
for index in queue:
|
|
desc = global_block.desc._insert_op(
|
|
min_output_index
|
|
)
|
|
desc.copy_from(global_block.ops[index].desc)
|
|
global_block.desc._remove_op(
|
|
index + 1, index + 2
|
|
)
|
|
global_block.ops[index].desc = desc
|
|
insert_op = global_block.ops.pop(index)
|
|
input_state = input_indexes.pop(index)
|
|
output_state = output_indexes.pop(index)
|
|
global_block.ops.insert(
|
|
min_output_index, insert_op
|
|
)
|
|
input_indexes.insert(
|
|
min_output_index, input_state
|
|
)
|
|
output_indexes.insert(
|
|
min_output_index, output_state
|
|
)
|
|
min_output_index = min_output_index + 1
|
|
|
|
assert global_block.desc.op_size() == len(
|
|
global_block.ops
|
|
)
|
|
for i in range(len(global_block.ops)):
|
|
assert (
|
|
global_block.desc.op(i)
|
|
== global_block.ops[i].desc
|
|
)
|
|
|
|
for param, ops in pull_sparse_ops.items():
|
|
all_ops = program.global_block().ops
|
|
|
|
inputs = [
|
|
program.global_block().vars[op.input("Ids")[0]]
|
|
for op in ops
|
|
]
|
|
|
|
w = program.global_block().vars[ops[0].input("W")[0]]
|
|
|
|
if w.name not in varname2tables.keys():
|
|
raise ValueError(
|
|
f"can not find variable {w.name}, please check your configuration"
|
|
)
|
|
|
|
table_id = varname2tables[w.name]
|
|
|
|
padding_idx = ops[0].attr("padding_idx")
|
|
is_distributed = ops[0].attr("is_distributed")
|
|
op_type = ops[0].type
|
|
|
|
outputs = [
|
|
program.global_block().vars[op.output("Out")[0]]
|
|
for op in ops
|
|
]
|
|
|
|
dag_check_up_and_reorder(program, inputs, outputs)
|
|
op_idxs = [all_ops.index(op) for op in ops]
|
|
|
|
for idx in op_idxs[::-1]:
|
|
program.global_block()._remove_op(idx)
|
|
|
|
inputs_idxs = [-1] * len(inputs)
|
|
outputs_idxs = [len(program.global_block().ops) + 1] * len(
|
|
outputs
|
|
)
|
|
|
|
for idx, op in enumerate(program.global_block().ops):
|
|
for i in range(0, len(op.output_names)):
|
|
outs = op.output(op.output_names[i])
|
|
for in_id, in_var in enumerate(inputs):
|
|
if in_var.name in outs:
|
|
inputs_idxs[in_id] = max(
|
|
idx, inputs_idxs[in_id]
|
|
)
|
|
for i in range(0, len(op.input_names)):
|
|
ins = op.input(op.input_names[i])
|
|
for out_id, out_var in enumerate(outputs):
|
|
if out_var.name in ins:
|
|
outputs_idxs[out_id] = min(
|
|
idx, outputs_idxs[out_id]
|
|
)
|
|
|
|
if min(outputs_idxs) - max(inputs_idxs) >= 1:
|
|
distributed_idx = max(inputs_idxs) + 1
|
|
|
|
program.global_block()._insert_op(
|
|
index=distributed_idx,
|
|
type="distributed_lookup_table",
|
|
inputs={"Ids": inputs, 'W': w},
|
|
outputs={"Outputs": outputs},
|
|
attrs={
|
|
"is_distributed": is_distributed,
|
|
"padding_idx": padding_idx,
|
|
"table_id": table_id,
|
|
"is_test": True,
|
|
"lookup_table_version": op_type,
|
|
},
|
|
)
|
|
else:
|
|
raise ValueError(
|
|
"something wrong with Fleet, submit a issue is recommended"
|
|
)
|
|
|
|
pull_sparse_ops = _get_pull_sparse_ops(program)
|
|
warnings.warn(
|
|
"lookup_table will be forced to test mode when use DistributedInfer"
|
|
)
|
|
_pull_sparse_fuse(program, pull_sparse_ops)
|
|
return program
|
|
|
|
covert_program = distributed_ops_pass(main_program)
|
|
return covert_program
|