| 432 | |
| 433 | |
| 434 | class BaseTransformer(torch.nn.Module): |
| 435 | def __init__(self, |
| 436 | num_layers, |
| 437 | vocab_size, |
| 438 | hidden_size, |
| 439 | num_attention_heads, |
| 440 | max_sequence_length, |
| 441 | embedding_dropout_prob=0, |
| 442 | attention_dropout_prob=0, |
| 443 | output_dropout_prob=0, |
| 444 | drop_path=0, |
| 445 | checkpoint_activations=False, |
| 446 | checkpoint_num_layers=1, |
| 447 | checkpoint_skip_layers=0, |
| 448 | layernorm_epsilon=1.0e-5, |
| 449 | init_method_std=0.02, |
| 450 | inner_hidden_size=None, |
| 451 | hidden_size_per_attention_head=None, |
| 452 | cross_hidden_size_per_attention_head=None, |
| 453 | layernorm_order='pre', |
| 454 | parallel_output=False, |
| 455 | is_decoder=False, |
| 456 | cross_attn_hidden_size=None, |
| 457 | use_bias=True, |
| 458 | use_qkv_bias=False, |
| 459 | num_multi_query_heads=0, |
| 460 | cross_num_multi_query_heads=0, |
| 461 | row_parallel_linear_final_bias=True, |
| 462 | activation_func=gelu, |
| 463 | is_gated_mlp=False, |
| 464 | is_rotary_emb=False, |
| 465 | num_experts=1, |
| 466 | layernorm=LayerNorm, |
| 467 | init_method=None, |
| 468 | use_final_layernorm=True, |
| 469 | hooks={}, |
| 470 | params_dtype=torch.float, |
| 471 | skip_init=False, |
| 472 | device=torch.device('cpu') |
| 473 | ): |
| 474 | super(BaseTransformer, self).__init__() |
| 475 | |
| 476 | # recording parameters |
| 477 | self.hidden_size = hidden_size |
| 478 | self.inner_hidden_size = inner_hidden_size |
| 479 | self.hidden_size_per_attention_head = hidden_size_per_attention_head |
| 480 | self.cross_hidden_size_per_attention_head = cross_hidden_size_per_attention_head |
| 481 | self.is_decoder = is_decoder |
| 482 | self.cross_attn_hidden_size = cross_attn_hidden_size |
| 483 | self.cross_num_multi_query_heads = cross_num_multi_query_heads |
| 484 | if not is_decoder and cross_attn_hidden_size is not None: |
| 485 | print('warning: cross_attn_hidden_size is set but is_decoder is False') |
| 486 | self.use_bias = use_bias |
| 487 | self.use_qkv_bias = use_qkv_bias |
| 488 | self.num_multi_query_heads = num_multi_query_heads |
| 489 | self.is_gated_mlp = is_gated_mlp |
| 490 | self.is_rotary_emb = is_rotary_emb |
| 491 | self.num_experts = num_experts |