| import torch | |
| import torch.nn as nn | |
| class NeighborExchange(nn.Module): | |
| def __init__(self, config: MeshConfig): | |
| super().__init__() | |
| self.config = config | |
| self.num_experts_x = config.mesh_grid_size[0] | |
| self.num_experts_y = config.mesh_grid_size[1] | |
| self.num_experts = self.num_experts_x * self.num_experts_y | |
| # Define parameters for neighbor communication. | |
| # A simple approach: a learned linear combination of neighbor features. | |
| # We can define a weight for each potential neighbor direction (e.g., up, down, left, right). | |
| # For a 2x2 grid, each expert has 2 or 3 neighbors. | |
| # A more general approach is a linear layer that takes concatenated neighbor features. | |
| # Let's use a linear layer to transform the aggregated neighbor information. | |
| # The input size to this layer will be the sum of hidden sizes of all potential neighbors | |
| # multiplied by the hidden size, but that's too complex. | |
| # A simpler approach: a linear layer per direction, or a single layer after aggregating. | |
| # Let's define a linear layer to process the information received from neighbors. | |
| # The input size is the hidden size (from neighbors), output size is hidden size | |
| # This layer will transform the aggregated neighbor features before adding to the expert's own output. | |
| self.exchange_projection = nn.Linear(config.hidden_size, config.hidden_size) # Projects aggregated neighbor info | |
| # Optional: Learned weights for different neighbor directions | |
| # self.neighbor_weights = nn.Parameter(torch.ones(4)) # Example for 4 directions (N, S, E, W) | |
| def forward(self, expert_outputs, expert_indices=None): | |
| # expert_outputs shape: (batch_size, sequence_length, num_experts, hidden_size) | |
| # expert_indices shape: (batch_size, sequence_length, k) - indices of selected experts (not directly used for neighbor exchange in this simple model) | |
| if not self.config.neighbor_exchange_enabled: | |
| return expert_outputs | |
| batch_size, seq_length, num_experts, hidden_size = expert_outputs.shape | |
| # Reshape expert_outputs to reflect the grid structure (batch_size, seq_length, grid_x, grid_y, hidden_size) | |
| reshaped_outputs = expert_outputs.view(batch_size, seq_length, self.num_experts_x, self.num_experts_y, hidden_size) | |
| # Create a tensor to store the aggregated neighbor information for each expert | |
| aggregated_neighbor_info = torch.zeros_like(reshaped_outputs) | |
| # Implement neighbor exchange logic | |
| # Iterate through each expert in the grid | |
| for i in range(self.num_experts_x): | |
| for j in range(self.num_experts_y): | |
| current_expert_output = reshaped_outputs[:, :, i, j, :] | |
| neighbor_info = torch.zeros_like(current_expert_output) # Accumulate info from neighbors | |
| # Define neighbor directions (example: up, down, left, right) | |
| neighbors = [] | |
| if i > 0: # Up neighbor | |
| neighbors.append(reshaped_outputs[:, :, i-1, j, :]) | |
| if i < self.num_experts_x - 1: # Down neighbor | |
| neighbors.append(reshaped_outputs[:, :, i+1, j, :]) | |
| if j > 0: # Left neighbor | |
| neighbors.append(reshaped_outputs[:, :, i, j-1, :]) | |
| if j < self.num_experts_y - 1: # Right neighbor | |
| neighbors.append(reshaped_outputs[:, :, i, j+1, :]) | |
| # Aggregate information from neighbors (simple average as an example) | |
| if neighbors: | |
| # Stack neighbors along a new dimension and take the mean | |
| neighbor_stack = torch.stack(neighbors, dim=-2) # shape (batch, seq, num_neighbors, hidden) | |
| aggregated_info = torch.mean(neighbor_stack, dim=-2) # shape (batch, seq, hidden) | |
| neighbor_info = aggregated_info # Use the aggregated info | |
| # Apply the exchange projection to the aggregated neighbor information | |
| transformed_neighbor_info = self.exchange_projection(neighbor_info) | |
| # Store the transformed neighbor info for the current expert's position | |
| aggregated_neighbor_info[:, :, i, j, :] = transformed_neighbor_info | |
| # Reshape aggregated_neighbor_info back to (batch_size, sequence_length, num_experts, hidden_size) | |
| aggregated_neighbor_info = aggregated_neighbor_info.view(batch_size, seq_length, num_experts, hidden_size) | |
| # Combine expert outputs with aggregated neighbor information (additive combination) | |
| exchanged_expert_outputs = expert_outputs + aggregated_neighbor_info | |
| return exchanged_expert_outputs | |