| import torch |
| import torch.nn.functional as F |
|
|
|
|
| def rotate_half(x): |
| """Rotates half the hidden dims of the input.""" |
| x1 = x[..., : x.shape[-1] // 2] |
| x2 = x[..., x.shape[-1] // 2 :] |
| return torch.cat((-x2, x1), dim=-1) |
|
|
|
|
| def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): |
| cos = cos.unsqueeze(unsqueeze_dim) |
| sin = sin.unsqueeze(unsqueeze_dim) |
| q_embed = (q * cos) + (rotate_half(q) * sin) |
| k_embed = (k * cos) + (rotate_half(k) * sin) |
| return q_embed, k_embed |
|
|
|
|
| def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: |
| batch, num_key_value_heads, slen, head_dim = hidden_states.shape |
| if n_rep == 1: |
| return hidden_states |
| hidden_states = hidden_states[:, :, None, :, :].expand( |
| batch, num_key_value_heads, n_rep, slen, head_dim |
| ) |
| return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) |
|
|
|
|
| def _get_unpad_data(attention_mask): |
| seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) |
| indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten() |
| max_seqlen_in_batch = seqlens_in_batch.max().item() |
| cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) |
| return indices, cu_seqlens, max_seqlen_in_batch |
|
|
|
|
| def load_balancing_loss_func( |
| gate_logits, num_experts=None, top_k=2, attention_mask=None |
| ): |
| if gate_logits is None or not isinstance(gate_logits, tuple): |
| return 0 |
|
|
| if isinstance(gate_logits, tuple): |
| compute_device = gate_logits[0].device |
| concatenated_gate_logits = torch.cat( |
| [layer_gate.to(compute_device) for layer_gate in gate_logits], dim=0 |
| ) |
|
|
| routing_weights = torch.nn.functional.softmax(concatenated_gate_logits, dim=-1) |
| _, selected_experts = torch.topk(routing_weights, top_k, dim=-1) |
| expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts) |
|
|
| if attention_mask is None: |
| |
| tokens_per_expert = torch.mean(expert_mask.float(), dim=(0, 1)) |
| router_prob_per_expert = torch.mean(routing_weights, dim=(0, 1)) |
| else: |
| batch_size, sequence_length = attention_mask.shape |
| num_hidden_layers = concatenated_gate_logits.shape[0] // ( |
| batch_size * sequence_length |
| ) |
|
|
| expert_attention_mask = ( |
| attention_mask[None, :, :, None, None] |
| .expand( |
| (num_hidden_layers, batch_size, sequence_length, top_k, num_experts) |
| ) |
| .reshape(-1, top_k, num_experts) |
| .to(compute_device) |
| ) |
|
|
| tokens_per_expert = torch.sum( |
| expert_mask.float() * expert_attention_mask, dim=(0, 1) |
| ) / torch.sum(expert_attention_mask, dim=(0, 1)) |
|
|
| router_per_expert_attention_mask = ( |
| attention_mask[None, :, :, None] |
| .expand((num_hidden_layers, batch_size, sequence_length, num_experts)) |
| .reshape(-1, num_experts) |
| .to(compute_device) |
| ) |
|
|
| router_prob_per_expert = torch.sum( |
| routing_weights * router_per_expert_attention_mask, dim=0 |
| ) / torch.sum(router_per_expert_attention_mask, dim=0) |
|
|
| overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert) |
| return overall_loss * num_experts |
|
|
|
|
| def count_mem_blocks_in_config(config): |
| num_blocks = 0 |
| for val in config.layers_block_type: |
| name = val.split("_")[0] if "_" in val else val |
| if name == "HelixSeekDelta": |
| num_blocks += 1 |
| return num_blocks |
|
|
|
|
| def layer_type_list(config): |
| ll = [] |
| for i, val in enumerate(config.layers_block_type): |
| name = val.split("_")[0] if "_" in val else val |
| if name == "HelixSeekDelta": |
| ll.append(i) |
| return ll |
|
|
