distribute_transpiler.py 22.7 KB
Newer Older
D
dzhwinter 已提交
1
#   Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserve.
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 17 18 19
import framework
from framework import Program, default_main_program, Parameter, Variable
import optimizer
from layer_helper import LayerHelper
T
typhoonzero 已提交
20
from distributed_spliter import *
T
typhoonzero 已提交
21
import math
22
from . import core
T
done  
typhoonzero 已提交
23 24


T
typhoonzero 已提交
25 26 27 28 29 30
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 已提交
31

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


36 37 38 39
def same_or_split_var(p_name, var_name):
    return p_name == var_name or p_name.startswith(var_name + ".block")


T
typhoonzero 已提交
40 41 42 43 44
def split_dense_variable(var_list,
                         pserver_count,
                         min_block_size=1024,
                         max_block_size=1048576):
    """
45
        We may need to split dense tensor to one or more blocks and put
T
typhoonzero 已提交
46 47
        them equally onto parameter server. One block is a sub-tensor
        aligned by dim[0] of the tensor.
48

T
typhoonzero 已提交
49 50
        We need to have a minimal block size so that the calculations in
        the parameter server side can gain better performance. By default
51 52
        minimum block size is 1024. The max block size is used to prevent
        very large blocks that may cause send error.
T
typhoonzero 已提交
53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70
    """
    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
71
        # update split_count after aligning
T
typhoonzero 已提交
72 73 74 75 76 77 78 79 80
        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 已提交
81 82 83 84 85 86 87 88 89
class DistributeTranspiler:
    def transpile(self,
                  optimize_ops,
                  params_grads,
                  program=None,
                  pservers="127.0.0.1:6174",
                  trainers=1,
                  split_method=round_robin):
        """
90 91
            Transpile the program to distributed data-parallelism programs.
            The main_program will be transformed to use a remote parameter server
T
done  
typhoonzero 已提交
92
            to do parameter optimization. And the optimization graph will be put
93
            into a parameter server program.
T
done  
typhoonzero 已提交
94

95
            Use different methods to split trainable variables to different
T
done  
typhoonzero 已提交
96 97 98 99 100
            parameter servers.

            :param optimize_ops: op list of optimization, should be the
                                 return value of Optimizer.minimize
            :type optimize_ops: list
101
            :param program: program to optimize, default is default_main_program
T
done  
typhoonzero 已提交
102 103 104 105
            :param pservers: parameter server endpoints like "m1:6174,m2:6174"
            :type pservers: string
            :return: return a list of programs
        """
T
typhoonzero 已提交
106
        assert (callable(split_method))
T
done  
typhoonzero 已提交
107 108
        if program is None:
            program = default_main_program()
T
typhoonzero 已提交
109
        self.program = program
T
done  
typhoonzero 已提交
110
        self.trainers = trainers
T
typhoonzero 已提交
111
        self.optimize_ops = optimize_ops
T
typhoonzero 已提交
112
        # steps to transpile:
113
        # 1. split variable to multiple blocks, aligned by product(dim[1:]) (width).
T
typhoonzero 已提交
114 115 116
        # 2. modify trainer program add split_op to each Grad.
        # 3. append send_op to trainer.
        # 4. append concat_op to trainer to update local weights.
117
        # 5. create new program for parameter server.
T
typhoonzero 已提交
118
        # 6. create parameter server program by split_method generated endpoint->VarBlock
T
typhoonzero 已提交
119

T
typhoonzero 已提交
120
        pserver_endpoints = pservers.split(",")
T
typhoonzero 已提交
121 122

        # step1
T
typhoonzero 已提交
123 124
        param_list = [pg[0] for pg in params_grads]
        grad_list = [pg[1] for pg in params_grads]
T
typhoonzero 已提交
125
        # TODO: add split selected rows support
T
typhoonzero 已提交
126 127
        grad_blocks = split_dense_variable(grad_list, len(pserver_endpoints))
        param_blocks = split_dense_variable(param_list, len(pserver_endpoints))
T
typhoonzero 已提交
128
        # step2
T
typhoonzero 已提交
129
        grad_var_mapping = self._append_split_op(program, grad_blocks)
T
typhoonzero 已提交
130 131 132

        # step3
        send_inputs = []
T
typhoonzero 已提交
133
        send_outputs = []
T
typhoonzero 已提交
134 135 136 137
        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 已提交
138 139
        param_var_mapping = self._create_vars_from_blocklist(program,
                                                             param_blocks)
T
typhoonzero 已提交
140 141 142
        for b in param_blocks:
            varname, block_id, _ = b.split(":")
            send_outputs.append(param_var_mapping[varname][int(block_id)])
143 144
        # let send_op know which endpoint to send which var to, eplist has the same
        # order as send_inputs.
T
typhoonzero 已提交
145
        eplist = split_method(send_inputs, pserver_endpoints)
146
        # create mapping of endpoint -> split var to create pserver side program
T
typhoonzero 已提交
147 148 149 150 151 152 153 154
        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 已提交
155

156
        # create send_op
T
typhoonzero 已提交
157 158 159 160 161
        send_op = program.global_block().append_op(
            type="send",
            inputs={"X": send_inputs},
            outputs={"Out": send_outputs},
            attrs={"endpoints": pserver_endpoints,
T
typhoonzero 已提交
162 163 164
                   "epmap": eplist})
        # step4
        for varname, splited_var in param_var_mapping.iteritems():
T
typhoonzero 已提交
165 166
            if len(splited_var) <= 1:
                continue
T
typhoonzero 已提交
167 168 169
            orig_param = program.global_block().vars[varname]
            concat = program.global_block().append_op(
                type="concat",
T
typhoonzero 已提交
170
                inputs={"X": splited_var},
T
typhoonzero 已提交
171
                outputs={"Out": [orig_param]},
T
typhoonzero 已提交
172
                attrs={"axis": 0})
T
typhoonzero 已提交
173 174

    def _create_vars_from_blocklist(self, program, block_list):
175
        # Create respective variables using the block_list
T
typhoonzero 已提交
176
        block_map = dict()
T
typhoonzero 已提交
177
        var_mapping = dict()
T
typhoonzero 已提交
178 179 180 181 182 183 184
        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():
            orig_var = program.global_block().vars[varname]
T
typhoonzero 已提交
185 186 187 188
            var_mapping[varname] = []
            if len(splited) == 1:
                var_mapping[varname] = [orig_var]
                continue
T
typhoonzero 已提交
189 190 191 192
            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 已提交
193

T
typhoonzero 已提交
194
            for i, block in enumerate(splited):
T
typhoonzero 已提交
195
                size = block[1]
T
typhoonzero 已提交
196 197 198 199
                rows = size / orig_dim1_flatten
                splited_shape = [rows]
                if len(orig_shape) >= 2:
                    splited_shape.extend(orig_shape[1:])
T
typhoonzero 已提交
200 201 202 203
                var = program.global_block().create_var(
                    name="%s.block%d" % (varname, i),
                    psersistable=False,
                    dtype=orig_var.dtype,
T
typhoonzero 已提交
204
                    shape=splited_shape)  # flattend splited var
T
typhoonzero 已提交
205
                var_mapping[varname].append(var)
T
typhoonzero 已提交
206
        return var_mapping
T
done  
typhoonzero 已提交
207 208 209 210 211 212 213 214 215

    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,
216
            # HACK: let all param in pserver be persistable so the child
T
typhoonzero 已提交
217 218
            # program in recv can get them
            persistable=True)
T
done  
typhoonzero 已提交
219

T
typhoonzero 已提交
220
    def _append_split_op(self, program, gradblocks):
221
        # Split variables that need to be split and append respective ops
T
typhoonzero 已提交
222 223
        var_mapping = self._create_vars_from_blocklist(program, gradblocks)
        for varname, splited_vars in var_mapping.iteritems():
T
typhoonzero 已提交
224 225
            # variable that don't need to split have empty splited_vars
            if len(splited_vars) <= 1:
T
typhoonzero 已提交
226
                continue
T
typhoonzero 已提交
227
            orig_var = program.global_block().vars[varname]
T
typhoonzero 已提交
228
            if orig_var.type == core.VarDesc.VarType.SELECTED_ROWS:
229 230 231 232 233 234 235 236
                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 已提交
237
            elif orig_var.type == core.VarDesc.VarType.LOD_TENSOR:
238 239 240 241 242 243 244 245 246 247 248 249
                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 已提交
250
        return var_mapping
T
done  
typhoonzero 已提交
251

T
typhoonzero 已提交
252
    def get_trainer_program(self):
T
typhoonzero 已提交
253
        # remove optimize ops and add a send op to main_program
T
typhoonzero 已提交
254 255
        self.program.global_block().delete_ops(self.optimize_ops)
        return self.program
T
typhoonzero 已提交
256

T
done  
typhoonzero 已提交
257
    def _create_var_for_trainers(self, block, var, trainers):
258
        # For each trainer, create the necessary variables
T
done  
typhoonzero 已提交
259 260 261 262 263 264 265 266 267 268
        var_list = []
        for i in xrange(trainers):
            var_each = block.create_var(
                name="%s.trainer_%d" % (var.name, i),
                psersistable=var.persistable,
                dtype=var.dtype,
                shape=var.shape)
            var_list.append(var_each)
        return var_list

T
typhoonzero 已提交
269 270 271 272
    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
273
        Param and Grad is split to multiple servers.
T
typhoonzero 已提交
274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295
        """
        # 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 已提交
296 297 298 299 300 301 302 303 304 305 306 307 308 309
    def _is_op_on_pserver(self, endpoint, all_ops, idx):
        """
        Recursively check if the op need to run on current server.
        Assume that ops are in the execution order.
        """
        param_names = [
            p.name for p in self.param_grad_ep_mapping[endpoint]["params"]
        ]
        op = all_ops[idx]
        if op.inputs.has_key("Param"):
            if op.inputs["Param"].name in param_names:
                return True
            else:
                for n in param_names:
310 311
                    if same_or_split_var(n, op.inputs[
                            "Param"].name) and n != op.inputs["Param"].name:
T
typhoonzero 已提交
312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334
                        return True
                return False
        else:
            j = idx - 1
            while j >= 0:
                prev_op = all_ops[j]
                prev_output_names = [o.name for o in prev_op.outputs.values()]
                prev_input_names = [o.name for o in prev_op.inputs.values()]
                found1 = False
                found2 = False
                for _, v in op.inputs.iteritems():
                    if v.name in prev_output_names:
                        found1 = self._is_op_on_pserver(endpoint, all_ops, j)
                # later ops may produce output for prev op's next batch use.
                for _, v in op.outputs.iteritems():
                    if v.name in prev_input_names:
                        found2 = self._is_op_on_pserver(endpoint, all_ops, j)
                if found1 or found2:
                    return True
                j -= 1
            return False

    def _append_pserver_ops(self, program, pserver_program, opt_op, endpoint):
T
typhoonzero 已提交
335
        new_inputs = dict()
T
typhoonzero 已提交
336 337
        # update param/grad shape first, then other inputs like
        # moment can use the updated shape
T
typhoonzero 已提交
338 339 340 341
        for key, var in opt_op.inputs.iteritems():
            if key == "Grad":
                grad_block = None
                for g in self.param_grad_ep_mapping[endpoint]["grads"]:
342
                    if same_or_split_var(g.name, var.name):
T
typhoonzero 已提交
343 344 345 346 347 348
                        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 已提交
349
                merged_var = program.global_block().create_var(
T
typhoonzero 已提交
350 351 352 353 354 355
                    name=grad_block.name,
                    persistable=grad_block.persistable,
                    dtype=grad_block.dtype,
                    shape=grad_block.shape)
                # append merging ops if trainers > 1
                if self.trainers > 1:
T
done  
typhoonzero 已提交
356
                    vars2merge = self._create_var_for_trainers(
T
typhoonzero 已提交
357 358
                        program.global_block(), grad_block, self.trainers)
                    program.global_block().append_op(
T
done  
typhoonzero 已提交
359 360 361
                        type="sum",
                        inputs={"X": vars2merge},
                        outputs={"Out": merged_var})
T
typhoonzero 已提交
362
                    program.global_block().append_op(
T
done  
typhoonzero 已提交
363 364 365 366
                        type="scale",
                        inputs={"X": merged_var},
                        outputs={"Out": merged_var},
                        attrs={"scale": 1.0 / float(self.trainers)})
T
typhoonzero 已提交
367 368 369 370 371
                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"]:
372
                    if same_or_split_var(p.name, var.name):
T
typhoonzero 已提交
373 374 375 376
                        param_block = p
                        break
                if not param_block:
                    return
T
typhoonzero 已提交
377
                tmpvar = program.global_block().create_var(
T
typhoonzero 已提交
378
                    name=param_block.name,
T
typhoonzero 已提交
379
                    persistable=True,
T
typhoonzero 已提交
380 381
                    dtype=param_block.dtype,
                    shape=param_block.shape)
T
typhoonzero 已提交
382

T
typhoonzero 已提交
383
                new_inputs[key] = tmpvar
T
typhoonzero 已提交
384 385 386 387 388 389 390 391 392 393 394 395 396 397

        for key, var in opt_op.inputs.iteritems():
            if key in ["Param", "Grad"]:
                continue
            # 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)
            tmpvar = program.global_block().create_var(
                name=var.name,
                persistable=var.persistable,
                dtype=var.dtype,
                shape=new_shape)
            new_inputs[key] = tmpvar
T
typhoonzero 已提交
398 399 400 401 402 403 404 405
            # create var in pserver program global block.
            # TODO(typhoonzero): put blocks in one program to avoid create two
            # variables.
            pserver_program.global_block().create_var(
                name=var.name,
                persistable=var.persistable,
                dtype=var.dtype,
                shape=new_shape)
T
typhoonzero 已提交
406

407
        # change output's ParamOut variable
T
typhoonzero 已提交
408
        opt_op.outputs["ParamOut"] = new_inputs["Param"]
T
typhoonzero 已提交
409
        program.global_block().append_op(
T
typhoonzero 已提交
410 411 412 413 414
            type=opt_op.type,
            inputs=new_inputs,
            outputs=opt_op.outputs,
            attrs=opt_op.attrs)

T
typhoonzero 已提交
415
    def _append_pserver_non_opt_ops(self, program, pserver_program, opt_op):
416
        # Append the ops for parameters that do not need to be optimized/updated
T
typhoonzero 已提交
417
        for _, var in opt_op.inputs.iteritems():
T
typhoonzero 已提交
418
            program.global_block().create_var(
T
typhoonzero 已提交
419 420 421 422
                name=var.name,
                persistable=var.persistable,
                dtype=var.dtype,
                shape=var.shape)
T
typhoonzero 已提交
423 424 425 426 427
            pserver_program.global_block().create_var(
                name=var.name,
                persistable=var.persistable,
                dtype=var.dtype,
                shape=var.shape)
T
typhoonzero 已提交
428
        program.global_block().append_op(
T
typhoonzero 已提交
429
            type=opt_op.type,
T
typhoonzero 已提交
430
            inputs=opt_op.inputs,
T
typhoonzero 已提交
431 432 433
            outputs=opt_op.outputs,
            attrs=opt_op.attrs)

434
    def get_pserver_program(self, endpoint):
T
typhoonzero 已提交
435
        """
436
        Get pserver side program using the endpoint
T
typhoonzero 已提交
437 438 439 440 441 442 443 444 445

        NOTE: assume blocks of the same variable is not distributed
        on the same pserver, only change param/grad varnames for
        trainers to fetch. For each pserver endpoint, server side
        program must be a sub-set of the original optimization program.
        """
        # step5
        pserver_program = Program()
        for v in self.param_grad_ep_mapping[endpoint]["params"]:
T
typhoonzero 已提交
446
            self._clone_var(pserver_program.global_block(), v)
T
typhoonzero 已提交
447 448 449 450 451 452 453 454 455 456 457 458 459
        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.
            pserver_program.global_block().create_var(
                name=v.name, persistable=True, dtype=v.dtype, shape=v.shape)
            for trainer_id in xrange(self.trainers):
                print("create variable for program: %s.trainer_%d" %
                      (v.name, trainer_id))
                pserver_program.global_block().create_var(
                    name="%s.trainer_%d" % (v.name, trainer_id),
                    persistable=True,
                    dtype=v.dtype,
                    shape=v.shape)
T
typhoonzero 已提交
460 461
        # step6
        optimize_sub_program = Program()
462
        # Iterate through the ops and append ops as needed
463 464 465
        for idx, opt_op in enumerate(self.optimize_ops):
            is_op_on_pserver = self._is_op_on_pserver(endpoint,
                                                      self.optimize_ops, idx)
T
typhoonzero 已提交
466 467
            if not is_op_on_pserver:
                continue
T
typhoonzero 已提交
468
            if opt_op.inputs.has_key("Grad"):
T
typhoonzero 已提交
469 470
                self._append_pserver_ops(optimize_sub_program, pserver_program,
                                         opt_op, endpoint)
T
typhoonzero 已提交
471
            else:
T
typhoonzero 已提交
472 473
                self._append_pserver_non_opt_ops(optimize_sub_program,
                                                 pserver_program, opt_op)
474
        # Append the recv op
T
done  
typhoonzero 已提交
475 476
        pserver_program.global_block().append_op(
            type="recv",
T
typhoonzero 已提交
477 478
            inputs={"RX": self.param_grad_ep_mapping[endpoint]["grads"]
                    },  # grads to recv
T
done  
typhoonzero 已提交
479 480
            outputs={},
            attrs={
481
                "OptimizeBlock": optimize_sub_program.global_block(),
T
done  
typhoonzero 已提交
482
                "endpoint": endpoint,
T
typhoonzero 已提交
483 484 485 486 487 488 489 490
                "ParamList": [
                    p.name
                    for p in self.param_grad_ep_mapping[endpoint]["params"]
                ],
                "GradList": [
                    p.name
                    for p in self.param_grad_ep_mapping[endpoint]["grads"]
                ],
T
typhoonzero 已提交
491
                "Fanin": self.trainers
T
done  
typhoonzero 已提交
492 493 494
            })
        pserver_program.sync_with_cpp()
        return pserver_program
T
typhoonzero 已提交
495

T
typhoonzero 已提交
496
    def get_startup_program(self, endpoint, pserver_program):
T
typhoonzero 已提交
497 498 499
        """
        Get startup program for current parameter server.
        Modify operator input variables if there are variables that
500
        were split to several blocks.
T
typhoonzero 已提交
501 502 503 504 505 506 507 508
        """
        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
509
                if same_or_split_var(pname, varname) and varname != pname:
T
typhoonzero 已提交
510 511 512
                    return pname, splited_param.shape
            return "", []

Y
update  
yi.wu 已提交
513 514
        # 1. create vars in pserver program to startup program
        pserver_vars = pserver_program.global_block().vars
T
typhoonzero 已提交
515
        created_var_map = dict()
Y
update  
yi.wu 已提交
516
        for _, var in pserver_vars.iteritems():
T
typhoonzero 已提交
517 518
            tmpvar = s_prog.global_block().create_var(
                name=var.name,
T
typhoonzero 已提交
519
                persistable=var.persistable,
T
typhoonzero 已提交
520 521 522 523 524 525 526
                dtype=var.dtype,
                shape=var.shape)
            created_var_map[var.name] = tmpvar

        # 2. rename op outputs
        for op in orig_s_prog.global_block().ops:
            new_outputs = dict()
Y
update  
yi.wu 已提交
527 528
            # do not append startup op if var is not on this pserver
            op_on_pserver = False
T
typhoonzero 已提交
529 530 531
            for key, var in op.outputs.iteritems():
                newname, _ = _get_splited_name_and_shape(var.name)
                if newname:
Y
update  
yi.wu 已提交
532
                    op_on_pserver = True
T
typhoonzero 已提交
533
                    new_outputs[key] = created_var_map[newname]
Y
update  
yi.wu 已提交
534
                elif var.name in pserver_vars:
T
typhoonzero 已提交
535
                    op_on_pserver = True
Y
update  
yi.wu 已提交
536 537
                    new_outputs[key] = pserver_vars[var.name]

T
typhoonzero 已提交
538
            if op_on_pserver:
T
typhoonzero 已提交
539 540 541
                if op.type in [
                        "gaussian_random", "fill_constant", "uniform_random"
                ]:
T
typhoonzero 已提交
542
                    op.attrs["shape"] = new_outputs["Out"].shape
T
typhoonzero 已提交
543 544 545 546 547 548
                s_prog.global_block().append_op(
                    type=op.type,
                    inputs=op.inputs,
                    outputs=new_outputs,
                    attrs=op.attrs)
        return s_prog