mirror of
https://github.com/huggingface/text-generation-inference.git
synced 2025-04-21 14:52:20 +00:00
291 lines
10 KiB
Python
291 lines
10 KiB
Python
import torch
|
|
import torch.distributed
|
|
|
|
from torch import nn
|
|
from typing import Optional, List, Tuple, Any
|
|
from transformers.configuration_utils import PretrainedConfig
|
|
from transformers.modeling_outputs import CausalLMOutputWithPast
|
|
import torch.nn.functional as F
|
|
|
|
from text_generation_server.utils.layers import (
|
|
TensorParallelColumnLinear,
|
|
TensorParallelEmbedding,
|
|
)
|
|
|
|
|
|
class MambaConfig(PretrainedConfig):
|
|
def __init__(
|
|
self,
|
|
vocab_size=50280,
|
|
d_model=768,
|
|
n_layer=32,
|
|
layer_norm_epsilon=1e-5,
|
|
tie_word_embeddings=False,
|
|
pad_token_id=0,
|
|
bos_token_id=1,
|
|
eos_token_id=2,
|
|
**kwargs,
|
|
):
|
|
self.vocab_size = vocab_size
|
|
self.n_layer = n_layer
|
|
self.layer_norm_epsilon = layer_norm_epsilon
|
|
self.d_model = d_model
|
|
self.d_inner = d_model * 2
|
|
self.d_conv = 4
|
|
|
|
super().__init__(
|
|
pad_token_id=pad_token_id,
|
|
bos_token_id=bos_token_id,
|
|
eos_token_id=eos_token_id,
|
|
tie_word_embeddings=tie_word_embeddings,
|
|
**kwargs,
|
|
)
|
|
|
|
|
|
class MambaBlock(nn.Module):
|
|
def __init__(self, prefix, config, weights):
|
|
super().__init__()
|
|
|
|
# TODO: use model config to set the dt_rank instead of hardcoding it
|
|
self.dt_rank = (config.d_model + 15) // 16
|
|
|
|
# TODO: improve how we load the conv1d weights
|
|
# explore a transposed conv1d that avoids the need for
|
|
# a transpose during inference
|
|
self.conv1 = nn.Conv1d(
|
|
config.d_inner,
|
|
config.d_inner,
|
|
kernel_size=config.d_conv,
|
|
groups=config.d_inner,
|
|
padding=config.d_conv - 1,
|
|
)
|
|
self.conv1.weight = nn.Parameter(weights.get_tensor(f"{prefix}.conv1d.weight"))
|
|
self.conv1.bias = nn.Parameter(weights.get_tensor(f"{prefix}.conv1d.bias"))
|
|
|
|
self.dt_proj = TensorParallelColumnLinear.load(
|
|
config=config,
|
|
prefix=f"{prefix}.dt_proj",
|
|
weights=weights,
|
|
bias=True,
|
|
)
|
|
self.in_proj = TensorParallelColumnLinear.load(
|
|
config=config,
|
|
prefix=f"{prefix}.in_proj",
|
|
weights=weights,
|
|
bias=False,
|
|
)
|
|
self.x_proj = TensorParallelColumnLinear.load(
|
|
config=config,
|
|
prefix=f"{prefix}.x_proj",
|
|
weights=weights,
|
|
bias=False,
|
|
)
|
|
self.out_proj = TensorParallelColumnLinear.load(
|
|
config=config,
|
|
prefix=f"{prefix}.out_proj",
|
|
weights=weights,
|
|
bias=False,
|
|
)
|
|
|
|
# TODO: improve how we load the weights
|
|
self.A_log = nn.Parameter(weights.get_tensor(f"{prefix}.A_log"))
|
|
self.D = nn.Parameter(weights.get_tensor(f"{prefix}.D"))
|
|
|
|
def selective_scan(
|
|
self, input_tensor, delta, a_tensor, b_tensor, c_tensor, d_tensor
|
|
):
|
|
batch_size, sequence_length, input_dim = input_tensor.shape
|
|
num_cols = a_tensor.shape[1]
|
|
|
|
# TODO: revisit this math to avoid the transposes when possible
|
|
# reshape and process delta
|
|
delta = delta.transpose(1, 2).view((batch_size, input_dim, sequence_length, 1))
|
|
exp_delta_a = (delta * a_tensor.view((1, input_dim, 1, num_cols))).exp()
|
|
|
|
# calc involving delta, b_tensor, and input_tensor
|
|
delta_b_input = (
|
|
delta
|
|
* b_tensor.view((batch_size, 1, sequence_length, num_cols))
|
|
* input_tensor.transpose(1, 2).view(
|
|
(batch_size, input_dim, sequence_length, 1)
|
|
)
|
|
)
|
|
|
|
# init output tensor
|
|
output_tensor = torch.zeros(
|
|
(batch_size, input_dim, num_cols),
|
|
dtype=exp_delta_a.dtype,
|
|
device=exp_delta_a.device,
|
|
)
|
|
|
|
# iterate over sequence_length
|
|
output_sequence = []
|
|
for i in range(sequence_length):
|
|
multiplier = exp_delta_a[:, :, i]
|
|
output_tensor = (multiplier * output_tensor) + delta_b_input[:, :, i]
|
|
y = output_tensor.matmul(c_tensor[:, i, :].unsqueeze(2)).squeeze(2)
|
|
output_sequence.append(y)
|
|
|
|
stacked_output = torch.stack(output_sequence, 1)
|
|
return stacked_output + input_tensor * d_tensor
|
|
|
|
def ssm(self, hidden_states):
|
|
_input_dim, num_cols = self.A_log.shape
|
|
negative_exponential_a = self.A_log.exp().neg()
|
|
d_matrix = self.D
|
|
projected_hidden_states = self.x_proj(hidden_states)
|
|
|
|
# narrow operations for delta, b, and c
|
|
delta = projected_hidden_states.narrow(-1, 0, self.dt_rank)
|
|
b_matrix = projected_hidden_states.narrow(-1, self.dt_rank, num_cols)
|
|
c_matrix = projected_hidden_states.narrow(-1, self.dt_rank + num_cols, num_cols)
|
|
|
|
# process delta
|
|
delta = self.dt_proj(delta)
|
|
delta = torch.log(torch.exp(delta) + 1)
|
|
|
|
# apply selective scan
|
|
selective_scan_output = self.selective_scan(
|
|
hidden_states, delta, negative_exponential_a, b_matrix, c_matrix, d_matrix
|
|
)
|
|
return selective_scan_output
|
|
|
|
def forward(self, index, hidden_states, past_transformed_state):
|
|
sequence_length = hidden_states.shape[1]
|
|
|
|
# minimal amount of new work on single hidden state (previous hidden state are cached)
|
|
only_last = hidden_states[:, -1, :]
|
|
projected_only_last = self.in_proj(only_last)
|
|
transformed_only_last, residual_only_last = torch.chunk(
|
|
projected_only_last, 2, dim=-1
|
|
)
|
|
|
|
if past_transformed_state is not None:
|
|
# build a new transformed_states tensor with past_transformed_state and transformed_only_last
|
|
new_transformed_states = torch.cat(
|
|
[past_transformed_state, transformed_only_last.unsqueeze(1)], dim=1
|
|
)
|
|
transformed_states = new_transformed_states
|
|
residual_states = residual_only_last
|
|
else:
|
|
# prefilling the cache with the last transformed state
|
|
projected_states = self.in_proj(hidden_states)
|
|
split_states = torch.chunk(projected_states, 2, dim=-1)
|
|
transformed_states, residual_states = split_states
|
|
|
|
# NOTE: we need the past hidden states to produce the correct output
|
|
# therefore we cannot simply compute the most recent and append it as we
|
|
# did for the transformed states
|
|
|
|
# TODO: avoid the transpose by using a transposed conv1d
|
|
# apply convolution and narrowing operation
|
|
conv_output = (
|
|
self.conv1(transformed_states.transpose(1, 2))
|
|
.narrow(-1, 0, sequence_length)
|
|
.transpose(1, 2)
|
|
)
|
|
|
|
# apply silu (Swish) activation function
|
|
activated_transformed = F.silu(conv_output)
|
|
activated_residual = F.silu(residual_states)
|
|
|
|
# Subsequent operations
|
|
output = self.ssm(activated_transformed)
|
|
combined_output = output * activated_residual
|
|
|
|
return self.out_proj(combined_output), transformed_states
|
|
|
|
|
|
# TODO: prefer a more optimized implementation of RMSNorm if possible
|
|
class RMSNorm(nn.Module):
|
|
def __init__(self, num_features, eps=1e-8):
|
|
super().__init__()
|
|
self.num_features = num_features
|
|
self.eps = eps
|
|
self.scale = nn.Parameter(torch.ones(num_features))
|
|
|
|
def forward(self, x):
|
|
rms = torch.sqrt(torch.mean(x**2, dim=-1, keepdim=True) + self.eps)
|
|
x = x / rms
|
|
return self.scale * x
|
|
|
|
|
|
class ResidualBlock(nn.Module):
|
|
def __init__(self, layer_id, config, weights):
|
|
super().__init__()
|
|
self.layer_id = layer_id
|
|
self.mamba_block = MambaBlock(
|
|
prefix=f"{layer_id}.mixer", config=config, weights=weights
|
|
)
|
|
self.layer_norm = RMSNorm(config.d_model, eps=config.layer_norm_epsilon)
|
|
self.layer_norm.scale = nn.Parameter(
|
|
weights.get_tensor(f"{layer_id}.norm.weight")
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
index,
|
|
hidden_states,
|
|
past_transformed_state,
|
|
):
|
|
residual = hidden_states
|
|
hidden_states = self.layer_norm(hidden_states)
|
|
attn_outputs, transformed_states = self.mamba_block(
|
|
index, hidden_states, past_transformed_state
|
|
)
|
|
hidden_states = residual + attn_outputs
|
|
return hidden_states, transformed_states
|
|
|
|
|
|
class MambaModel(nn.Module):
|
|
def __init__(self, config, weights):
|
|
super().__init__()
|
|
self.tp_rank = weights.process_group.rank()
|
|
self.tp_world_size = weights.process_group.size()
|
|
self.embed_tokens = TensorParallelEmbedding(
|
|
prefix="backbone.embedding", weights=weights
|
|
)
|
|
self.blocks = nn.ModuleList(
|
|
[
|
|
ResidualBlock(f"backbone.layers.{layer_id}", config, weights)
|
|
for layer_id in range(config.n_layer)
|
|
]
|
|
)
|
|
|
|
# TODO: avoid hardcoded sizes and improve how we load the weights
|
|
self.norm_f = RMSNorm(config.d_model, eps=config.layer_norm_epsilon)
|
|
self.norm_f.scale = nn.Parameter(weights.get_tensor(f"backbone.norm_f.weight"))
|
|
# use the same weights for the embedding and the final layer norm
|
|
self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
|
|
self.lm_head.weight = nn.Parameter(
|
|
self.embed_tokens.weight[: config.vocab_size, :]
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
input_ids: torch.LongTensor,
|
|
past_input_ids: Optional[List[Tuple[torch.FloatTensor]]] = None,
|
|
past_transformed_states: Optional[List[Tuple[torch.FloatTensor]]] = None,
|
|
) -> Tuple[torch.Tensor, List[Tuple[torch.Tensor, torch.Tensor]]]:
|
|
# NOTE: we need all input_ids to compute the correct embeddings
|
|
if past_input_ids is not None:
|
|
input_ids = past_input_ids
|
|
|
|
hidden_states = self.embed_tokens(input_ids)
|
|
|
|
past_transformed_states = (
|
|
[None] * len(self.blocks)
|
|
if past_transformed_states is None
|
|
else past_transformed_states
|
|
)
|
|
|
|
for index, block in enumerate(self.blocks):
|
|
hidden_states, transformed_states = block(
|
|
index, hidden_states, past_transformed_states[index]
|
|
)
|
|
past_transformed_states[index] = transformed_states
|
|
|
|
final_hidden_states = self.norm_f(hidden_states)
|
|
after_lm_head = self.lm_head(final_hidden_states)
|
|
return after_lm_head, input_ids, past_transformed_states
|