| 35 | class ChatGLMDecoderLayer(Module): |
| 36 | |
| 37 | def __init__(self, config: ChatGLMConfig, layer_idx: int): |
| 38 | super().__init__() |
| 39 | self.layer_idx = layer_idx |
| 40 | self.config = config |
| 41 | self.chatglm_version = config.chatglm_version |
| 42 | |
| 43 | hidden_size = config.hidden_size |
| 44 | dtype = config.dtype |
| 45 | tp_group = config.mapping.tp_group |
| 46 | tp_size = config.mapping.tp_size |
| 47 | tp_rank = config.mapping.tp_rank |
| 48 | layernorm_epsilon = config.norm_epsilon |
| 49 | |
| 50 | self.apply_residual_connection_post_layernorm = config.apply_residual_connection_post_layernorm |
| 51 | self.alpha = (2 * config.num_hidden_layers)**0.5 |
| 52 | norm_cls = RmsNorm if config.rmsnorm else LayerNorm |
| 53 | |
| 54 | if config.chatglm_version == 'glm': |
| 55 | attention_mask_type = AttentionMaskType.bidirectionalglm |
| 56 | elif config.chatglm_version == 'chatglm': |
| 57 | attention_mask_type = AttentionMaskType.bidirectional |
| 58 | elif config.chatglm_version in GLM_ARCH2_VERSIONS: |
| 59 | attention_mask_type = AttentionMaskType.causal |
| 60 | |
| 61 | self.input_layernorm = norm_cls( |
| 62 | normalized_shape=hidden_size, |
| 63 | eps=layernorm_epsilon, |
| 64 | elementwise_affine=True, |
| 65 | dtype=dtype, |
| 66 | ) |
| 67 | |
| 68 | layers_range = config.mapping.pp_layers(config.num_hidden_layers) |
| 69 | local_layer_idx = layer_idx - layers_range[0] |
| 70 | self.attention = Attention( |
| 71 | local_layer_idx=local_layer_idx, |
| 72 | hidden_size=hidden_size, |
| 73 | num_attention_heads=config.num_attention_heads, |
| 74 | num_kv_heads=config.num_key_value_heads, |
| 75 | max_position_embeddings=config.max_position_embeddings, |
| 76 | num_layers=config.num_hidden_layers, |
| 77 | apply_query_key_layer_scaling=config.apply_query_key_layer_scaling, |
| 78 | attention_mask_type=attention_mask_type, |
| 79 | bias=config.add_qkv_bias, |
| 80 | dense_bias=config.add_bias_linear, |
| 81 | dtype=config.dtype, |
| 82 | position_embedding_type=config.position_embedding_type, |
| 83 | rotary_embedding_base=config.rotary_base, |
| 84 | rotary_embedding_scaling=config.rotary_scaling, |
| 85 | rotary_embedding_percentage=config.rotary_pct, |
| 86 | tp_group=tp_group, |
| 87 | tp_size=tp_size, |
| 88 | tp_rank=tp_rank, |
| 89 | quant_mode=config.quant_mode, |
| 90 | q_scaling=1.0, |
| 91 | cross_attention=False, |
| 92 | relative_attention=False, |
| 93 | max_distance=0, |
| 94 | num_buckets=0, |