grad_tensor_holder.cc 5.7 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
// Copyright (c) 2021 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.

#include "paddle/fluid/eager/grad_tensor_holder.h"
16
#include "paddle/fluid/imperative/gradient_accumulator.h"
17

18
#include "paddle/fluid/eager/api/generated/eager_generated/forwards/dygraph_functions.h"
19
#include "paddle/fluid/framework/convert_utils.h"
20
#include "paddle/fluid/framework/var_type.h"
21
#include "paddle/phi/kernels/funcs/math_function.h"
22 23 24

namespace egr {

25 26 27 28 29
void GradTensorHolder::SetBufferSlotRankZeros(size_t slot_id, size_t rank) {
  buffer_[slot_id][rank] =
      paddle::experimental::zeros_like(buffer_[slot_id][rank]);
}

30 31 32
void GradTensorHolder::CopyValueFromTensor(
    size_t slot_id, size_t rank, const paddle::experimental::Tensor& t,
    bool fill_one) {
33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52
  // TODO(jiabin): We need to deal with empty input_buffer with slot size not
  // empty;
  PADDLE_ENFORCE(slot_id < buffer_.size(),
                 paddle::platform::errors::Fatal(
                     "Invalid slot_id for GradTensorHolder::add() "
                     "which exceeds size of buffer"));
  VLOG(6) << "Add Tensor for buffer_ slot: " << slot_id
          << ", size: " << buffer_[slot_id].size();
  if (buffer_[slot_id].empty()) {
    VLOG(6) << "Pass add Tensor for buffer_ slot: " << slot_id
            << " since its buffer_ is empty ";
    return;
  }
  PADDLE_ENFORCE(
      rank < buffer_[slot_id].size(),
      paddle::platform::errors::Fatal(
          "Invalid rank for GradTensorHolder::add() which exceeds size "
          "of buffer slot %d, got slot size is: %d rank is: %d",
          slot_id, buffer_[slot_id].size(), rank));
  if (!fill_one) {
53 54
    paddle::experimental::Tensor& buffer_tensor = buffer_[slot_id][rank];
    if ((!buffer_tensor.defined() || !buffer_tensor.initialized())) {
55 56 57 58
      // Perform deep copy here
      buffer_tensor.copy_(t, t.inner_place(), false);
      buffer_tensor.set_autograd_meta(t.mutable_autograd_meta());

59
    } else {
60 61 62
      PADDLE_THROW(paddle::platform::errors::Fatal(
          "Cannot copy grad_tensors' value to grad tensor holders,"
          "input buffer has already been initialized."));
63 64 65 66 67
    }
  } else {
    // Create new tensor->impl and fill it with 1.0
    if (t.defined()) {
      // Fill 1.0
H
hong 已提交
68
      buffer_[slot_id][rank] = paddle::experimental::ones_like(t, t.dtype());
69 70 71 72
    }
  }
}

73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140
void GradTensorHolder::add(size_t slot_id, size_t rank,
                           const paddle::experimental::Tensor& t) {
  // TODO(jiabin): We need to deal with empty input_buffer with slot size not
  // empty;
  PADDLE_ENFORCE(slot_id < buffer_.size(),
                 paddle::platform::errors::Fatal(
                     "Invalid slot_id for GradTensorHolder::add() "
                     "which exceeds size of buffer"));
  VLOG(6) << "Add Tensor for buffer_ slot: " << slot_id
          << ", size: " << buffer_[slot_id].size();
  if (buffer_[slot_id].empty()) {
    VLOG(6) << "Pass add Tensor for buffer_ slot: " << slot_id
            << " since its buffer_ is empty ";
    return;
  }
  PADDLE_ENFORCE(
      rank < buffer_[slot_id].size(),
      paddle::platform::errors::Fatal(
          "Invalid rank for GradTensorHolder::add() which exceeds size "
          "of buffer slot %d, got slot size is: %d rank is: %d",
          slot_id, buffer_[slot_id].size(), rank));

  paddle::experimental::Tensor& buffer_tensor = buffer_[slot_id][rank];
  // TODO(jiabin): Code bellow is ugly to divide which inner var we used,
  // remove framework::Variable
  // related code later.
  // This if statement is trying to test neither phi::Tensor nor
  // framework::Variable is initialized.
  if ((!buffer_tensor.defined() || !buffer_tensor.initialized())) {
    // Simply copy tensor->impl
    buffer_tensor = t;
  } else {
    // Accumulation
    PADDLE_ENFORCE_EQ(t.initialized(), true,
                      paddle::platform::errors::Fatal(
                          "We can only accumulate initialized tensor, but we "
                          "got tensor: %s is empty please check you network "
                          "and make sure it creates grads.",
                          t.name()));
    if (t.is_dense_tensor()) {
      if (buffer_tensor.is_dense_tensor()) {
        buffer_tensor = add_final_state_dygraph_function(t, buffer_tensor);

      } else {
        // TODO(jiabin): Support Other TensorBase later
        // TODO(zhanlve): Replace SelectedRowsAddTensor with
        // add_dygraph_function once it's supported
        paddle::experimental::Tensor new_buffer(
            std::make_shared<phi::DenseTensor>(), "tmp_accumulator");
        paddle::imperative::SelectedRowsAddTensor(buffer_tensor, t,
                                                  &new_buffer);
        buffer_tensor.set_impl(new_buffer.impl());
      }
    } else {
      // TODO(jiabin): Support Other TensorBase later
      // TODO(zhanlve): Replace SelectedRowsAddTensor with add_dygraph_function
      // once it's supported
      if (buffer_tensor.is_dense_tensor()) {
        paddle::imperative::SelectedRowsAddToTensor(t, &buffer_tensor);
      } else {
        buffer_tensor =
            std::move(*paddle::imperative::SelectedRowsMerge<
                      paddle::experimental::Tensor>(t, buffer_tensor));
      }
    }
  }
}

141
}  // namespace egr