| import torch |
| import torch.distributed as dist |
| from typing import Union, Tuple, Optional |
|
|
| def solution( |
| gate_logits: Union[torch.Tensor, Tuple[torch.Tensor, ...]], |
| num_experts: int, |
| top_k: int = 2, |
| attention_mask: Optional[torch.Tensor] = None, |
| ) -> torch.Tensor: |
| if isinstance(gate_logits, (tuple, list)): |
| compute_device = gate_logits[0].device |
| concatenated_gate_logits = torch.cat( |
| [layer_gate.to(compute_device) for layer_gate in gate_logits], dim=0 |
| ) |
| else: |
| compute_device = gate_logits.device |
| concatenated_gate_logits = gate_logits |
|
|
| 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) |
| router_prob_per_expert = torch.mean(routing_weights, dim=0) |
| 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) / torch.sum(expert_attention_mask, dim=0) |
| |
| 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.unsqueeze(0)) |
| overall_loss = overall_loss * num_experts |
|
|
| if dist.is_available() and dist.is_initialized(): |
| dist.all_reduce(overall_loss, op=dist.ReduceOp.SUM) |
| overall_loss = overall_loss / dist.get_world_size() |
|
|
| return overall_loss |
|
|
