Merge branch 'master' into dev_fix_task_graph_builder

Oneflow-Inc · Dec 26, 2024 · cbd61e9 · cbd61e9
2 parents dcf232a + f7fa76f
commit cbd61e9
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 0 deletions.
diff --git a/oneflow/user/ops/group_norm_op.cpp b/oneflow/user/ops/group_norm_op.cpp
@@ -19,9 +19,12 @@ limitations under the License.
 
 namespace oneflow {
 
+DEFINE_ENV_BOOL(ONEFLOW_GROUP_NORM_USE_FP16_DIRECTLY, false);
+
 namespace {
 
 oneflow::DataType InferGnParamDataType(const DataType x_data_type) {
+  if (EnvBool<ONEFLOW_GROUP_NORM_USE_FP16_DIRECTLY>()) { return x_data_type; }
   return (x_data_type == DataType::kFloat16 || x_data_type == DataType::kBFloat16)
              ? DataType::kFloat
              : x_data_type;

diff --git a/oneflow/user/ops/layer_norm_op.cpp b/oneflow/user/ops/layer_norm_op.cpp
@@ -18,6 +18,8 @@ limitations under the License.
 
 namespace oneflow {
 
+DEFINE_ENV_BOOL(ONEFLOW_LAYER_NORM_PARAM_KEEP_DIM, false);
+
 namespace {
 
 int64_t ShiftNegativeAxisIfNeed(const Shape& shape, int64_t axis) {
@@ -31,6 +33,11 @@ Shape InferBnParamShape(const Shape& x_shape, const int64_t begin_norm_axis) {
   DimVector bn_param_shape_dim_vec;
   bn_param_shape_dim_vec.insert(bn_param_shape_dim_vec.end(), x_shape.dim_vec().cbegin(),
                                 x_shape.dim_vec().cbegin() + begin_norm_axis);
+  if (EnvBool<ONEFLOW_LAYER_NORM_PARAM_KEEP_DIM>()) {
+    while (bn_param_shape_dim_vec.size() < x_shape.dim_vec().size()) {
+      bn_param_shape_dim_vec.push_back(1);
+    }
+  }
   const Shape bn_param_shape(bn_param_shape_dim_vec);
   return bn_param_shape;
 }