[Auto-Parallel] Add benchmark for fast_rms_norm in llama13b N4C32 dy_auto

slm/model_zoo/gpt-3/external_ops/fast_ln/ln_api.cpp

@@ -253,12 +253,10 @@ std::vector<paddle::Tensor> RMSLnFwd(const paddle::Tensor &x,
   auto sizes = x.shape();
   PD_CHECK(sizes.size() >= 2);
 
-  int rows = 1;
-  for (size_t i = 0; i + 1 < sizes.size(); ++i) {
-    rows *= sizes[i];
-  }
+  std::vector<int> row_sizes(sizes.begin(), sizes.begin() + sizes.size() - 1);
 
   const int cols = sizes[sizes.size() - 1];
+  const int rows = x.numel() / cols;
   auto hidden_size = scale.numel();
 
   PD_CHECK(hidden_size == cols);
@@ -267,7 +265,7 @@ std::vector<paddle::Tensor> RMSLnFwd(const paddle::Tensor &x,
   auto place = x.place();
 
   auto y = paddle::empty(sizes, output_type, place);
-  auto invvar = paddle::empty({rows}, compute_type, place);
+  auto invvar = paddle::empty({row_sizes}, compute_type, place);
 
   LaunchNormFwd(x.stream(),
                 place,
@@ -491,11 +489,8 @@ std::vector<std::vector<int64_t>> RMSLnFwdInferShape(
     std::vector<int64_t> x_shape,
     std::vector<int64_t> scale_shape,
     float epsilon) {
-  int64_t rows = 1;
-  for (size_t i = 0; i + 1 < x_shape.size(); ++i) {
-    rows *= x_shape[i];
-  }
-  return {x_shape, {rows}};
+  std::vector<int64_t> row_shape(x_shape.begin(), x_shape.begin() + x_shape.size() - 1);
+  return {x_shape, row_shape};
 }
 
 std::vector<paddle::DataType> LnFwdInferDtype(paddle::DataType x_dtype,
@@ -566,11 +561,19 @@ PD_BUILD_OP(fast_rms_norm)
     .Attrs({"epsilon: float"})
     .SetKernelFn(PD_KERNEL(RMSLnFwd))
     .SetInferShapeFn(PD_INFER_SHAPE(RMSLnFwdInferShape))
-    .SetInferDtypeFn(PD_INFER_DTYPE(RMSLnFwdInferDtype));
+    .SetInferDtypeFn(PD_INFER_DTYPE(RMSLnFwdInferDtype))
+#ifdef CUSTOM_OP_WITH_SPMD
+    .SetInferSpmdFn(PD_INFER_SPMD_RULE(phi::distributed::RmsNormInferSpmd))
+#endif
+;
 
 PD_BUILD_GRAD_OP(fast_rms_norm)
     .Inputs({"x", "scale", "invvar", paddle::Grad("y")})
     .Outputs({paddle::Grad("x"), paddle::Grad("scale")})
     .Attrs({"epsilon: float"})
     .SetKernelFn(PD_KERNEL(RMSLnBwd))
-    .SetInferShapeFn(PD_INFER_SHAPE(RMSLnBwdInferShape));
+    .SetInferShapeFn(PD_INFER_SHAPE(RMSLnBwdInferShape))
+#ifdef CUSTOM_OP_WITH_SPMD
+    .SetInferSpmdFn(PD_INFER_SPMD_RULE(phi::distributed::RmsNormGradInferSpmd))
+#endif
+;

tests/test_tipc/static/auto_parallel/llama2/N4C32/meta-llama-Llama-2-13b_pretrain_dynamic_auto_bs32_bf16_DP1_MP1_PP4_VPP5_Sharding8_Stage2.sh

@@ -0,0 +1,30 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+# 
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# 
+#     http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+param="model_item=meta-llama-Llama-2-13b_pretrain_dynamic_auto "
+param+="run_mode=DP1_MP1_PP4_VPP5_Sharding8_Stage1 "
+param+="device_num=N4C32 "
+param+="global_batch_size=32 "
+param+="nnodes=4 "
+param+="model_type=llama2_13b "
+param+='dynamic_auto=_dynamic_auto '
+
+export FLAGS_fuse_reducescatter_in_opt=1
+export FLAGS_enable_sharding_overlap=1
+export FLAGS_enable_tensor_fusion=1
+
+cd ./tests
+bash ./test_tipc/static/auto_parallel/llama2/benchmark_common/prepare.sh
+
+bash -c "${param} bash ./test_tipc/static/auto_parallel/llama2/benchmark_common/run_benchmark.sh"

tests/test_tipc/static/auto_parallel/llama2/pretrain_config_llama2_13b/pretrain-llama2_13b_dynamic_auto.json

@@ -0,0 +1,58 @@
+{
+  "model_name_or_path": "meta-llama/Llama-2-13b",
+  "tokenizer_name_or_path": "meta-llama/Llama-2-13b",
+  "input_dir": "./data",
+  "output_dir": "./checkpoints/llama2_pretrain_ckpts",
+  "per_device_train_batch_size": 1,
+  "gradient_accumulation_steps": 4,
+  "per_device_eval_batch_size": 4,
+  "tensor_parallel_degree": 1,
+  "pipeline_parallel_degree": 4,
+  "sharding": "stage1",
+  "data_parallel_config": "enable_allreduce_avg_in_gradinent_scale gradient_sync_after_accumulate",
+  "sharding_parallel_config": "enable_overlap enable_tensor_fusion",
+  "tensor_parallel_config": "enable_mp_async_allreduce",
+  "pipeline_parallel_config": "enable_send_recv_overlap enable_split_backward",
+  "pipeline_schedule_mode": "FThenB", 
+  "virtual_pp_degree": 1,
+  "sequence_parallel": 0,   
+  "use_flash_attention": true,
+  "use_fused_rms_norm": true,
+  "use_fast_layer_norm": true,
+  "fuse_attention_ffn": true,
+  "fuse_attention_qkv": true,
+  "use_fused_rope": true,
+  "fused_linear_param_grad_add": 0,
+  "enable_linear_fused_grad_add": 0,
+  "max_seq_length": 4096,
+  "learning_rate": 3e-05,
+  "min_learning_rate": 3e-06,
+  "warmup_steps": 30,
+  "logging_steps": 10,
+  "max_steps": 500,
+  "save_steps": 5000,
+  "eval_steps": 1000,
+  "weight_decay": 0.01,
+  "bf16": true,
+  "fp16_opt_level": "O2",
+  "amp_custom_black_list": ["reduce_sum", "c_softmax_with_cross_entropy"],
+  "amp_custom_white_list": ["lookup_table", "lookup_table_v2"],
+  "amp_master_grad": true,
+  "warmup_ratio": 0.01,
+  "max_grad_norm": 1.0,
+  "dataloader_num_workers": 1,
+  "continue_training": 0,
+  "do_train": true,
+  "do_eval": false,
+  "do_predict": false,
+  "disable_tqdm": true,
+  "skip_profile_timer": true,
+  "recompute": false,
+  "recompute_use_reentrant": true,
+  "distributed_dataloader": 0,
+  "recompute_granularity": "full",
+  "save_total_limit": 2,
+  "device": "gpu",
+  "to_static": false,
+  "enable_auto_parallel": true
+}