| 20 | class Phi3DecoderLayer(Module): |
| 21 | |
| 22 | def __init__(self, config: PretrainedConfig, layer_idx: int): |
| 23 | super().__init__() |
| 24 | self.config = config |
| 25 | self.layer_idx = layer_idx |
| 26 | tp_group = config.mapping.tp_group |
| 27 | tp_size = config.mapping.tp_size |
| 28 | |
| 29 | attention_mask_type = AttentionMaskType.causal |
| 30 | block_sparse_attn_params = BlockSparseAttnParams() |
| 31 | q_scaling = 1.0 |
| 32 | self.gegelu_limit = None |
| 33 | |
| 34 | self.small_variant = config.architecture == "Phi3SmallForCausalLM" |
| 35 | self.moe_variant = config.architecture == "PhiMoEForCausalLM" |
| 36 | if self.small_variant: |
| 37 | self.gegelu_limit = config.gegelu_limit |
| 38 | |
| 39 | # MuP uses norm_factor=attention_head_size (rather than sqrt(attention_head_size)) |
| 40 | # We achieve this using q_scaling = sqrt(attention_head_size) |
| 41 | hidden_size = config.hidden_size |
| 42 | num_attention_heads = config.num_attention_heads |
| 43 | attention_head_size = hidden_size / num_attention_heads |
| 44 | q_scaling = attention_head_size**.5 |
| 45 | |
| 46 | block_sparse = ((layer_idx + 1) % |
| 47 | config.dense_attention_every_n_layers) != 0 |
| 48 | attention_mask_type = AttentionMaskType.blocksparse if block_sparse else AttentionMaskType.causal |
| 49 | |
| 50 | block_sparse_attn_params = BlockSparseAttnParams( |
| 51 | config.blocksparse_block_size, |
| 52 | config.blocksparse_homo_head_pattern, |
| 53 | config.blocksparse_num_local_blocks, |
| 54 | config.blocksparse_vertical_stride) |
| 55 | |
| 56 | if self.small_variant or self.moe_variant: |
| 57 | self.input_layernorm = LayerNorm( |
| 58 | normalized_shape=config.hidden_size, |
| 59 | dtype=config.dtype, |
| 60 | eps=config.norm_epsilon) |
| 61 | self.post_layernorm = LayerNorm(normalized_shape=config.hidden_size, |
| 62 | dtype=config.dtype, |
| 63 | eps=config.norm_epsilon) |
| 64 | else: |
| 65 | self.input_layernorm = RmsNorm(normalized_shape=config.hidden_size, |
| 66 | eps=config.norm_epsilon, |
| 67 | dtype=config.dtype) |
| 68 | self.post_layernorm = RmsNorm(normalized_shape=config.hidden_size, |
| 69 | eps=config.norm_epsilon, |
| 70 | dtype=config.dtype) |
| 71 | |
| 72 | layers_range = config.mapping.pp_layers(config.num_hidden_layers) |
| 73 | local_layer_idx = layer_idx - layers_range[0] |
| 74 | position_embedding_type = PositionEmbeddingType.rope_gpt_neox |
| 75 | |
| 76 | rope_scaling_short_factors, rope_scaling_long_factors = None, None |
| 77 | rope_scaling_short_mscale, rope_scaling_long_mscale = None, None |
| 78 | original_max_position_embeddings = config.max_position_embeddings |
| 79 | |