[PHI]fix momentum dtype infer (#51353)

* fix momentum dtype infer

* fix momentum datatype

* fix on cpu

* add momentum

paddle/phi/infermeta/multiary.cc

@@ -2172,10 +2172,15 @@ void MomentumInferMeta(const MetaTensor& param,
 
   auto param_dim = param.dims();
   param_out->set_dims(param_dim);
+  auto MPType = (param.dtype() == phi::DataType::FLOAT16 ||
+                 param.dtype() == phi::DataType::BFLOAT16)
+                    ? phi::DataType::FLOAT32
+                    : param.dtype();
   velocity_out->set_dims(param_dim);
-
+  velocity_out->set_dtype(MPType);
   if (master_param_out) {
     master_param_out->set_dims(param_dim);
+    master_param_out->set_dtype(MPType);
   }
 }
 

paddle/phi/kernels/gpu/momentum_kernel.cu

@@ -24,7 +24,10 @@ PD_REGISTER_KERNEL(momentum,
                    phi::MomentumDenseKernel,
                    float,
                    double,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16) {
+  kernel->OutputAt(1).SetDataType(phi::DataType::UNDEFINED);
+  kernel->OutputAt(2).SetDataType(phi::DataType::UNDEFINED);
+}
 
 PD_REGISTER_KERNEL(momentum_dense_param_sparse_grad,
                    GPU,
@@ -32,4 +35,7 @@ PD_REGISTER_KERNEL(momentum_dense_param_sparse_grad,
                    phi::MomentumSparseKernel,
                    float,
                    double,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16) {
+  kernel->OutputAt(1).SetDataType(phi::DataType::UNDEFINED);
+  kernel->OutputAt(2).SetDataType(phi::DataType::UNDEFINED);
+}

paddle/phi/kernels/xpu/momentum_kernel.cc

@@ -69,4 +69,7 @@ PD_REGISTER_KERNEL(momentum,
                    ALL_LAYOUT,
                    phi::MomentumDenseKernel,
                    float,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16) {
+  kernel->OutputAt(1).SetDataType(phi::DataType::UNDEFINED);
+  kernel->OutputAt(2).SetDataType(phi::DataType::UNDEFINED);
+}