| 244 | |
| 245 | |
| 246 | class MixtralDecoderLayer(nn.Module): |
| 247 | |
| 248 | def __init__( |
| 249 | self, |
| 250 | config: MixtralConfig, |
| 251 | attention_backend: str, |
| 252 | linear_method: Optional[LinearMethodBase] = None, |
| 253 | ) -> None: |
| 254 | super().__init__() |
| 255 | self.hidden_size = config.hidden_size |
| 256 | # Requires transformers > 4.32.0 |
| 257 | rope_theta = getattr(config, "rope_theta", 10000) |
| 258 | self.self_attn = MixtralAttention( |
| 259 | hidden_size=self.hidden_size, |
| 260 | num_heads=config.num_attention_heads, |
| 261 | max_position=config.max_position_embeddings, |
| 262 | num_kv_heads=config.num_key_value_heads, |
| 263 | rope_theta=rope_theta, |
| 264 | sliding_window=config.sliding_window, |
| 265 | linear_method=linear_method, |
| 266 | attention_backend=attention_backend) |
| 267 | self.block_sparse_moe = MixtralMoE(config=config, |
| 268 | linear_method=linear_method) |
| 269 | self.input_layernorm = RMSNorm(config.hidden_size, |
| 270 | eps=config.rms_norm_eps) |
| 271 | self.post_attention_layernorm = RMSNorm(config.hidden_size, |
| 272 | eps=config.rms_norm_eps) |
| 273 | |
| 274 | def forward( |
| 275 | self, |
| 276 | positions: torch.Tensor, |
| 277 | hidden_states: torch.Tensor, |
| 278 | kv_cache: KVCache, |
| 279 | input_metadata: InputMetadata, |
| 280 | residual: Optional[torch.Tensor], |
| 281 | ) -> torch.Tensor: |
| 282 | # Self Attention |
| 283 | if residual is None: |
| 284 | residual = hidden_states |
| 285 | hidden_states = self.input_layernorm(hidden_states) |
| 286 | else: |
| 287 | hidden_states, residual = self.input_layernorm( |
| 288 | hidden_states, residual) |
| 289 | hidden_states = self.self_attn( |
| 290 | positions=positions, |
| 291 | hidden_states=hidden_states, |
| 292 | kv_cache=kv_cache, |
| 293 | input_metadata=input_metadata, |
| 294 | ) |
| 295 | |
| 296 | # Fully Connected |
| 297 | hidden_states, residual = self.post_attention_layernorm( |
| 298 | hidden_states, residual) |
| 299 | hidden_states = self.block_sparse_moe(hidden_states) |
| 300 | return hidden_states, residual |
| 301 | |
| 302 | |
| 303 | class MixtralModel(nn.Module): |