@@ -301,7 +301,7 @@ void InplacePass::TryInplaceOpInputOutput(ir::Node* op,
301
301
302
302
// 3. if output has been memory optimize by python(fluid.memory_optmize()).
303
303
// this candidate can not be inplaced. Will be deprecated in the future.
304
- if (view_.ReusedInPythonMemOpt (out_node->Name ())) {
304
+ if (view_.InSkipSet (out_node->Name ())) {
305
305
VLOG (4 ) << string::Sprintf (
306
306
" Skiped %s => %s reused previous memory block in python memory "
307
307
" optmize,"
@@ -385,7 +385,7 @@ void GraphView::Build(ir::Graph* g) {
385
385
// resolve data harzards depends on the var nodes in right order.
386
386
ops_ = SortOpLikeDescOrder (*g);
387
387
388
- // track the nodes which reused previous node in Python memory optimize.
388
+ // 1. track the nodes which reused previous node in Python memory optimize.
389
389
// these node can not be inplaced, otherwise may generate a circle in graph.
390
390
std::unordered_set<std::string> all_vars;
391
391
for (auto & node : g->Nodes ()) {
@@ -399,11 +399,28 @@ void GraphView::Build(ir::Graph* g) {
399
399
}
400
400
}
401
401
}
402
+
403
+ // 2. track the nodes which used by parameter server.
404
+ // these node can not be inplaced, otherwise trainer
405
+ // pserver can not find each other name.
406
+ for (auto & node : g->Nodes ()) {
407
+ if (!node->IsOp ()) continue ;
408
+ if (node->Name () == " send" ) {
409
+ for (auto & in : node->inputs ) {
410
+ dup_nodes_.emplace (in->Name ());
411
+ }
412
+ }
413
+ if (node->Name () == " recv" ) {
414
+ for (auto & out : node->outputs ) {
415
+ dup_nodes_.emplace (out->Name ());
416
+ }
417
+ }
418
+ }
402
419
}
403
420
404
421
const std::vector<ir::Node*>& GraphView::AllOps () { return ops_; }
405
422
406
- bool GraphView::ReusedInPythonMemOpt (const std::string& var) const {
423
+ bool GraphView::InSkipSet (const std::string& var) const {
407
424
return dup_nodes_.count (var);
408
425
}
409
426
0 commit comments