From da67a91c14575139b6a2c72dba196ec4db3bd29f Mon Sep 17 00:00:00 2001 From: yao_yf Date: Tue, 6 Jul 2021 12:46:38 +0800 Subject: [PATCH] add pangu gpu --- model_zoo/official/nlp/pangu_alpha/predict.py | 18 ++------ .../scripts/run_distribute_train_gpu.sh | 38 ++++++++++++++++ .../nlp/pangu_alpha/src/pangu_alpha.py | 43 ++++++++++--------- .../official/nlp/pangu_alpha/src/utils.py | 5 +++ model_zoo/official/nlp/pangu_alpha/train.py | 23 +++++----- 5 files changed, 78 insertions(+), 49 deletions(-) create mode 100644 model_zoo/official/nlp/pangu_alpha/scripts/run_distribute_train_gpu.sh diff --git a/model_zoo/official/nlp/pangu_alpha/predict.py b/model_zoo/official/nlp/pangu_alpha/predict.py index a30c96d9ee0..78bff0fd512 100644 --- a/model_zoo/official/nlp/pangu_alpha/predict.py +++ b/model_zoo/official/nlp/pangu_alpha/predict.py @@ -37,33 +37,21 @@ def load_model(args_opt): r""" The main function for load model """ - device_id = int(os.getenv("DEVICE_ID")) - rank_id_str = os.getenv('RANK_ID', '0') - rank_id = int( - rank_id_str[rank_id_str.rfind('-') + - 1:]) - print('rank_id:{}'.format(rank_id), "rank_id str:{}".format(rank_id_str)) - device_id = int(os.getenv('DEVICE_ID')) - local_rank = rank_id - print('local_rank:{}, device id:{} start to run...'.format(local_rank, device_id), flush=True) # Set execution mode context.set_context(save_graphs=False, mode=context.GRAPH_MODE, - device_target="Ascend", - device_id=device_id) + device_target=args_opt.device_target) context.set_context(variable_memory_max_size="30GB") # Set parallel context if args_opt.distribute == "true": D.init() device_num = D.get_group_size() rank = D.get_rank() - print("device_id is {}, rank_id is {}, device_num is {}".format( - device_id, rank, device_num)) + print("rank_id is {}, device_num is {}".format(rank, device_num)) context.reset_auto_parallel_context() context.set_auto_parallel_context( parallel_mode=ParallelMode.SEMI_AUTO_PARALLEL, gradients_mean=False, - device_num=device_num, full_batch=True, loss_repeated_mean=True, enable_parallel_optimizer=False, @@ -75,7 +63,7 @@ def load_model(args_opt): else: rank = 0 device_num = 1 - + print('local_rank:{}, start to run...'.format(rank), flush=True) use_past = False if args_opt.export: use_past = True diff --git a/model_zoo/official/nlp/pangu_alpha/scripts/run_distribute_train_gpu.sh b/model_zoo/official/nlp/pangu_alpha/scripts/run_distribute_train_gpu.sh new file mode 100644 index 00000000000..3fd78dd5d90 --- /dev/null +++ b/model_zoo/official/nlp/pangu_alpha/scripts/run_distribute_train_gpu.sh @@ -0,0 +1,38 @@ +#!/bin/bash +# Copyright 2020 Huawei Technologies Co., Ltd +# +# 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. +# ============================================================================ + +echo "==============================================================================================================" +echo "Please run the script as: " +echo "bash run_distributed_train_gpu.sh RANK_SIZE HOSTFILE DATASET MODE" +echo "for example: bash run_distributed_train_gpu.sh 16 hostfile_16p /mass_dataset/train_data/ 2.6B" +echo "It is better to use absolute path." +echo "==============================================================================================================" + +script_self=$(readlink -f "$0") +self_path=$(dirname "${script_self}") +RANK_SIZE=$1 +HOSTFILE=$2 +DATASET=$3 +MODE=$4 + +mpirun --allow-run-as-root -x PATH -x LD_LIBRARY_PATH -x PYTHONPATH -x NCCL_DEBUG -x GLOG_v -n $RANK_SIZE --hostfile $HOSTFILE --output-filename log_output --merge-stderr-to-stdout \ + python -s ${self_path}/../train.py \ + --distribute=true \ + --device_num=$RANK_SIZE \ + --device_target="GPU" \ + --data_url=$DATASET \ + --mode=$MODE \ + --run_type=train > train_log.txt 2>&1 & diff --git a/model_zoo/official/nlp/pangu_alpha/src/pangu_alpha.py b/model_zoo/official/nlp/pangu_alpha/src/pangu_alpha.py index 0806e287c40..269d0d7b731 100644 --- a/model_zoo/official/nlp/pangu_alpha/src/pangu_alpha.py +++ b/model_zoo/official/nlp/pangu_alpha/src/pangu_alpha.py @@ -76,6 +76,13 @@ class Dropout(nn.Cell): def extend_repr(self): return 'keep_prob={}, dtype={}'.format(self.keep_prob, self.dtype) + def shard(self, strategy): + if self.is_ascend: + self.dropout_gen_mask.shard(strategy) + self.dropout_do_mask.shard(strategy) + else: + self.dropout.shard(strategy) + class LayerNorm(nn.Cell): r""" A self-defined layer norm operation using reduce sum and reduce mean @@ -147,7 +154,7 @@ class Mapping(nn.Cell): return output -class Mapping_output(nn.Cell): +class MappingOutput(nn.Cell): """ A mapping function with a 3d input Args: @@ -161,7 +168,7 @@ class Mapping_output(nn.Cell): output: Tensor, a 3d tensor after projection """ def __init__(self, config, input_size, output_size, scale=1.0): - super(Mapping_output, self).__init__() + super(MappingOutput, self).__init__() self.output_size = output_size self.input_size = input_size self.weight = Parameter(initializer(Normal(sigma=0.02 * scale), @@ -203,14 +210,13 @@ class Output(nn.Cell): input_size = config.embedding_size output_size = config.embedding_size * config.expand_ratio # Project to expand_ratio*embedding_size - self.mapping = Mapping_output(config, input_size, output_size) + self.mapping = MappingOutput(config, input_size, output_size) # Project back to embedding_size self.projection = Mapping(config, output_size, input_size, scale) self.activation = nn.GELU() self.activation.gelu.shard(((config.dp, 1, config.mp),)) self.dropout = Dropout(1 - config.dropout_rate) - self.dropout.dropout_gen_mask.shard(((config.dp, 1, 1),)) - self.dropout.dropout_do_mask.shard(((config.dp, 1, 1),)) + self.dropout.shard(((config.dp, 1, 1),)) def construct(self, x): # [bs, seq_length, expand_ratio*embedding_size] @@ -282,13 +288,9 @@ class Attention(nn.Cell): self.coeff = Tensor(self.coeff) self.use_past = config.use_past self.dropout = Dropout(1 - config.dropout_rate) - self.dropout.dropout_gen_mask.shard(((config.dp, 1, 1),)) - self.dropout.dropout_do_mask.shard(((config.dp, 1, 1),)) + self.dropout.shard(((config.dp, 1, 1),)) self.prob_dropout = Dropout(1 - config.dropout_rate) - self.prob_dropout.dropout_gen_mask.shard( - ((config.dp, config.mp, 1, 1),)) - self.prob_dropout.dropout_do_mask.shard( - ((config.dp, config.mp, 1, 1),)) + self.prob_dropout.shard(((config.dp, config.mp, 1, 1),)) self.softmax = nn.Softmax() self.softmax.softmax.shard(((config.dp, config.mp, 1),)) self.expand_dims = P.ExpandDims().shard(((config.dp, 1, 1),)) @@ -631,12 +633,12 @@ class Decoder(nn.Cell): output = self.add(x, mlp_logit) return output, layer_present -class EmbeddingCell(nn.Cell): +class Embedding(nn.Cell): """ - EmbeddingCell + Embedding """ def __init__(self, config): - super(EmbeddingCell, self).__init__() + super(Embedding, self).__init__() self.word_embedding = EmbeddingLookup().set_comm_fusion(1) if config.word_emb_dp: self.word_embedding.gather.shard(((1, 1), (config.dp, 1))) @@ -669,8 +671,7 @@ class EmbeddingCell(nn.Cell): self.position_embedding.expand.shard(((config.dp, 1),)) self.add = P.TensorAdd().shard(((config.dp, 1, 1), (config.dp, 1, 1))) self.dropout = Dropout(1 - config.dropout_rate) - self.dropout.dropout_gen_mask.shard(((config.dp, 1, 1),)) - self.dropout.dropout_do_mask.shard(((config.dp, 1, 1),)) + self.dropout.shard(((config.dp, 1, 1),)) self.use_past = config.use_past self.is_first_iteration = True @@ -686,12 +687,12 @@ class EmbeddingCell(nn.Cell): return hidden_states -class MaskCell(nn.Cell): +class Mask(nn.Cell): """ - MaskCell + Mask """ def __init__(self, config): - super(MaskCell, self).__init__() + super(Mask, self).__init__() self.dtype = config.compute_dtype self.expand_dims = P.ExpandDims().shard(((config.dp, 1, 1),)) def construct(self, attention_mask): @@ -871,10 +872,10 @@ class PanguAlphaEmbedding(nn.Cell): """ def __init__(self, config): super(PanguAlphaEmbedding, self).__init__() - self.embedding = EmbeddingCell(config) + self.embedding = Embedding(config) if config.stage_num > 1: self.embedding.pipeline_stage = 0 - self.mask = MaskCell(config) + self.mask = Mask(config) def construct(self, input_ids, input_mask, table, input_position, attention_mask, valid_index=None): """ diff --git a/model_zoo/official/nlp/pangu_alpha/src/utils.py b/model_zoo/official/nlp/pangu_alpha/src/utils.py index 355cc071778..90d9fad9794 100644 --- a/model_zoo/official/nlp/pangu_alpha/src/utils.py +++ b/model_zoo/official/nlp/pangu_alpha/src/utils.py @@ -397,6 +397,11 @@ def get_args(inference=False): default="2.6B", choices=["200B", "13B", "2.6B", "self_define"], help="The scale of the model parameters") + parser.add_argument("--device_target", + type=str, + default="Ascend", + choices=["Ascend", "GPU"], + help="The running device") parser.add_argument("--strategy_load_ckpt_path", type=str, default="", diff --git a/model_zoo/official/nlp/pangu_alpha/train.py b/model_zoo/official/nlp/pangu_alpha/train.py index 777154b0fcd..4bb5b676cc8 100644 --- a/model_zoo/official/nlp/pangu_alpha/train.py +++ b/model_zoo/official/nlp/pangu_alpha/train.py @@ -44,12 +44,13 @@ class LossCallBack(Callback): If the loss in NAN or INF terminating training. """ - def __init__(self, dataset_size=-1, local_rank=0, has_trained_epoch=0, has_trained_step=0): + def __init__(self, dataset_size=-1, local_rank=0, has_trained_epoch=0, has_trained_step=0, micro_size=1): super(LossCallBack, self).__init__() self._dataset_size = dataset_size self.local_rank = local_rank self.has_trained_epoch = has_trained_epoch self.has_trained_step = has_trained_step + self.micro_size = micro_size print("load has trained epoch :{} and step: {}".format(has_trained_epoch, has_trained_step), flush=True) def step_end(self, run_context): @@ -63,9 +64,10 @@ class LossCallBack(Callback): if percent == 0: epoch_num -= 1 date = time.asctime(time.localtime(time.time())) + loss_value = cb_params.net_outputs[0].asnumpy() / self.micro_size print("time: {} local_rank: {}, epoch: {}, step: {}, output is {}, overflow is {}, scale is {}". format(date, int(self.local_rank), int(epoch_num) + int(self.has_trained_epoch), - cb_params.cur_step_num + int(self.has_trained_step), cb_params.net_outputs[0].asnumpy(), + cb_params.cur_step_num + int(self.has_trained_step), loss_value, cb_params.net_outputs[1].asnumpy(), cb_params.net_outputs[2].asnumpy())) @@ -78,25 +80,21 @@ def run_train(args_opt): r""" The main training process. """ - device_id = int(os.getenv('DEVICE_ID')) # Set execution mode context.set_context(mode=context.GRAPH_MODE, - device_target="Ascend", - device_id=device_id) + device_target=args_opt.device_target) context.set_context(variable_memory_max_size="30GB") # Set parallel context if args_opt.distribute == "true": D.init() device_num = D.get_group_size() rank = D.get_rank() - print("device_id is {}, rank_id is {}, device_num is {}".format( - device_id, rank, device_num)) + print("rank_id is {}, device_num is {}".format(rank, device_num)) context.reset_auto_parallel_context() context.set_auto_parallel_context( parallel_mode=ParallelMode.SEMI_AUTO_PARALLEL, gradients_mean=False, - device_num=device_num, full_batch=bool(args_opt.full_batch), enable_parallel_optimizer=bool(args_opt.optimizer_shard)) set_algo_parameters(elementwise_op_strategy_follow=True) @@ -105,7 +103,7 @@ def run_train(args_opt): else: rank = 0 device_num = 1 - + context.set_context(save_graphs=False, save_graphs_path="./graphs_of_device_id_" + str(rank)) # copy data from the cloud to the /cache/Data cache_url = '/cache/Data/' if args_opt.offline: @@ -194,18 +192,17 @@ def run_train_pipeline(args_opt): r""" The main training process in pipeline. """ - device_id = int(os.getenv("DEVICE_ID")) - context.set_context(save_graphs=False, mode=context.GRAPH_MODE, device_target="Ascend", device_id=device_id) + context.set_context(save_graphs=False, mode=context.GRAPH_MODE, device_target=args_opt.device_target) context.set_context(variable_memory_max_size="31GB") if args_opt.distribute == "true": D.init() device_num = D.get_group_size() rank_id = D.get_rank() + print("rank_id is {}, device_num is {}".format(rank_id, device_num)) context.reset_auto_parallel_context() context.set_auto_parallel_context( parallel_mode=ParallelMode.SEMI_AUTO_PARALLEL, gradients_mean=False, - device_num=device_num, full_batch=bool(args_opt.full_batch), loss_repeated_mean=True, enable_parallel_optimizer=bool(args_opt.optimizer_shard), @@ -281,7 +278,7 @@ def run_train_pipeline(args_opt): step_per_epoch = ds.get_dataset_size() callback_size = args_opt.sink_size actual_epoch_num = int(epoch_num * step_per_epoch / callback_size) - callback = [TimeMonitor(callback_size), LossCallBack(callback_size, rank_id, config.stage_num)] + callback = [TimeMonitor(callback_size), LossCallBack(callback_size, rank_id, config.stage_num, config.micro_size)] loss_scale_value = math.pow(2, 32) update_cell = DynamicLossScaleUpdateCell(loss_scale_value=loss_scale_value, scale_factor=2, scale_window=1000) pangu_alpha_with_grads = PanguAlphaTrainPipelineWithLossScaleCell(