|
|
|
|
@ -31,6 +31,7 @@ static constexpr char kParallelScopes[] = "parallel_scopes";
|
|
|
|
|
static constexpr char kParallelBlock[] = "sub_block";
|
|
|
|
|
|
|
|
|
|
using LoDTensor = framework::LoDTensor;
|
|
|
|
|
using SelectedRows = framework::SelectedRows;
|
|
|
|
|
|
|
|
|
|
static void SplitTensorAndMoveTensorToScopes(
|
|
|
|
|
const framework::Scope &scope, std::vector<framework::Scope *> *sub_scopes,
|
|
|
|
|
@ -64,6 +65,30 @@ static void SplitTensorAndMoveTensorToScopes(
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
inline void CopyOrShare(const framework::Variable &src,
|
|
|
|
|
const platform::Place &dst_place,
|
|
|
|
|
framework::Variable *dst) {
|
|
|
|
|
if (src.IsType<LoDTensor>()) {
|
|
|
|
|
if (src.Get<LoDTensor>().place() == dst_place) {
|
|
|
|
|
dst->GetMutable<LoDTensor>()->ShareDataWith(src.Get<LoDTensor>());
|
|
|
|
|
} else {
|
|
|
|
|
Copy(src.Get<LoDTensor>(), dst_place, dst->GetMutable<LoDTensor>());
|
|
|
|
|
}
|
|
|
|
|
} else if (src.IsType<SelectedRows>()) {
|
|
|
|
|
auto &src_sr = src.Get<SelectedRows>();
|
|
|
|
|
auto *dst_sr = dst->GetMutable<SelectedRows>();
|
|
|
|
|
dst_sr->set_rows(src_sr.rows());
|
|
|
|
|
dst_sr->set_height(src_sr.height());
|
|
|
|
|
if (src_sr.value().place() == dst_place) {
|
|
|
|
|
dst_sr->mutable_value()->ShareDataWith(src_sr.value());
|
|
|
|
|
} else {
|
|
|
|
|
Copy(src_sr.value(), dst_place, dst_sr->mutable_value());
|
|
|
|
|
}
|
|
|
|
|
} else {
|
|
|
|
|
PADDLE_THROW("Expect LoDTensor/SelectedRows, get %s", src.Type().name());
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
void WaitOnPlace(const platform::Place place) {
|
|
|
|
|
platform::DeviceContextPool &pool = platform::DeviceContextPool::Instance();
|
|
|
|
|
auto &dev_ctx = *pool.Get(place);
|
|
|
|
|
@ -210,30 +235,30 @@ class ParallelDoGradOp : public framework::OperatorBase {
|
|
|
|
|
}
|
|
|
|
|
WaitOnPlaces(places);
|
|
|
|
|
|
|
|
|
|
// merge grad
|
|
|
|
|
AccumulateGrad(scope, place, sub_scopes, places);
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
void AccumulateGrad(const framework::Scope &scope,
|
|
|
|
|
const platform::Place &place,
|
|
|
|
|
const std::vector<framework::Scope *> &sub_scopes,
|
|
|
|
|
const platform::PlaceList &places) const {
|
|
|
|
|
for (auto &s : Outputs(framework::GradVarName(kParameters))) {
|
|
|
|
|
auto &result = sub_scopes[0]->FindVar(s)->Get<LoDTensor>();
|
|
|
|
|
std::string tmp_name;
|
|
|
|
|
auto *tmp = sub_scopes[0]->Var(&tmp_name)->GetMutable<LoDTensor>();
|
|
|
|
|
auto *tmp = sub_scopes[0]->Var(&tmp_name);
|
|
|
|
|
|
|
|
|
|
for (size_t i = 1; i < sub_scopes.size(); ++i) {
|
|
|
|
|
auto &tensor_to_merge = sub_scopes[i]->FindVar(s)->Get<LoDTensor>();
|
|
|
|
|
if (!(places[i] == places[0])) {
|
|
|
|
|
framework::Copy(tensor_to_merge, places[0], tmp);
|
|
|
|
|
WaitOnPlace(places[0]);
|
|
|
|
|
} else {
|
|
|
|
|
tmp->ShareDataWith(tensor_to_merge);
|
|
|
|
|
}
|
|
|
|
|
CopyOrShare(*sub_scopes[i]->FindVar(s), places[0], tmp);
|
|
|
|
|
WaitOnPlace(places[0]);
|
|
|
|
|
|
|
|
|
|
auto sum_op = framework::OpRegistry::CreateOp(
|
|
|
|
|
"sum", {{"X", {s, tmp_name}}}, {{"Out", {s}}},
|
|
|
|
|
framework::AttributeMap{});
|
|
|
|
|
VLOG(3) << sum_op->DebugStringEx(sub_scopes[0]);
|
|
|
|
|
sum_op->Run(*sub_scopes[0], places[0]);
|
|
|
|
|
WaitOnPlace(places[0]);
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
VLOG(3) << result;
|
|
|
|
|
framework::Copy(result, place, scope.FindVar(s)->GetMutable<LoDTensor>());
|
|
|
|
|
CopyOrShare(*sub_scopes[0]->FindVar(s), place, scope.FindVar(s));
|
|
|
|
|
}
|
|
|
|
|
WaitOnPlaces(places);
|
|
|
|
|
}
|
|
|
|
|
@ -289,7 +314,7 @@ class ParallelDoGradOpShapeInference : public framework::InferShapeBase {
|
|
|
|
|
|
|
|
|
|
PADDLE_ENFORCE(ctx->HasInputs(kParameters));
|
|
|
|
|
PADDLE_ENFORCE(ctx->HasOutputs(framework::GradVarName(kParameters)));
|
|
|
|
|
PADDLE_ENFORCE(ctx->HasInput(kInputs));
|
|
|
|
|
PADDLE_ENFORCE(ctx->HasInputs(kInputs));
|
|
|
|
|
|
|
|
|
|
for (auto &s : output) {
|
|
|
|
|
PADDLE_ENFORCE(ctx->HasInputs(s));
|
|
|
|
|
|