fleet_base.py 37.6 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
#   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import print_function
16
import copy
17
import warnings
18
import paddle
19
import os
20
from paddle.fluid.framework import dygraph_only
21
from paddle.fluid import compiler
22
from .role_maker import UserDefinedRoleMaker, PaddleCloudRoleMaker, RoleMakerBase
23
from .strategy_compiler import StrategyCompiler
24
from .distributed_strategy import DistributedStrategy
25 26
from .meta_optimizer_factory import MetaOptimizerFactory
from .runtime_factory import RuntimeFactory
27
from paddle.fluid.wrapped_decorator import wrap_decorator
28
from paddle.fluid.dygraph import parallel_helper
29

30

31 32 33 34 35 36 37 38 39 40 41 42
def _inited_runtime_handler_(func):
    def __impl__(*args, **kwargs):
        cls = args[0]

        if cls._runtime_handle is None:
            raise ValueError("Fleet can not find suitable runtime handler")

        return func(*args, **kwargs)

    return __impl__


43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58
def _is_non_distributed_check_(func):
    def __impl__(*args, **kwargs):
        cls = args[0]

        if cls._role_maker is not None and cls._role_maker._is_non_distributed(
        ) is True:
            warnings.warn(
                "%s() function doesn't work when use non_distributed fleet." %
                (func.__name__))
            return

        return func(*args, **kwargs)

    return __impl__


59
inited_runtime_handler = wrap_decorator(_inited_runtime_handler_)
60
is_non_distributed_check = wrap_decorator(_is_non_distributed_check_)
61 62


63 64 65
class Fleet(object):
    """
    Unified API for distributed training of PaddlePaddle
66
    Please reference the https://github.com/PaddlePaddle/FleetX for details
67 68 69 70 71


    Returns:
        Fleet: A Fleet instance

72
    Example for collective training:
1
123malin 已提交
73

74 75
        .. code-block:: python

1
123malin 已提交
76 77
            import paddle
            paddle.enable_static()
78
            import paddle.distributed.fleet as fleet
79 80 81

            fleet.init(is_collective=True)

82 83 84
            strategy = fleet.DistributedStrategy()
            optimizer = paddle.optimizer.SGD(learning_rate=0.001)
            optimizer = fleet.distributed_optimizer(optimizer, strategy=strategy)
85 86 87 88 89 90 91 92

            # do distributed training


    Example for parameter server training:

        .. code-block:: python

1
123malin 已提交
93 94
            import paddle
            paddle.enable_static()
95 96
            import paddle.distributed.fleet as fleet
            strategy = fleet.DistributedStrategy()
S
ShenLiang 已提交
97
            fleet.init(strategy=strategy)
98

99
            optimizer = paddle.optimizer.SGD(learning_rate=0.001)
100
            optimizer = fleet.distributed_optimizer(optimizer)
101

102 103
            if fleet.is_first_worker():
                print("this is first worker")
104

105 106
            print("current node index: {}".format(fleet.worker_index()))
            print("total number of worker num: {}".format(fleet.worker_num()))
107

108 109 110
            if fleet.is_worker():
                print("this is worker")
            print("worker endpoints: {}".format(fleet.worker_endpoints(to_string=True)))
111

112 113
            print("server num: {}".format(fleet.server_num()))
            print("server endpoints: {}".format(fleet.server_endpoints(to_string=True)))
114

115 116 117
            if fleet.is_server():
                print("this is server")
            fleet.stop_worker()
118 119


120 121 122
    """

    def __init__(self):
123
        self._role_maker = None
124
        self.strategy_compiler = None
125
        self._is_collective = False
126
        self._runtime_handle = None
D
Dong Daxiang 已提交
127 128
        self._util = None
        self._context = {}
129

130
    def init(self, role_maker=None, is_collective=False, strategy=None):
131 132 133
        """
        Initialize role_maker in Fleet.

134 135 136 137 138 139 140 141 142 143 144
        This function is responsible for the distributed architecture
        what you want to run your code behind.

        Args:
            role_maker (RoleMakerBase, optional): A ``RoleMakerBase`` containing the configuration
                of environment variables related to distributed training.If you did not initialize 
                the rolemaker by yourself, it will be automatically initialized to PaddleRoleMaker.
                The default value is None.
            is_collective (Boolean, optional): A ``Boolean`` variable determines whether the program 
                runs on the CPU or GPU. False means set distributed training using CPU, and True means
                GPU.The default value is False.The default value is False.
145 146 147 148
            strategy (DistributedStrategy): Extra properties for distributed training. 
                For details, please refer to paddle.distributed.fleet.DistributedStrategy. Default: None.


149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170
        Returns:
            None

        Examples1:

            .. code-block:: python

                import paddle.distributed.fleet as fleet
                fleet.init()

        Examples2:

            .. code-block:: python

                import paddle.distributed.fleet as fleet
                fleet.init(is_collective=True)

        Examples3:

            .. code-block:: python

                import paddle.distributed.fleet as fleet
1
123malin 已提交
171
                role = fleet.PaddleCloudRoleMaker()
172
                fleet.init(role)
173

174 175 176 177 178 179
        Examples4:

            .. code-block:: python

                import paddle.distributed.fleet as fleet
                strategy = fleet.DistributedStrategy()
S
ShenLiang 已提交
180
                fleet.init(strategy=strategy)
181

182
        """
S
ShenLiang 已提交
183 184 185
        if strategy is None:
            strategy = DistributedStrategy()
        self._user_defined_strategy = copy.deepcopy(strategy)
186 187

        if role_maker is None:
188 189 190 191 192 193
            if isinstance(is_collective, bool):
                self._is_collective = is_collective
                self._role_maker = PaddleCloudRoleMaker(
                    is_collective=self._is_collective)
            else:
                raise ValueError(
194 195
                    "`is_collective` should be instance of `bool`, but got {}".
                    format(type(is_collective)))
196
        else:
197 198 199 200 201 202
            if isinstance(role_maker, RoleMakerBase):
                self._role_maker = role_maker
            else:
                raise ValueError(
                    "`role_maker` should be subclass of `RoleMakerBase`, but got {}".
                    format(type(role_maker)))
203
        self._role_maker._generate_role()
204

205 206 207
        import paddle.distributed.fleet as fleet
        fleet.util._set_role_maker(self._role_maker)

208
        self.strategy_compiler = StrategyCompiler()
209 210 211 212 213 214 215 216 217

        if self._role_maker._is_non_distributed() and self._is_collective:
            if paddle.fluid.core.is_compiled_with_cuda():
                gpus_num = paddle.fluid.core.get_cuda_device_count()
                if gpus_num != 1:
                    raise ValueError(
                        "CUDA_VISIBLE_DEVICES shoule be set only 1 card if you use `python` to launch fleet program."
                    )

218
        if paddle.fluid.framework.in_dygraph_mode():
219 220
            if self.worker_num() == 1:
                return
221 222 223 224
            if parallel_helper._is_parallel_ctx_initialized():
                warnings.warn(
                    "The dygraph parallel environment has been initialized.")
            else:
225 226 227 228 229 230 231 232 233
                # FLAGS_nccl_nrings is used for dynamic graph multi-stream communication
                if "FLAGS_nccl_nrings" in os.environ:
                    warnings.warn(
                        "You have set the environment variable FLAGS_nccl_nrings "
                        "outside the program, so the nccl_comm_num in "
                        "DistributedStrategy will not take effect here.")
                else:
                    os.environ["FLAGS_nccl_nrings"] = str(
                        self._user_defined_strategy.nccl_comm_num)
234
                paddle.distributed.init_parallel_env()
235 236 237 238 239 240 241 242

    def is_first_worker(self):
        """
        Check whether the node is the first instance of worker.

        Returns:
            bool: True if this is the first node of worker,
                  False if not.
243

244 245 246 247 248 249 250 251
        Examples:

            .. code-block:: python

                import paddle.distributed.fleet as fleet
                fleet.init()
                fleet.is_first_worker()

252
        """
253
        return self._role_maker._is_first_worker()
254 255 256 257 258 259 260

    def worker_index(self):
        """
        Get current worker index.

        Returns:
            int: node id
261 262 263 264

        Examples:

            .. code-block:: python
1
123malin 已提交
265

266 267 268 269
                import paddle.distributed.fleet as fleet
                fleet.init()
                fleet.worker_index()

270
        """
271
        return self._role_maker._worker_index()
272 273 274 275 276 277 278

    def worker_num(self):
        """
        Get current total worker number.

        Returns:
            int: worker numbers
1
123malin 已提交
279

280
        Examples:
1
123malin 已提交
281

282 283 284 285 286 287
            .. code-block:: python

                import paddle.distributed.fleet as fleet
                fleet.init()
                fleet.worker_num()

288
        """
289
        return self._role_maker._worker_num()
290

G
gongweibao 已提交
291 292 293 294 295 296
    def node_num(self):
        return self._role_maker._get_node_num()

    def rank_in_node(self):
        return self._role_maker._get_rank_in_node()

297 298 299 300 301 302 303
    def is_worker(self):
        """
        Check whether the node is an instance of worker.

        Returns:
            bool: True if this is a node of worker,
                  False if not.
304 305

        Examples:
1
123malin 已提交
306

307 308 309 310 311 312
            .. code-block:: python

                import paddle.distributed.fleet as fleet
                fleet.init()
                fleet.is_worker()

313
        """
314
        return self._role_maker._is_worker()
315 316 317

    def worker_endpoints(self, to_string=False):
        """
318
        Get current worker endpoints, such as ["127.0.0.1:1001", "127.0.0.1:1002"].
319 320 321

        Returns:
            list/string: server endpoints
322 323

        Examples:
1
123malin 已提交
324

325 326 327 328 329 330
            .. code-block:: python

                import paddle.distributed.fleet as fleet
                fleet.init()
                fleet.worker_endpoints()

331 332
        """
        if to_string:
333
            return ",".join(self._role_maker._get_trainer_endpoints())
334
        else:
335
            return self._role_maker._get_trainer_endpoints()
336 337 338 339 340 341 342

    def server_num(self):
        """
        Get current total worker number.

        Returns:
            int: server number
343 344

        Examples:
1
123malin 已提交
345

346
            .. code-block:: python
1
123malin 已提交
347 348 349 350

                import paddle.distributed.fleet as fleet
                fleet.init()
                fleet.server_num()
351
        """
352
        return len(self._role_maker._get_pserver_endpoints())
353 354 355 356 357 358 359

    def server_index(self):
        """
        Get current server index.

        Returns:
            int: node id
360 361

        Examples:
1
123malin 已提交
362

363 364 365 366 367 368
            .. code-block:: python

                import paddle.distributed.fleet as fleet
                fleet.init()
                fleet.server_index()

369
        """
370
        return self._role_maker._server_index()
371 372 373 374 375 376 377

    def server_endpoints(self, to_string=False):
        """
        Get current server endpoints, such as ["127.0.0.1:1001", "127.0.0.1:1002"].

        Returns:
            list/string: server endpoints
378 379

        Examples:
1
123malin 已提交
380

381 382 383 384 385 386
            .. code-block:: python

                import paddle.distributed.fleet as fleet
                fleet.init()
                fleet.server_endpoints()

387
        """
388

389
        if to_string:
390
            return ",".join(self._role_maker._get_pserver_endpoints())
391
        else:
392
            return self._role_maker._get_pserver_endpoints()
393 394 395 396 397 398 399 400

    def is_server(self):
        """
        Check whether the node is an instance of server.

        Returns:
            bool: True if this is a node of server,
                  False if not.
401 402 403 404

        Examples:

            .. code-block:: python
1
123malin 已提交
405

406 407 408 409
                import paddle.distributed.fleet as fleet
                fleet.init()
                fleet.is_server()

410
        """
411
        return self._role_maker._is_server(
412
        ) or self._role_maker._is_heter_worker()
413 414 415

    def barrier_worker(self):
        """
416 417 418 419
        barrier all workers

        Returns:
            None
420
        """
421
        self._role_maker._barrier("worker")
422

423
    @is_non_distributed_check
424
    @inited_runtime_handler
425 426
    def init_worker(self):
        """
427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444
        initialize `Communicator` for parameter server training.


        Returns:
            None

        Examples:

            .. code-block:: python

                import paddle.distributed.fleet as fleet
                fleet.init()

                # build net
                # fleet.distributed_optimizer(...)

                fleet.init_worker()

445 446 447
        """
        self._runtime_handle._init_worker()

448
    @is_non_distributed_check
449
    @inited_runtime_handler
450
    def init_server(self, *args, **kwargs):
451
        """
452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470
        init_server executor to initialize startup program,
        if the `args` is not empty, it will run load_persistables for increment training.


        Returns:
            None

        Examples:

            .. code-block:: python

                import paddle.distributed.fleet as fleet
                fleet.init()

                # build net
                # fleet.distributed_optimizer(...)

                fleet.init_server()

471
        """
472
        self._runtime_handle._init_server(*args, **kwargs)
473

474
    @is_non_distributed_check
475
    @inited_runtime_handler
476 477
    def run_server(self):
        """
478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495
        run server will run pserver main program with executor.

        Returns:
            None

        Examples:

            .. code-block:: python

                import paddle.distributed.fleet as fleet
                fleet.init()

                # build net
                # fleet.distributed_optimizer(...)

                if fleet.is_server():
                    fleet.init_server()

496 497 498
        """
        self._runtime_handle._run_server()

499
    @is_non_distributed_check
500
    @inited_runtime_handler
501 502
    def stop_worker(self):
        """
503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519
        stop `Communicator` and give training complete notice to parameter server.

        Returns:
            None

        Examples:

            .. code-block:: python

                import paddle.distributed.fleet as fleet
                fleet.init()

                # build net
                # fleet.distributed_optimizer(...)

                fleet.init_server()

520 521 522
        """
        self._runtime_handle._stop_worker()

523 524 525 526 527 528 529
    def save_inference_model(self,
                             executor,
                             dirname,
                             feeded_var_names,
                             target_vars,
                             main_program=None,
                             export_for_deployment=True):
530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549
        """
        save inference model for inference.

        Returns:
            None

        Examples:

            .. code-block:: python

                import paddle.distributed.fleet as fleet
                fleet.init()

                # build net
                # fleet.distributed_optimizer(...)

                fleet.init_server()

        """

550 551 552 553
        self._runtime_handle._save_inference_model(
            executor, dirname, feeded_var_names, target_vars, main_program,
            export_for_deployment)

554
    def save_persistables(self, executor, dirname, main_program=None, mode=0):
555 556
        """

1
123malin 已提交
557
        saves all persistable tensors from :code:`main_program` to
558 559
        the folder :code:`dirname`. You can refer to

1
123malin 已提交
560 561
        The :code:`dirname` is used to specify the folder where persistable tensors
        are going to be saved. If you would like to save tensors in separate
562 563 564
        files, set :code:`filename` None.

        Args:
1
123malin 已提交
565
            executor(Executor): The executor to run for saving persistable tensors.
566 567 568 569 570
                                You can refer to :ref:`api_guide_executor_en` for
                                more details.

            dirname(str, optional): The saving directory path.
                                When you need to save the parameter to the memory, set it to None.
1
123malin 已提交
571
            main_program(Program, optional): The program whose persistbale tensors will
572 573 574 575 576 577 578 579 580 581
                                             be saved. Default: None.


        Returns:
            None

        Examples:

            .. code-block:: text

1
123malin 已提交
582 583
                import paddle
                paddle.enable_static()
584 585 586 587 588 589 590
                import paddle.distributed.fleet as fleet

                fleet.init()

                # build net
                # fleet.distributed_optimizer(...)

1
123malin 已提交
591 592
                exe = paddle.static.Executor(paddle.CPUPlace())
                fleet.save_persistables(exe, "dirname", paddle.static.default_main_program())
593 594 595

        """

596 597
        self._runtime_handle._save_persistables(executor, dirname, main_program,
                                                mode)
598

599
    def distributed_optimizer(self, optimizer, strategy=None):
600
        """
601 602 603 604 605 606 607
        Optimizer for distributed training.

        For the distributed training, this method would rebuild a new instance of DistributedOptimizer.
        Which has basic Optimizer function and special features for distributed training.

        Args:
            optimizer(Optimizer): The executor to run for init server.
608 609 610 611 612
            strategy(DistributedStrategy): Extra properties for distributed optimizer. 
                It is recommended to use DistributedStrategy in fleet.init(). The strategy
                here is for compatibility. If the strategy in fleet.distributed_optimizer() 
                is not None, then it will overwrite the DistributedStrategy in fleet.init(), 
                which will take effect in distributed training.
613

614
        Returns:
615
            Fleet: instance of fleet.
616 617

        Examples:
618

619
            .. code-block:: python
620

1
123malin 已提交
621
                import paddle
622
                import paddle.distributed.fleet as fleet
1
123malin 已提交
623
                fleet.init(is_collective=True)
624 625 626 627
                strategy = fleet.DistributedStrategy()
                optimizer = paddle.optimizer.SGD(learning_rate=0.001)
                optimizer = fleet.distributed_optimizer(optimizer, strategy=strategy)

628 629
        """
        self.user_defined_optimizer = optimizer
630

631 632
        if strategy is not None:
            warnings.warn(
S
ShenLiang 已提交
633 634 635 636
                "It is recommended to use DistributedStrategy "
                "in fleet.init(). The strategy here is only for compatibility. "
                "If the strategy in fleet.distributed_optimizer() is "
                "not None, then it will overwrite the DistributedStrategy in fleet.init(), "
637 638
                "which will take effect in distributed training.")
            self._user_defined_strategy = copy.deepcopy(strategy)
D
Dong Daxiang 已提交
639 640

        self._context = {}
641 642
        return self

643
    @dygraph_only
644
    def distributed_model(self, model):
645
        """
646 647 648 649 650 651 652
        Return distributed data parallel model (Only work in dygraph mode)

        Args:
            model (Layer): the user-defind model which inherits Layer.

        Returns:
            distributed data parallel model which inherits Layer.
653 654

        Examples:
655

656 657
            .. code-block:: python

658 659 660 661 662 663 664 665 666
                import paddle
                import paddle.nn as nn
                from paddle.distributed import fleet

                class LinearNet(nn.Layer):
                    def __init__(self):
                        super(LinearNet, self).__init__()
                        self._linear1 = nn.Linear(10, 10)
                        self._linear2 = nn.Linear(10, 1)
667

668 669
                    def forward(self, x):
                        return self._linear2(self._linear1(x))
670

1
123malin 已提交
671
                # 1. initialize fleet environment
672 673
                fleet.init(is_collective=True)

1
123malin 已提交
674
                # 2. create layer & optimizer
675 676 677 678 679
                layer = LinearNet()
                loss_fn = nn.MSELoss()
                adam = paddle.optimizer.Adam(
                    learning_rate=0.001, parameters=layer.parameters())

1
123malin 已提交
680
                # 3. get data_parallel model using fleet
681 682 683
                adam = fleet.distributed_optimizer(adam)
                dp_layer = fleet.distributed_model(layer)

1
123malin 已提交
684
                # 4. run layer
685 686 687 688 689 690 691 692 693 694 695 696
                inputs = paddle.randn([10, 10], 'float32')
                outputs = dp_layer(inputs)
                labels = paddle.randn([10, 1], 'float32')
                loss = loss_fn(outputs, labels)

                print("loss:", loss.numpy())

                loss.backward()

                adam.step()
                adam.clear_grad()

697

698 699
        """
        assert model is not None
700 701
        self.model = paddle.DataParallel(
            model,
702 703 704
            comm_buffer_size=self._user_defined_strategy.fuse_grad_size_in_MB,
            last_comm_buffer_size=self._user_defined_strategy.
            last_comm_group_size_MB)
705 706 707 708 709 710
        return self.model

    @dygraph_only
    def state_dict(self):
        """
        Get state dict information from optimizer.
711
        (Only work in dygraph mode)
712 713 714 715 716 717 718

        Returns: 
            state_dict(dict) : dict contains all the Tensor used by optimizer

        Examples:
            .. code-block:: python

719 720 721 722 723
                import numpy as np
                import paddle
                from paddle.distributed import fleet

                fleet.init(is_collective=True)
724

725
                value = np.arange(26).reshape(2, 13).astype("float32")
1
123malin 已提交
726
                a = paddle.to_tensor(value)
727

728 729
                layer = paddle.nn.Linear(13, 5)
                adam = paddle.optimizer.Adam(learning_rate=0.01, parameters=layer.parameters())
730

731 732 733
                adam = fleet.distributed_optimizer(adam)
                dp_layer = fleet.distributed_model(layer)
                state_dict = adam.state_dict()
734 735 736 737 738 739 740 741
        """
        # imitate target optimizer retrieval
        return self.user_defined_optimizer.state_dict()

    @dygraph_only
    def set_state_dict(self, state_dict):
        """
        Load optimizer state dict.
742
        (Only work in dygraph mode)
743 744 745 746

        Args: 
            state_dict(dict) : Dict contains all the Tensor needed by optimizer

747 748
        Returns:
            None
749 750 751 752

        Examples:
            .. code-block:: python

753 754 755
                import numpy as np
                import paddle
                from paddle.distributed import fleet
756

757 758 759
                fleet.init(is_collective=True)

                value = np.arange(26).reshape(2, 13).astype("float32")
1
123malin 已提交
760
                a = paddle.to_tensor(value)
761

762 763
                layer = paddle.nn.Linear(13, 5)
                adam = paddle.optimizer.Adam(learning_rate=0.01, parameters=layer.parameters())
764

765 766 767
                adam = fleet.distributed_optimizer(adam)
                dp_layer = fleet.distributed_model(layer)
                state_dict = adam.state_dict()
1
123malin 已提交
768 769 770
                paddle.save(state_dict, "paddle_dy")
                para_state_dict = paddle.load("paddle_dy")
                adam.set_state_dict(para_state_dict)
771 772 773 774 775 776 777 778
        """
        # imitate target optimizer retrieval
        return self.user_defined_optimizer.set_state_dict(state_dict)

    @dygraph_only
    def set_lr(self, value):
        """
        Set the value of the learning rate manually in the optimizer. 
779
        (Only work in dygraph mode)
780

781 782 783
        Args:
            value (float|Tensor): the value of learning rate

784 785
        Returns: 
            None 
786 787 788 789

        Examples:
            .. code-block:: python

790 791 792
                import numpy as np
                import paddle
                from paddle.distributed import fleet
793

794
                fleet.init(is_collective=True)
795

796
                value = np.arange(26).reshape(2, 13).astype("float32")
1
123malin 已提交
797
                a = paddle.to_tensor(value)
798

799 800
                layer = paddle.nn.Linear(13, 5)
                adam = paddle.optimizer.Adam(learning_rate=0.01, parameters=layer.parameters())
801

802 803 804 805 806 807 808 809 810 811 812 813 814 815
                adam = fleet.distributed_optimizer(adam)
                dp_layer = fleet.distributed_model(layer)

                lr_list = [0.2, 0.3, 0.4, 0.5, 0.6]
                for i in range(5):
                    adam.set_lr(lr_list[i])
                    lr = adam.get_lr()
                    print("current lr is {}".format(lr))
                # Print:
                #    current lr is 0.2
                #    current lr is 0.3
                #    current lr is 0.4
                #    current lr is 0.5
                #    current lr is 0.6
816 817 818 819 820 821 822 823
        """
        # imitate target optimizer retrieval
        return self.user_defined_optimizer.set_lr(value)

    @dygraph_only
    def get_lr(self):
        """
        Get current step learning rate.
824
        (Only work in dygraph mode)
825 826 827 828 829

        Returns:
            float: The learning rate of the current step.

        Examples:
1
123malin 已提交
830

831 832
            .. code-block:: python

833 834 835 836 837
                import numpy as np
                import paddle
                from paddle.distributed import fleet

                fleet.init(is_collective=True)
838

839
                value = np.arange(26).reshape(2, 13).astype("float32")
1
123malin 已提交
840
                a = paddle.to_tensor(value)
841

842 843
                layer = paddle.nn.Linear(13, 5)
                adam = paddle.optimizer.Adam(learning_rate=0.01, parameters=layer.parameters())
844

845 846
                adam = fleet.distributed_optimizer(adam)
                dp_layer = fleet.distributed_model(layer)
847

848 849
                lr = adam.get_lr()
                print(lr) # 0.01
850 851 852 853 854 855 856 857
        """
        # imitate target optimizer retrieval
        return self.user_defined_optimizer.get_lr()

    @dygraph_only
    def step(self):
        """
        Execute the optimizer once.
858
        (Only work in dygraph mode)
859

860 861
        Returns:
            None
862 863

        Examples:
1
123malin 已提交
864

865 866
            .. code-block:: python

867 868 869
                import paddle
                import paddle.nn as nn
                from paddle.distributed import fleet
870

871 872 873 874 875
                class LinearNet(nn.Layer):
                    def __init__(self):
                        super(LinearNet, self).__init__()
                        self._linear1 = nn.Linear(10, 10)
                        self._linear2 = nn.Linear(10, 1)
876

877 878
                    def forward(self, x):
                        return self._linear2(self._linear1(x))
879

1
123malin 已提交
880
                # 1. initialize fleet environment
881 882
                fleet.init(is_collective=True)

1
123malin 已提交
883
                # 2. create layer & optimizer
884 885 886 887 888
                layer = LinearNet()
                loss_fn = nn.MSELoss()
                adam = paddle.optimizer.Adam(
                    learning_rate=0.001, parameters=layer.parameters())

1
123malin 已提交
889
                # 3. get data_parallel model using fleet
890 891 892
                adam = fleet.distributed_optimizer(adam)
                dp_layer = fleet.distributed_model(layer)

1
123malin 已提交
893
                # 4. run layer
894 895 896 897 898 899 900 901 902 903 904 905 906 907 908 909 910 911 912 913
                inputs = paddle.randn([10, 10], 'float32')
                outputs = dp_layer(inputs)
                labels = paddle.randn([10, 1], 'float32')
                loss = loss_fn(outputs, labels)

                print("loss:", loss.numpy())

                loss.backward()

                adam.step()
                adam.clear_grad()


        """
        # imitate target optimizer retrieval
        return self.user_defined_optimizer.step()

    @dygraph_only
    def clear_grad(self):
        """
914 915
        Clear the gradients of all optimized parameters for model.
        (Only work in dygraph mode)
916

917 918
        Returns: 
            None
919 920

        Examples:
1
123malin 已提交
921

922 923
            .. code-block:: python

924 925 926
                import paddle
                import paddle.nn as nn
                from paddle.distributed import fleet
927

928 929 930 931 932
                class LinearNet(nn.Layer):
                    def __init__(self):
                        super(LinearNet, self).__init__()
                        self._linear1 = nn.Linear(10, 10)
                        self._linear2 = nn.Linear(10, 1)
933

934 935
                    def forward(self, x):
                        return self._linear2(self._linear1(x))
936

1
123malin 已提交
937
                # 1. initialize fleet environment
938 939
                fleet.init(is_collective=True)

1
123malin 已提交
940
                # 2. create layer & optimizer
941 942 943 944 945
                layer = LinearNet()
                loss_fn = nn.MSELoss()
                adam = paddle.optimizer.Adam(
                    learning_rate=0.001, parameters=layer.parameters())

1
123malin 已提交
946
                # 3. get data_parallel model using fleet
947 948 949
                adam = fleet.distributed_optimizer(adam)
                dp_layer = fleet.distributed_model(layer)

1
123malin 已提交
950
                # 4. run layer
951 952 953 954 955 956 957 958 959 960 961 962 963 964 965 966
                inputs = paddle.randn([10, 10], 'float32')
                outputs = dp_layer(inputs)
                labels = paddle.randn([10, 1], 'float32')
                loss = loss_fn(outputs, labels)

                print("loss:", loss.numpy())

                loss.backward()

                adam.step()
                adam.clear_grad()

        """
        # imitate target optimizer retrieval
        return self.user_defined_optimizer.clear_grad()

D
Dong Daxiang 已提交
967 968 969 970 971 972 973 974 975
    def _final_strategy(self):
        if "valid_strategy" not in self._context:
            print(
                "WARNING: You may need to call minimize function before this function is called"
            )
            return {}
        else:
            return self._context["valid_strategy"]

976 977 978 979 980 981 982 983 984 985 986 987 988 989 990 991 992 993
    def _get_applied_meta_list(self):
        if "applied_meta_list" not in self._context:
            print(
                "WARNING: You may need to call minimize function before _get_applied_meta_list called"
            )
            return []
        else:
            return self._context["applied_meta_list"]

    def _get_applied_graph_list(self):
        if "applied_graph_list" not in self._context:
            print(
                "WARNING: You may need to call minimize function before _get_applied_graph_list called"
            )
            return []
        else:
            return self._context["applied_graph_list"]

994 995 996 997 998 999 1000 1001 1002
    def minimize(self,
                 loss,
                 startup_program=None,
                 parameter_list=None,
                 no_grad_set=None):
        """
        Add distributed operations to minimize ``loss`` by updating ``parameter_list``.

        Args:
1
123malin 已提交
1003
            loss (Tensor): A ``Tensor`` containing the value to minimize.
1004 1005 1006
            startup_program (Program, optional): :ref:`api_fluid_Program` for
                initializing parameters in ``parameter_list``. The default value
                is None, at this time :ref:`api_fluid_default_startup_program` will be used.
1
123malin 已提交
1007
            parameter_list (Iterable, optional): Iterable of ``Tensor`` or ``Tensor.name`` to update
1008 1009
                to minimize ``loss``. The default value is None, at this time all parameters
                will be updated.
1
123malin 已提交
1010
            no_grad_set (set, optional): Set of ``Tensor``  or ``Tensor.name`` that don't need
1011 1012 1013 1014
                to be updated. The default value is None.

        Returns:
            tuple: tuple (optimize_ops, params_grads), A list of operators appended
1
123malin 已提交
1015
            by minimize and a list of (param, grad) tensor pairs, param is
1016
            ``Parameter``, grad is the gradient value corresponding to the parameter.
1017 1018
            The returned tuple can be passed to ``fetch_list`` in ``Executor.run()`` to
            indicate program pruning. If so, the program will be pruned by ``feed`` and
1019 1020 1021
            ``fetch_list`` before run, see details in ``Executor``.

        Examples:
1
123malin 已提交
1022

1023
            .. code-block:: python
1024

1025
                import paddle
1
123malin 已提交
1026
                paddle.enable_static()
1027
                import paddle.distributed.fleet as fleet
1
123malin 已提交
1028 1029 1030 1031 1032 1033 1034 1035 1036 1037 1038
                import paddle.nn.functional as F

                hid_dim = 10
                label_dim = 2
                input_x = paddle.static.data(name='x', shape=[None, 13], dtype='float32')
                input_y = paddle.static.data(name='y', shape=[None, 1], dtype='int64')
                fc_1 = paddle.static.nn.fc(x=input_x, size=hid_dim, activation='tanh')
                fc_2 = paddle.static.nn.fc(x=fc_1, size=hid_dim, activation='tanh')
                prediction = paddle.static.nn.fc(x=[fc_2], size=label_dim, activation='softmax')
                cost = F.cross_entropy(input=prediction, label=input_y)
                avg_cost = paddle.mean(x=cost)
1039

1
123malin 已提交
1040
                fleet.init(is_collective=True)
1041 1042 1043 1044
                strategy = fleet.DistributedStrategy()
                optimizer = paddle.optimizer.SGD(learning_rate=0.001)
                optimizer = fleet.distributed_optimizer(optimizer, strategy=strategy)
                optimizer.minimize(avg_cost)
1045

1046
                # for more examples, please reference https://github.com/PaddlePaddle/FleetX
1047 1048

        """
D
Dong Daxiang 已提交
1049 1050 1051
        context = {}
        context["user_defined_strategy"] = copy.deepcopy(
            self._user_defined_strategy)
1052 1053 1054
        if paddle.fluid.framework.in_dygraph_mode():
            # imitate target optimizer retrieval
            target_opt = self.user_defined_optimizer
D
Dong Daxiang 已提交
1055
            self._context = context
1056 1057
            return target_opt.minimize(loss)

1058 1059
        # cache original feed forward program
        self.origin_main_program = loss.block.program
1060 1061
        context["origin_main_program"] = self.origin_main_program
        context["loss"] = loss
1062 1063
        if startup_program == None:
            self.origin_startup_program = \
1064 1065
                paddle.static.default_startup_program().clone(for_test=False)
            startup_program = paddle.static.default_startup_program()
1066 1067 1068
        else:
            self.origin_startup_program = \
                startup_program.clone(for_test=False)
1069

1070 1071
        context["origin_startup_program"] = startup_program
        context["role_maker"] = self._role_maker
1072 1073 1074 1075 1076

        # compile time
        distributed_optimizer_list = \
            MetaOptimizerFactory()._get_valid_meta_optimizers(
                self.user_defined_optimizer)
D
Dong Daxiang 已提交
1077

D
Dong Daxiang 已提交
1078 1079 1080
        context["user_defined_strategy"] = copy.deepcopy(
            self._user_defined_strategy)
        copy_user_defined_strategy = copy.deepcopy(self._user_defined_strategy)
1081 1082 1083 1084 1085 1086

        # trigger the auto-parallel in very strict condition
        # strategy = DistributedStrategy()
        # strategy.auto = True
        # optimizer = paddle.optimizer.SGD(learning_rate=0.1)
        # optimizer = fleet.distributed_optimizer(optimizer, strategy)
D
Dong Daxiang 已提交
1087
        if copy_user_defined_strategy._is_strict_auto():
1088 1089
            # turn on all the strategy for each optimizer
            for opt in distributed_optimizer_list:
D
Dong Daxiang 已提交
1090
                opt._enable_strategy(copy_user_defined_strategy, context)
1091

1092 1093
        valid_optimizer_list = []
        valid_graph_optimizer_list = []
D
Dong Daxiang 已提交
1094
        can_not_apply_optimizer_list = []
1095 1096 1097 1098
        # recall meta optimizers for ranking
        for opt in distributed_optimizer_list:
            opt._set_basic_info(loss, self._role_maker,
                                self.user_defined_optimizer,
D
Dong Daxiang 已提交
1099
                                copy_user_defined_strategy)
1100 1101
            if opt._can_apply() and not opt._is_graph_out():
                valid_optimizer_list.append(opt)
D
Dong Daxiang 已提交
1102
            elif opt._can_apply() and opt._is_graph_out():
1103
                valid_graph_optimizer_list.append(opt)
D
Dong Daxiang 已提交
1104 1105
            else:
                can_not_apply_optimizer_list.append(opt)
1106
        # combine recalled meta optimizers to be a valid meta optimizer
D
Dong Daxiang 已提交
1107
        meta_optimizer, graph_optimizer = \
1108 1109
            self.strategy_compiler.generate_optimizer(
                loss, self._role_maker, self.user_defined_optimizer,
D
Dong Daxiang 已提交
1110
                copy_user_defined_strategy, valid_optimizer_list,
1111
                valid_graph_optimizer_list)
D
Dong Daxiang 已提交
1112

D
Dong Daxiang 已提交
1113
        valid_strategy = self.strategy_compiler._get_valid_strategy(
D
Dong Daxiang 已提交
1114 1115 1116
            copy_user_defined_strategy, can_not_apply_optimizer_list)

        context["valid_strategy"] = copy.deepcopy(valid_strategy)
1117

1118 1119 1120 1121 1122 1123
        applied_meta_list = self.strategy_compiler._get_applied_meta_list()
        applied_graph_list = self.strategy_compiler._get_applied_graph_list()

        context['applied_meta_list'] = applied_meta_list
        context['applied_graph_list'] = applied_graph_list

D
Dong Daxiang 已提交
1124
        self._context = context
1125

D
Dong Daxiang 已提交
1126
        self.valid_strategy = valid_strategy
1127
        self.valid_strategy._enable_env()
D
Dong Daxiang 已提交
1128

1129 1130
        optimize_ops = []
        params_grads = []
1131

1132 1133 1134 1135 1136 1137 1138 1139 1140
        if self._role_maker._is_non_distributed() and not self._is_collective:
            if self._runtime_handle is None:
                self._runtime_handle = RuntimeFactory()._create_runtime(context)

            compiled_program = compiler.CompiledProgram(
                self.origin_main_program).with_data_parallel(
                    loss_name=loss.name, share_vars_from=None)
            loss.block.program._graph = compiled_program
            return self.user_defined_optimizer.minimize(
M
MRXLT 已提交
1141
                loss, startup_program, parameter_list, no_grad_set=no_grad_set)
1142

1143 1144
        if meta_optimizer:
            optimize_ops, params_grads = meta_optimizer.minimize(
M
MRXLT 已提交
1145
                loss, startup_program, parameter_list, no_grad_set=no_grad_set)
1146

1147
            default_program = paddle.static.default_main_program()
1148 1149 1150 1151

            if id(default_program) != id(loss.block.program):
                paddle.fluid.framework.switch_main_program(loss.block.program)

1152 1153
        else:
            optimize_ops, params_grads = self.user_defined_optimizer.minimize(
M
MRXLT 已提交
1154
                loss, startup_program, parameter_list, no_grad_set=no_grad_set)
1155

1156 1157
        context["program_optimize_ops"] = optimize_ops
        context["program_params_grads"] = params_grads
1158

1159
        if graph_optimizer:
D
Dong Daxiang 已提交
1160
            optimize_ops, params_grads = graph_optimizer.minimize(
M
MRXLT 已提交
1161
                loss, startup_program, parameter_list, no_grad_set=no_grad_set)
1162 1163 1164 1165
            # since we do not encourage users to use graph operations
            # if a graph optimizer takes effect, mostly
            # optimizers_ops and params_grads are None
            # i.e. users can not modify current computation graph anymore
1166 1167 1168
            context["graph_optimize_ops"] = optimize_ops
            context["graph_optimize_grads"] = params_grads

1169
        if self._runtime_handle is None:
1170
            self._runtime_handle = RuntimeFactory()._create_runtime(context)
1171

1172 1173
        import paddle.distributed.fleet as fleet
        fleet.util._set_strategy(context["valid_strategy"])
1174 1175

        return optimize_ops, params_grads