reducer.h 4.9 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18
// 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.

#pragma once

#include <algorithm>
#include <iostream>
19
#include <map>
20
#include <memory>
21
#include <queue>
22 23
#include <string>
#include <unordered_map>
24
#include <unordered_set>
25 26
#include <utility>
#include <vector>
27

28
#include "paddle/fluid/framework/data_type.h"
29 30
#include "paddle/fluid/framework/tensor.h"
#include "paddle/fluid/framework/variable.h"
31
#include "paddle/fluid/operators/math/math_function.h"
32

33 34 35 36 37 38 39 40 41 42 43
namespace paddle {
namespace platform {
class DeviceContext;
}  // namespace platform

namespace imperative {
class ParallelContext;
class VarBase;
class VariableWrapper;
}  // namespace imperative
}  // namespace paddle
44 45 46 47

namespace paddle {
namespace imperative {

48
#if (defined PADDLE_WITH_NCCL) || (defined PADDLE_WITH_XPU_BKCL)
49 50 51 52 53 54 55 56 57 58 59 60 61
class Group {
 public:
  // Here, we use dense_contents_ & sparse_contents_ to
  // achieve the tensor fuse. When is_sparse_ is true, sparse_contents_ work,
  // conversely, dense_contents_ works. It is mutex relationship.
  framework::Variable dense_contents_;
  framework::Variable* sparse_contents_ = nullptr;
  bool is_sparse_ = false;

  // for concat kernel
  std::vector<framework::Tensor> dense_tensors_;

  std::vector<size_t> length_;
62 63

  int64_t all_length_{0};
64 65 66 67 68 69 70 71 72 73 74
  // Global indices of participating variables in the group
  std::vector<size_t> variable_indices_;

  // Number of params that haven't been ready. When it is 0, it means
  // the group is ready.
  size_t pending_ = -1;

  // external message of group
  framework::proto::VarType::Type dtype_;

  // context is used to select the stream for concat
75
  void ConcatTensors(const platform::DeviceContext& context);
76 77

  // context is used to select the stream for split
78
  void SplitTensors(const platform::DeviceContext& context);
79 80

  friend std::ostream& operator<<(std::ostream&, const Group&);
81 82
};

83
struct VariableLocator {
84 85 86 87 88 89 90 91 92 93 94
  // record the index in groups_
  size_t group_index;
  size_t inside_group_index;
};

class Reducer {
 public:
  explicit Reducer(
      const std::vector<std::shared_ptr<imperative::VarBase>>& vars,
      const std::vector<std::vector<size_t>>& group_indices,
      const std::vector<bool>& is_sparse_gradient,
95
      std::shared_ptr<imperative::ParallelContext> parallel_ctx,
96
      const std::vector<size_t>& group_size_limits, bool find_unused_vars);
97 98 99 100 101

  virtual ~Reducer() {}

  void InitializeGroups(const std::vector<std::vector<size_t>>& group_indices);

102 103
  void InitializeDenseGroups(const std::vector<size_t>& variable_indices_,
                             Group* p_group);
104

105
  void PrepareDeps(const std::unordered_set<GradOpNode*>& init_nodes);
106

107 108
  void PrepareForBackward(
      const std::vector<std::shared_ptr<imperative::VarBase>>& outputs);
109

110
  void AddDistHook(size_t var_index);
111

112
  void MarkVarReady(const size_t var_index, const bool is_used_var);
113 114 115 116 117

  void MarkGroupReady(size_t group_index);

  void FinalizeBackward();

118 119
  std::vector<std::vector<size_t>> RebuildGruops();

120 121
  inline bool NeedRebuildGroup() { return !has_rebuilt_group_; }

122 123 124 125 126 127 128 129 130 131
 private:
  std::vector<std::shared_ptr<imperative::VarBase>> vars_;
  std::vector<std::vector<size_t>> group_indices_;
  static std::shared_ptr<Reducer> s_instance_;
  std::vector<Group> groups_;
  size_t next_group_ = 0;
  platform::Place place_;
  std::once_flag once_flag_;
  std::vector<bool> is_sparse_gradient_;
  std::shared_ptr<imperative::ParallelContext> parallel_ctx_;
132
  std::vector<VariableLocator> variable_locators_;
133

134
  int nrings_ = 1;
135 136

  // Following variables are to help rebuild group
137 138
  // TODO(shenliang03): Support rebuild in the future.
  bool has_rebuilt_group_{true};
139 140 141
  std::vector<std::shared_ptr<imperative::VarBase>> rebuild_vars_;
  std::vector<int64_t> rebuild_var_indices_;
  const std::vector<size_t> group_size_limits_;
142 143 144 145 146 147 148 149

  // Following variables are to help unused vars
  std::unordered_map<GradOpNode*, size_t> node_deps_;
  std::unordered_map<VariableWrapper*, size_t> var_index_map_;
  std::vector<size_t> unused_vars_;
  bool has_marked_unused_vars_{false};
  bool find_unused_vars_{false};
  bool all_group_ready_{false};
150 151 152 153 154
};

std::vector<std::vector<size_t>> AssignGroupBySize(
    const std::vector<std::shared_ptr<imperative::VarBase>>& tensors,
    const std::vector<bool>& is_sparse_gradient,
155 156
    const std::vector<size_t>& group_size_limits,
    const std::vector<int64_t>& tensor_indices = {});
157 158 159 160
#endif

}  // namespace imperative
}  // namespace paddle