mirror of
https://github.com/xai-org/grok-1.git
synced 2024-11-24 12:39:54 +03:00
402 lines
15 KiB
Python
402 lines
15 KiB
Python
from configuration_grok_1 import Grok1Config
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from transformers import PreTrainedModel
|
|
from transformers.activations import ACT2FN
|
|
from typing import Optional, NamedTuple, List
|
|
|
|
class TiedWeightEmbedding(nn.Embedding):
|
|
"""Module for tied weight embedding."""
|
|
|
|
def __init__(
|
|
self,
|
|
config: Grok1Config,
|
|
):
|
|
super().__init__(
|
|
num_embeddings=config.vocab_size,
|
|
embedding_dim=config.hidden_size,
|
|
padding_idx=config.pad_token_id,
|
|
)
|
|
|
|
def decode(
|
|
self,
|
|
inputs: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
return torch.matmul(inputs, self.weight.T)
|
|
|
|
class Gating(nn.Module):
|
|
"""Gating module for spare MoE expert selection."""
|
|
|
|
def __init__(
|
|
self,
|
|
config: Grok1Config,
|
|
):
|
|
super().__init__()
|
|
self.router_weights = nn.Linear(config.hidden_size, config.num_local_experts, bias=False)
|
|
|
|
def forward(
|
|
self,
|
|
inputs: torch.Tensor,
|
|
padding_mask: Optional[torch.Tensor] = None,
|
|
):
|
|
routing_logits = self.router_weights(inputs)
|
|
routing_probs = F.softmax(routing_logits, dim=-1, dtype=torch.float32)
|
|
|
|
if padding_mask is not None:
|
|
# [batch * seq, expert]
|
|
routing_probs = routing_probs * padding_mask.view(-1).unsqueeze(-1)
|
|
|
|
# Note routing_probs is using float32.
|
|
return routing_probs, routing_logits
|
|
|
|
class MLPExpert(nn.Module):
|
|
"""MLP expert module for sparse MoE."""
|
|
|
|
def __init__(
|
|
self,
|
|
config: Grok1Config,
|
|
):
|
|
super().__init__()
|
|
self.w1 = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
|
|
self.v = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
|
|
self.dense = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
|
|
self.act_fn = ACT2FN[config.hidden_act]
|
|
|
|
def forward(
|
|
self,
|
|
inputs: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
h_w1 = self.act_fn(self.w1(inputs))
|
|
h_v = self.v(inputs)
|
|
h_dense = self.dense(h_w1 * h_v)
|
|
return h_dense
|
|
|
|
class SparseMoEMLP(nn.Module):
|
|
"""Sparse MoE MLP module."""
|
|
|
|
def __init__(
|
|
self,
|
|
config: Grok1Config,
|
|
):
|
|
super().__init__()
|
|
self.num_experts = config.num_local_experts
|
|
self.top_k = config.num_experts_per_tok
|
|
self.gating = Gating(config)
|
|
self.experts = nn.ModuleList([MLPExpert(config) for _ in range(self.num_experts)])
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
padding_mask: Optional[torch.Tensor] = None,
|
|
) -> torch.Tensor:
|
|
batch_size, sequence_length, hidden_dim = hidden_states.shape
|
|
hidden_states = hidden_states.view(-1, hidden_dim)
|
|
|
|
# Get routing probabilities and selected experts.
|
|
routing_probs, routing_logits = self.gating(hidden_states, padding_mask)
|
|
routing_probs, selected_experts = torch.topk(routing_probs, self.top_k, dim=-1)
|
|
routing_probs = routing_probs / routing_probs.sum(dim=-1, keepdim=True)
|
|
# Now routing_probs is using the hidden_states' dtype instead of float32.
|
|
routing_probs = routing_probs.to(hidden_states.dtype)
|
|
|
|
# Initialize output hidden states.
|
|
final_hidden_states = torch.zeros(
|
|
(batch_size * sequence_length, hidden_dim),
|
|
dtype=hidden_states.dtype,
|
|
device=hidden_states.device,
|
|
)
|
|
|
|
# Create expert mask.
|
|
expert_mask = F.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0)
|
|
|
|
# Loop over experts and compute their contributions.
|
|
for expert_idx in range(self.num_experts):
|
|
expert_layer = self.experts[expert_idx]
|
|
idx, top_x = torch.where(expert_mask[expert_idx])
|
|
|
|
if top_x.shape[0] == 0:
|
|
continue
|
|
|
|
top_x_list = top_x.tolist()
|
|
idx_list = idx.tolist()
|
|
|
|
current_state = hidden_states[None, top_x_list].reshape(-1, hidden_dim)
|
|
current_hidden_states = expert_layer(current_state) * routing_probs[top_x_list, idx_list, None]
|
|
|
|
final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype))
|
|
|
|
final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim)
|
|
return final_hidden_states, routing_logits
|
|
|
|
def rotate_half(x: torch.Tensor) -> torch.Tensor:
|
|
x1, x2 = x.chunk(2, dim=-1)
|
|
return torch.cat((-x2, x1), dim=-1)
|
|
|
|
class RotaryPositionalEmbedding(nn.Module):
|
|
|
|
def __init__(
|
|
self,
|
|
dim: int,
|
|
base_exponent: int = int(1e4),
|
|
):
|
|
super().__init__()
|
|
self.dim = dim
|
|
self.base_exponent = base_exponent
|
|
assert self.dim % 2 == 0, "Embedding dimension must be even for rotary embeddings."
|
|
|
|
def forward(
|
|
self,
|
|
x: torch.Tensor,
|
|
seq_dim: int,
|
|
offset: torch.Tensor,
|
|
const_position: Optional[int] = None,
|
|
t: Optional[torch.Tensor] = None,
|
|
) -> torch.Tensor:
|
|
# Compute the per-dimension frequencies.
|
|
dtype = x.dtype
|
|
exponents = torch.arange(0, self.dim, 2, dtype=torch.float32, device=x.device)
|
|
inv_freq = (self.base_exponent ** (exponents / self.dim)).reciprocal()
|
|
|
|
if not isinstance(offset, torch.Tensor):
|
|
offset = torch.tensor(offset, dtype=torch.float32, device=x.device)
|
|
if offset.dim() == 0:
|
|
# Offset can be a scalar or one offset per batch element.
|
|
offset = offset.unsqueeze(0)
|
|
|
|
# Compute the per-element phase (to pass into sin and cos).
|
|
if const_position is not None:
|
|
t = const_position * torch.ones(
|
|
(1, x.shape[seq_dim]),
|
|
dtype=torch.float32,
|
|
device=x.device,
|
|
)
|
|
elif t is None:
|
|
t = torch.arange(x.shape[seq_dim], dtype=torch.float32, device=x.device)
|
|
t = t.unsqueeze(0) + offset.unsqueeze(-1)
|
|
|
|
phase = torch.einsum("bi,j->bij", t, inv_freq)
|
|
phase = torch.cat([phase, phase], dim=-1)[:, :, None, :]
|
|
|
|
x_rotated = x * phase.cos() + rotate_half(x) * phase.sin()
|
|
|
|
return x_rotated.to(dtype)
|
|
|
|
class RMSNorm(nn.Module):
|
|
|
|
def __init__(self, hidden_size, eps=1e-5):
|
|
super().__init__()
|
|
self.weight = nn.Parameter(torch.ones(hidden_size))
|
|
self.eps = eps
|
|
|
|
def forward(self, hidden_states):
|
|
input_dtype = hidden_states.dtype
|
|
hidden_states = hidden_states.to(torch.float32)
|
|
variance = hidden_states.pow(2).mean(-1, keepdim=True)
|
|
hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
|
|
return self.weight * hidden_states.to(input_dtype)
|
|
|
|
class KVMemory(NamedTuple):
|
|
k: Optional[torch.Tensor]
|
|
v: Optional[torch.Tensor]
|
|
step: Optional[torch.Tensor]
|
|
|
|
def init_layer_memories(
|
|
batch_size: int,
|
|
sequence_len: int,
|
|
num_kv_heads: int,
|
|
key_size: int,
|
|
num_layers: int,
|
|
step: Optional[torch.Tensor] = None,
|
|
dtype: torch.dtype = torch.bfloat16,
|
|
device: str = "cpu",
|
|
):
|
|
if step is None:
|
|
step = torch.zeros(batch_size, dtype=torch.int32, device=device)
|
|
return [
|
|
KVMemory(
|
|
k=torch.zeros(batch_size, sequence_len, num_kv_heads, key_size, dtype=dtype, device=device),
|
|
v=torch.zeros(batch_size, sequence_len, num_kv_heads, key_size, dtype=dtype, device=device),
|
|
step=step,
|
|
)
|
|
for _ in range(num_layers)
|
|
]
|
|
|
|
class MultiHeadAttention(nn.Module):
|
|
|
|
def __init__(self, config: Grok1Config):
|
|
super().__init__()
|
|
self.num_q_heads = config.num_attention_heads
|
|
self.num_kv_heads = config.num_key_value_heads
|
|
self.key_size = config.hidden_size // config.num_attention_heads
|
|
self.value_size = self.key_size
|
|
self.attn_output_multiplier = config.attn_output_multiplier
|
|
|
|
self.q_proj = nn.Linear(config.hidden_size, self.num_q_heads * self.key_size, bias=False)
|
|
self.k_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.key_size, bias=False)
|
|
self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.value_size, bias=False)
|
|
self.out_proj = nn.Linear(self.num_q_heads * self.value_size, config.hidden_size, bias=False)
|
|
self.rotary_pos_emb = RotaryPositionalEmbedding(self.key_size, base_exponent=config.rope_theta)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
mask: Optional[torch.Tensor] = None,
|
|
layer_memory: Optional[KVMemory] = None,
|
|
):
|
|
batch_size, seq_len, _ = hidden_states.shape
|
|
|
|
query = self.q_proj(hidden_states).view(batch_size, seq_len, self.num_q_heads, self.key_size)
|
|
key = self.k_proj(hidden_states).view(batch_size, seq_len, self.num_kv_heads, self.key_size)
|
|
value = self.v_proj(hidden_states).view(batch_size, seq_len, self.num_kv_heads, self.value_size)
|
|
|
|
query = self.rotary_pos_emb(query, seq_dim=1, offset=layer_memory.step if layer_memory else 0)
|
|
key = self.rotary_pos_emb(key, seq_dim=1, offset=layer_memory.step if layer_memory else 0)
|
|
|
|
if layer_memory:
|
|
key = torch.cat([layer_memory.k, key], dim=1)
|
|
value = torch.cat([layer_memory.v, value], dim=1)
|
|
new_step = layer_memory.step + seq_len
|
|
memory_mask = torch.arange(key.shape[1], device=key.device) < new_step[:, None]
|
|
memory_mask = memory_mask[:, None, None, :]
|
|
if mask is not None:
|
|
mask = mask * memory_mask
|
|
else:
|
|
mask = memory_mask
|
|
new_memory = KVMemory(k=key, v=value, step=new_step)
|
|
else:
|
|
new_memory = None
|
|
|
|
query = query.view(batch_size, seq_len, self.num_kv_heads, self.num_q_heads // self.num_kv_heads, self.key_size)
|
|
attn_logits = torch.einsum("...thHd,...Thd->...hHtT", query, key).to(torch.float32)
|
|
attn_logits *= self.attn_output_multiplier
|
|
max_attn_val = torch.tensor(30.0, dtype=attn_logits.dtype, device=attn_logits.device)
|
|
attn_logits = max_attn_val * torch.tanh(attn_logits / max_attn_val)
|
|
|
|
if mask is not None:
|
|
mask = mask[:, :, None, :, :]
|
|
attn_logits = torch.where(mask, attn_logits, torch.full_like(attn_logits, float("-inf")))
|
|
|
|
attn_weights = F.softmax(attn_logits, dim=-1).to(query.dtype)
|
|
|
|
attn = torch.einsum("...hHtT,...Thd->...thHd", attn_weights, value)
|
|
attn = attn.reshape(batch_size, seq_len, -1)
|
|
attn = self.out_proj(attn)
|
|
|
|
return attn, new_memory
|
|
|
|
class Decoder(nn.Module):
|
|
|
|
def __init__(self, config: Grok1Config):
|
|
super().__init__()
|
|
self.num_layers = config.num_hidden_layers
|
|
self.attention = MultiHeadAttention(config)
|
|
self.norm1 = RMSNorm(config.hidden_size)
|
|
self.norm2 = RMSNorm(config.hidden_size)
|
|
self.norm3 = RMSNorm(config.hidden_size)
|
|
self.norm4 = RMSNorm(config.hidden_size)
|
|
|
|
if config.num_local_experts > 1:
|
|
self.mlp = SparseMoEMLP(config)
|
|
else:
|
|
self.mlp = MLPExpert(config)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
mask: Optional[torch.Tensor] = None,
|
|
padding_mask: Optional[torch.Tensor] = None,
|
|
layer_memory: Optional[KVMemory] = None,
|
|
):
|
|
residual = hidden_states
|
|
hidden_states = self.norm1(hidden_states)
|
|
attn_output, new_memory = self.attention(hidden_states, mask, layer_memory)
|
|
attn_output = self.norm2(attn_output)
|
|
hidden_states = residual + attn_output
|
|
|
|
residual = hidden_states
|
|
hidden_states = self.norm3(hidden_states)
|
|
if isinstance(self.mlp, SparseMoEMLP):
|
|
mlp_output, routing_logits = self.mlp(hidden_states, padding_mask)
|
|
else:
|
|
mlp_output = self.mlp(hidden_states)
|
|
routing_logits = None
|
|
mlp_output = self.norm4(mlp_output)
|
|
hidden_states = residual + mlp_output
|
|
|
|
return hidden_states, new_memory, routing_logits
|
|
|
|
class Grok1PreTrainedModel(PreTrainedModel):
|
|
config_class = Grok1Config
|
|
base_model_prefix = "model"
|
|
supports_gradient_checkpointing = True
|
|
_no_split_modules = ["Decoder"]
|
|
_skip_keys_device_placement = "past_key_values"
|
|
_supports_flash_attn_2 = False
|
|
_supports_sdpa = False
|
|
_supports_cache_class = True
|
|
|
|
def _init_weights(self, module):
|
|
std = self.config.initializer_range
|
|
if isinstance(module, nn.Linear):
|
|
module.weight.data.normal_(mean=0.0, std=std)
|
|
if module.bias is not None:
|
|
module.bias.data.zero_()
|
|
elif isinstance(module, nn.Embedding):
|
|
module.weight.data.normal_(mean=0.0, std=std)
|
|
if module.padding_idx is not None:
|
|
module.weight.data[module.padding_idx].zero_()
|
|
|
|
class Grok1ForCausalLM(Grok1PreTrainedModel):
|
|
|
|
def __init__(self, config: Grok1Config):
|
|
super().__init__(config)
|
|
self.embedding = TiedWeightEmbedding(self.config)
|
|
self.layers = nn.ModuleList([Decoder(self.config) for _ in range(self.config.num_hidden_layers)])
|
|
self.norm = RMSNorm(self.config.hidden_size)
|
|
|
|
def forward(
|
|
self,
|
|
input_ids: torch.Tensor,
|
|
attention_mask: Optional[torch.Tensor] = None,
|
|
memory: Optional[List[KVMemory]] = None,
|
|
last_hid_only: bool = False,
|
|
length: Optional[torch.Tensor] = None,
|
|
):
|
|
batch_size, seq_len = input_ids.shape
|
|
if attention_mask is None:
|
|
attention_mask = torch.ones_like(input_ids, dtype=torch.bool)
|
|
else:
|
|
attention_mask = attention_mask.bool()
|
|
|
|
padding_mask = attention_mask.view(batch_size, seq_len)
|
|
causal_mask = torch.tril(torch.ones(1, 1, seq_len, seq_len, dtype=torch.bool, device=input_ids.device))
|
|
mask = padding_mask[:, None, None, :] * causal_mask
|
|
|
|
hidden_states = self.embedding(input_ids) * self.config.embedding_multiplier_scale
|
|
kv_memories = []
|
|
|
|
for i, layer in enumerate(self.layers):
|
|
layer_memory = memory[i] if memory else None
|
|
hidden_states, new_memory, routing_logits = layer(
|
|
hidden_states,
|
|
mask,
|
|
padding_mask,
|
|
layer_memory,
|
|
)
|
|
kv_memories.append(new_memory)
|
|
|
|
hidden_states = self.norm(hidden_states)
|
|
|
|
if last_hid_only:
|
|
last_step = torch.maximum(torch.sum(padding_mask, dim=1) - 1, torch.tensor(0, device=input_ids.device))
|
|
hidden_states = hidden_states[torch.arange(batch_size, device=input_ids.device), last_step]
|
|
elif length is not None:
|
|
last_step = torch.maximum(length - 1, torch.tensor(0, device=input_ids.device))
|
|
hidden_states = hidden_states[torch.arange(batch_size, device=input_ids.device), last_step]
|
|
hidden_states = hidden_states.unsqueeze(1)
|
|
|
|
logits = self.embedding.decode(hidden_states) * torch.tensor(self.config.output_multiplier_scale, dtype=hidden_states.dtype, device=hidden_states.device)
|
|
|
|
return logits, kv_memories
|