127 lines
4.4 KiB
Python
127 lines
4.4 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.utils.checkpoint as cp
|
|
from transformers.activations import ACT2FN
|
|
|
|
from wall_x.model.core.ops import permute, unpermute
|
|
|
|
|
|
class TokenTypeRouter(nn.Module):
|
|
def __init__(self, num_experts: int):
|
|
super().__init__()
|
|
self.num_experts = num_experts
|
|
|
|
def forward(self, token_types: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Route tokens to experts based on token_type.
|
|
|
|
Args:
|
|
token_types (torch.Tensor): Tensor of shape (batch_size, seq_length) containing each token type.
|
|
|
|
Returns:
|
|
experts_indices (torch.Tensor): Tensor of shape (batch_size, seq_length) containing each assigned expert index.
|
|
"""
|
|
# Simple rule: assign by token_type modulo the expert count
|
|
experts_indices = token_types % self.num_experts
|
|
return experts_indices
|
|
|
|
|
|
class BlockSparseMLP(nn.Module):
|
|
def __init__(self, config, use_selective_recompute: bool = False):
|
|
super().__init__()
|
|
self.hidden_size = config["hidden_size"]
|
|
self.intermediate_size = config["intermediate_size"]
|
|
self.hidden_act = config["hidden_act"]
|
|
|
|
self.use_selective_recompute = use_selective_recompute
|
|
|
|
self.gate_up_proj = nn.Linear(
|
|
self.hidden_size, 2 * self.intermediate_size, bias=False
|
|
)
|
|
self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
|
|
|
|
self.act_fn = ACT2FN[self.hidden_act]
|
|
|
|
# FIXME: The full MLP is recomputed for now; recomputing only activations can be optimized later.
|
|
def _full_mlp(self, hidden_state):
|
|
gate_up_out = self.gate_up_proj(hidden_state)
|
|
gate_out, up_out = gate_up_out.split(
|
|
[self.intermediate_size, self.intermediate_size], dim=-1
|
|
)
|
|
|
|
act_out = self.act_fn(gate_out) * up_out
|
|
return self.down_proj(act_out)
|
|
|
|
def forward(self, hidden_state):
|
|
if self.use_selective_recompute:
|
|
# Checkpoint-recompute the whole expert MLP
|
|
return cp.checkpoint(
|
|
self._full_mlp,
|
|
hidden_state,
|
|
use_reentrant=False,
|
|
)
|
|
else:
|
|
return self._full_mlp(hidden_state)
|
|
|
|
|
|
class SparseMoeBlock(nn.Module):
|
|
def __init__(self, config, num_experts: int, use_selective_recompute: bool = False):
|
|
super().__init__()
|
|
self.num_experts = num_experts
|
|
self.use_selective_recompute = use_selective_recompute
|
|
|
|
# Pass use_selective_recompute to each expert
|
|
self.experts = nn.ModuleList(
|
|
[
|
|
BlockSparseMLP(
|
|
config.experts[i], use_selective_recompute=use_selective_recompute
|
|
)
|
|
for i in range(num_experts)
|
|
]
|
|
)
|
|
|
|
if not hasattr(config, "dim_inputs") or not config.dim_inputs:
|
|
raise ValueError("config.dim_inputs must be set")
|
|
|
|
self.dim_inputs = config.dim_inputs
|
|
self.permuted = config.mot_opt
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
experts_indices: torch.Tensor,
|
|
start_indices: torch.Tensor,
|
|
end_indices: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
|
|
if self.permuted:
|
|
permuted_inputs = hidden_states
|
|
else:
|
|
batch_size, seq_length, hidden_dim = hidden_states.shape
|
|
|
|
flat_hidden = hidden_states.reshape(-1, hidden_dim)
|
|
experts_indices = experts_indices.reshape(-1)
|
|
probs = torch.ones_like(experts_indices, dtype=torch.float32).reshape(-1, 1)
|
|
permuted_inputs, row_id_map = permute(flat_hidden, experts_indices)
|
|
|
|
# buffer
|
|
final_output = torch.zeros_like(permuted_inputs)
|
|
|
|
# Expert forward, including selective recompute
|
|
for expert_idx, expert in enumerate(self.experts):
|
|
start, end = start_indices[expert_idx], end_indices[expert_idx]
|
|
if start == end:
|
|
continue
|
|
|
|
dim_input = self.dim_inputs[expert_idx]
|
|
expert_input = permuted_inputs[start:end, :dim_input]
|
|
|
|
partial_output = expert(expert_input)
|
|
final_output[start:end, :dim_input] = partial_output[:, :dim_input]
|
|
|
|
if self.permuted:
|
|
return final_output
|
|
else:
|
|
final_output = unpermute(final_output, row_id_map, probs)
|
|
return final_output.reshape(batch_size, seq_length, hidden_dim)
|