| 29 | class FalconDecoderLayer(Module): |
| 30 | |
| 31 | def __init__(self, config: FalconConfig, layer_idx: int): |
| 32 | super().__init__() |
| 33 | self.layer_idx = layer_idx |
| 34 | self.config = config |
| 35 | |
| 36 | hidden_size = config.hidden_size |
| 37 | dtype = config.dtype |
| 38 | tp_group = config.mapping.tp_group |
| 39 | tp_size = config.mapping.tp_size |
| 40 | tp_rank = config.mapping.tp_rank |
| 41 | layernorm_epsilon = config.norm_epsilon |
| 42 | |
| 43 | self.input_layernorm = LayerNorm(normalized_shape=hidden_size, |
| 44 | eps=layernorm_epsilon, |
| 45 | dtype=dtype) |
| 46 | |
| 47 | self.new_decoder_architecture = config.new_decoder_architecture |
| 48 | self.parallel_attn = config.parallel_attention |
| 49 | self.num_ln_in_parallel_attn = config.num_ln_in_parallel_attn |
| 50 | if self.num_ln_in_parallel_attn is None and self.new_decoder_architecture: |
| 51 | self.num_ln_in_parallel_attn = 2 |
| 52 | if self.is_parallel_attention: |
| 53 | # Not to apply allreduce inside the Attention/MLP layers. |
| 54 | # allreduce applies after those layer. |
| 55 | tp_group = None |
| 56 | layers_range = config.mapping.pp_layers(config.num_hidden_layers) |
| 57 | local_layer_idx = layer_idx - layers_range[0] |
| 58 | self.attention = Attention( |
| 59 | local_layer_idx=local_layer_idx, |
| 60 | hidden_size=hidden_size, |
| 61 | num_attention_heads=config.num_attention_heads, |
| 62 | num_kv_heads=config.num_key_value_heads, |
| 63 | max_position_embeddings=config.max_position_embeddings, |
| 64 | attention_mask_type=AttentionMaskType.causal, |
| 65 | dtype=dtype, |
| 66 | tp_group=tp_group, |
| 67 | tp_size=tp_size, |
| 68 | tp_rank=tp_rank, |
| 69 | bias=config.bias, |
| 70 | position_embedding_type=config.position_embedding_type, |
| 71 | rotary_embedding_base=config.rotary_base, |
| 72 | quant_mode=config.quantization.quant_mode, |
| 73 | ) |
| 74 | |
| 75 | mlp_hidden_size = hidden_size * 4 if config.intermediate_size is None else config.intermediate_size |
| 76 | |
| 77 | if self.new_decoder_architecture and self.num_ln_in_parallel_attn == 2: |
| 78 | # Layernorm before MLP. |
| 79 | self.mlp_layernorm = LayerNorm(normalized_shape=hidden_size, |
| 80 | eps=layernorm_epsilon, |
| 81 | dtype=dtype) |
| 82 | else: |
| 83 | self.mlp_layernorm = None |
| 84 | self.mlp = MLP( |
| 85 | hidden_size=hidden_size, |
| 86 | ffn_hidden_size=mlp_hidden_size, |
| 87 | hidden_act=config.hidden_act, |
| 88 | dtype=dtype, |