| 44 | class LLaMADecoderLayer(Module): |
| 45 | |
| 46 | def __init__(self, config: LLaMAConfig, layer_idx: int): |
| 47 | super().__init__() |
| 48 | self.layer_idx = layer_idx |
| 49 | layer_idx += config.layer_idx_offset |
| 50 | self.config = config |
| 51 | self.mapping = config.mapping |
| 52 | |
| 53 | if (self.config.use_input_layernorm_in_first_layer |
| 54 | and self.layer_idx == 0) or self.layer_idx > 0: |
| 55 | self.input_layernorm = RmsNorm(normalized_shape=config.hidden_size, |
| 56 | eps=config.norm_epsilon, |
| 57 | dtype=config.dtype) |
| 58 | |
| 59 | layers_range = config.mapping.pp_layers(config.num_hidden_layers) |
| 60 | self.local_layer_idx = layer_idx - layers_range[0] |
| 61 | self.is_last_local_layer = layer_idx == layers_range[-1] |
| 62 | self.attention = Attention( |
| 63 | local_layer_idx=self.local_layer_idx, |
| 64 | hidden_size=config.hidden_size, |
| 65 | attention_head_size=config.head_size, |
| 66 | num_attention_heads=config.num_attention_heads, |
| 67 | num_kv_heads=config.num_key_value_heads, |
| 68 | max_position_embeddings=config.max_position_embeddings, |
| 69 | dtype=config.dtype, |
| 70 | attention_mask_type=AttentionMaskType.causal, |
| 71 | bias=config.attn_bias, |
| 72 | position_embedding_type=PositionEmbeddingType.rope_gpt_neox, |
| 73 | rotary_embedding_base=config.rotary_base, |
| 74 | rotary_embedding_scaling=config.rotary_scaling, |
| 75 | tp_group=config.mapping.tp_group, |
| 76 | tp_size=config.mapping.tp_size, |
| 77 | tp_rank=config.mapping.tp_rank, |
| 78 | q_scaling=1.0 / config.attention_multiplier, |
| 79 | quant_mode=config.quant_mode, |
| 80 | cp_group=config.mapping.cp_group, |
| 81 | cp_size=config.mapping.cp_size, |
| 82 | cp_rank=config.mapping.cp_rank) |
| 83 | |
| 84 | mlp_hidden_size = config.hidden_size * 4 if config.intermediate_size is None else config.intermediate_size |
| 85 | |
| 86 | ClsMLP = GatedMLP |
| 87 | mlp_kwargs = {} |
| 88 | if config.moe.has_moe(): |
| 89 | ClsMLP = MOE |
| 90 | mlp_kwargs = { |
| 91 | "moe_config": config.moe, |
| 92 | "mapping": config.mapping, |
| 93 | } |
| 94 | self.mlp = ClsMLP(hidden_size=config.hidden_size, |
| 95 | ffn_hidden_size=mlp_hidden_size, |
| 96 | hidden_act=config.hidden_act, |
| 97 | dtype=config.dtype, |
| 98 | bias=config.mlp_bias, |
| 99 | tp_group=config.mapping.tp_group, |
| 100 | tp_size=config.mapping.tp_size, |
| 101 | quant_mode=config.quant_mode, |
| 102 | **mlp_kwargs) |
| 103 | |