r"""The base backbone of the PanGuAlpha model
| 290 | |
| 291 | |
| 292 | class PanguAlpha_Model(Cell): |
| 293 | r"""The base backbone of the PanGuAlpha model""" |
| 294 | |
| 295 | def __init__(self, config): |
| 296 | super(PanguAlpha_Model, self).__init__() |
| 297 | self.is_pipeline = config.parallel_config.pipeline_stage > 1 |
| 298 | self.embedding = EmbeddingLayer(config) |
| 299 | self.config = config |
| 300 | self.layernorm = _LayerNorm((config.hidden_size,)).to_float( |
| 301 | mstype.float32 |
| 302 | ) |
| 303 | if config.parallel_config.pipeline_stage > 1: |
| 304 | self.layernorm.set_comm_fusion(2) |
| 305 | else: |
| 306 | self.layernorm.set_comm_fusion( |
| 307 | config.parallel_config.gradient_aggregation_group |
| 308 | ) |
| 309 | self.layernorm.shard(((config.parallel_config.data_parallel, 1),)) |
| 310 | self.layernorm.pipeline_stage = ( |
| 311 | config.parallel_config.pipeline_stage - 1 |
| 312 | ) |
| 313 | # Configure the shard configure of the Embedding layer |
| 314 | self.embedding.pipeline_stage = 0 |
| 315 | self.num_layers = config.num_layers |
| 316 | if config.use_moe: |
| 317 | moe_config = MoEConfig( |
| 318 | expert_num=config.parallel_config.data_parallel |
| 319 | * config.per_dp_dim_expert_num |
| 320 | ) |
| 321 | else: |
| 322 | moe_config = MoEConfig(expert_num=1) |
| 323 | # The shard setting of Transformer is set within the class StackedTransformer |
| 324 | self.blocks = TransformerEncoder(num_layers=config.num_layers - 1, |
| 325 | batch_size=config.batch_size, |
| 326 | hidden_size=config.hidden_size, |
| 327 | ffn_hidden_size=config.ffn_hidden_size, |
| 328 | num_heads=config.num_heads, |
| 329 | seq_length=config.seq_length, |
| 330 | attention_dropout_rate=config.dropout_rate, |
| 331 | hidden_dropout_rate=config.dropout_rate, |
| 332 | lambda_func=set_parallel_configure_for_layer, |
| 333 | hidden_act="fast_gelu", |
| 334 | param_init_type=config.param_init_type, |
| 335 | use_past=config.use_past, |
| 336 | parallel_config=config.parallel_config, |
| 337 | moe_config=moe_config, |
| 338 | softmax_compute_type=config.softmax_compute_type).blocks |
| 339 | for block in self.blocks: |
| 340 | block.attention.dense1.bias.parallel_optimizer = False |
| 341 | block.attention.dense2.bias.parallel_optimizer = False |
| 342 | block.attention.dense3.bias.parallel_optimizer = False |
| 343 | block.output.mapping.bias.parallel_optimizer = False |
| 344 | copied_parallel_config = copy.deepcopy(config.parallel_config) |
| 345 | copied_parallel_config.vocab_emb_dp = True |
| 346 | self.top_query_embedding = VocabEmbedding(vocab_size=config.seq_length, |
| 347 | embedding_size=config.hidden_size, |
| 348 | param_init=initializer("normal", |
| 349 | [config.seq_length, config.hidden_size], |