distribute_transpiler.py 33.1 KB
Newer Older
1
#   Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
D
dzhwinter 已提交
2
#
D
dzhwinter 已提交
3 4 5
# 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
D
dzhwinter 已提交
6
#
D
dzhwinter 已提交
7
#     http://www.apache.org/licenses/LICENSE-2.0
D
dzhwinter 已提交
8
#
D
dzhwinter 已提交
9 10 11 12 13 14
# 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.

T
typhoonzero 已提交
15
from __future__ import print_function
T
done  
typhoonzero 已提交
16
import framework
17
from framework import Program, default_main_program, default_startup_program, Parameter, Variable
T
done  
typhoonzero 已提交
18 19
import optimizer
from layer_helper import LayerHelper
20
import distributed_splitter as splitter
T
typhoonzero 已提交
21
import math
22
from . import core
23
import debuger
T
done  
typhoonzero 已提交
24 25


T
typhoonzero 已提交
26 27 28 29 30 31
class VarBlock:
    def __init__(self, varname, offset, size):
        self.varname = varname
        # NOTE: real offset is offset * size
        self.offset = offset
        self.size = size
T
done  
typhoonzero 已提交
32

T
typhoonzero 已提交
33 34
    def __str__(self):
        return "%s:%d:%d" % (self.varname, self.offset, self.size)
T
done  
typhoonzero 已提交
35 36


37 38
class UnionFind(object):
    """ Union-find data struct.
39

40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87
    Union-find is a data struct that keeps track of a set of elements partitioned
    into a number of disjoint (non-overlapping) subsets.

    Reference:
    https://en.wikipedia.org/wiki/Disjoint-set_data_structure

    Args:
      elements(list): The initialize element list.
    """

    def __init__(self, elementes=None):
        self._parents = []  # index -> parent index
        self._index = {}  # element -> index
        self._curr_idx = 0
        if not elementes:
            elementes = []
        for ele in elementes:
            self._parents.append(self._curr_idx)
            self._index.update({ele: self._curr_idx})
            self._curr_idx += 1

    def find(self, x):
        # Find the root index of given element x,
        # execute the path compress while findind the root index
        if not x in self._index:
            return -1
        idx = self._index[x]
        while idx != self._parents[idx]:
            t = self._parents[idx]
            self._parents[idx] = self._parents[t]
            idx = t
        return idx

    def union(self, x, y):
        # Union two given element
        x_root = self.find(x)
        y_root = self.find(y)

        if x_root == y_root:
            return
        self._parents[x_root] = y_root

    def is_connected(self, x, y):
        # If two given elements have the same root index,
        # then they are connected.
        return self.find(x) == self.find(y)


88 89 90 91
def same_or_split_var(p_name, var_name):
    return p_name == var_name or p_name.startswith(var_name + ".block")


T
typhoonzero 已提交
92 93 94 95 96
def split_dense_variable(var_list,
                         pserver_count,
                         min_block_size=1024,
                         max_block_size=1048576):
    """
97
        We may need to split dense tensor to one or more blocks and put
T
typhoonzero 已提交
98 99
        them equally onto parameter server. One block is a sub-tensor
        aligned by dim[0] of the tensor.
100

T
typhoonzero 已提交
101 102
        We need to have a minimal block size so that the calculations in
        the parameter server side can gain better performance. By default
103 104
        minimum block size is 1024. The max block size is used to prevent
        very large blocks that may cause send error.
T
typhoonzero 已提交
105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122
    """
    blocks = []
    for var in var_list:
        split_count = pserver_count
        var_numel = reduce(lambda x, y: x * y, var.shape)
        max_pserver_count = int(math.floor(var_numel / float(min_block_size)))
        if max_pserver_count == 0:
            max_pserver_count = 1
        if max_pserver_count < pserver_count:
            split_count = max_pserver_count
        block_size = int(math.ceil(var_numel / float(split_count)))

        if len(var.shape) >= 2:
            # align by dim1(width)
            dim1 = reduce(lambda x, y: x * y, var.shape[1:])
            remains = block_size % dim1
            if remains != 0:
                block_size += dim1 - remains
123
        # update split_count after aligning
T
typhoonzero 已提交
124 125 126 127 128 129 130 131 132
        split_count = int(math.ceil(var_numel / float(block_size)))
        for block_id in xrange(split_count):
            curr_block_size = min(block_size, var_numel - (
                (block_id) * block_size))
            block = VarBlock(var.name, block_id, curr_block_size)
            blocks.append(str(block))
    return blocks


T
done  
typhoonzero 已提交
133 134 135 136
class DistributeTranspiler:
    def transpile(self,
                  optimize_ops,
                  params_grads,
T
typhoonzero 已提交
137
                  trainer_id,
T
done  
typhoonzero 已提交
138 139 140
                  program=None,
                  pservers="127.0.0.1:6174",
                  trainers=1,
141
                  split_method=splitter.round_robin):
T
done  
typhoonzero 已提交
142
        """
143 144
            Transpile the program to distributed data-parallelism programs.
            The main_program will be transformed to use a remote parameter server
T
done  
typhoonzero 已提交
145
            to do parameter optimization. And the optimization graph will be put
146
            into a parameter server program.
T
done  
typhoonzero 已提交
147

148
            Use different methods to split trainable variables to different
T
done  
typhoonzero 已提交
149 150
            parameter servers.

T
typhoonzero 已提交
151 152 153 154 155 156 157 158 159 160 161 162 163 164 165
            Steps to transpile trainer:
            1. split variable to multiple blocks, aligned by product(dim[1:]) (width).
            2. rename splited grad variables to add trainer_id suffix ".trainer_%d".
            3. modify trainer program add split_op to each grad variable.
            4. append send_op to send splited variables to server and fetch
               params(splited blocks or origin param) from server.
            5. append concat_op to merge splited blocks to update local weights.

            Steps to transpile pserver:
            1. create new program for parameter server.
            2. create params and grad variables that assigned to current server instance.
            3. create a sub-block in the server side program
            4. append ops that should run on current server instance.
            5. add listen_and_serv op

T
done  
typhoonzero 已提交
166
            :param optimize_ops: op list of optimization, should be the
167
                                    return value of Optimizer.minimize
T
done  
typhoonzero 已提交
168
            :type optimize_ops: list
T
typhoonzero 已提交
169 170 171 172
            :param params_grads: list of tuple(weight, gradient)
            :type params_grads: list
            :param trainer_id: one unique id for each trainer in a job.
            :type trainer_id: int
T
typhoonzero 已提交
173
            :param program: program to transpile, default is default_main_program
T
typhoonzero 已提交
174
            :type program: Program
T
done  
typhoonzero 已提交
175 176
            :param pservers: parameter server endpoints like "m1:6174,m2:6174"
            :type pservers: string
T
typhoonzero 已提交
177 178 179 180 181
            :param trainers: total number of workers/trainers in the job
            :type trainers: int
            :param split_method: A function to determin how to split variables
                to different servers equally.
            :type split_method: function
T
done  
typhoonzero 已提交
182
        """
T
typhoonzero 已提交
183
        assert (callable(split_method))
T
done  
typhoonzero 已提交
184 185
        if program is None:
            program = default_main_program()
186
        self.program = program
T
done  
typhoonzero 已提交
187
        self.trainers = trainers
188
        self.optimize_ops = optimize_ops
T
typhoonzero 已提交
189 190 191 192
        # TODO(typhoonzero): currently trainer_id is fetched from cluster system
        # like Kubernetes, we should port this to use etcd later when developing
        # fluid distributed training with fault-tolerance.
        self.trainer_id = trainer_id
T
typhoonzero 已提交
193
        pserver_endpoints = pservers.split(",")
T
typhoonzero 已提交
194 195

        # step1
T
typhoonzero 已提交
196 197 198 199
        param_list = [pg[0] for pg in params_grads]
        grad_list = [pg[1] for pg in params_grads]
        grad_blocks = split_dense_variable(grad_list, len(pserver_endpoints))
        param_blocks = split_dense_variable(param_list, len(pserver_endpoints))
T
typhoonzero 已提交
200
        # step2
T
typhoonzero 已提交
201
        grad_var_mapping = self._append_split_op(program, grad_blocks)
T
typhoonzero 已提交
202 203
        # step3
        send_inputs = []
T
typhoonzero 已提交
204
        send_outputs = []
T
typhoonzero 已提交
205 206 207 208
        for b in grad_blocks:  # append by order
            varname, block_id, _ = b.split(":")
            send_inputs.append(grad_var_mapping[varname][int(block_id)])

T
typhoonzero 已提交
209 210
        param_var_mapping = self._create_vars_from_blocklist(program,
                                                             param_blocks)
T
typhoonzero 已提交
211 212 213
        for b in param_blocks:
            varname, block_id, _ = b.split(":")
            send_outputs.append(param_var_mapping[varname][int(block_id)])
214 215
        # let send_op know which endpoint to send which var to, eplist has the same
        # order as send_inputs.
T
typhoonzero 已提交
216
        eplist = split_method(send_inputs, pserver_endpoints)
217
        # create mapping of endpoint -> split var to create pserver side program
218 219 220 221 222 223 224 225
        self.param_grad_ep_mapping = dict()
        for i, ep in enumerate(eplist):
            param = send_outputs[i]
            grad = send_inputs[i]
            if not self.param_grad_ep_mapping.has_key(ep):
                self.param_grad_ep_mapping[ep] = {"params": [], "grads": []}
            self.param_grad_ep_mapping[ep]["params"].append(param)
            self.param_grad_ep_mapping[ep]["grads"].append(grad)
T
typhoonzero 已提交
226

227 228
        rpc_client_var = program.global_block().create_var(
            name="RPC_CLIENT_VAR",
T
typhoonzero 已提交
229
            persistable=True,
230
            type=core.VarDesc.VarType.RAW)
231

232
        # create send_op
T
typhoonzero 已提交
233
        program.global_block().append_op(
T
typhoonzero 已提交
234 235
            type="send",
            inputs={"X": send_inputs},
236 237
            outputs={"Out": send_outputs,
                     "RPCClient": rpc_client_var},
T
typhoonzero 已提交
238
            attrs={"endpoints": pserver_endpoints,
T
typhoonzero 已提交
239 240 241
                   "epmap": eplist})
        # step4
        for varname, splited_var in param_var_mapping.iteritems():
T
typhoonzero 已提交
242 243
            if len(splited_var) <= 1:
                continue
T
typhoonzero 已提交
244
            orig_param = program.global_block().vars[varname]
T
typhoonzero 已提交
245
            program.global_block().append_op(
T
typhoonzero 已提交
246
                type="concat",
T
typhoonzero 已提交
247
                inputs={"X": splited_var},
248
                outputs={"Out": [orig_param]},
T
typhoonzero 已提交
249
                attrs={"axis": 0})
T
typhoonzero 已提交
250

T
typhoonzero 已提交
251 252 253
    def get_trainer_program(self):
        # remove optimize ops and add a send op to main_program
        self.program.global_block().delete_ops(self.optimize_ops)
254 255
        # FIXME(typhoonzero): serialize once will fix error occurs when clone.
        self.program.__str__()
T
typhoonzero 已提交
256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278
        return self.program

    def get_pserver_program(self, endpoint):
        """
        Get pserver side program using the endpoint.
        NOTE: assume blocks of the same variable is not distributed
        on the same pserver, only change param/grad varnames for
        trainers to fetch.
        """
        # step1
        pserver_program = Program()
        # step2
        recv_inputs = []
        for v in self.param_grad_ep_mapping[endpoint]["params"]:
            self._clone_var(pserver_program.global_block(), v)
        for v in self.param_grad_ep_mapping[endpoint]["grads"]:
            # create vars for each trainer in global scope, so
            # we don't need to create them when grad arrives.
            # change client side var name to origin name by
            # removing ".trainer_%d" suffix
            suff_idx = v.name.find(".trainer_")
            if suff_idx >= 0:
                orig_var_name = v.name[:suff_idx]
T
typhoonzero 已提交
279 280 281
            else:
                orig_var_name = v.name
            single_trainer_var = pserver_program.global_block().create_var(
T
typhoonzero 已提交
282 283
                name=orig_var_name,
                persistable=True,
284
                type=v.type,
T
typhoonzero 已提交
285 286
                dtype=v.dtype,
                shape=v.shape)
T
typhoonzero 已提交
287 288 289 290 291 292 293 294 295 296 297
            if self.trainers > 1:
                for trainer_id in xrange(self.trainers):
                    var = pserver_program.global_block().create_var(
                        name="%s.trainer_%d" % (orig_var_name, trainer_id),
                        persistable=False,
                        type=v.type,
                        dtype=v.dtype,
                        shape=v.shape)
                    recv_inputs.append(var)
            else:
                recv_inputs.append(single_trainer_var)
298

T
typhoonzero 已提交
299 300 301 302 303 304 305
        # step3
        optimize_block = pserver_program.create_block(0)
        # step 4
        # Create a union-find data struct from optimize ops,
        # If two ops are connected, we could add these two ops
        # into one set.
        ufind = self._create_ufind(self.optimize_ops)
306
        # step 4.2
T
typhoonzero 已提交
307 308 309 310 311 312 313 314
        # Iterate through the ops and append optimize op which
        # located on current pserver
        opt_op_on_pserver = []
        for _, op in enumerate(self.optimize_ops):
            if self._is_opt_op(op) and self._is_opt_op_on_pserver(endpoint, op):
                opt_op_on_pserver.append(op)
        # step 4.3
        # Iterate through the ops, and if an op and the optimize ops
315
        # which located on current pserver are in one set, then
T
typhoonzero 已提交
316
        # append it into the sub program.
T
typhoonzero 已提交
317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345

        # We try to put optimization program run parallelly, assume
        # optimization program always looks like:
        #
        # prevop -> prevop -> opt op -> following op -> following op; ->
        # prevop -> prevop -> opt op -> following op -> following op; ->
        # global op -> global op
        #
        # we put operators that can run parallelly to many program blocks.
        # in above example, we seperate ops by the ";". Global ops must run
        # after all the optimize ops finished.

        global_ops = []
        # HACK: optimization global ops only used to scale beta1 and beta2
        # replace it with dependency engine.
        for op in self.optimize_ops:
            if op.type == "scale":
                for in_name in op.input_arg_names:
                    if in_name.startswith("beta1_pow_acc") or\
                        in_name.startswith("beta2_pow_acc"):
                        global_ops.append(op)

        def __append_optimize_op__(op, block):
            if self._is_opt_op(op):
                self._append_pserver_ops(block, op, endpoint,
                                         default_main_program())
            else:
                self._append_pserver_non_opt_ops(block, op)

346 347 348 349 350 351 352 353 354
        append_block = optimize_block
        # append lr decay ops to the child block if exits
        lr_ops = self._get_lr_ops()
        if len(lr_ops) > 0:
            for _, op in enumerate(lr_ops):
                self._append_pserver_non_opt_ops(append_block, op)

            append_block = pserver_program.create_block(append_block.idx)

T
typhoonzero 已提交
355
        # append op to the current block
356
        per_opt_block = append_block
T
typhoonzero 已提交
357 358 359 360 361 362
        for _, opt_op in enumerate(opt_op_on_pserver):
            for _, op in enumerate(self.optimize_ops):
                # optimizer is connected to itself
                if ufind.is_connected(op, opt_op) and \
                    op not in global_ops:
                    __append_optimize_op__(op, per_opt_block)
363
            per_opt_block = pserver_program.create_block(append_block.idx)
T
typhoonzero 已提交
364 365 366 367 368 369 370 371 372 373 374 375 376

        # append global ops
        for glb_op in global_ops:
            __append_optimize_op__(glb_op, per_opt_block)

        # NOT USED: single block version:
        #
        # for _, op in enumerate(self.optimize_ops):
        #     for _, opt_op in enumerate(opt_op_on_pserver):
        #         if ufind.is_connected(op, opt_op):
        #             __append_optimize_op__(glb_op, optimize_block)
        #             break

T
typhoonzero 已提交
377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410
        # step5 append the listen_and_serv op
        pserver_program.global_block().append_op(
            type="listen_and_serv",
            inputs={'X': recv_inputs},
            outputs={},
            attrs={
                "OptimizeBlock": optimize_block,
                "endpoint": endpoint,
                "Fanin": self.trainers
            })
        pserver_program.sync_with_cpp()
        return pserver_program

    def get_startup_program(self, endpoint, pserver_program):
        """
        Get startup program for current parameter server.
        Modify operator input variables if there are variables that
        were split to several blocks.
        """
        s_prog = Program()
        orig_s_prog = framework.default_startup_program()
        params = self.param_grad_ep_mapping[endpoint]["params"]

        def _get_splited_name_and_shape(varname):
            for idx, splited_param in enumerate(params):
                pname = splited_param.name
                if same_or_split_var(pname, varname) and varname != pname:
                    return pname, splited_param.shape
            return "", []

        # 1. create vars in pserver program to startup program
        pserver_vars = pserver_program.global_block().vars
        created_var_map = dict()
        for _, var in pserver_vars.iteritems():
T
update  
typhoonzero 已提交
411
            tmpvar = s_prog.global_block().clone_variable(var)
T
typhoonzero 已提交
412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452
            created_var_map[var.name] = tmpvar

        # 2. rename op outputs
        for op in orig_s_prog.global_block().ops:
            new_inputs = dict()
            new_outputs = dict()
            # do not append startup op if var is not on this pserver
            op_on_pserver = False
            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]]

            # most startup program ops have no inputs
            new_inputs = self._get_input_map_from_op(pserver_vars, op)

            if op_on_pserver:
                if op.type in [
                        "gaussian_random", "fill_constant", "uniform_random"
                ]:
                    op.attrs["shape"] = new_outputs["Out"].shape
                s_prog.global_block().append_op(
                    type=op.type,
                    inputs=new_inputs,
                    outputs=new_outputs,
                    attrs=op.attrs)
        return s_prog

    # ====================== private transpiler functions =====================
    def _create_vars_from_blocklist(self,
                                    program,
                                    block_list,
                                    add_trainer_suffix=False):
        """
        NOTE: only grads need to be named for different trainers, use
              add_trainer_suffix to rename the grad vars.
        """
T
typhoonzero 已提交
453
        block_map = dict()
T
typhoonzero 已提交
454
        var_mapping = dict()
T
typhoonzero 已提交
455 456 457 458 459 460
        for block_str in block_list:
            varname, offset, size = block_str.split(":")
            if not block_map.has_key(varname):
                block_map[varname] = []
            block_map[varname].append((long(offset), long(size)))
        for varname, splited in block_map.iteritems():
T
typhoonzero 已提交
461
            orig_var = program.global_block().var(varname)
T
typhoonzero 已提交
462
            if len(splited) == 1:
T
typhoonzero 已提交
463 464 465 466 467 468 469 470 471
                if add_trainer_suffix:
                    new_var_name = "%s.trainer_%d" % \
                        (orig_var.name, self.trainer_id)
                    program.global_block().rename_var(varname, new_var_name)
                    var_mapping[varname] = \
                        [program.global_block().var(new_var_name)]
                else:
                    var_mapping[varname] = \
                        [program.global_block().var(orig_var.name)]
T
typhoonzero 已提交
472
                continue
T
typhoonzero 已提交
473 474

            var_mapping[varname] = []
T
typhoonzero 已提交
475 476 477 478
            orig_shape = orig_var.shape
            orig_dim1_flatten = 1
            if len(orig_shape) >= 2:
                orig_dim1_flatten = reduce(lambda x, y: x * y, orig_shape[1:])
T
typhoonzero 已提交
479

T
typhoonzero 已提交
480
            for i, block in enumerate(splited):
T
typhoonzero 已提交
481
                size = block[1]
T
typhoonzero 已提交
482 483 484 485
                rows = size / orig_dim1_flatten
                splited_shape = [rows]
                if len(orig_shape) >= 2:
                    splited_shape.extend(orig_shape[1:])
T
typhoonzero 已提交
486 487 488 489 490 491 492
                new_var_name = ""
                if add_trainer_suffix:
                    new_var_name = "%s.block%d.trainer_%d" % \
                        (varname, i, self.trainer_id)
                else:
                    new_var_name = "%s.block%d" % \
                        (varname, i)
T
typhoonzero 已提交
493
                var = program.global_block().create_var(
T
typhoonzero 已提交
494 495
                    name=new_var_name,
                    persistable=False,
T
typhoonzero 已提交
496
                    dtype=orig_var.dtype,
497
                    type=orig_var.type,
T
typhoonzero 已提交
498
                    shape=splited_shape)  # flattend splited var
T
typhoonzero 已提交
499
                var_mapping[varname].append(var)
T
typhoonzero 已提交
500
            program.global_block().sync_with_cpp()
T
typhoonzero 已提交
501
        return var_mapping
T
done  
typhoonzero 已提交
502 503 504 505 506 507 508 509 510

    def _clone_var(self, block, var):
        assert isinstance(var, Variable)
        return block.create_var(
            name=var.name,
            shape=var.shape,
            dtype=var.dtype,
            type=var.type,
            lod_level=var.lod_level,
T
typhoonzero 已提交
511
            persistable=True)
T
done  
typhoonzero 已提交
512

T
typhoonzero 已提交
513
    def _append_split_op(self, program, gradblocks):
514
        # Split variables that need to be split and append respective ops
T
typhoonzero 已提交
515 516 517
        add_suffix = False
        if self.trainers > 1:
            add_suffix = True
T
typhoonzero 已提交
518
        var_mapping = self._create_vars_from_blocklist(
T
typhoonzero 已提交
519
            program, gradblocks, add_trainer_suffix=add_suffix)
T
typhoonzero 已提交
520
        for varname, splited_vars in var_mapping.iteritems():
T
typhoonzero 已提交
521 522
            # variable that don't need to split have empty splited_vars
            if len(splited_vars) <= 1:
T
typhoonzero 已提交
523
                continue
T
typhoonzero 已提交
524
            orig_var = program.global_block().vars[varname]
T
typhoonzero 已提交
525
            if orig_var.type == core.VarDesc.VarType.SELECTED_ROWS:
526 527 528 529 530 531 532 533
                height_sections = []
                for v in splited_vars:
                    height_sections.append(v.shape[0])
                program.global_block().append_op(
                    type="split_selected_rows",
                    inputs={"X": orig_var},
                    outputs={"Out": splited_vars},
                    attrs={"height_sections": height_sections})
T
typhoonzero 已提交
534
            elif orig_var.type == core.VarDesc.VarType.LOD_TENSOR:
535 536 537 538 539 540 541 542 543 544 545 546
                sections = []
                for v in splited_vars:
                    sections.append(v.shape[0])
                program.global_block().append_op(
                    type="split",
                    inputs={"X": orig_var},
                    outputs={"Out": splited_vars},
                    attrs={"sections": sections}  # assume split evenly
                )
            else:
                AssertionError("Variable type should be in set "
                               "[LOD_TENSOR, SELECTED_ROWS]")
T
typhoonzero 已提交
547
        return var_mapping
T
done  
typhoonzero 已提交
548

T
typhoonzero 已提交
549 550 551 552
    def _get_optimizer_input_shape(self, op_type, varkey, orig_shape,
                                   param_shape):
        """
        Returns the shape for optimizer inputs that need to be reshaped when
553
        Param and Grad is split to multiple servers.
T
typhoonzero 已提交
554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575
        """
        # HACK(typhoonzero): Should use functions of corresponding optimizer in
        # optimizer.py to get the shape, do not  bind this in the transpiler.
        if op_type == "adam":
            if varkey in ["Moment1", "Moment2"]:
                return param_shape
        elif op_type == "adagrad":
            if varkey == "Moment":
                return param_shape
        elif op_type == "adamax":
            if varkey in ["Moment", "InfNorm"]:
                return param_shape
        elif op_type == "momentum":
            if varkey == "Velocity":
                return param_shape
        elif op_type == "":
            if varkey == "Moment":
                return param_shape
        elif op_type == "sgd":
            pass
        return orig_shape

T
typhoonzero 已提交
576 577 578 579 580
    def _orig_varname(self, varname):
        suff_idx = varname.find(".trainer_")
        orig_var_name = ""
        if suff_idx >= 0:
            orig_var_name = varname[:suff_idx]
T
typhoonzero 已提交
581 582
        else:
            orig_var_name = varname
T
typhoonzero 已提交
583 584
        return orig_var_name

585 586
    def _append_pserver_ops(self, optimize_block, opt_op, endpoint,
                            origin_program):
587
        program = optimize_block.program
588
        pserver_block = program.global_block()
589
        new_inputs = dict()
T
typhoonzero 已提交
590 591
        # update param/grad shape first, then other inputs like
        # moment can use the updated shape
592
        for key in opt_op.input_names:
593 594 595
            if key == "Grad":
                grad_block = None
                for g in self.param_grad_ep_mapping[endpoint]["grads"]:
T
typhoonzero 已提交
596
                    if same_or_split_var(
T
typhoonzero 已提交
597 598
                            self._orig_varname(g.name),
                            self._orig_varname(opt_op.input(key)[0])):
599 600 601 602 603 604
                        grad_block = g
                        break
                if not grad_block:
                    # do not append this op if current endpoint
                    # is not dealing with this grad block
                    return
T
typhoonzero 已提交
605 606
                merged_var = \
                    pserver_block.vars[self._orig_varname(grad_block.name)]
607
                if self.trainers > 1:
T
typhoonzero 已提交
608 609 610 611 612 613
                    vars2merge = []
                    for i in xrange(self.trainers):
                        per_trainer_name = "%s.trainer_%d" % \
                        (self._orig_varname(grad_block.name), i)
                        vars2merge.append(pserver_block.vars[per_trainer_name])

614
                    optimize_block.append_op(
T
done  
typhoonzero 已提交
615 616 617
                        type="sum",
                        inputs={"X": vars2merge},
                        outputs={"Out": merged_var})
618 619 620 621 622 623
                    if not merged_var.type == core.VarDesc.VarType.SELECTED_ROWS:
                        optimize_block.append_op(
                            type="scale",
                            inputs={"X": merged_var},
                            outputs={"Out": merged_var},
                            attrs={"scale": 1.0 / float(self.trainers)})
624 625 626 627 628
                new_inputs[key] = merged_var
            elif key == "Param":
                # param is already created on global program
                param_block = None
                for p in self.param_grad_ep_mapping[endpoint]["params"]:
T
typhoonzero 已提交
629
                    if same_or_split_var(p.name, opt_op.input(key)[0]):
630 631 632 633
                        param_block = p
                        break
                if not param_block:
                    return
634
                tmpvar = pserver_block.create_var(
635
                    name=param_block.name,
T
typhoonzero 已提交
636
                    persistable=True,
637 638 639
                    dtype=param_block.dtype,
                    shape=param_block.shape)
                new_inputs[key] = tmpvar
640 641 642
            elif key == "LearningRate":
                # leraning rate variable has already be created by non-optimize op,
                # don't create it once again.
643 644 645 646 647 648 649 650 651 652 653
                lr_varname = opt_op.input(key)[0]
                if pserver_block.vars.has_key(lr_varname):
                    new_inputs[key] = pserver_block.vars[opt_op.input(key)[0]]
                else:
                    origin_var = origin_program.global_block().vars[lr_varname]
                    tmpvar = pserver_block.create_var(
                        name=origin_var.name,
                        persistable=origin_var.persistable,
                        dtype=origin_var.dtype,
                        shape=origin_var.shape)
                    new_inputs[key] = tmpvar
T
typhoonzero 已提交
654

655
        for key in opt_op.input_names:
656 657
            new_shape = None
            if key in ["Param", "Grad", "LearningRate"]:
T
typhoonzero 已提交
658
                continue
659
            var = self.program.global_block().vars[opt_op.input(key)[0]]
T
typhoonzero 已提交
660 661 662 663
            # update accumulator variable shape
            param_shape = new_inputs["Param"].shape
            new_shape = self._get_optimizer_input_shape(opt_op.type, key,
                                                        var.shape, param_shape)
664
            tmpvar = pserver_block.create_var(
T
typhoonzero 已提交
665 666 667 668 669
                name=var.name,
                persistable=var.persistable,
                dtype=var.dtype,
                shape=new_shape)
            new_inputs[key] = tmpvar
T
typhoonzero 已提交
670

671
        # change output's ParamOut variable
672 673
        outputs = self._get_output_map_from_op(self.program.global_block().vars,
                                               opt_op)
674
        outputs["ParamOut"] = new_inputs["Param"]
675

676
        optimize_block.append_op(
677 678
            type=opt_op.type,
            inputs=new_inputs,
679
            outputs=outputs,
680 681
            attrs=opt_op.attrs)

682 683
    def _append_pserver_non_opt_ops(self, optimize_block, opt_op):
        program = optimize_block.program
684
        # Append the ops for parameters that do not need to be optimized/updated
685 686
        inputs = self._get_input_map_from_op(self.program.global_block().vars,
                                             opt_op)
687 688 689 690
        for varlist in inputs.itervalues():
            if not isinstance(varlist, list):
                varlist = [varlist]

691
            for var in varlist:
692 693
                if not program.global_block().vars.has_key(var.name):
                    program.global_block().create_var(
694 695 696 697 698 699 700 701
                        name=var.name,
                        persistable=var.persistable,
                        dtype=var.dtype,
                        shape=var.shape)

        outputs = self._get_output_map_from_op(self.program.global_block().vars,
                                               opt_op)

702 703 704 705 706
        for varlist in outputs.itervalues():
            if not isinstance(varlist, list):
                varlist = [varlist]

            for var in varlist:
T
update  
typhoonzero 已提交
707
                program.global_block().clone_variable(var)
708

709
        optimize_block.append_op(
710
            type=opt_op.type,
711 712
            inputs=inputs,
            outputs=outputs,
713 714
            attrs=opt_op.attrs)

715 716 717 718
    def _is_op_connected(self, op1, op2):
        # If one op's input is another op's output or
        # one op's output is another op's input, we say
        # the two operator is connected.
T
typhoonzero 已提交
719 720 721 722 723 724 725 726 727 728 729 730 731
        def _append_inname_remove_beta(varname_list):
            op_input_names = []
            for in_name in varname_list:
                # HACK: remove beta1 and beta2 to avoid let all
                # ops connected.
                if in_name.startswith("beta2_pow_acc") or \
                    in_name.startswith("beta1_pow_acc"):
                    continue
                else:
                    op_input_names.append(in_name)
            return op_input_names

        op1_input_names = _append_inname_remove_beta(op1.desc.input_arg_names())
732 733
        op1_output_names = op1.desc.output_arg_names()

T
typhoonzero 已提交
734
        op2_input_names = _append_inname_remove_beta(op2.desc.input_arg_names())
735
        op2_output_names = op2.desc.output_arg_names()
736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754

        if set(op1_output_names) & set(op2_input_names) or \
           set(op1_input_names) & set(op2_output_names):
            return True
        return False

    def _create_ufind(self, optimize_ops):
        # Create a unit find data struct by optimize ops
        ufind = UnionFind(optimize_ops)
        for i in xrange(len(optimize_ops)):
            for j in xrange(i, len(optimize_ops)):
                op1 = optimize_ops[i]
                op2 = optimize_ops[j]
                if self._is_op_connected(op1, op2):
                    ufind.union(op1, op2)
        return ufind

    def _is_opt_op(self, op):
        # NOTE: It's a HACK implement.
755
        # optimize op: SGDOptimize, MomentumOptimizer, AdamOptimizer and etc...
T
typhoonzero 已提交
756 757
        if "Param" in op.input_names and \
            "LearningRate" in op.input_names:
758 759 760 761 762 763 764
            return True
        return False

    def _is_opt_op_on_pserver(self, endpoint, op):
        param_names = [
            p.name for p in self.param_grad_ep_mapping[endpoint]["params"]
        ]
T
typhoonzero 已提交
765
        if op.input("Param")[0] in param_names:
766 767 768
            return True
        else:
            for n in param_names:
769
                param = op.input("Param")[0]
T
typhoonzero 已提交
770
                if same_or_split_var(n, param) and n != param:
771 772 773 774
                    return True
            return False
        return False

775 776 777 778 779 780 781 782 783 784 785 786 787 788 789 790 791 792 793 794 795 796 797
    def _get_input_map_from_op(self, varmap, op):
        iomap = dict()
        for key in op.input_names:
            vars = []
            for varname in op.input(key):
                vars.append(varmap[varname])
            if len(vars) == 1:
                iomap[key] = vars[0]
            else:
                iomap[key] = vars
        return iomap

    def _get_output_map_from_op(self, varmap, op):
        iomap = dict()
        for key in op.output_names:
            vars = []
            for varname in op.output(key):
                vars.append(varmap[varname])
            if len(vars) == 1:
                iomap[key] = vars[0]
            else:
                iomap[key] = vars
        return iomap
798 799 800 801 802 803 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819 820 821 822 823 824 825 826 827

    def _get_lr_ops(self):
        lr_ops = []
        # find learning rate variables by optimize op
        lr_vars = set()
        for op in self.optimize_ops:
            if self._is_opt_op(op):
                lr_vars.add(op.input("LearningRate")[0])

        find_ops = []
        # find ops which output is lr var
        block = self.program.global_block()
        for op in block.ops:
            if set(op.output_arg_names) & lr_vars:
                find_ops.append(op)
        # make a union find struct by the ops in default_main_program
        ufind = UnionFind(block.ops)
        for op1 in block.ops:
            for op2 in block.ops:
                # NOTE: we need to skip all optimize ops, since it is connected
                # with forward/backward ops and lr ops, we only need the lr ops.
                if op1 != op2 and self._is_op_connected(op1, op2) and \
                    not self._is_opt_op(op1) and not self._is_opt_op(op2):
                    ufind.union(op1, op2)
        # find all ops which is related with lr var
        for op1 in block.ops:
            for op2 in find_ops:
                if ufind.is_connected(op1, op2):
                    lr_ops.append(op1)
        return lr_ops
反馈
建议
客服 返回
顶部