diff --git a/model_zoo/official/nlp/pangu_alpha/README.md b/model_zoo/official/nlp/pangu_alpha/README.md
index f863875e45e..3aa82373b85 100644
--- a/model_zoo/official/nlp/pangu_alpha/README.md
+++ b/model_zoo/official/nlp/pangu_alpha/README.md
@@ -151,22 +151,23 @@ The above command involves some `args` described below:
 - MICRO_SIZE: The number of micro batches in pipeline parallel mode. It should large than `stage_num`.
 - PER_BATCH: The batch size for each data parallel-way. default 8.
 - RANK_START: The start of rank_id in current machines, it helps to set the rank_id for each machine in multi-machine scenario.
+- LOCAL_DEVICE_NUM: The device number of the local machine.
 
 The following command will launch he program will train 2.6B model with the following command:
 
 ```bash
 # run distributed training example in one ascend machine
 
-bash scripts/run_distribute_train.sh /path/dataset /path/hccl.json 8 fp32 2.6B 1 1 8 0
+bash scripts/run_distribute_train.sh /path/dataset /path/hccl.json 8 fp32 2.6B 1 1 8 0 8
 ```
 
 ```bash
 # run distributed training example in two ascend machine
 
 # machine A
-bash scripts/run_distribute_train.sh /path/dataset /path/hccl.json 16 fp32 2.6B 2 4 8 0
+bash scripts/run_distribute_train.sh /path/dataset /path/hccl.json 16 fp32 2.6B 2 4 8 0 8
 # machine B
-bash scripts/run_distribute_train.sh /path/dataset /path/hccl.json 16 fp32 2.6B 2 4 8 8
+bash scripts/run_distribute_train.sh /path/dataset /path/hccl.json 16 fp32 2.6B 2 4 8 8 8
 ```
 
 For distributed training, an hccl configuration file with JSON format needs to be created in advance.
diff --git a/model_zoo/official/nlp/pangu_alpha/scripts/run_distribute_train.sh b/model_zoo/official/nlp/pangu_alpha/scripts/run_distribute_train.sh
index 0330acd8449..6127d2f8235 100644
--- a/model_zoo/official/nlp/pangu_alpha/scripts/run_distribute_train.sh
+++ b/model_zoo/official/nlp/pangu_alpha/scripts/run_distribute_train.sh
@@ -16,13 +16,14 @@
 
 echo "=============================================================================================================="
 echo "Please run the script as: "
-echo "bash run_distributed_train.sh DATA_DIR RANK_TABLE_FILE DEVICE_NUM TYPE MODE"
+echo "bash run_distributed_train.sh DATA_DIR RANK_TABLE_FILE DEVICE_NUM TYPE MODE STAGE_NUM MICRO_SIZE"
+echo "PER_BATCH RANK_START RANK_START LOCAL_DEVICE_NUM"
 echo "for example:"
 echo "#######no pipeline#######"
-echo "bash run_distributed_train.sh /path/dataset /path/hccl.json 8 fp32 2.6B 1 1 16 0"
+echo "bash run_distributed_train.sh /path/dataset /path/hccl.json 8 fp32 2.6B 1 1 16 0 8"
 echo "#######pipeline#######"
-echo "bash run_distributed_train.sh /path/dataset /path/hccl.json 16 fp32 2.6B 2 4 16 0"
-echo "bash run_distributed_train.sh /path/dataset /path/hccl.json 16 fp32 2.6B 2 4 16 8"
+echo "bash run_distributed_train.sh /path/dataset /path/hccl.json 16 fp32 2.6B 2 4 16 0 8"
+echo "bash run_distributed_train.sh /path/dataset /path/hccl.json 16 fp32 2.6B 2 4 16 8 8"
 echo "It is better to use absolute path."
 echo "=============================================================================================================="
 
@@ -36,8 +37,9 @@ STAGE_NUM=$6
 MICRO_SIZE=$7
 PER_BATCH=$8
 RANK_START=$9
+LOCAL_DEVICE_NUM=${10}
 
-for((i=0;i<${RANK_SIZE};i++));
+for((i=0;i<${LOCAL_DEVICE_NUM};i++));
 do
     rm ${ROOT_PATH}/device$i/ -rf
     mkdir ${ROOT_PATH}/device$i
diff --git a/model_zoo/official/nlp/pangu_alpha/train.py b/model_zoo/official/nlp/pangu_alpha/train.py
index a1e06a8e266..fd2a83a3784 100644
--- a/model_zoo/official/nlp/pangu_alpha/train.py
+++ b/model_zoo/official/nlp/pangu_alpha/train.py
@@ -102,7 +102,7 @@ def run_train(args_opt):
     """
     # Set execution mode
     context.set_context(mode=context.GRAPH_MODE, device_target=args_opt.device_target)
-    context.set_context(variable_memory_max_size="30GB")
+    context.set_context(variable_memory_max_size="31GB")
     # Set parallel context
     if args_opt.distribute == "true":
         D.init()