mode_npu_gather_v2 (#34194)

paddle/fluid/operators/lookup_table_v2_op_npu.cc

@@ -40,6 +40,9 @@ class LookupTableV2NPUKernel : public framework::OpKernel<T> {
         platform::errors::InvalidArgument("npu only accept LoDTensor"));
     output_t->mutable_data<T>(ctx.GetPlace());
 
+    std::vector<int> ids;
+    TensorToVector(*ids_t, ctx.device_context(), &ids);
+
     NpuOpRunner runner;
     runner.SetType("GatherV2")
         .AddInput(*table_t)

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):