diff --git a/README.md b/README.md index 9be8d09..aa4b00d 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,15 @@ A minimal Transformer implementation using NumPy with **KV Cache** support for e ## Files -- `transformer.py` - Core Transformer model implementation with KV Cache +- `transformer/` - Transformer model package (按架构组件拆分) + - `__init__.py` - 包初始化，导出所有组件 + - `utils.py` - 激活函数和掩码工具 + - `layers.py` - 基础层（Linear, LayerNorm） + - `attention.py` - 多头注意力机制 + - `feedforward.py` - 前馈网络 + - `block.py` - Transformer编码器块 + - `model.py` - SimpleTransformer模型 + - `inference.py` - 推理引擎 - `example.py` - Usage examples (training & inference) ## Usage diff --git a/transformer.py b/transformer.py deleted file mode 100644 index 9542c7b..0000000 --- a/transformer.py +++ /dev/null @@ -1,438 +0,0 @@ -import numpy as np - - -def softmax(x, axis=-1): - """Softmax激活函数，用于将 logits 转换为概率分布""" - # 减去最大值防止数值溢出 - e_x = np.exp(x - np.max(x, axis=axis, keepdims=True)) - return e_x / np.sum(e_x, axis=axis, keepdims=True) - - -def relu(x): - """ReLU激活函数""" - return np.maximum(0, x) - - -class Linear: - """全连接层（线性层）""" - def __init__(self, in_features, out_features): - # 使用He初始化权重 - self.weight = np.random.randn(in_features, out_features) * np.sqrt(2.0 / in_features) - # 偏置初始化为0 - self.bias = np.zeros(out_features) - - def forward(self, x): - """前向传播：y = xW + b""" - return x @ self.weight + self.bias - - -class LayerNorm: - """层归一化（Layer Normalization）""" - def __init__(self, d_model, eps=1e-6): - # 缩放参数 - self.gamma = np.ones(d_model) - # 偏移参数 - self.beta = np.zeros(d_model) - # 防止除零的小常数 - self.eps = eps - - def forward(self, x): - """前向传播：对最后一个维度进行归一化""" - mean = np.mean(x, axis=-1, keepdims=True) - std = np.std(x, axis=-1, keepdims=True) - return self.gamma * (x - mean) / (std + self.eps) + self.beta - - -class MultiHeadAttention: - """多头注意力机制（支持KV cache）""" - def __init__(self, d_model, num_heads): - self.d_model = d_model # 模型维度 - self.num_heads = num_heads # 注意力头数 - self.d_k = d_model // num_heads # 每个头的维度 - - # 定义Q、K、V的线性变换层 - self.W_q = Linear(d_model, d_model) # Query变换 - self.W_k = Linear(d_model, d_model) # Key变换 - self.W_v = Linear(d_model, d_model) # Value变换 - self.W_o = Linear(d_model, d_model) # 输出变换 - - # KV cache：用于推理时缓存历史Key和Value - self.k_cache = None # 缓存的Key [batch_size, num_heads, seq_len, d_k] - self.v_cache = None # 缓存的Value [batch_size, num_heads, seq_len, d_k] - - def clear_cache(self): - """清除KV cache""" - self.k_cache = None - self.v_cache = None - - def forward(self, Q, K, V, mask=None, use_cache=False): - """ - 前向传播 - Q: Query张量 [batch_size, seq_len, d_model] - K: Key张量 [batch_size, seq_len, d_model] - V: Value张量 [batch_size, seq_len, d_model] - mask: 注意力掩码 - use_cache: 是否使用KV cache（推理时设为True） - """ - batch_size = Q.shape[0] - seq_len = Q.shape[1] - - # 线性变换并分割成多头 - Q = self.W_q.forward(Q).reshape(batch_size, seq_len, self.num_heads, self.d_k).transpose(0, 2, 1, 3) - K = self.W_k.forward(K).reshape(batch_size, seq_len, self.num_heads, self.d_k).transpose(0, 2, 1, 3) - V = self.W_v.forward(V).reshape(batch_size, seq_len, self.num_heads, self.d_k).transpose(0, 2, 1, 3) - - # KV cache逻辑 - if use_cache: - if self.k_cache is None: - # Prefill阶段：首次计算，缓存完整的K和V - self.k_cache = K - self.v_cache = V - else: - # Decode阶段：拼接历史缓存和新的K、V - self.k_cache = np.concatenate([self.k_cache, K], axis=2) - self.v_cache = np.concatenate([self.v_cache, V], axis=2) - # 使用完整的K和V进行注意力计算 - K = self.k_cache - V = self.v_cache - - # 计算缩放点积注意力分数 - # Q: [batch, heads, q_len, d_k] - # K: [batch, heads, kv_len, d_k]（decode时kv_len > q_len） - attn_scores = Q @ K.transpose(0, 1, 3, 2) / np.sqrt(self.d_k) - - # 应用掩码（如果提供） - if mask is not None: - attn_scores = np.where(mask == 0, -1e9, attn_scores) - - # Softmax得到注意力权重 - attn_probs = softmax(attn_scores, axis=-1) - - # 加权求和 - attn_output = attn_probs @ V - - # 拼接多头输出 - attn_output = attn_output.transpose(0, 2, 1, 3).reshape(batch_size, seq_len, self.d_model) - - # 最终线性变换 - output = self.W_o.forward(attn_output) - return output - - -class FeedForward: - """前馈神经网络（两层全连接网络）""" - def __init__(self, d_model, d_ff): - # 第一层：扩展维度 - self.linear1 = Linear(d_model, d_ff) - # 第二层：恢复维度 - self.linear2 = Linear(d_ff, d_model) - - def forward(self, x): - """前向传播：线性 -> ReLU -> 线性""" - return self.linear2.forward(relu(self.linear1.forward(x))) - - -class TransformerBlock: - """Transformer编码器块""" - def __init__(self, d_model, num_heads, d_ff): - # 多头自注意力层 - self.attention = MultiHeadAttention(d_model, num_heads) - # 前馈网络层 - self.feed_forward = FeedForward(d_model, d_ff) - # 两个层归一化 - self.norm1 = LayerNorm(d_model) - self.norm2 = LayerNorm(d_model) - - def forward(self, x, mask=None, use_cache=False): - """ - 前向传播（残差连接 + 层归一化） - 1. 自注意力 -> 残差连接 -> 层归一化 - 2. 前馈网络 -> 残差连接 -> 层归一化 - x: 输入张量 [batch_size, seq_len, d_model] - mask: 注意力掩码 - use_cache: 是否使用KV cache - """ - # 自注意力子层（Q=K=V，自注意力） - attn_output = self.attention.forward(x, x, x, mask, use_cache=use_cache) - x = self.norm1.forward(x + attn_output) - - # 前馈网络子层 - ff_output = self.feed_forward.forward(x) - x = self.norm2.forward(x + ff_output) - return x - - def clear_cache(self): - """清除该层的KV cache""" - self.attention.clear_cache() - - -class SimpleTransformer: - """简化的Transformer编码器模型（支持prefill/decode分离）""" - def __init__(self, vocab_size, d_model, num_heads, num_layers, d_ff, max_seq_len): - self.d_model = d_model - self.vocab_size = vocab_size - # 词嵌入层 - self.embedding = np.random.randn(vocab_size, d_model) * 0.01 - # 位置编码 - self.positional_encoding = self._create_positional_encoding(max_seq_len, d_model) - # Transformer编码器层堆叠 - self.transformer_blocks = [ - TransformerBlock(d_model, num_heads, d_ff) for _ in range(num_layers) - ] - - def _create_positional_encoding(self, max_seq_len, d_model): - """ - 创建正弦余弦位置编码 - 使用不同频率的正弦和余弦函数生成位置信息 - """ - pe = np.zeros((max_seq_len, d_model)) - position = np.arange(0, max_seq_len).reshape(-1, 1).astype(float) - div_term = np.exp(np.arange(0, d_model, 2).astype(float) * -(np.log(10000.0) / d_model)) - pe[:, 0::2] = np.sin(position * div_term) # 偶数维度用正弦 - pe[:, 1::2] = np.cos(position * div_term) # 奇数维度用余弦 - return pe - - def clear_cache(self): - """清除所有层的KV cache""" - for block in self.transformer_blocks: - block.clear_cache() - - def prefill(self, x, mask=None): - """ - Prefill阶段：处理完整的输入序列 - x: 输入序列 [batch_size, seq_len] - mask: 注意力掩码 - 返回: 输出张量 [batch_size, seq_len, d_model] - """ - seq_len = x.shape[1] - - # 词嵌入 + 缩放 - x = self.embedding[x] * np.sqrt(self.d_model) - - # 加上位置编码 - x = x + self.positional_encoding[:seq_len, :] - - # 通过所有Transformer块，使用KV cache - for block in self.transformer_blocks: - x = block.forward(x, mask, use_cache=True) - - return x - - def decode(self, x, position): - """ - Decode阶段：处理单个新token - x: 新token [batch_size, 1] - position: 当前token在序列中的位置 - 返回: 输出张量 [batch_size, 1, d_model] - """ - # 词嵌入 + 缩放 - x = self.embedding[x] * np.sqrt(self.d_model) - - # 加上位置编码（使用当前位置） - x = x + self.positional_encoding[position:position+1, :] - - # 通过所有Transformer块，使用KV cache - for block in self.transformer_blocks: - x = block.forward(x, use_cache=True) - - return x - - def forward(self, x, mask=None): - """ - 前向传播（不使用KV cache，用于训练） - x: 输入序列 [batch_size, seq_len] - mask: 注意力掩码 - """ - seq_len = x.shape[1] - - # 词嵌入 + 缩放 - x = self.embedding[x] * np.sqrt(self.d_model) - - # 加上位置编码 - x = x + self.positional_encoding[:seq_len, :] - - # 通过所有Transformer块 - for block in self.transformer_blocks: - x = block.forward(x, mask) - - return x - - def count_parameters(self): - """统计模型参数数量""" - count = 0 - count += self.embedding.size # 嵌入层参数 - for block in self.transformer_blocks: - # 注意力层参数 - count += block.attention.W_q.weight.size + block.attention.W_q.bias.size - count += block.attention.W_k.weight.size + block.attention.W_k.bias.size - count += block.attention.W_v.weight.size + block.attention.W_v.bias.size - count += block.attention.W_o.weight.size + block.attention.W_o.bias.size - # 前馈网络参数 - count += block.feed_forward.linear1.weight.size + block.feed_forward.linear1.bias.size - count += block.feed_forward.linear2.weight.size + block.feed_forward.linear2.bias.size - # 层归一化参数 - count += block.norm1.gamma.size + block.norm1.beta.size - count += block.norm2.gamma.size + block.norm2.beta.size - return count - - -def create_padding_mask(seq, pad_idx=0): - """ - 创建填充掩码 - 用于忽略序列中的填充位置（padding tokens） - """ - return (seq != pad_idx).astype(float)[:, np.newaxis, np.newaxis, :] - - -def create_look_ahead_mask(size): - """ - 创建前瞻掩码（下三角掩码） - 用于解码器中，防止位置i看到i之后的信息 - """ - mask = np.triu(np.ones((size, size)), k=1).astype(bool) - return (~mask).astype(float) - - -class InferenceEngine: - """推理引擎：管理prefill和decode阶段""" - def __init__(self, model): - """ - 初始化推理引擎 - model: SimpleTransformer模型实例 - """ - self.model = model - self.generated_tokens = [] # 已生成的token列表 - self.current_position = 0 # 当前位置 - - def reset(self): - """重置推理状态""" - self.model.clear_cache() - self.generated_tokens = [] - self.current_position = 0 - - def prefill(self, input_ids): - """ - Prefill阶段：处理完整的输入序列 - input_ids: 输入token序列 [batch_size, seq_len] - 返回: 下一个token的logits [batch_size, vocab_size] - """ - self.reset() - batch_size = input_ids.shape[0] - seq_len = input_ids.shape[1] - - # 创建前瞻掩码（防止看到未来信息） - mask = create_look_ahead_mask(seq_len) - - # 执行prefill前向传播（使用KV cache） - output = self.model.prefill(input_ids, mask) - - # 获取最后一个位置的输出（用于预测下一个token） - last_output = output[:, -1, :] # [batch_size, d_model] - - # 简单的logits计算：使用嵌入矩阵的转置作为输出投影 - # logits = last_output @ embedding.T - logits = last_output @ self.model.embedding.T # [batch_size, vocab_size] - - # 更新状态 - self.current_position = seq_len - - return logits - - def decode_step(self, input_token): - """ - Decode阶段：处理单个新token - input_token: 新token [batch_size, 1] - 返回: 下一个token的logits [batch_size, vocab_size] - """ - # 执行decode前向传播（使用KV cache） - output = self.model.decode(input_token, self.current_position) - - # 获取输出（单token，所以直接取[:, 0, :]） - last_output = output[:, 0, :] # [batch_size, d_model] - - # 计算logits - logits = last_output @ self.model.embedding.T # [batch_size, vocab_size] - - # 更新位置 - self.current_position += 1 - - return logits - - def generate(self, input_ids, max_new_tokens=50, temperature=1.0): - """ - 自回归生成文本 - input_ids: 输入token序列 [batch_size, seq_len] - max_new_tokens: 最大生成token数 - temperature: 温度参数（控制随机性） - 返回: 完整的生成序列 [batch_size, seq_len + max_new_tokens] - """ - batch_size = input_ids.shape[0] - - # 保存原始输入 - generated = input_ids.copy() - - # Prefill阶段 - logits = self.prefill(input_ids) - - # 采样第一个生成的token - next_token = self._sample(logits, temperature) - generated = np.concatenate([generated, next_token], axis=1) - - # Decode阶段：逐个生成token - for i in range(max_new_tokens - 1): - logits = self.decode_step(next_token) - next_token = self._sample(logits, temperature) - generated = np.concatenate([generated, next_token], axis=1) - - # 检查是否生成了结束符（这里用0作为结束符） - if np.all(next_token == 0): - break - - return generated - - def _sample(self, logits, temperature=1.0): - """ - 从logits中采样 - logits: [batch_size, vocab_size] - temperature: 温度参数 - 返回: 采样的token [batch_size, 1] - """ - # 应用温度缩放 - logits = logits / temperature - - # 计算概率分布 - probs = softmax(logits, axis=-1) - - # 按概率采样 - batch_size = logits.shape[0] - next_tokens = np.zeros((batch_size, 1), dtype=int) - - for i in range(batch_size): - # 使用numpy的random.choice采样 - next_tokens[i, 0] = np.random.choice(logits.shape[1], p=probs[i]) - - return next_tokens - - -if __name__ == "__main__": - # 模型超参数 - vocab_size = 1000 # 词汇表大小 - d_model = 512 # 模型维度 - num_heads = 8 # 注意力头数 - num_layers = 6 # Transformer层数 - d_ff = 2048 # 前馈网络隐藏层维度 - max_seq_len = 100 # 最大序列长度 - - # 创建模型 - model = SimpleTransformer(vocab_size, d_model, num_heads, num_layers, d_ff, max_seq_len) - - # 创建随机输入 - batch_size = 2 - seq_len = 10 - x = np.random.randint(0, vocab_size, (batch_size, seq_len)) - - # 前向传播 - output = model.forward(x) - print(f"输入形状: {x.shape}") - print(f"输出形状: {output.shape}") - print(f"模型参数量: {model.count_parameters():,}") \ No newline at end of file diff --git a/transformer/__init__.py b/transformer/__init__.py new file mode 100644 index 0000000..9ba48c0 --- /dev/null +++ b/transformer/__init__.py @@ -0,0 +1,33 @@ +""" +Transformer模型实现包 +按架构组件拆分： +- utils: 激活函数和掩码工具 +- layers: 基础层（Linear, LayerNorm） +- attention: 多头注意力机制 +- feedforward: 前馈网络 +- block: Transformer块 +- model: SimpleTransformer模型 +- inference: 推理引擎 +""" + +from .utils import softmax, relu, create_padding_mask, create_look_ahead_mask +from .layers import Linear, LayerNorm +from .attention import MultiHeadAttention +from .feedforward import FeedForward +from .block import TransformerBlock +from .model import SimpleTransformer +from .inference import InferenceEngine + +__all__ = [ + 'softmax', + 'relu', + 'create_padding_mask', + 'create_look_ahead_mask', + 'Linear', + 'LayerNorm', + 'MultiHeadAttention', + 'FeedForward', + 'TransformerBlock', + 'SimpleTransformer', + 'InferenceEngine' +] \ No newline at end of file diff --git a/transformer/attention.py b/transformer/attention.py new file mode 100644 index 0000000..2f0673f --- /dev/null +++ b/transformer/attention.py @@ -0,0 +1,84 @@ +""" +多头注意力机制模块 +包含MultiHeadAttention类，支持KV cache +""" + +import numpy as np +from .layers import Linear +from .utils import softmax + + +class MultiHeadAttention: + """多头注意力机制（支持KV cache）""" + def __init__(self, d_model, num_heads): + self.d_model = d_model # 模型维度 + self.num_heads = num_heads # 注意力头数 + self.d_k = d_model // num_heads # 每个头的维度 + + # 定义Q、K、V的线性变换层 + self.W_q = Linear(d_model, d_model) # Query变换 + self.W_k = Linear(d_model, d_model) # Key变换 + self.W_v = Linear(d_model, d_model) # Value变换 + self.W_o = Linear(d_model, d_model) # 输出变换 + + # KV cache：用于推理时缓存历史Key和Value + self.k_cache = None # 缓存的Key [batch_size, num_heads, seq_len, d_k] + self.v_cache = None # 缓存的Value [batch_size, num_heads, seq_len, d_k] + + def clear_cache(self): + """清除KV cache""" + self.k_cache = None + self.v_cache = None + + def forward(self, Q, K, V, mask=None, use_cache=False): + """ + 前向传播 + Q: Query张量 [batch_size, seq_len, d_model] + K: Key张量 [batch_size, seq_len, d_model] + V: Value张量 [batch_size, seq_len, d_model] + mask: 注意力掩码 + use_cache: 是否使用KV cache（推理时设为True） + """ + batch_size = Q.shape[0] + seq_len = Q.shape[1] + + # 线性变换并分割成多头 + Q = self.W_q.forward(Q).reshape(batch_size, seq_len, self.num_heads, self.d_k).transpose(0, 2, 1, 3) + K = self.W_k.forward(K).reshape(batch_size, seq_len, self.num_heads, self.d_k).transpose(0, 2, 1, 3) + V = self.W_v.forward(V).reshape(batch_size, seq_len, self.num_heads, self.d_k).transpose(0, 2, 1, 3) + + # KV cache逻辑 + if use_cache: + if self.k_cache is None: + # Prefill阶段：首次计算，缓存完整的K和V + self.k_cache = K + self.v_cache = V + else: + # Decode阶段：拼接历史缓存和新的K、V + self.k_cache = np.concatenate([self.k_cache, K], axis=2) + self.v_cache = np.concatenate([self.v_cache, V], axis=2) + # 使用完整的K和V进行注意力计算 + K = self.k_cache + V = self.v_cache + + # 计算缩放点积注意力分数 + # Q: [batch, heads, q_len, d_k] + # K: [batch, heads, kv_len, d_k]（decode时kv_len > q_len） + attn_scores = Q @ K.transpose(0, 1, 3, 2) / np.sqrt(self.d_k) + + # 应用掩码（如果提供） + if mask is not None: + attn_scores = np.where(mask == 0, -1e9, attn_scores) + + # Softmax得到注意力权重 + attn_probs = softmax(attn_scores, axis=-1) + + # 加权求和 + attn_output = attn_probs @ V + + # 拼接多头输出 + attn_output = attn_output.transpose(0, 2, 1, 3).reshape(batch_size, seq_len, self.d_model) + + # 最终线性变换 + output = self.W_o.forward(attn_output) + return output \ No newline at end of file diff --git a/transformer/block.py b/transformer/block.py new file mode 100644 index 0000000..107b368 --- /dev/null +++ b/transformer/block.py @@ -0,0 +1,42 @@ +""" +Transformer块模块 +包含TransformerBlock类 +""" + +from .attention import MultiHeadAttention +from .feedforward import FeedForward +from .layers import LayerNorm + + +class TransformerBlock: + """Transformer编码器块""" + def __init__(self, d_model, num_heads, d_ff): + # 多头自注意力层 + self.attention = MultiHeadAttention(d_model, num_heads) + # 前馈网络层 + self.feed_forward = FeedForward(d_model, d_ff) + # 两个层归一化 + self.norm1 = LayerNorm(d_model) + self.norm2 = LayerNorm(d_model) + + def forward(self, x, mask=None, use_cache=False): + """ + 前向传播（残差连接 + 层归一化） + 1. 自注意力 -> 残差连接 -> 层归一化 + 2. 前馈网络 -> 残差连接 -> 层归一化 + x: 输入张量 [batch_size, seq_len, d_model] + mask: 注意力掩码 + use_cache: 是否使用KV cache + """ + # 自注意力子层（Q=K=V，自注意力） + attn_output = self.attention.forward(x, x, x, mask, use_cache=use_cache) + x = self.norm1.forward(x + attn_output) + + # 前馈网络子层 + ff_output = self.feed_forward.forward(x) + x = self.norm2.forward(x + ff_output) + return x + + def clear_cache(self): + """清除该层的KV cache""" + self.attention.clear_cache() \ No newline at end of file diff --git a/transformer/feedforward.py b/transformer/feedforward.py new file mode 100644 index 0000000..6541bdb --- /dev/null +++ b/transformer/feedforward.py @@ -0,0 +1,20 @@ +""" +前馈网络模块 +包含FeedForward类 +""" + +from .layers import Linear +from .utils import relu + + +class FeedForward: + """前馈神经网络（两层全连接网络）""" + def __init__(self, d_model, d_ff): + # 第一层：扩展维度 + self.linear1 = Linear(d_model, d_ff) + # 第二层：恢复维度 + self.linear2 = Linear(d_ff, d_model) + + def forward(self, x): + """前向传播：线性 -> ReLU -> 线性""" + return self.linear2.forward(relu(self.linear1.forward(x))) \ No newline at end of file diff --git a/transformer/inference.py b/transformer/inference.py new file mode 100644 index 0000000..14a0417 --- /dev/null +++ b/transformer/inference.py @@ -0,0 +1,128 @@ +""" +推理引擎模块 +包含InferenceEngine类，管理prefill和decode阶段 +""" + +import numpy as np +from .utils import softmax, create_look_ahead_mask + + +class InferenceEngine: + """推理引擎：管理prefill和decode阶段""" + def __init__(self, model): + """ + 初始化推理引擎 + model: SimpleTransformer模型实例 + """ + self.model = model + self.generated_tokens = [] # 已生成的token列表 + self.current_position = 0 # 当前位置 + + def reset(self): + """重置推理状态""" + self.model.clear_cache() + self.generated_tokens = [] + self.current_position = 0 + + def prefill(self, input_ids): + """ + Prefill阶段：处理完整的输入序列 + input_ids: 输入token序列 [batch_size, seq_len] + 返回: 下一个token的logits [batch_size, vocab_size] + """ + self.reset() + batch_size = input_ids.shape[0] + seq_len = input_ids.shape[1] + + # 创建前瞻掩码（防止看到未来信息） + mask = create_look_ahead_mask(seq_len) + + # 执行prefill前向传播（使用KV cache） + output = self.model.prefill(input_ids, mask) + + # 获取最后一个位置的输出（用于预测下一个token） + last_output = output[:, -1, :] # [batch_size, d_model] + + # 简单的logits计算：使用嵌入矩阵的转置作为输出投影 + # logits = last_output @ embedding.T + logits = last_output @ self.model.embedding.T # [batch_size, vocab_size] + + # 更新状态 + self.current_position = seq_len + + return logits + + def decode_step(self, input_token): + """ + Decode阶段：处理单个新token + input_token: 新token [batch_size, 1] + 返回: 下一个token的logits [batch_size, vocab_size] + """ + # 执行decode前向传播（使用KV cache） + output = self.model.decode(input_token, self.current_position) + + # 获取输出（单token，所以直接取[:, 0, :]） + last_output = output[:, 0, :] # [batch_size, d_model] + + # 计算logits + logits = last_output @ self.model.embedding.T # [batch_size, vocab_size] + + # 更新位置 + self.current_position += 1 + + return logits + + def generate(self, input_ids, max_new_tokens=50, temperature=1.0): + """ + 自回归生成文本 + input_ids: 输入token序列 [batch_size, seq_len] + max_new_tokens: 最大生成token数 + temperature: 温度参数（控制随机性） + 返回: 完整的生成序列 [batch_size, seq_len + max_new_tokens] + """ + batch_size = input_ids.shape[0] + + # 保存原始输入 + generated = input_ids.copy() + + # Prefill阶段 + logits = self.prefill(input_ids) + + # 采样第一个生成的token + next_token = self._sample(logits, temperature) + generated = np.concatenate([generated, next_token], axis=1) + + # Decode阶段：逐个生成token + for i in range(max_new_tokens - 1): + logits = self.decode_step(next_token) + next_token = self._sample(logits, temperature) + generated = np.concatenate([generated, next_token], axis=1) + + # 检查是否生成了结束符（这里用0作为结束符） + if np.all(next_token == 0): + break + + return generated + + def _sample(self, logits, temperature=1.0): + """ + 从logits中采样 + logits: [batch_size, vocab_size] + temperature: 温度参数 + 返回: 采样的token [batch_size, 1] + """ + # 应用温度缩放 + logits = logits / temperature + + # 计算概率分布 + probs = softmax(logits, axis=-1) + + # 按概率采样 + batch_size = logits.shape[0] + next_tokens = np.zeros((batch_size, 1), dtype=int) + + for i in range(batch_size): + # 使用numpy的random.choice采样 + next_tokens[i, 0] = np.random.choice(logits.shape[1], p=probs[i]) + + return next_tokens \ No newline at end of file diff --git a/transformer/layers.py b/transformer/layers.py new file mode 100644 index 0000000..c56f1c1 --- /dev/null +++ b/transformer/layers.py @@ -0,0 +1,36 @@ +""" +基础层模块 +包含Linear和LayerNorm层 +""" + +import numpy as np + + +class Linear: + """全连接层（线性层）""" + def __init__(self, in_features, out_features): + # 使用He初始化权重 + self.weight = np.random.randn(in_features, out_features) * np.sqrt(2.0 / in_features) + # 偏置初始化为0 + self.bias = np.zeros(out_features) + + def forward(self, x): + """前向传播：y = xW + b""" + return x @ self.weight + self.bias + + +class LayerNorm: + """层归一化（Layer Normalization）""" + def __init__(self, d_model, eps=1e-6): + # 缩放参数 + self.gamma = np.ones(d_model) + # 偏移参数 + self.beta = np.zeros(d_model) + # 防止除零的小常数 + self.eps = eps + + def forward(self, x): + """前向传播：对最后一个维度进行归一化""" + mean = np.mean(x, axis=-1, keepdims=True) + std = np.std(x, axis=-1, keepdims=True) + return self.gamma * (x - mean) / (std + self.eps) + self.beta \ No newline at end of file diff --git a/transformer/model.py b/transformer/model.py new file mode 100644 index 0000000..f612e00 --- /dev/null +++ b/transformer/model.py @@ -0,0 +1,117 @@ +""" +Transformer模型模块 +包含SimpleTransformer类 +""" + +import numpy as np +from .block import TransformerBlock + + +class SimpleTransformer: + """简化的Transformer编码器模型（支持prefill/decode分离）""" + def __init__(self, vocab_size, d_model, num_heads, num_layers, d_ff, max_seq_len): + self.d_model = d_model + self.vocab_size = vocab_size + # 词嵌入层 + self.embedding = np.random.randn(vocab_size, d_model) * 0.01 + # 位置编码 + self.positional_encoding = self._create_positional_encoding(max_seq_len, d_model) + # Transformer编码器层堆叠 + self.transformer_blocks = [ + TransformerBlock(d_model, num_heads, d_ff) for _ in range(num_layers) + ] + + def _create_positional_encoding(self, max_seq_len, d_model): + """ + 创建正弦余弦位置编码 + 使用不同频率的正弦和余弦函数生成位置信息 + """ + pe = np.zeros((max_seq_len, d_model)) + position = np.arange(0, max_seq_len).reshape(-1, 1).astype(float) + div_term = np.exp(np.arange(0, d_model, 2).astype(float) * -(np.log(10000.0) / d_model)) + pe[:, 0::2] = np.sin(position * div_term) # 偶数维度用正弦 + pe[:, 1::2] = np.cos(position * div_term) # 奇数维度用余弦 + return pe + + def clear_cache(self): + """清除所有层的KV cache""" + for block in self.transformer_blocks: + block.clear_cache() + + def prefill(self, x, mask=None): + """ + Prefill阶段：处理完整的输入序列 + x: 输入序列 [batch_size, seq_len] + mask: 注意力掩码 + 返回: 输出张量 [batch_size, seq_len, d_model] + """ + seq_len = x.shape[1] + + # 词嵌入 + 缩放 + x = self.embedding[x] * np.sqrt(self.d_model) + + # 加上位置编码 + x = x + self.positional_encoding[:seq_len, :] + + # 通过所有Transformer块，使用KV cache + for block in self.transformer_blocks: + x = block.forward(x, mask, use_cache=True) + + return x + + def decode(self, x, position): + """ + Decode阶段：处理单个新token + x: 新token [batch_size, 1] + position: 当前token在序列中的位置 + 返回: 输出张量 [batch_size, 1, d_model] + """ + # 词嵌入 + 缩放 + x = self.embedding[x] * np.sqrt(self.d_model) + + # 加上位置编码（使用当前位置） + x = x + self.positional_encoding[position:position+1, :] + + # 通过所有Transformer块，使用KV cache + for block in self.transformer_blocks: + x = block.forward(x, use_cache=True) + + return x + + def forward(self, x, mask=None): + """ + 前向传播（不使用KV cache，用于训练） + x: 输入序列 [batch_size, seq_len] + mask: 注意力掩码 + """ + seq_len = x.shape[1] + + # 词嵌入 + 缩放 + x = self.embedding[x] * np.sqrt(self.d_model) + + # 加上位置编码 + x = x + self.positional_encoding[:seq_len, :] + + # 通过所有Transformer块 + for block in self.transformer_blocks: + x = block.forward(x, mask) + + return x + + def count_parameters(self): + """统计模型参数数量""" + count = 0 + count += self.embedding.size # 嵌入层参数 + for block in self.transformer_blocks: + # 注意力层参数 + count += block.attention.W_q.weight.size + block.attention.W_q.bias.size + count += block.attention.W_k.weight.size + block.attention.W_k.bias.size + count += block.attention.W_v.weight.size + block.attention.W_v.bias.size + count += block.attention.W_o.weight.size + block.attention.W_o.bias.size + # 前馈网络参数 + count += block.feed_forward.linear1.weight.size + block.feed_forward.linear1.bias.size + count += block.feed_forward.linear2.weight.size + block.feed_forward.linear2.bias.size + # 层归一化参数 + count += block.norm1.gamma.size + block.norm1.beta.size + count += block.norm2.gamma.size + block.norm2.beta.size + return count \ No newline at end of file diff --git a/transformer/utils.py b/transformer/utils.py new file mode 100644 index 0000000..1305f0a --- /dev/null +++ b/transformer/utils.py @@ -0,0 +1,40 @@ +""" +工具函数模块 +包含激活函数和掩码创建函数 +""" + +import numpy as np + + +def softmax(x, axis=-1): + """Softmax激活函数，用于将 logits 转换为概率分布""" + # 减去最大值防止数值溢出 + e_x = np.exp(x - np.max(x, axis=axis, keepdims=True)) + return e_x / np.sum(e_x, axis=axis, keepdims=True) + + +def relu(x): + """ReLU激活函数""" + return np.maximum(0, x) + + +def create_padding_mask(seq, pad_idx=0): + """ + 创建填充掩码 + 用于忽略序列中的填充位置（padding tokens） + seq: 输入序列 [batch_size, seq_len] + pad_idx: 填充token的索引 + 返回: 掩码 [batch_size, 1, 1, seq_len] + """ + return (seq != pad_idx).astype(float)[:, np.newaxis, np.newaxis, :] + + +def create_look_ahead_mask(size): + """ + 创建前瞻掩码（下三角掩码） + 用于解码器中，防止位置i看到i之后的信息 + size: 序列长度 + 返回: 掩码 [size, size] + """ + mask = np.triu(np.ones((size, size)), k=1).astype(bool) + return (~mask).astype(float) \ No newline at end of file