fix one_embedding adam (#7974)

* fix one_embedding adam * fix tidy * fix normal Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Oneflow-Inc · Apr 24, 2022 · a1f23bb · a1f23bb
1 parent ef6dbb1
commit a1f23bb
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/oneflow/user/kernels/one_embedding_kernels.cu b/oneflow/user/kernels/one_embedding_kernels.cu
@@ -222,7 +222,7 @@ __global__ void InitValueKernel(uint64_t seed, one::CUDAGeneratorState* cuda_gen
       } else if (initializer.type == InitializerType::kNormal) {
         const float mean = initializer.normal_param.mean;
         const float std = initializer.normal_param.std;
-        value = (curand_normal(&state) + mean) / std;
+        value = curand_normal(&state) * std + mean;
       } else {
         __trap();
       }

diff --git a/oneflow/user/ops/one_embedding_ops.cpp b/oneflow/user/ops/one_embedding_ops.cpp
@@ -263,7 +263,8 @@ Maybe<void> CheckDataType(user_op::InferContext* ctx) {
     user_op::InferContext* ctx) {
   JUST(CheckDataShape(ctx));
   const Shape& unique_embeddings_shape = ctx->InputShape("unique_embeddings", 0);
-  CHECK_EQ_OR_RETURN(unique_embeddings_shape.At(1), 2 * ctx->InputShape("embedding_grad", 0).At(1));
+  CHECK_EQ_OR_RETURN(unique_embeddings_shape.At(1), 2 * ctx->InputShape("embedding_grad", 0).At(1))
+      << "please adjust size_factor of MultiTableEmbedding's store_options to 2";
   *ctx->OutputShape("updated_unique_embeddings", 0) = unique_embeddings_shape;
   return Maybe<void>::Ok();
 }
@@ -293,7 +294,8 @@ Maybe<void> CheckDataType(user_op::InferContext* ctx) {
 /* static */ Maybe<void> AdamEmbeddingUpdateOp::InferLogicalTensorDesc(user_op::InferContext* ctx) {
   JUST(CheckDataShape(ctx));
   const Shape& unique_embeddings_shape = ctx->InputShape("unique_embeddings", 0);
-  CHECK_EQ_OR_RETURN(unique_embeddings_shape.At(1), 2 * ctx->InputShape("embedding_grad", 0).At(1));
+  CHECK_EQ_OR_RETURN(unique_embeddings_shape.At(1), 3 * ctx->InputShape("embedding_grad", 0).At(1))
+      << "please adjust size_factor of MultiTableEmbedding's store_options to 3";
   *ctx->OutputShape("updated_unique_embeddings", 0) = unique_embeddings_shape;
   return Maybe<void>::Ok();
 }