From 8b59f5e01bf60f8b73033931180b79431cbdaed4 Mon Sep 17 00:00:00 2001 From: Baibaifan <39549453+Baibaifan@users.noreply.github.com> Date: Fri, 16 Jul 2021 15:50:08 +0800 Subject: [PATCH] mode_npu_gather_v2 (#34194) --- paddle/fluid/operators/lookup_table_v2_op_npu.cc | 3 +++ .../distributed/fleet/meta_optimizers/sharding_optimizer.py | 4 ++-- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/paddle/fluid/operators/lookup_table_v2_op_npu.cc b/paddle/fluid/operators/lookup_table_v2_op_npu.cc index 2a8f4746234..c75ea537216 100644 --- a/paddle/fluid/operators/lookup_table_v2_op_npu.cc +++ b/paddle/fluid/operators/lookup_table_v2_op_npu.cc @@ -40,6 +40,9 @@ class LookupTableV2NPUKernel : public framework::OpKernel { platform::errors::InvalidArgument("npu only accept LoDTensor")); output_t->mutable_data(ctx.GetPlace()); + std::vector ids; + TensorToVector(*ids_t, ctx.device_context(), &ids); + NpuOpRunner runner; runner.SetType("GatherV2") .AddInput(*table_t) diff --git a/python/paddle/distributed/fleet/meta_optimizers/sharding_optimizer.py b/python/paddle/distributed/fleet/meta_optimizers/sharding_optimizer.py index a74f923dea4..878614ca152 100755 --- a/python/paddle/distributed/fleet/meta_optimizers/sharding_optimizer.py +++ b/python/paddle/distributed/fleet/meta_optimizers/sharding_optimizer.py @@ -374,8 +374,8 @@ class ShardingOptimizer(MetaOptimizerBase): 'w') as f: f.writelines(str(main_block.program)) - if core.is_compiled_with_cuda(): - self._wait() + # GPU and NPU need to wait server ready + self._wait() return optimize_ops, params_grads def _init_comm(self): -- GitLab