| 16 | class CohereDecoderLayer(Module): |
| 17 | |
| 18 | def __init__(self, config: CohereConfig, layer_idx: int): |
| 19 | super().__init__() |
| 20 | self.layer_idx = layer_idx |
| 21 | self.config = config |
| 22 | |
| 23 | self.input_layernorm = LayerNorm(normalized_shape=config.hidden_size, |
| 24 | eps=config.norm_epsilon, |
| 25 | bias=False, |
| 26 | dtype=config.dtype) |
| 27 | |
| 28 | layers_range = config.mapping.pp_layers(config.num_hidden_layers) |
| 29 | self.local_layer_idx = layer_idx - layers_range[0] |
| 30 | self.attention = Attention( |
| 31 | local_layer_idx=self.local_layer_idx, |
| 32 | hidden_size=config.hidden_size, |
| 33 | attention_head_size=config.head_size, |
| 34 | num_attention_heads=config.num_attention_heads, |
| 35 | num_kv_heads=config.num_key_value_heads, |
| 36 | max_position_embeddings=config.max_position_embeddings, |
| 37 | dtype=config.dtype, |
| 38 | attention_mask_type=AttentionMaskType.causal, |
| 39 | bias=config.attn_bias, |
| 40 | position_embedding_type=PositionEmbeddingType.rope_gptj, |
| 41 | rotary_embedding_base=config.rotary_base, |
| 42 | tp_group=config.mapping.tp_group, |
| 43 | tp_size=config.mapping.tp_size, |
| 44 | tp_rank=config.mapping.tp_rank, |
| 45 | qk_layernorm=config.qk_layernorm, |
| 46 | layernorm_share=False, |
| 47 | eps=config.norm_epsilon, |
| 48 | quant_mode=config.quant_mode) |
| 49 | |
| 50 | self.mlp = GatedMLP(hidden_size=config.hidden_size, |
| 51 | ffn_hidden_size=config.intermediate_size, |
| 52 | hidden_act=config.hidden_act, |
| 53 | dtype=config.dtype, |
| 54 | bias=False, |
| 55 | tp_group=config.mapping.tp_group, |
| 56 | tp_size=config.mapping.tp_size, |
| 57 | quant_mode=config.quant_mode) |
| 58 | |
| 59 | def forward(self, |
| 60 | hidden_states, |