| 208 | |
| 209 | |
| 210 | class TeleChatDecoderLayer(nn.Module): |
| 211 | |
| 212 | def __init__( |
| 213 | self, |
| 214 | config: PretrainedConfig, |
| 215 | cache_config: Optional[CacheConfig] = None, |
| 216 | quant_config: Optional[QuantizationConfig] = None, |
| 217 | ) -> None: |
| 218 | super().__init__() |
| 219 | self.hidden_size = config.hidden_size |
| 220 | rope_theta = getattr(config, "rope_theta", 10000) |
| 221 | rope_scaling = getattr(config, "rope_scaling", None) |
| 222 | if rope_scaling is not None and getattr( |
| 223 | config, "original_max_position_embeddings", None): |
| 224 | rope_scaling["original_max_position_embeddings"] = ( |
| 225 | config.original_max_position_embeddings) |
| 226 | max_position_embeddings = getattr(config, "max_position_embeddings", |
| 227 | 8192) |
| 228 | # Support abacusai/Smaug-72B-v0.1 with attention_bias |
| 229 | # Support internlm/internlm-7b with bias |
| 230 | attention_bias = getattr(config, "attention_bias", False) or getattr( |
| 231 | config, "bias", False) |
| 232 | self.self_attention = TeleChatAttention( |
| 233 | config, |
| 234 | hidden_size=self.hidden_size, |
| 235 | num_heads=config.num_attention_heads, |
| 236 | num_kv_heads=getattr(config, "num_key_value_heads", |
| 237 | config.num_attention_heads), |
| 238 | rope_theta=rope_theta, |
| 239 | rope_scaling=rope_scaling, |
| 240 | max_position_embeddings=max_position_embeddings, |
| 241 | quant_config=quant_config, |
| 242 | bias=attention_bias, |
| 243 | cache_config=cache_config, |
| 244 | ) |
| 245 | self.mlp = TeleChatMLP( |
| 246 | hidden_size=self.hidden_size, |
| 247 | intermediate_size=config.intermediate_size, |
| 248 | hidden_act=config.hidden_act, |
| 249 | quant_config=quant_config, |
| 250 | bias=getattr(config, "mlp_bias", False), |
| 251 | ) |
| 252 | self.input_layernorm = RMSNorm(config.hidden_size, |
| 253 | eps=config.rms_norm_eps) |
| 254 | self.post_attention_layernorm = RMSNorm(config.hidden_size, |
| 255 | eps=config.rms_norm_eps) |
| 256 | self.apply_residual_connection_post_layernorm = config.apply_residual_connection_post_layernorm |
| 257 | |
| 258 | |
| 259 | def forward( |
| 260 | self, |
| 261 | positions: torch.Tensor, |
| 262 | hidden_states: torch.Tensor, |
| 263 | kv_cache: torch.Tensor, |
| 264 | attn_metadata: AttentionMetadata, |
| 265 | ) -> Tuple[torch.Tensor, torch.Tensor]: |
| 266 | residual = hidden_states |
| 267 | layernorm_output = self.input_layernorm(hidden_states) |