| 161 | from functools import partial |
| 162 | |
| 163 | def mlp_forward_default(self, hidden_states, expert_id=-1, **kw_args): |
| 164 | if self.transformer.num_experts == 1 or expert_id > -1: |
| 165 | self = self.transformer.layers[kw_args['layer_id']].mlp |
| 166 | suffix = f"_{expert_id}" if expert_id > 0 else "" |
| 167 | if self.is_gated_mlp: |
| 168 | intermediate_parallel = getattr(self, "dense_h_to_4h"+suffix)(hidden_states) |
| 169 | gated_intermediate_parallel = getattr(self, "dense_h_to_4h_gate"+suffix)(hidden_states) |
| 170 | intermediate_parallel = self.activation_func(gated_intermediate_parallel) * intermediate_parallel |
| 171 | output = getattr(self, "dense_4h_to_h"+suffix)(intermediate_parallel) |
| 172 | else: |
| 173 | intermediate_parallel = getattr(self, "dense_h_to_4h"+suffix)(hidden_states) |
| 174 | intermediate_parallel = self.activation_func(intermediate_parallel) |
| 175 | output = getattr(self, "dense_4h_to_h"+suffix)(intermediate_parallel) |
| 176 | return output |
| 177 | else: |
| 178 | mlp_forward = self.hooks.get('mlp_forward', partial(mlp_forward_default, self)) |
| 179 | routing_forward = self.hooks.get('routing_forward', partial(routing_forward_default, self)) |
| 180 | self = self.transformer.layers[kw_args['layer_id']].mlp |
| 181 | fwd_weight, fwd_idx = routing_forward(hidden_states, **kw_args) |
| 182 | |
| 183 | # Adapted from mixtral-8x7b https://github.com/huggingface/transformers/blob/main/src/transformers/models/mixtral/modeling_mixtral.py |
| 184 | batch_size, sequence_length, hidden_dim = hidden_states.shape |
| 185 | hidden_states = hidden_states.view(-1, hidden_dim) |
| 186 | final_hidden_states = torch.zeros( |
| 187 | (batch_size * sequence_length, hidden_dim), dtype=hidden_states.dtype, device=hidden_states.device |
| 188 | ) |
| 189 | # One hot encode the selected experts to create an expert mask |
| 190 | # this will be used to easily index which expert is going to be sollicitated |
| 191 | expert_mask = torch.nn.functional.one_hot(fwd_idx, num_classes=self.num_experts).permute(2, 1, 0) |
| 192 | # Loop over all available experts in the model and perform the computation on each expert |
| 193 | for expert_idx in range(self.num_experts): |
| 194 | idx, top_x = torch.where(expert_mask[expert_idx]) |
| 195 | if top_x.shape[0] == 0: |
| 196 | continue |
| 197 | # in torch it is faster to index using lists than torch tensors |
| 198 | top_x_list = top_x.tolist() |
| 199 | idx_list = idx.tolist() |
| 200 | # Index the correct hidden states and compute the expert hidden state for |
| 201 | # the current expert. We need to make sure to multiply the output hidden |
| 202 | # states by `routing_weights` on the corresponding tokens (top-1 and top-2) |
| 203 | current_state = hidden_states[top_x_list] # I don't know why using hidden_states[None, top_x_list].reshape(-1, hidden_dim) |
| 204 | current_hidden_states = mlp_forward(current_state, expert_id=expert_idx, **kw_args) * fwd_weight[top_x_list, idx_list, None] |
| 205 | # However `index_add_` only support torch tensors for indexing so we'll use |
| 206 | # the `top_x` tensor here. |
| 207 | final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype)) |
| 208 | output = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim) |
| 209 | return output |
| 210 | |
| 211 | def word_embedding_forward_default(self, input_ids, output_cross_layer, **kw_args): |
| 212 | return self.transformer.word_embeddings(input_ids) |