Initial commit
This commit is contained in:
@@ -0,0 +1,27 @@
|
||||
from khaosz.core.tokenizer import BpeTokenizer
|
||||
from khaosz.core.transformer import Transformer, TransformerConfig
|
||||
from khaosz.core.parameter import ParameterLoader, ModelParameter, Checkpoint
|
||||
from khaosz.core.generator import (
|
||||
TextGenerator,
|
||||
ChatGenerator,
|
||||
StreamGenerator,
|
||||
BatchGenerator,
|
||||
RetrievalGenerator,
|
||||
EmbeddingEncoder
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Transformer",
|
||||
"TransformerConfig",
|
||||
"BpeTokenizer",
|
||||
"ParameterLoader",
|
||||
"ModelParameter",
|
||||
"Checkpoint",
|
||||
"TextGenerator",
|
||||
"ChatGenerator",
|
||||
"StreamGenerator",
|
||||
"BatchGenerator",
|
||||
"RetrievalGenerator",
|
||||
"EmbeddingEncoder"
|
||||
]
|
||||
@@ -0,0 +1,568 @@
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from typing import List, Tuple, Union, Optional, Generator, Self
|
||||
from khaosz.core.parameter import ModelParameter
|
||||
|
||||
|
||||
def build_prompt(query: str, history: Optional[List[Tuple[str, str]]] = None) -> str:
|
||||
"""
|
||||
Build prompt for query and history
|
||||
|
||||
Args:
|
||||
query(str): query string
|
||||
history(Optional[List[Tuple[str, str]]]): history list of query and response
|
||||
|
||||
Returns:
|
||||
str: prompt string
|
||||
|
||||
"""
|
||||
prompt_parts = []
|
||||
|
||||
if history is None:
|
||||
history = []
|
||||
|
||||
for his_query, his_response in history:
|
||||
prompt_parts.append(f"<|user|> {his_query} <|system|> <bos>{his_response}<eos>")
|
||||
|
||||
if query is not None:
|
||||
prompt_parts.append(f"<|user|> {query} <|system|> <bos>")
|
||||
|
||||
return "\n".join(prompt_parts)
|
||||
|
||||
def pad_sequence(ids_list: List[List[int]], max_ids_len: int, pad_id: int) -> List[List[int]]:
|
||||
"""
|
||||
Pad a list of sequences to a fixed length.
|
||||
|
||||
Args:
|
||||
ids_list (List[List[int]]): A list of sequences.
|
||||
max_ids_len (int): The maximum length of sequences.
|
||||
pad_id (int): The id to pad sequences.
|
||||
|
||||
Returns:
|
||||
List[List[int]]: A list of padded sequences.
|
||||
|
||||
"""
|
||||
new_ids_list = []
|
||||
for ids in ids_list:
|
||||
pad_len = max_ids_len - len(ids)
|
||||
padded_seq = [pad_id] * pad_len + ids
|
||||
new_ids_list.append(padded_seq)
|
||||
|
||||
return new_ids_list
|
||||
|
||||
def apply_sampling_strategies(
|
||||
logits: Tensor,
|
||||
temperature: float,
|
||||
top_k: int,
|
||||
top_p: float,
|
||||
filter_value: float = -float("inf")
|
||||
) -> Tensor:
|
||||
"""
|
||||
Apply sampling strategies to the logits tensor.
|
||||
|
||||
Args:
|
||||
logits (Tensor): The logits tensor.
|
||||
temperature (float): The temperature parameter.
|
||||
top_k (int): The top-k parameter.
|
||||
top_p (float): The top-p parameter.
|
||||
filter_value (float, optional): The filter value. Defaults to -float("inf").
|
||||
|
||||
Returns:
|
||||
Tensor: The sampled logits tensor.
|
||||
|
||||
"""
|
||||
|
||||
if temperature != 1.0:
|
||||
logits = logits / temperature
|
||||
|
||||
if top_k > 0:
|
||||
top_k = min(top_k, logits.size(-1))
|
||||
indices_to_remove = logits < torch.topk(logits, top_k, dim=-1)[0][..., -1, None]
|
||||
logits[indices_to_remove] = filter_value
|
||||
|
||||
if top_p < 1.0:
|
||||
sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
|
||||
cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
|
||||
|
||||
sorted_indices_to_remove = cumulative_probs > top_p
|
||||
sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
|
||||
sorted_indices_to_remove[..., 0] = 0
|
||||
|
||||
indices_to_remove = torch.zeros_like(logits, dtype=torch.bool)
|
||||
indices_to_remove.scatter_(
|
||||
dim=1,
|
||||
index=sorted_indices,
|
||||
src=sorted_indices_to_remove
|
||||
)
|
||||
|
||||
logits[indices_to_remove] = filter_value
|
||||
|
||||
return logits
|
||||
|
||||
|
||||
class KVCacheManager:
|
||||
def __init__(
|
||||
self,
|
||||
num_layers: int,
|
||||
batch_size: int,
|
||||
max_len: int,
|
||||
num_heads: int,
|
||||
head_dim: int,
|
||||
device: torch.device = "cuda",
|
||||
dtype: torch.dtype = torch.bfloat16
|
||||
):
|
||||
self.num_layers = num_layers
|
||||
self.batch_size = batch_size
|
||||
self.max_len = max_len
|
||||
self.num_heads = num_heads
|
||||
self.head_dim = head_dim
|
||||
self.device = device
|
||||
self.dtype = dtype
|
||||
|
||||
self._kv_cache: List[Tuple[Tensor, Tensor]] = None
|
||||
self._seq_mask: Tensor = None
|
||||
self._initialize()
|
||||
|
||||
def _initialize(self):
|
||||
self._kv_cache = []
|
||||
for _ in range(self.num_layers):
|
||||
k_cache = torch.zeros(
|
||||
(self.batch_size, self.max_len, self.num_heads, self.head_dim),
|
||||
device=self.device, dtype=self.dtype
|
||||
)
|
||||
v_cache = torch.zeros(
|
||||
(self.batch_size, self.max_len, self.num_heads, self.head_dim),
|
||||
device=self.device, dtype=self.dtype
|
||||
)
|
||||
self._kv_cache.append((k_cache, v_cache))
|
||||
|
||||
self._seq_mask = torch.ones(
|
||||
(self.batch_size, self.max_len),
|
||||
device=self.device, dtype=torch.bool
|
||||
)
|
||||
|
||||
def update(self, active_mask: Tensor):
|
||||
for i in range(self.num_layers):
|
||||
k_cache, v_cache = self._kv_cache[i]
|
||||
new_k_cache, new_v_cache = k_cache[active_mask], v_cache[active_mask]
|
||||
self._kv_cache[i] = (new_k_cache, new_v_cache)
|
||||
|
||||
self._seq_mask = self._seq_mask[active_mask]
|
||||
|
||||
def reset(self, full_reset=False):
|
||||
if full_reset:
|
||||
self._kv_cache = None
|
||||
self._seq_mask = None
|
||||
else:
|
||||
self._initialize()
|
||||
|
||||
def set_seq_mask(self, input_ids: Tensor, pad_id: int):
|
||||
batch_size, seq_len = input_ids.shape
|
||||
bool_mask = (input_ids != pad_id)
|
||||
self._seq_mask[: batch_size, : seq_len] = bool_mask
|
||||
|
||||
def get_kvcache(self) -> List[Tuple[Tensor, Tensor]]:
|
||||
return self._kv_cache
|
||||
|
||||
def get_seq_mask(self) -> Tensor:
|
||||
return self._seq_mask
|
||||
|
||||
|
||||
class GeneratorCore:
|
||||
def __init__(self, parameter: ModelParameter):
|
||||
self.model = parameter.model
|
||||
self.tokenizer = parameter.tokenizer
|
||||
self.config = parameter.config
|
||||
|
||||
def compute_logits(
|
||||
self,
|
||||
input_ids: Tensor,
|
||||
attn_mask: Optional[Tensor] = None,
|
||||
kv_caches: Optional[List[Tuple[Tensor, Tensor]]] = None,
|
||||
start_pos: int = 0
|
||||
) -> Tuple[Tensor, int]:
|
||||
with torch.inference_mode():
|
||||
outputs = self.model(input_ids, attn_mask, kv_caches, start_pos)
|
||||
logits = outputs["logits"][:, -1, :]
|
||||
cache_increase = input_ids.size(-1)
|
||||
|
||||
return logits, cache_increase
|
||||
|
||||
def to(self, *args, **kargs) -> Self:
|
||||
self.model.to(*args, **kargs)
|
||||
return self
|
||||
|
||||
|
||||
class EmbeddingEncoderCore:
|
||||
def __init__(self, parameter: ModelParameter):
|
||||
self.model = parameter.model
|
||||
self.tokenizer = parameter.tokenizer
|
||||
self.config = parameter.config
|
||||
|
||||
def encode(self, sentence: Union[str, List[str]]) -> Union[Tensor, List[Tensor]]:
|
||||
with_batch = isinstance(sentence, list)
|
||||
ids = self.tokenizer.encode(sentence)
|
||||
batch_ids = ids if with_batch else [ids]
|
||||
max_model_len = self.config.m_len
|
||||
|
||||
all_fragments = []
|
||||
fragment_origin_idx = []
|
||||
|
||||
for i, seq in enumerate(batch_ids):
|
||||
if len(seq) > max_model_len:
|
||||
fragments = [seq[j:j+max_model_len] for j in range(0, len(seq), max_model_len)]
|
||||
all_fragments.extend(fragments)
|
||||
fragment_origin_idx.extend([i] * len(fragments))
|
||||
else:
|
||||
all_fragments.append(seq)
|
||||
fragment_origin_idx.append(i)
|
||||
|
||||
#if empty fragments
|
||||
if not all_fragments or not ids:
|
||||
return [] if with_batch else torch.tensor([])
|
||||
|
||||
device = next(self.model.parameters()).device
|
||||
max_len = min(max(len(seq) for seq in all_fragments), max_model_len)
|
||||
|
||||
padded_ids = []
|
||||
masks = []
|
||||
for seq in all_fragments:
|
||||
pad_len = max_len - len(seq)
|
||||
padded_seq = seq + [self.tokenizer.pad_id] * pad_len
|
||||
mask = [token_id != self.tokenizer.pad_id for token_id in padded_seq]
|
||||
padded_ids.append(padded_seq)
|
||||
masks.append(mask)
|
||||
|
||||
input_tensor = torch.tensor(padded_ids, device=device, dtype=torch.long)
|
||||
seq_mask = torch.tensor(masks, device=device, dtype=torch.bool)
|
||||
|
||||
with torch.inference_mode():
|
||||
outputs = self.model(input_tensor, seq_mask)["hidden_states"]
|
||||
# [num_fragments, seq_len, hidden_size]
|
||||
fragment_embs = torch.mul(outputs, seq_mask.unsqueeze(-1))
|
||||
|
||||
sentence_embs: List[Tensor] = []
|
||||
for i in range(len(batch_ids)):
|
||||
indices = [idx for idx, orig_idx in enumerate(fragment_origin_idx) if orig_idx == i]
|
||||
if indices is not None:
|
||||
sum_frags = torch.sum(fragment_embs[indices, :, :], dim=1) # [frags, hidden_size]
|
||||
length = torch.sum(seq_mask[indices, :], dim=1).unsqueeze(1) # [frags, 1]
|
||||
emb = torch.sum(sum_frags / length, dim=0) # [frags, hidden_size]
|
||||
sentence_embs.append(emb.flatten())
|
||||
|
||||
if with_batch:
|
||||
return [emb.flatten() for emb in sentence_embs]
|
||||
else:
|
||||
return sentence_embs[0].flatten()
|
||||
|
||||
def to(self, *args, **kargs) -> Self:
|
||||
self.model.to(*args, **kargs)
|
||||
return self
|
||||
|
||||
|
||||
class TextGenerator(GeneratorCore):
|
||||
def __init__(self, parameter: ModelParameter):
|
||||
super().__init__(parameter)
|
||||
|
||||
def generate(
|
||||
self,
|
||||
query: str,
|
||||
temperature: float,
|
||||
top_k: int,
|
||||
top_p: float,
|
||||
) -> str:
|
||||
assert temperature >= 0.0
|
||||
assert top_k >= 0
|
||||
assert top_p >= 0.0 and top_p <= 1.0
|
||||
|
||||
device = next(self.model.parameters()).device
|
||||
cache_manager = KVCacheManager(
|
||||
num_layers=self.config.n_layer,
|
||||
batch_size=1,
|
||||
max_len=self.config.m_len,
|
||||
num_heads=self.config.n_kvhead,
|
||||
head_dim=self.config.n_dim // self.config.n_head,
|
||||
device=device,
|
||||
)
|
||||
|
||||
ids = self.tokenizer.encode(query)
|
||||
input_ids = torch.tensor([ids], device=device, dtype=torch.long)
|
||||
|
||||
start_cache_pos = len(ids)
|
||||
cur_cache_pos = 0
|
||||
self.model.eval()
|
||||
|
||||
while len(ids) < self.config.m_len:
|
||||
kv_caches = cache_manager.get_kvcache()
|
||||
logits, cache_increase = self.compute_logits(
|
||||
input_ids,
|
||||
kv_caches=kv_caches,
|
||||
start_pos=cur_cache_pos
|
||||
)
|
||||
logits = apply_sampling_strategies(logits, temperature, top_k, top_p)
|
||||
probs = torch.softmax(logits, dim=-1)
|
||||
next_token_id = torch.multinomial(probs, num_samples=1)
|
||||
|
||||
input_ids = next_token_id
|
||||
ids.append(next_token_id.item())
|
||||
cur_cache_pos += cache_increase
|
||||
|
||||
if next_token_id.item() in self.tokenizer.stop_ids:
|
||||
break
|
||||
|
||||
response = self.tokenizer.decode(ids[start_cache_pos:])
|
||||
|
||||
return response
|
||||
|
||||
|
||||
|
||||
class ChatGenerator(GeneratorCore):
|
||||
def __init__(self, parameter: ModelParameter):
|
||||
super().__init__(parameter)
|
||||
|
||||
def generate(
|
||||
self,
|
||||
query: str,
|
||||
history: List[Tuple[str, str]],
|
||||
temperature: float,
|
||||
top_k: int,
|
||||
top_p: float,
|
||||
) -> str:
|
||||
|
||||
assert temperature >= 0.0
|
||||
assert top_k >= 0
|
||||
assert top_p >= 0.0 and top_p <= 1.0
|
||||
|
||||
if history is None:
|
||||
history = []
|
||||
|
||||
device = next(self.model.parameters()).device
|
||||
cache_manager = KVCacheManager(
|
||||
num_layers=self.config.n_layer,
|
||||
batch_size=1,
|
||||
max_len=self.config.m_len,
|
||||
num_heads=self.config.n_kvhead,
|
||||
head_dim=self.config.n_dim // self.config.n_head,
|
||||
device=device,
|
||||
)
|
||||
ids = self.tokenizer.encode(build_prompt(query, history))
|
||||
input_ids = torch.tensor([ids], device=device, dtype=torch.long)
|
||||
cpy_history = history.copy()
|
||||
|
||||
start_cache_pos = len(ids)
|
||||
cur_cache_pos = 0
|
||||
self.model.eval()
|
||||
|
||||
|
||||
while len(ids) < self.config.m_len:
|
||||
kv_caches = cache_manager.get_kvcache()
|
||||
logits, cache_increase = self.compute_logits(
|
||||
input_ids,
|
||||
kv_caches=kv_caches,
|
||||
start_pos=cur_cache_pos
|
||||
)
|
||||
logits = apply_sampling_strategies(logits, temperature, top_k, top_p)
|
||||
probs = torch.softmax(logits, dim=-1)
|
||||
next_token_id = torch.multinomial(probs, num_samples=1)
|
||||
|
||||
input_ids = next_token_id
|
||||
ids.append(next_token_id.item())
|
||||
cur_cache_pos += cache_increase
|
||||
|
||||
if next_token_id.item() in self.tokenizer.stop_ids:
|
||||
break
|
||||
|
||||
response = self.tokenizer.decode(ids[start_cache_pos:])
|
||||
cpy_history.append((query, response))
|
||||
|
||||
return response, cpy_history
|
||||
|
||||
|
||||
class StreamGenerator(GeneratorCore):
|
||||
def __init__(self, parameter: ModelParameter):
|
||||
super().__init__(parameter)
|
||||
|
||||
def generate(
|
||||
self,
|
||||
query: str,
|
||||
history: List[Tuple[str, str]],
|
||||
temperature: float,
|
||||
top_k: int,
|
||||
top_p: float,
|
||||
) -> Generator[Tuple[str, List[Tuple[str, str]]], None, None]:
|
||||
|
||||
assert temperature >= 0.0
|
||||
assert top_k >= 0
|
||||
assert top_p >= 0.0 and top_p <= 1.0
|
||||
|
||||
if history is None:
|
||||
history = []
|
||||
|
||||
device = next(self.model.parameters()).device
|
||||
cache_manager = KVCacheManager(
|
||||
num_layers=self.config.n_layer,
|
||||
batch_size=1,
|
||||
max_len=self.config.m_len,
|
||||
num_heads=self.config.n_kvhead,
|
||||
head_dim=self.config.n_dim // self.config.n_head,
|
||||
device=device,
|
||||
)
|
||||
ids = self.tokenizer.encode(build_prompt(query, history))
|
||||
input_ids = torch.tensor([ids], device=device, dtype=torch.long)
|
||||
cpy_history = history.copy()
|
||||
|
||||
start_cache_pos = len(ids)
|
||||
cur_cache_pos = 0
|
||||
self.model.eval()
|
||||
|
||||
|
||||
while len(ids) < self.config.m_len:
|
||||
kv_caches = cache_manager.get_kvcache()
|
||||
logits, cache_increase = self.compute_logits(
|
||||
input_ids,
|
||||
kv_caches=kv_caches,
|
||||
start_pos=cur_cache_pos
|
||||
)
|
||||
logits = apply_sampling_strategies(logits, temperature, top_k, top_p)
|
||||
probs = torch.softmax(logits, dim=-1)
|
||||
next_token_id = torch.multinomial(probs, num_samples=1)
|
||||
|
||||
input_ids = next_token_id
|
||||
ids.append(next_token_id.item())
|
||||
cur_cache_pos += cache_increase
|
||||
|
||||
response = self.tokenizer.decode(ids[start_cache_pos:])
|
||||
yield response, cpy_history + [(query, response)]
|
||||
|
||||
if next_token_id.item() in self.tokenizer.stop_ids:
|
||||
yield response + "\n", cpy_history + [(query, response)]
|
||||
break
|
||||
|
||||
|
||||
class BatchGenerator(GeneratorCore):
|
||||
def __init__(self, parameter: ModelParameter):
|
||||
super().__init__(parameter)
|
||||
|
||||
def generate(
|
||||
self,
|
||||
queries: List[str],
|
||||
histories: List[List[Tuple[str, str]]],
|
||||
temperature: float,
|
||||
top_k: int,
|
||||
top_p: float
|
||||
) -> List[str]:
|
||||
|
||||
assert temperature >= 0.0
|
||||
assert top_k >= 0
|
||||
assert top_p >= 0.0 and top_p <= 1.0
|
||||
|
||||
batch_size = len(queries)
|
||||
if histories is None:
|
||||
histories = [[] for _ in range(batch_size)]
|
||||
|
||||
prompts = [build_prompt(query, history) for query, history in zip(queries, histories)]
|
||||
ids_list = [self.tokenizer.encode(prompt) for prompt in prompts]
|
||||
max_ids_len = max(len(ids) for ids in ids_list)
|
||||
ids_list = pad_sequence(ids_list, max_ids_len, self.tokenizer.pad_id)
|
||||
|
||||
device = next(self.model.parameters()).device
|
||||
cache_manager = KVCacheManager(
|
||||
num_layers=self.config.n_layer,
|
||||
batch_size=batch_size,
|
||||
max_len=self.config.m_len,
|
||||
num_heads=self.config.n_kvhead,
|
||||
head_dim=self.config.n_dim // self.config.n_head,
|
||||
device=device,
|
||||
)
|
||||
|
||||
input_tensor = torch.tensor(ids_list, device=device, dtype=torch.long)
|
||||
cache_manager.set_seq_mask(input_tensor, self.tokenizer.pad_id)
|
||||
activate_task_mask = [True] * batch_size
|
||||
|
||||
start_cache_pos = max_ids_len
|
||||
cur_cache_pos = 0
|
||||
|
||||
while max_ids_len < self.config.m_len and sum(activate_task_mask) != 0:
|
||||
kv_caches = cache_manager.get_kvcache()
|
||||
attn_mask =cache_manager.get_seq_mask()
|
||||
|
||||
logits, cache_increase = self.compute_logits(
|
||||
input_tensor,
|
||||
attn_mask=attn_mask,
|
||||
kv_caches=kv_caches,
|
||||
start_pos=cur_cache_pos
|
||||
)
|
||||
|
||||
cur_cache_pos += cache_increase
|
||||
logits = apply_sampling_strategies(logits, temperature, top_k, top_p)
|
||||
probs = torch.softmax(logits, dim=-1)
|
||||
next_token_id = torch.multinomial(probs, num_samples=1)
|
||||
|
||||
active_mask = []
|
||||
c_ids = 0
|
||||
|
||||
for i in range(batch_size):
|
||||
if activate_task_mask[i]:
|
||||
token = next_token_id[c_ids, :].item()
|
||||
ids_list[i].append(token)
|
||||
c_ids += 1
|
||||
|
||||
is_active = not token in self.tokenizer.stop_ids
|
||||
activate_task_mask[i] = is_active
|
||||
active_mask.append(is_active)
|
||||
|
||||
active_mask = torch.tensor(active_mask, device=device, dtype=torch.bool)
|
||||
cache_manager.update(active_mask)
|
||||
input_tensor = next_token_id[active_mask, :]
|
||||
|
||||
max_ids_len += 1
|
||||
|
||||
|
||||
responses = [str()] * batch_size
|
||||
for i in range(batch_size):
|
||||
responses[i] = self.tokenizer.decode(ids_list[i][start_cache_pos:])
|
||||
histories[i].append((queries[i], responses[i]))
|
||||
|
||||
return responses
|
||||
|
||||
|
||||
|
||||
class RetrievalGenerator(GeneratorCore):
|
||||
def __init__(self, retriever_parameter: ModelParameter):
|
||||
super().__init__(retriever_parameter)
|
||||
|
||||
def generate(
|
||||
self,
|
||||
retrieved: List[str],
|
||||
query: str,
|
||||
history: List[Tuple[str, str]],
|
||||
temperature: float,
|
||||
top_k: int,
|
||||
top_p: float,
|
||||
) -> str:
|
||||
assert temperature >= 0.0
|
||||
assert top_k >= 0
|
||||
assert top_p >= 0.0 and top_p <= 1.0
|
||||
|
||||
if history is None:
|
||||
history = []
|
||||
|
||||
retrieved = "\n".join([f"{idx + 1}. {key}" for idx, key in enumerate(retrieved)]) if retrieved else ""
|
||||
retrieved_query = f"{retrieved}<eos>\n\n根据以上内容回答: {query}" if retrieved else query
|
||||
parameter = ModelParameter(self.model, self.tokenizer, self.config)
|
||||
|
||||
return ChatGenerator(parameter).generate(
|
||||
retrieved_query,
|
||||
history,
|
||||
temperature=temperature,
|
||||
top_k=top_k,
|
||||
top_p=top_p,
|
||||
)
|
||||
|
||||
class EmbeddingEncoder(EmbeddingEncoderCore):
|
||||
def __init__(self, parameter: ModelParameter):
|
||||
super().__init__(parameter)
|
||||
|
||||
def encode(self, sentence: Union[str, List[str]]) -> Union[Tensor, List[Tensor]]:
|
||||
return super().encode(sentence)
|
||||
|
||||
@@ -0,0 +1,238 @@
|
||||
import pickle as pkl
|
||||
import matplotlib.pyplot as plt
|
||||
import safetensors.torch as st
|
||||
import torch.nn as nn
|
||||
import torch.optim as optim
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional, Self, Union
|
||||
from pathlib import Path
|
||||
|
||||
from khaosz.core.tokenizer import BpeTokenizer
|
||||
from khaosz.core.transformer import TransformerConfig, Transformer
|
||||
|
||||
|
||||
class BaseModelIO:
|
||||
"""Base class for model I/O operations."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: Optional[nn.Module] = None,
|
||||
tokenizer: Optional[BpeTokenizer] = None,
|
||||
config: Optional[TransformerConfig] = None
|
||||
):
|
||||
self.model = model
|
||||
self.tokenizer = tokenizer or BpeTokenizer()
|
||||
self.config = config or TransformerConfig()
|
||||
|
||||
def _get_file_paths(self, directory: Union[str, Path]) -> dict[str, Path]:
|
||||
"""Get standardized file paths for model components."""
|
||||
dir_path = Path(directory)
|
||||
return {
|
||||
"model": dir_path / "model.safetensors",
|
||||
"config": dir_path / "config.json",
|
||||
"tokenizer": dir_path / "tokenizer.json"
|
||||
}
|
||||
|
||||
def save_components(self, save_dir: Union[str, Path]):
|
||||
"""Save core model components."""
|
||||
paths = self._get_file_paths(save_dir)
|
||||
paths["model"].parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
if self.model is not None:
|
||||
st.save_file(self.model.state_dict(), str(paths["model"]))
|
||||
self.config.save(str(paths["config"]))
|
||||
self.tokenizer.save(str(paths["tokenizer"]))
|
||||
|
||||
def load_components(self, load_dir: Union[str, Path]) -> Self:
|
||||
"""Load core model components."""
|
||||
paths = self._get_file_paths(load_dir)
|
||||
|
||||
self.config.load(str(paths["config"]))
|
||||
self.tokenizer.load(str(paths["tokenizer"]))
|
||||
|
||||
if paths["model"].exists():
|
||||
state_dict = st.load_file(str(paths["model"]))
|
||||
if self.model is None:
|
||||
self.model = Transformer(self.config)
|
||||
self.model.load_state_dict(state_dict)
|
||||
|
||||
return self
|
||||
|
||||
def to(self, *args, **kwargs) -> Self:
|
||||
"""Move model to device."""
|
||||
if self.model is not None:
|
||||
self.model.to(*args, **kwargs)
|
||||
return self
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelParameter(BaseModelIO):
|
||||
"""Container for model parameters with serialization capabilities."""
|
||||
|
||||
model: Optional[nn.Module] = field(
|
||||
default=None,
|
||||
metadata={"help": "Transformer model."}
|
||||
)
|
||||
tokenizer: BpeTokenizer = field(
|
||||
default_factory=BpeTokenizer,
|
||||
metadata={"help": "Tokenizer for the model."}
|
||||
)
|
||||
config: TransformerConfig = field(
|
||||
default_factory=TransformerConfig,
|
||||
metadata={"help": "Transformer model configuration."}
|
||||
)
|
||||
|
||||
def save(self, save_dir: Union[str, Path]):
|
||||
"""Save model parameters."""
|
||||
self.save_components(save_dir)
|
||||
|
||||
def load(self, load_dir: Union[str, Path]) -> Self:
|
||||
"""Load model parameters."""
|
||||
return self.load_components(load_dir)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Checkpoint(BaseModelIO):
|
||||
"""Extended model parameters with training state."""
|
||||
|
||||
model: Optional[nn.Module] = field(
|
||||
default=None,
|
||||
metadata={"help": "Transformer model."}
|
||||
)
|
||||
tokenizer: BpeTokenizer = field(
|
||||
default_factory=BpeTokenizer,
|
||||
metadata={"help": "Tokenizer for the model."}
|
||||
)
|
||||
config: TransformerConfig = field(
|
||||
default_factory=TransformerConfig,
|
||||
metadata={"help": "Transformer model configuration."}
|
||||
)
|
||||
loss_list: list[float] = field(
|
||||
default_factory=list,
|
||||
metadata={"help": "List of training losses."}
|
||||
)
|
||||
current_iter: int = field(
|
||||
default=0,
|
||||
metadata={"help": "Current training iteration."}
|
||||
)
|
||||
optimizer: Optional[optim.Optimizer] = field(
|
||||
default=None,
|
||||
metadata={"help": "Optimizer state."}
|
||||
)
|
||||
|
||||
def __post_init__(self):
|
||||
# Ensure current_iter matches loss list length if not explicitly set
|
||||
if self.current_iter == 0 and self.loss_list:
|
||||
self.current_iter = len(self.loss_list)
|
||||
|
||||
def _get_training_paths(self, directory: Union[str, Path]) -> dict[str, Path]:
|
||||
"""Get file paths for training-specific files."""
|
||||
paths = self._get_file_paths(directory)
|
||||
paths.update({
|
||||
"loss_list": paths["model"].parent / "loss.pkl",
|
||||
"loss_plot": paths["model"].parent / "loss.png",
|
||||
"optimizer": paths["model"].parent / "optimizer.pkl"
|
||||
})
|
||||
return paths
|
||||
|
||||
def save_training_state(self, save_dir: Union[str, Path]):
|
||||
"""Save training-specific state."""
|
||||
paths = self._get_training_paths(save_dir)
|
||||
|
||||
# Save loss plot
|
||||
self._plot_loss(str(paths["loss_plot"]))
|
||||
|
||||
# Save loss list
|
||||
with open(str(paths["loss_list"]), "wb") as f:
|
||||
pkl.dump(self.loss_list, f)
|
||||
|
||||
# Save optimizer state
|
||||
if self.optimizer is not None:
|
||||
with open(str(paths["optimizer"]), "wb") as f:
|
||||
pkl.dump(self.optimizer.state_dict(), f)
|
||||
|
||||
def load_training_state(self, load_dir: Union[str, Path]) -> Self:
|
||||
"""Load training-specific state."""
|
||||
paths = self._get_training_paths(load_dir)
|
||||
|
||||
# Load loss list
|
||||
if paths["loss_list"].exists():
|
||||
with open(str(paths["loss_list"]), "rb") as f:
|
||||
self.loss_list = pkl.load(f)
|
||||
self.current_iter = len(self.loss_list)
|
||||
|
||||
# Load optimizer state
|
||||
if paths["optimizer"].exists() and self.optimizer is not None:
|
||||
with open(str(paths["optimizer"]), "rb") as f:
|
||||
optim_state = pkl.load(f)
|
||||
self.optimizer.load_state_dict(optim_state)
|
||||
|
||||
return self
|
||||
|
||||
def _plot_loss(self, save_path: str):
|
||||
"""Plot and save loss curve."""
|
||||
if not self.loss_list:
|
||||
return
|
||||
|
||||
plt.figure(figsize=(10, 6))
|
||||
plt.plot(self.loss_list)
|
||||
plt.title(f"Training Loss - Iteration {self.current_iter}")
|
||||
plt.xlabel("Batch")
|
||||
plt.ylabel("Loss")
|
||||
plt.grid(True)
|
||||
plt.savefig(save_path, dpi=300, bbox_inches="tight")
|
||||
plt.close()
|
||||
|
||||
def save(self, save_dir: Union[str, Path]):
|
||||
"""Save complete checkpoint."""
|
||||
self.save_components(save_dir)
|
||||
self.save_training_state(save_dir)
|
||||
|
||||
def load(self, load_dir: Union[str, Path]) -> Self:
|
||||
"""Load complete checkpoint."""
|
||||
self.load_components(load_dir)
|
||||
self.load_training_state(load_dir)
|
||||
return self
|
||||
|
||||
|
||||
class ParameterLoader:
|
||||
"""Factory class for loading model parameters or checkpoints."""
|
||||
|
||||
@staticmethod
|
||||
def load(load_dir: Union[str, Path]) -> Union[ModelParameter, Checkpoint]:
|
||||
"""Load either ModelParameter or Checkpoint based on directory contents."""
|
||||
load_dir = Path(load_dir)
|
||||
|
||||
# Check for training-specific files
|
||||
loss_file = load_dir / "loss.pkl"
|
||||
has_training_data = loss_file.exists()
|
||||
|
||||
# Create appropriate instance
|
||||
if has_training_data:
|
||||
checkpoint = Checkpoint()
|
||||
checkpoint.load(str(load_dir))
|
||||
return checkpoint
|
||||
else:
|
||||
params = ModelParameter()
|
||||
params.load(str(load_dir))
|
||||
return params
|
||||
|
||||
@staticmethod
|
||||
def create_checkpoint(
|
||||
model: nn.Module,
|
||||
tokenizer: BpeTokenizer,
|
||||
config: TransformerConfig,
|
||||
loss_list: Optional[list[float]] = None,
|
||||
optimizer: Optional[optim.Optimizer] = None
|
||||
) -> Checkpoint:
|
||||
"""Convenience method to create a training checkpoint."""
|
||||
return Checkpoint(
|
||||
model=model,
|
||||
tokenizer=tokenizer,
|
||||
config=config,
|
||||
loss_list=loss_list or [],
|
||||
optimizer=optimizer
|
||||
)
|
||||
|
||||
|
||||
@@ -0,0 +1,111 @@
|
||||
from tokenizers import Tokenizer, Encoding
|
||||
from tokenizers import decoders, processors, normalizers, pre_tokenizers
|
||||
from tokenizers.models import BPE
|
||||
from tokenizers.trainers import BpeTrainer
|
||||
from typing import List, Union
|
||||
|
||||
|
||||
class BpeTokenizer:
|
||||
def __init__(self, path=None):
|
||||
self._control_tokens = ["<bos>", "<eos>", "<pad>"]
|
||||
self._special_tokens = ["<|user|>", "<|system|>"]
|
||||
model = BPE()
|
||||
tokenizer = Tokenizer(model)
|
||||
tokenizer.normalizer = normalizers.Sequence([
|
||||
normalizers.NFC()
|
||||
])
|
||||
tokenizer.pre_tokenizer = pre_tokenizers.Sequence([
|
||||
pre_tokenizers.Punctuation(behavior="isolated"),
|
||||
pre_tokenizers.Metaspace(prepend_scheme="never"),
|
||||
pre_tokenizers.Split(pattern=r"(\d+|[a-zA-Z]+|(?:'s|'t|'re|'ve|'m|'ll|'d))", behavior="isolated"),
|
||||
pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=False)
|
||||
])
|
||||
tokenizer.decoder = decoders.Sequence([
|
||||
decoders.ByteLevel(),
|
||||
decoders.Metaspace(prepend_scheme="never")
|
||||
])
|
||||
tokenizer.post_processor = processors.Sequence([
|
||||
processors.ByteLevel(trim_offsets=False)
|
||||
])
|
||||
self._tokenizer = tokenizer
|
||||
|
||||
if path is not None:
|
||||
self._tokenizer = Tokenizer.from_file(path)
|
||||
|
||||
def _prepare_trainer(self, vocab_size: int, min_freq: int, reserved_token_size: int) -> tuple:
|
||||
assert reserved_token_size > len(self._special_tokens)
|
||||
reserved_tokens = [f"<|rsv{i:02d}|>" for i in range(reserved_token_size - len(self._special_tokens))]
|
||||
detail_vocab_size = vocab_size - (len(reserved_tokens) + len(self._special_tokens))
|
||||
|
||||
alphabet = pre_tokenizers.ByteLevel.alphabet()
|
||||
min_size = len(alphabet) + len(self._control_tokens)
|
||||
assert detail_vocab_size > min_size
|
||||
|
||||
trainer = BpeTrainer(
|
||||
vocab_size=detail_vocab_size,
|
||||
min_frequency=min_freq,
|
||||
limit_alphabet=detail_vocab_size // 4,
|
||||
max_token_length=18,
|
||||
special_tokens=self._control_tokens,
|
||||
show_progress=True,
|
||||
initial_alphabet=alphabet,
|
||||
)
|
||||
|
||||
return trainer, detail_vocab_size, reserved_tokens
|
||||
|
||||
def train(self, files, vocab_size, min_freq, reserved_token_size=100):
|
||||
trainer, _, reserved_tokens = self._prepare_trainer(
|
||||
vocab_size=vocab_size,
|
||||
min_freq=min_freq,
|
||||
reserved_token_size=reserved_token_size
|
||||
)
|
||||
self._tokenizer.train(files=files, trainer=trainer)
|
||||
self._tokenizer.add_special_tokens(self._special_tokens + reserved_tokens)
|
||||
|
||||
def train_from_iterator(self, iterator, vocab_size, min_freq, reserved_token_size=100):
|
||||
trainer, _, reserved_tokens = self._prepare_trainer(
|
||||
vocab_size=vocab_size,
|
||||
min_freq=min_freq,
|
||||
reserved_token_size=reserved_token_size
|
||||
)
|
||||
self._tokenizer.train_from_iterator(iterator=iterator, trainer=trainer)
|
||||
self._tokenizer.add_special_tokens(self._special_tokens + reserved_tokens)
|
||||
|
||||
def save(self, path):
|
||||
self._tokenizer.save(path)
|
||||
|
||||
def load(self, path):
|
||||
self._tokenizer = Tokenizer.from_file(path)
|
||||
|
||||
def encode(self, tokens: Union[str, List[str]], out_ids: bool=True, add_special_tokens: bool=False) -> List:
|
||||
if isinstance(tokens, str):
|
||||
encoded: Encoding = self._tokenizer.encode(tokens, add_special_tokens=add_special_tokens)
|
||||
return encoded.ids if out_ids else encoded.tokens
|
||||
elif isinstance(tokens, list):
|
||||
encoded_list: List[Encoding] = self._tokenizer.encode_batch(tokens, add_special_tokens=add_special_tokens)
|
||||
return [encoded.ids if out_ids else encoded.tokens for encoded in encoded_list]
|
||||
|
||||
def decode(self, tokens: List[int], skip_special_tokens: bool=True) -> str:
|
||||
return self._tokenizer.decode(tokens, skip_special_tokens=skip_special_tokens)
|
||||
|
||||
def __len__(self) -> int:
|
||||
return self._tokenizer.get_vocab_size()
|
||||
|
||||
@property
|
||||
def stop_ids(self) -> List[int]:
|
||||
stop_ids = []
|
||||
for token in self._control_tokens:
|
||||
stop_ids.append(self._tokenizer.token_to_id(token))
|
||||
return stop_ids
|
||||
|
||||
@property
|
||||
def bos_id(self) -> int:
|
||||
return self._tokenizer.token_to_id("<bos>")
|
||||
|
||||
@property
|
||||
def eos_id(self) -> int:
|
||||
return self._tokenizer.token_to_id("<eos>")
|
||||
|
||||
@property
|
||||
def pad_id(self) -> int:
|
||||
return self._tokenizer.token_to_id("<pad>")
|
||||
@@ -0,0 +1,341 @@
|
||||
import json
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from torch import Tensor
|
||||
from torch.nn import init
|
||||
from dataclasses import asdict, dataclass
|
||||
from typing import List, Optional, Self, Tuple
|
||||
|
||||
|
||||
def repeat_kv(x: Tensor, n_rep: int) -> Tensor:
|
||||
"""
|
||||
Repeat k times along the dimension for attention heads.
|
||||
Args:
|
||||
x (Tensor): The input tensor.
|
||||
n_rep (int): The number of repetitions.
|
||||
Returns:
|
||||
Tensor: The repeated tensor.
|
||||
"""
|
||||
|
||||
bs, slen, n_heads, head_dim = x.shape
|
||||
if n_rep == 1:
|
||||
return x
|
||||
return (
|
||||
x[:, :, :, None, :]
|
||||
.expand(bs, slen, n_heads, n_rep, head_dim)
|
||||
.reshape(bs, slen, n_heads * n_rep, head_dim)
|
||||
)
|
||||
|
||||
def get_rotary_emb(
|
||||
dim: int,
|
||||
max_len: int,
|
||||
base: float = 10000,
|
||||
device: torch.device = "cuda",
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Get the rotary embedding for the given dimension and maximum length.
|
||||
Args:
|
||||
dim (int): The dimension of the input.
|
||||
max_len (int): The maximum length of the input.
|
||||
base (float, optional): The base for the frequency. Defaults to 10000.
|
||||
device (torch.device, optional): The device to use. Defaults to "cuda".
|
||||
Returns:
|
||||
Tensor: The rotary embedding tensor.
|
||||
"""
|
||||
|
||||
theta = base ** (-torch.arange(0, dim, 2, device=device).float() / dim)
|
||||
t = torch.arange(0, max_len, device=device).float()
|
||||
freqs = torch.outer(t, theta)
|
||||
freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
|
||||
|
||||
return freqs_cis
|
||||
|
||||
def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor:
|
||||
"""
|
||||
Apply rotary embedding to the input tensor.
|
||||
Args:
|
||||
x (Tensor): The input tensor.
|
||||
freqs_cis (Tensor): The rotary embedding tensor.
|
||||
Returns:
|
||||
Tensor: The output tensor.
|
||||
"""
|
||||
|
||||
dtype = x.dtype
|
||||
seq_len = x.size(1)
|
||||
|
||||
x_complex = torch.view_as_complex(x.view(*x.shape[:-1], -1, 2).float())
|
||||
freqs_cis = freqs_cis.reshape(1, seq_len, 1, -1)
|
||||
x_out = torch.view_as_real(x_complex * freqs_cis).flatten(3)
|
||||
|
||||
return x_out.to(dtype)
|
||||
|
||||
def create_attention_mask(
|
||||
seq_mask: Tensor,
|
||||
start_pos: int = 0,
|
||||
seq_len: int = 0,
|
||||
is_causal: bool = False,
|
||||
device: torch.device = "cuda",
|
||||
dtype: torch.dtype = torch.float32
|
||||
) -> Tensor:
|
||||
"""
|
||||
Create attention mask for GQA
|
||||
Args:
|
||||
seq_mask (Tensor): A tensor indicating whether each position is valid or not.
|
||||
start_pos (int): The starting position of the sequence.
|
||||
seq_len (int): The length of the sequence.
|
||||
is_causal (bool): Whether the attention is causal or not.
|
||||
device (torch.device): The device to use.
|
||||
Returns:
|
||||
Tensor: The attention mask tensor.
|
||||
"""
|
||||
|
||||
if start_pos != 0 and seq_mask is None:
|
||||
# for single prompt chat
|
||||
seq_mask = torch.ones((1, seq_len), dtype=torch.bool, device=device)
|
||||
|
||||
if seq_mask is None:
|
||||
return None
|
||||
|
||||
batch_size = seq_mask.size(0)
|
||||
seq_mask = seq_mask[:, :start_pos + seq_len].to(device=device, dtype=torch.bool)
|
||||
# (bsz, start_pos + seq_len)
|
||||
expanded_mask = seq_mask.unsqueeze(1).expand(batch_size, seq_len, start_pos + seq_len)
|
||||
# (bsz, seq_len, start_pos + seq_len)
|
||||
|
||||
if is_causal:
|
||||
causal_mask = torch.tril(
|
||||
torch.ones((seq_len, start_pos + seq_len), dtype=torch.bool, device=device),
|
||||
diagonal=start_pos
|
||||
)
|
||||
causal_mask = causal_mask.unsqueeze(0).expand(batch_size, seq_len, start_pos + seq_len)
|
||||
expanded_mask = expanded_mask & causal_mask
|
||||
|
||||
attention_mask = torch.zeros_like(expanded_mask, dtype=dtype, device=device)
|
||||
attention_mask = attention_mask.masked_fill_(~expanded_mask, -torch.finfo(dtype).max / 2).unsqueeze(1)
|
||||
# (bsz, 1, seq_len, seq_len + start_pos)
|
||||
|
||||
return attention_mask
|
||||
|
||||
|
||||
@dataclass
|
||||
class TransformerConfig:
|
||||
# basic config
|
||||
vocab_size: Optional[int] = None
|
||||
n_dim: Optional[int] = None
|
||||
n_head: Optional[int] = None
|
||||
n_layer: Optional[int] = None
|
||||
m_len: Optional[int] = None
|
||||
norm_eps: Optional[float] = None
|
||||
d_ffn: Optional[int] = None
|
||||
|
||||
# GQA
|
||||
n_kvhead: Optional[int] = None
|
||||
|
||||
|
||||
def load(self, config_path: str) -> Self:
|
||||
with open(config_path, 'r') as f:
|
||||
config: dict = json.load(f)
|
||||
for key, value in config.items():
|
||||
if hasattr(self, key):
|
||||
setattr(self, key, value)
|
||||
|
||||
return self
|
||||
|
||||
def save(self, config_path: str) -> None:
|
||||
config_dict = asdict(self)
|
||||
config_dict = {k: v for k, v in config_dict.items() if v is not None}
|
||||
with open(config_path, 'w') as f:
|
||||
json.dump(config_dict, f, indent=4)
|
||||
|
||||
|
||||
class Linear(nn.Module):
|
||||
def __init__(self, in_dim: int, out_dim: int, bias: bool=False):
|
||||
super().__init__()
|
||||
self.weight = nn.Parameter(torch.empty((out_dim, in_dim)))
|
||||
self.bias = nn.Parameter(torch.zeros(out_dim)) if bias else None
|
||||
init.normal_(self.weight, mean=0, std=0.006)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
return F.linear(x, self.weight, self.bias)
|
||||
|
||||
|
||||
class RMSNorm(nn.Module):
|
||||
def __init__(self, n_dim, norm_eps):
|
||||
super().__init__()
|
||||
self.weight = nn.Parameter(torch.ones(n_dim))
|
||||
self.norm_eps = norm_eps
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
dtype = x.dtype
|
||||
x = x.float()
|
||||
mean_square = torch.mean(torch.pow(x, 2), dim=-1, keepdim=True)
|
||||
norm = x * torch.rsqrt(mean_square + self.norm_eps)
|
||||
norm = norm.to(dtype)
|
||||
out = norm * self.weight
|
||||
return out
|
||||
|
||||
|
||||
class MLP(nn.Module):
|
||||
def __init__(self, n_dim: int, d_ffn: int):
|
||||
super().__init__()
|
||||
self.up = Linear(n_dim, d_ffn)
|
||||
self.gate = Linear(n_dim, d_ffn)
|
||||
self.down = Linear(d_ffn, n_dim)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
gated = self.up(x) * F.silu(self.gate(x))
|
||||
out = self.down(gated)
|
||||
return out
|
||||
|
||||
|
||||
class GQA(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
n_dim: int,
|
||||
n_head: int,
|
||||
n_kvhead: int,
|
||||
):
|
||||
super().__init__()
|
||||
assert n_dim % n_head == 0
|
||||
assert n_head % n_kvhead == 0
|
||||
|
||||
self.head_dim = n_dim // n_head
|
||||
self.n_dim = n_dim
|
||||
self.n_heads = n_head
|
||||
self.n_kvheads = n_kvhead
|
||||
self.n_rep = n_head // n_kvhead
|
||||
|
||||
self.q_proj = Linear(n_dim, n_head * self.head_dim)
|
||||
self.k_proj = Linear(n_dim, n_kvhead * self.head_dim)
|
||||
self.v_proj = Linear(n_dim, n_kvhead * self.head_dim)
|
||||
self.o_proj = Linear(n_dim, n_dim)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: Tensor,
|
||||
freqs_cis: Tensor,
|
||||
mask: Tensor = None,
|
||||
kv_cache: Optional[Tuple[Tensor, Tensor]] = None,
|
||||
start_pos: int = 0
|
||||
) -> Tensor:
|
||||
bsz, seq_len, _ = x.size()
|
||||
# x(bsz, seq_len, n_heads * head_dim) -> (bsz, seq_len, n_heads, head_dim)
|
||||
q = self._split_heads(self.q_proj(x), self.n_heads)
|
||||
k = self._split_heads(self.k_proj(x), self.n_kvheads)
|
||||
v = self._split_heads(self.v_proj(x), self.n_kvheads)
|
||||
q, k = apply_rotary_emb(q, freqs_cis), apply_rotary_emb(k, freqs_cis)
|
||||
|
||||
if kv_cache is not None:
|
||||
k_cache, v_cache = kv_cache
|
||||
|
||||
# copy to cache
|
||||
k_cache[:bsz, start_pos:start_pos + seq_len] = k
|
||||
v_cache[:bsz, start_pos:start_pos + seq_len] = v
|
||||
|
||||
# get cache
|
||||
k = k_cache[:bsz, :start_pos + seq_len]
|
||||
v = v_cache[:bsz, :start_pos + seq_len]
|
||||
|
||||
k, v = repeat_kv(k, self.n_rep), repeat_kv(v, self.n_rep)
|
||||
|
||||
# (bsz, seq_len, n_heads, head_dim) -> (bsz, n_heads, seq_len, head_dim)
|
||||
q, k, v = q.permute(0, 2, 1, 3), k.permute(0, 2, 1, 3), v.permute(0, 2, 1, 3)
|
||||
sdqa_out = F.scaled_dot_product_attention(q, k, v, mask, is_causal=(mask == None)).permute(0, 2, 1, 3)
|
||||
out = self.o_proj(sdqa_out.contiguous().view(bsz, seq_len, -1))
|
||||
|
||||
return out
|
||||
|
||||
def _split_heads(self, x: Tensor, n_heads) -> Tensor:
|
||||
batch_size, seq_len, _ = x.shape
|
||||
x = x.reshape(batch_size, seq_len, n_heads, self.head_dim)
|
||||
return x
|
||||
|
||||
|
||||
class DecoderBlock(nn.Module):
|
||||
def __init__(self, n_dim, n_head, d_ffn, n_kvhead, norm_eps):
|
||||
super().__init__()
|
||||
self.attention = GQA(n_dim, n_head, n_kvhead)
|
||||
self.norm_attn = RMSNorm(n_dim, norm_eps)
|
||||
self.ffn = MLP(n_dim, d_ffn)
|
||||
self.norm_ffn = RMSNorm(n_dim, norm_eps)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: Tensor,
|
||||
freqs_cis: Tensor,
|
||||
attention_mask: Optional[Tensor] = None,
|
||||
kv_cache: Optional[Tuple[Tensor, Tensor]] = None,
|
||||
start_pos: int = 0
|
||||
) -> Tensor:
|
||||
# attention
|
||||
attn_output = self.attention(
|
||||
self.norm_attn(x),
|
||||
freqs_cis,
|
||||
attention_mask,
|
||||
kv_cache,
|
||||
start_pos
|
||||
)
|
||||
x = attn_output + x
|
||||
|
||||
# feed forward
|
||||
x = self.ffn(self.norm_ffn(x)) + x
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class Transformer(nn.Module):
|
||||
def __init__(self, config: TransformerConfig):
|
||||
super().__init__()
|
||||
self.embedding = nn.Parameter(torch.empty(config.vocab_size, config.n_dim))
|
||||
self.layers = nn.ModuleList([
|
||||
DecoderBlock(
|
||||
config.n_dim,
|
||||
config.n_head,
|
||||
config.d_ffn,
|
||||
config.n_kvhead,
|
||||
config.norm_eps
|
||||
)
|
||||
for _ in range(config.n_layer)
|
||||
])
|
||||
self.norm = RMSNorm(config.n_dim, config.norm_eps)
|
||||
self.freq_cis = get_rotary_emb(config.n_dim // config.n_head, config.m_len)
|
||||
init.normal_(self.embedding, mean=0, std=0.02)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: Tensor,
|
||||
seq_mask: Optional[Tensor]=None,
|
||||
persistent_key_values: Optional[List[Tuple[Tensor, Tensor]]]=None,
|
||||
start_pos: int = 0
|
||||
) -> Tensor:
|
||||
assert input_ids.ndim == 2
|
||||
seq_len = input_ids.size(-1)
|
||||
x = F.embedding(input_ids, self.embedding)
|
||||
|
||||
self.freq_cis = self.freq_cis.to(x.device)
|
||||
freqs_cis = self.freq_cis[start_pos:start_pos+seq_len]
|
||||
has_kvcache = persistent_key_values is not None
|
||||
|
||||
attn_mask = create_attention_mask(
|
||||
seq_mask,
|
||||
start_pos=start_pos,
|
||||
seq_len=seq_len,
|
||||
is_causal=has_kvcache,
|
||||
device=x.device,
|
||||
dtype=x.dtype
|
||||
)
|
||||
|
||||
for i, layer in enumerate(self.layers):
|
||||
kv_cache = persistent_key_values[i] if persistent_key_values else None
|
||||
x = layer(x, freqs_cis, attn_mask, kv_cache, start_pos)
|
||||
|
||||
hidden_states = self.norm(x)
|
||||
logits = F.linear(hidden_states, self.embedding)
|
||||
|
||||
return {
|
||||
"logits": logits,
|
||||
"hidden_states": hidden_states
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user