From 035eca2485c14ecc654fa205af8e9205d49e1e37 Mon Sep 17 00:00:00 2001 From: huangxinjing Date: Sat, 4 Sep 2021 16:52:25 +0800 Subject: [PATCH] Fix performance degrades --- .../ccsrc/frontend/parallel/device_manager.h | 2 +- mindspore/parallel/nn/layers.py | 6 ++ mindspore/parallel/nn/transformer.py | 3 + .../nlp/pangu_alpha/src/pangu_alpha.py | 57 +++++++++++++++++-- 4 files changed, 62 insertions(+), 6 deletions(-) diff --git a/mindspore/ccsrc/frontend/parallel/device_manager.h b/mindspore/ccsrc/frontend/parallel/device_manager.h index 1d3d711d389..230ac9df622 100644 --- a/mindspore/ccsrc/frontend/parallel/device_manager.h +++ b/mindspore/ccsrc/frontend/parallel/device_manager.h @@ -35,7 +35,7 @@ namespace mindspore { namespace parallel { -#define MAX_DEVICE_NUM 1024 +#define MAX_DEVICE_NUM 4096 constexpr char HCCL_BACKEND[] = "hccl"; constexpr char NCCL_BACKEND[] = "nccl"; diff --git a/mindspore/parallel/nn/layers.py b/mindspore/parallel/nn/layers.py index 33d0044fa7f..d4a081a0242 100644 --- a/mindspore/parallel/nn/layers.py +++ b/mindspore/parallel/nn/layers.py @@ -359,6 +359,12 @@ class FixedSparseAttention(nn.Cell): Validator.check_positive_int(seq_length, "seq_length") Validator.check_positive_int(num_different_global_patterns, "num_different_global_patterns") dp, mp = parallel_config.data_parallel, parallel_config.model_parallel + if num_heads % mp != 0: + raise ValueError(f"The number of heads {num_heads} must be a " + f"multiple of parallel_config.model_parallel {mp}.") + if batch_size % dp != 0: + raise ValueError(f"The batch_size {batch_size} must be a " + f"multiple of parallel_config.data_parallel {parallel_config.data_parallel}.") self.seq_length = seq_length self.batch_size = batch_size self.hidden_size = size_per_head * num_heads diff --git a/mindspore/parallel/nn/transformer.py b/mindspore/parallel/nn/transformer.py index 0db195fb666..f81e0cf5c8b 100644 --- a/mindspore/parallel/nn/transformer.py +++ b/mindspore/parallel/nn/transformer.py @@ -598,6 +598,9 @@ class MultiHeadAttention(Cell): if num_heads % parallel_config.model_parallel != 0: raise ValueError(f"The number of heads {num_heads} must be a " f"multiple of parallel_config.model_parallel {parallel_config.model_parallel}.") + if batch_size % parallel_config.data_parallel != 0: + raise ValueError(f"The batch size {num_heads} must be a " + f"multiple of parallel_config.data_parallel {parallel_config.data_parallel}.") # Output layer self.projection = _Linear(in_channels=hidden_size, out_channels=hidden_size, 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 d8735f70316..2b461e9495f 100644 --- a/model_zoo/official/nlp/pangu_alpha/src/pangu_alpha.py +++ b/model_zoo/official/nlp/pangu_alpha/src/pangu_alpha.py @@ -14,6 +14,7 @@ # ============================================================================ """PanguAlpha model""" import os +import copy import numpy as np import mindspore.nn as nn import mindspore.common.dtype as mstype @@ -31,11 +32,15 @@ class EmbeddingLayer(nn.Cell): r"""Embedding layer of the PanGUAlpha Model""" def __init__(self, config): super(EmbeddingLayer, self).__init__() + # Only for the pipeline mode, the embedding needs to be row sliced. + copied_parallel_config = copy.deepcopy(config.parallel_config) + if copied_parallel_config.pipeline_stage > 1: + copied_parallel_config.vocab_emb_dp = False self.word_embedding = VocabEmbedding(vocab_size=config.vocab_size, embedding_size=config.hidden_size, param_init=initializer("normal", [config.vocab_size, config.hidden_size], dtype=config.param_init_type), - parallel_config=config.parallel_config.embedding_dp_mp_config) + parallel_config=copied_parallel_config.embedding_dp_mp_config) self.position_embedding = VocabEmbedding(vocab_size=config.seq_length, embedding_size=config.hidden_size, param_init=initializer("normal", @@ -180,6 +185,39 @@ class PanGuHead(Cell): return logits +def set_parallel_configure_for_layer(network, layer_id, offset, parallel_config, layers): + r""" + Default setting for the pipeline is: `(layer_id + offset) // (layers / pipeline_stage)`. + + + Args: + network(Cell) - Represents the transformer block + layer_id(int) - Means the layer index for the current module, counts from zero. + offset(int) - Means the layer_index needs a offset, if there are other modules in the net. + layers(int) - The total layers used for the model. + """ + # Used for the pipeline's stages setting + # As the final layer is not included here, so we need to manually add here. + # original: if set two stages, layers on two stages will be [15, 16+1] + # with 1 added, the layers on two stages will be [16, 15 +1] + pp_dis = max(int((layers + 1)/ parallel_config.pipeline_stage), 1) + # the pipeline stage must be in [0, parallel_config.pipeline_stage - 1] + pp_id = min((layer_id + offset) // pp_dis, parallel_config.pipeline_stage - 1) + network.pipeline_stage = pp_id + print(f"pipeline stage id is {pp_id}", flush=True) + + # Used for optimizer's fusion tag + dis = max(int(layers / parallel_config.gradient_aggregation_group), 1) + if parallel_config.pipeline_stage > 1: + # we give the fusion in pipeline mode a fixed value, otherwise the performance may become worse. + network.set_comm_fusion(2) + else: + network.set_comm_fusion(int((layer_id + offset) / dis) + 1) + # Used for enabling recomputation of the block + if parallel_config.recompute: + network.recompute() + + class PanguAlpha_Model(Cell): r"""The base backbone of the PanGuAlpha model""" def __init__(self, config): @@ -188,11 +226,13 @@ class PanguAlpha_Model(Cell): self.embedding = EmbeddingLayer(config) self.config = config self.layernorm = _LayerNorm((config.hidden_size,)).to_float(mstype.float32) - self.layernorm.set_comm_fusion(config.parallel_config.gradient_aggregation_group) + if config.parallel_config.pipeline_stage > 1: + self.layernorm.set_comm_fusion(2) + else: + self.layernorm.set_comm_fusion(config.parallel_config.gradient_aggregation_group) self.layernorm.shard(((config.parallel_config.data_parallel, 1, 1),)) self.layernorm.pipeline_stage = config.parallel_config.pipeline_stage - 1 # Configure the shard configure of the Embedding layer - self.embedding.set_comm_fusion(0) self.embedding.pipeline_stage = 0 self.num_layers = config.num_layers @@ -205,6 +245,7 @@ class PanguAlpha_Model(Cell): seq_length=config.seq_length, attention_dropout_rate=config.dropout_rate, hidden_dropout_rate=config.dropout_rate, + lambda_func=set_parallel_configure_for_layer, param_init_type=config.param_init_type, use_past=config.use_past, parallel_config=config.parallel_config).blocks @@ -216,7 +257,10 @@ class PanguAlpha_Model(Cell): dtype=config.param_init_type), parallel_config=config.parallel_config.embedding_dp_mp_config) self.top_query_embedding.pipeline_stage = config.parallel_config.pipeline_stage - 1 - self.top_query_embedding.set_comm_fusion(config.parallel_config.gradient_aggregation_group) + if config.parallel_config.pipeline_stage > 1: + self.top_query_embedding.set_comm_fusion(2) + else: + self.top_query_embedding.set_comm_fusion(config.parallel_config.gradient_aggregation_group) self.top_query_layer = QueryLayer(batch_size=config.batch_size, hidden_size=config.hidden_size, @@ -312,8 +356,11 @@ class PanguAlphaModel(nn.Cell): def __init__(self, config): super(PanguAlphaModel, self).__init__() # Network head to get logits over vocabulary + copied_parallel_config = copy.deepcopy(config.parallel_config) + if copied_parallel_config.pipeline_stage > 1: + copied_parallel_config.vocab_emb_dp = False self.head = PanGuHead(hidden_size=config.hidden_size, - parallel_config=config.parallel_config) + parallel_config=copied_parallel_config) self.head.pipeline_stage = config.parallel_config.pipeline_stage - 1 self.backbone = PanguAlpha_Model(config) self.backbone.embedding.word_embedding.embedding_table.add_pipeline_stage(self.head.pipeline_stage)