You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
195 lines
7.6 KiB
195 lines
7.6 KiB
// Copyright (c) 2018 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 <memory>
|
|
#include <set>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
#include "paddle/fluid/operators/distributed/parameter_recv.h"
|
|
|
|
#include "paddle/fluid/framework/lod_tensor.h"
|
|
#include "paddle/fluid/framework/scope.h"
|
|
#include "paddle/fluid/framework/selected_rows.h"
|
|
#include "paddle/fluid/framework/tensor.h"
|
|
|
|
#include "paddle/fluid/operators/distributed/distributed.h"
|
|
#include "paddle/fluid/operators/distributed/rpc_client.h"
|
|
#include "paddle/fluid/operators/distributed/variable_response.h"
|
|
#include "paddle/fluid/operators/distributed_ops/send_recv_util.h"
|
|
#include "paddle/fluid/operators/strided_memcpy.h"
|
|
|
|
namespace paddle {
|
|
namespace operators {
|
|
namespace distributed {
|
|
|
|
using LoDTensor = framework::LoDTensor;
|
|
using LoDTensor = framework::LoDTensor;
|
|
using SelectedRows = framework::SelectedRows;
|
|
using DDim = framework::DDim;
|
|
|
|
template <typename T>
|
|
void ParameterRecv<T>::operator()(const RpcContext &rpc_ctx,
|
|
const framework::Scope &scope) {
|
|
VLOG(2) << "ParameterRecv in " << rpc_ctx.var_name;
|
|
std::unique_ptr<framework::Scope> local_scope = scope.NewTmpScope();
|
|
|
|
platform::DeviceContextPool &pool = platform::DeviceContextPool::Instance();
|
|
auto &cpu_ctx = *pool.Get(platform::CPUPlace());
|
|
|
|
distributed::RPCClient *rpc_client =
|
|
distributed::RPCClient::GetInstance<RPCCLIENT_T>(rpc_ctx.trainer_id);
|
|
|
|
auto *recv_var = scope.FindVar(rpc_ctx.var_name);
|
|
|
|
// recv all vars to local scope
|
|
if (recv_var->IsType<framework::LoDTensor>() ||
|
|
recv_var->IsType<framework::SelectedRows>()) {
|
|
std::vector<distributed::VarHandlePtr> rets;
|
|
for (size_t i = 0; i < rpc_ctx.splited_var_names.size(); i++) {
|
|
auto &recv_var_name = rpc_ctx.splited_var_names[i];
|
|
local_scope->Var(recv_var_name);
|
|
VLOG(4) << "recv " << recv_var_name << " from " << rpc_ctx.epmap[i];
|
|
if (recv_var->IsType<framework::LoDTensor>()) {
|
|
// sparse param in recv_scope is LoDTensor
|
|
rets.push_back(rpc_client->AsyncGetVar(rpc_ctx.epmap[i], cpu_ctx,
|
|
*local_scope.get(),
|
|
recv_var_name, recv_var_name));
|
|
} else {
|
|
// sparse param in pserver_scope is SelectedRows
|
|
rets.push_back(rpc_client->AsyncGetVar(
|
|
rpc_ctx.epmap[i], cpu_ctx, *local_scope.get(), recv_var_name,
|
|
recv_var_name, recv_var_name));
|
|
}
|
|
}
|
|
for (size_t i = 0; i < rets.size(); i++) {
|
|
PADDLE_ENFORCE(rets[i]->Wait(), "internal error in RPCClient");
|
|
}
|
|
} else {
|
|
PADDLE_THROW("unsupported var type to recv!");
|
|
}
|
|
|
|
// concat recved tensor into one var
|
|
if (recv_var->IsType<framework::LoDTensor>()) {
|
|
size_t output_offset = 0;
|
|
size_t row_offset = 0;
|
|
framework::Tensor *recv_tensor =
|
|
recv_var->GetMutable<framework::LoDTensor>();
|
|
auto dev_ctx = paddle::platform::CPUDeviceContext();
|
|
int64_t recv_numel = 0;
|
|
for (auto &recv_var_name : rpc_ctx.splited_var_names) {
|
|
auto *recv_var = local_scope->FindVar(recv_var_name);
|
|
if (recv_var->IsType<framework::LoDTensor>()) {
|
|
auto &in = recv_var->Get<framework::LoDTensor>();
|
|
recv_numel += in.numel();
|
|
auto in_stride = framework::stride_numel(in.dims());
|
|
auto out_stride = framework::stride_numel(recv_tensor->dims());
|
|
StridedNumelCopyWithAxis<T>(
|
|
dev_ctx, 0, recv_tensor->data<T>() + output_offset, out_stride,
|
|
in.data<T>(), in_stride, in_stride[0]);
|
|
output_offset += in_stride[0];
|
|
} else if (recv_var->IsType<framework::SelectedRows>()) {
|
|
auto &recv_slr = recv_var->Get<framework::SelectedRows>();
|
|
auto &recv_dims = recv_tensor->dims();
|
|
int64_t width = recv_dims[1];
|
|
recv_numel += recv_slr.height() * width;
|
|
PADDLE_ENFORCE_EQ(recv_slr.value().dims()[1], width);
|
|
PADDLE_ENFORCE_EQ(recv_slr.value().dims()[0], recv_slr.rows().size());
|
|
VLOG(3) << "recv slr " << recv_var_name << " dims "
|
|
<< recv_slr.value().dims();
|
|
if (VLOG_IS_ON(3)) {
|
|
std::ostringstream sstream;
|
|
sstream << "[";
|
|
for (auto &row_id : recv_slr.rows()) {
|
|
sstream << row_id << ", ";
|
|
}
|
|
sstream << "]";
|
|
VLOG(3) << "recv_slr size: " << recv_slr.rows().size() << " "
|
|
<< sstream.str();
|
|
}
|
|
|
|
for (size_t i = 0; i < recv_slr.rows().size(); ++i) {
|
|
auto row_id = recv_slr.rows()[i] + row_offset;
|
|
PADDLE_ENFORCE_LT(row_id, recv_dims[0]);
|
|
memcpy(recv_tensor->data<T>() + row_id * width,
|
|
recv_slr.value().data<T>() + i * width, sizeof(T) * width);
|
|
}
|
|
row_offset += recv_slr.height();
|
|
} else {
|
|
PADDLE_THROW("unsupported recieved var type");
|
|
}
|
|
}
|
|
auto numel = recv_tensor->numel();
|
|
if (recv_numel != numel) {
|
|
LOG(FATAL) << "recv_numel: " << recv_numel << " acture numel: " << numel;
|
|
}
|
|
PADDLE_ENFORCE_EQ(recv_numel, numel);
|
|
} else if (recv_var->IsType<framework::SelectedRows>()) {
|
|
auto cpu_place = platform::CPUPlace();
|
|
auto *slr = recv_var->GetMutable<framework::SelectedRows>();
|
|
slr->mutable_rows()->clear();
|
|
slr->mutable_value()->mutable_data<float>({{}}, cpu_place);
|
|
int64_t width = 0;
|
|
int64_t height = 0;
|
|
std::vector<int64_t> new_rows{};
|
|
|
|
// trans sparse ids from local to global
|
|
std::vector<int64_t> abs_sections =
|
|
ToAbsoluteSection(rpc_ctx.height_sections);
|
|
|
|
for (size_t i = 0; i < rpc_ctx.splited_var_names.size(); i++) {
|
|
auto &recv_var_name = rpc_ctx.splited_var_names[i];
|
|
auto *var = local_scope->FindVar(recv_var_name);
|
|
auto *var_slr = var->GetMutable<framework::SelectedRows>();
|
|
auto *var_slr_row = var_slr->mutable_rows();
|
|
width = var_slr->mutable_value()->dims()[1];
|
|
height += var_slr->height();
|
|
auto row_offset = abs_sections[i];
|
|
VLOG(4) << "Recv split_var " << recv_var_name << " Row size "
|
|
<< var_slr_row->size();
|
|
for (size_t j = 0; j < var_slr_row->size(); j++) {
|
|
new_rows.push_back(row_offset + (*var_slr_row)[j]);
|
|
}
|
|
}
|
|
slr->set_rows(new_rows);
|
|
slr->set_height(height);
|
|
slr->mutable_value()->mutable_data<float>(
|
|
framework::make_ddim(
|
|
{static_cast<int64_t>(slr->mutable_rows()->size()), width}),
|
|
cpu_place);
|
|
auto *slr_data = slr->mutable_value()->data<float>();
|
|
|
|
size_t row_offset = 0;
|
|
for (auto &recv_var_name : rpc_ctx.splited_var_names) {
|
|
auto *var = local_scope->FindVar(recv_var_name);
|
|
auto *var_slr = var->GetMutable<framework::SelectedRows>();
|
|
auto *var_slr_row = var_slr->mutable_rows();
|
|
auto var_slr_row_size = var_slr_row->size();
|
|
auto *var_slr_data = var_slr->mutable_value()->data<float>();
|
|
|
|
memcpy(slr_data + row_offset * width, var_slr_data,
|
|
sizeof(float) * width * var_slr_row_size);
|
|
row_offset += var_slr_row_size;
|
|
}
|
|
}
|
|
|
|
VLOG(2) << "ParameterRecv out " << rpc_ctx.var_name;
|
|
}
|
|
|
|
template struct ParameterRecv<float>;
|
|
|
|
}; // namespace distributed
|
|
}; // namespace operators
|
|
}; // namespace paddle
|