In [1]:
# 全局常量
DIR = "../models/HuggingFace/Qwen3-0.6B"
In [2]:
# 处理tokenize逻辑
from tokenizers import Tokenizer
tokenizer = Tokenizer.from_file(DIR + "/tokenizer.json")
output = tokenizer.encode("测试一下tokenizer的效果")
print(output.ids)
tokenizer.decode(output.ids)
[81705, 100158, 85593, 105005]
Out[2]:
'测试一下tokenizer的效果'
In [ ]:
# 处理参数和超参数
from safetensors import safe_open
import json
import torch
# 常数和超参数
eos_token_id: int # 用于判断模型输出完的token id
rms_norm_eps: float # RMSNorm防止归一化分母为0的eps
rope_theta: int # RoPE的theta的底数部分
num_hidden_layers: int # transformer层数
num_attention_heads: int # Q的注意力头数
num_key_value_heads: int # KV的注意力头数
head_dim: int # 单头注意力维数
max_position_embeddings: int # 允许位置编码的最大序列长度
# 参数
E: torch.Tensor # 嵌入层
Gamma_finalRMSNorm: torch.Tensor # 最终RMSNorm的缩放系数
W_output: torch.Tensor # 最终投影到词表大小的线性层
Gamma_inputRMSNorm: list[torch.Tensor] = [] # transformer层中对输入的RMSNorm的缩放系数
Gamma_postATTNRMSNorm: list[torch.Tensor] = [] # transformer层中对自注意力结果的RMSNorm的缩放系数
W_FFNGate: list[torch.Tensor] = [] # transformer层中FFN的门控层
W_FFNUp: list[torch.Tensor] = [] # transformer层中FFN的升维层
W_FFNDown: list[torch.Tensor] = [] # transformer层中FFN的降维层
Gamma_QRMSNorm: list[torch.Tensor] = [] # transformer层中Q投影后、RoPE前的每个头的RMSNorm的缩放系数
Gamma_KRMSNorm: list[torch.Tensor] = [] # transformer层中K投影后、RoPE前的每个头的RMSNorm的缩放系数
W_Q: list[torch.Tensor] = [] # transformer层中Q投影矩阵
W_K: list[torch.Tensor] = [] # transformer层中K投影矩阵
W_V: list[torch.Tensor] = [] # transformer层中V投影矩阵
W_O: list[torch.Tensor] = [] # transformer层中O投影矩阵
# 读取常数和超参数
with open(DIR + "/config.json", "r", encoding="utf-8") as f:
data = json.load(f)
def readParam(name, t):
d = data.get(name)
if not isinstance(d, t):
raise ValueError(f"get {name} error")
return d
eos_token_id = readParam("eos_token_id", int) # 151645
rms_norm_eps = readParam("rms_norm_eps", float) # 1e-06
rope_theta = readParam("rope_theta", int) # 1000000
num_hidden_layers = readParam("num_hidden_layers", int) # 28
num_attention_heads = readParam("num_attention_heads", int) # 16
num_key_value_heads = readParam("num_key_value_heads", int) # 8
head_dim = readParam("head_dim", int) # 128
max_position_embeddings = readParam("max_position_embeddings", int) # 40960
# 读取参数
with safe_open(DIR + "/model.safetensors", framework="pt", device=0) as f:
E = f.get_tensor("model.embed_tokens.weight") # torch.Size([151936, 1024])
Gamma_finalRMSNorm = f.get_tensor("model.norm.weight") # torch.Size([1024])
W_output = f.get_tensor("lm_head.weight") # torch.Size([151936, 1024])
for i in range(num_hidden_layers):
Gamma_inputRMSNorm.append(f.get_tensor(f"model.layers.{i}.input_layernorm.weight")) # torch.Size([1024])
Gamma_postATTNRMSNorm.append(f.get_tensor(f"model.layers.{i}.post_attention_layernorm.weight")) # torch.Size([1024])
W_FFNGate.append(f.get_tensor(f"model.layers.{i}.mlp.gate_proj.weight")) # torch.Size([3072, 1024])
W_FFNUp.append(f.get_tensor(f"model.layers.{i}.mlp.up_proj.weight")) # torch.Size([3072, 1024])
W_FFNDown.append(f.get_tensor(f"model.layers.{i}.mlp.down_proj.weight")) # torch.Size([1024, 3072])
Gamma_QRMSNorm.append(f.get_tensor(f"model.layers.{i}.self_attn.q_norm.weight")) # torch.Size([128])
Gamma_KRMSNorm.append(f.get_tensor(f"model.layers.{i}.self_attn.k_norm.weight")) # torch.Size([128])
W_Q.append(f.get_tensor(f"model.layers.{i}.self_attn.q_proj.weight")) # torch.Size([2048, 1024])
W_K.append(f.get_tensor(f"model.layers.{i}.self_attn.k_proj.weight")) # torch.Size([1024, 1024])
W_V.append(f.get_tensor(f"model.layers.{i}.self_attn.v_proj.weight")) # torch.Size([1024, 1024])
W_O.append(f.get_tensor(f"model.layers.{i}.self_attn.o_proj.weight")) # torch.Size([1024, 2048])
model_dtype = E.dtype
In [4]:
# 工具函数
def RMSNorm(X, gamma):
return X / torch.sqrt(torch.mean(X[..., :] ** 2, -1, keepdim=True) + rms_norm_eps) * gamma
# interleaved: 交错配对theta[0,0,1,1,2,2,3,3]
# rope_theta_repeat2 = torch.pow(rope_theta, -torch.arange(head_dim // 2) * 2 / head_dim).repeat_interleave(2)
# rotate_half: 前后两半配对theta_0,1,2,3,0,1,2,3
rope_theta_repeat2 = torch.pow(rope_theta, -torch.arange(head_dim // 2) * 2 / head_dim).repeat(2)
rope_m_theta = torch.arange(max_position_embeddings).reshape(-1, 1) * rope_theta_repeat2.reshape(1, -1)
rope_sin_m_theta = torch.sin(rope_m_theta).cuda() # torch.Size([40960, 128])
rope_cos_m_theta = torch.cos(rope_m_theta).cuda() # torch.Size([40960, 128])
def RoPE(X, offset=0):
Ls = X.shape[-2]
d = X.shape[-1]
cos, sin = rope_cos_m_theta[offset : offset + Ls], rope_sin_m_theta[offset : offset + Ls]
# interleaved: 交错配对x[-1,0,-3,2]
# X_rot = torch.stack((-X[..., 1::2], X[..., 0::2]), dim=-1).flatten(-2)
# rotate_half: 前后两半配对x[-0,-1,-2,-3,0,1,2,3]
X_rot = torch.cat((-X[..., d // 2 :], X[..., : d // 2]), dim=-1)
return (X * cos + X_rot * sin).to(model_dtype)
# RoPE(torch.arange(2 * 3 * head_dim).reshape(2, 3, head_dim).cuda(), 0)
def SiLU(X):
return X / (1 + torch.exp(-X))
# SiLU(torch.arange(-10, 10).reshape(2, 10))
def softmax(X, temperature=1):
X_shifted = X - torch.max(X / temperature, dim=-1, keepdims=True).values
exp_X = torch.exp(X_shifted)
return exp_X / torch.sum(exp_X, dim=-1, keepdims=True)
# softmax(torch.arange(-10, 10).reshape(2, 10))
# 右上方三角被mask
def masked_softmax(X):
# X: (batch_size, num_heads, num_steps, num_steps)
maxlen = X.size(-1)
mask = (torch.arange(maxlen).reshape(-1, 1) >= torch.arange(maxlen)).to(X.device)
mask_not = ~mask
return softmax(X * mask + (mask_not * -1e6).to(X.dtype))
# masked_softmax(torch.arange(96).reshape(2, 3, 4, 4))
In [ ]:
# 核心计算函数与Transformer层
def FFN(X, gate, up, down):
return (SiLU(X @ gate.T) * (X @ up.T)) @ down.T
# 通过转置实现Q/K/V张量按注意力头的拆分与恢复，便于一次性计算所有注意力头
def transpose_qkv(X, num_heads):
# X_h: (batch_size, num_steps, num_heads, num_hiddens/num_heads)
X_h = X.reshape(*X.shape[:-1], num_heads, -1)
return X_h.transpose(-2, -3)
def transpose_output(X):
# X_t: (batch_size, num_steps, num_heads, num_hiddens/num_heads)
X_t = X.transpose(-2, -3)
return X_t.reshape(*X_t.shape[:-2], -1)
# print(transpose_qkv(torch.arange(80).reshape(2, 5, 8), 2).shape)
# print(transpose_output(transpose_qkv(torch.arange(80).reshape(2, 5, 8), 2)).shape)
def DotProductAttention(Q, K, V, offset=0):
KVScale = num_attention_heads // num_key_value_heads
# Q/K/V: (batch_size, num_heads, num_steps, num_hiddens/num_heads)
# scores/A: (batch_size, num_heads, num_steps, num_steps)
scores = Q @ torch.repeat_interleave(K.transpose(-1, -2), KVScale, -3) / torch.sqrt(torch.tensor(Q.shape[-1]))
A = softmax(scores) if offset else masked_softmax(scores)
return A @ torch.repeat_interleave(V, KVScale, -3)
KVCache: list[list[torch.Tensor]]
def ClearKVCache():
global KVCache
KVCache = [None] * num_hidden_layers
# KVCache中已经prefill或predict了多少个字符
def GetOffset():
global KVCache
return KVCache[0][0].size(-2)
# offset为0代表prefill阶段
def TransfomerLayer(layer, X, Gamma_inputRMSNorm, Gamma_postATTNRMSNorm, W_FFNGate, W_FFNUp, W_FFNDown, Gamma_QRMSNorm, Gamma_KRMSNorm, W_Q, W_K, W_V, W_O, offset=0):
# X: (batch_size, num_steps, num_hiddens)
X_inputRMSNorm = RMSNorm(X, Gamma_inputRMSNorm)
# 投影 -> QK RMSNorm -> RoPE -> attention
# Q/K/V: (batch_size, num_steps, num_heads, num_hiddens) -> (batch_size, num_heads, num_steps, num_hiddens/num_heads)
Q = transpose_qkv(X_inputRMSNorm @ W_Q.T, num_attention_heads)
K = transpose_qkv(X_inputRMSNorm @ W_K.T, num_key_value_heads)
V = transpose_qkv(X_inputRMSNorm @ W_V.T, num_key_value_heads)
Q_RMSNorm = RMSNorm(Q, Gamma_QRMSNorm)
K_RMSNorm = RMSNorm(K, Gamma_KRMSNorm)
Q_RoPE = RoPE(Q_RMSNorm, offset)
K_RoPE = RoPE(K_RMSNorm, offset)
# prefill阶段初始化KVCache
global KVCache
if offset == 0:
KVCache[layer] = [K_RoPE, V]
else:
KVCache[layer][0] = torch.cat((KVCache[layer][0], K_RoPE), dim=-2)
KVCache[layer][1] = torch.cat((KVCache[layer][1], V), dim=-2)
# V_attn: (batch_size, num_heads, num_steps, num_hiddens/num_heads)
V_attn = DotProductAttention(Q_RoPE, KVCache[layer][0], KVCache[layer][1], offset)
X2 = X + transpose_output(V_attn) @ W_O.T
# FFN前的RMSNorm -> FFN
# O: (batch_size, num_steps, num_hiddens)
O = X2 + FFN(RMSNorm(X2, Gamma_postATTNRMSNorm), W_FFNGate, W_FFNUp, W_FFNDown)
return O
In [6]:
# prefill与predict
def prefill(X):
ClearKVCache()
# X: (batch_size, num_steps)
# O: (batch_size, num_steps, num_hiddens)
O = E[X]
for i in range(num_hidden_layers):
O = TransfomerLayer(
i, O, Gamma_inputRMSNorm[i], Gamma_postATTNRMSNorm[i], W_FFNGate[i], W_FFNUp[i], W_FFNDown[i], Gamma_QRMSNorm[i], Gamma_KRMSNorm[i], W_Q[i], W_K[i], W_V[i], W_O[i]
)
return O
# 多层Transformer得到的结果经过最终的RMSNorm、线性层和softmax得到概率
def lastOutputToken(O, temperature=1):
O = RMSNorm(O, Gamma_finalRMSNorm)
Y = O @ W_output.T
P = softmax(Y, temperature)
return torch.argmax(P, dim=-1)[..., -1:].to(torch.long)
def predict(X, temperature=1):
offset = GetOffset()
# O: (batch_size, 1)
res = X.clone().cpu()
while offset < max_position_embeddings:
token_id = res[0][-1].numpy()
print(tokenizer.decode([token_id], skip_special_tokens=False), end="")
if token_id == eos_token_id:
break
O = E[X]
for i in range(num_hidden_layers):
O = TransfomerLayer(
i, O, Gamma_inputRMSNorm[i], Gamma_postATTNRMSNorm[i], W_FFNGate[i], W_FFNUp[i], W_FFNDown[i], Gamma_QRMSNorm[i], Gamma_KRMSNorm[i], W_Q[i], W_K[i], W_V[i], W_O[i], offset
)
X = lastOutputToken(O, temperature)
res = torch.cat((res, X.cpu()), dim=-1)
offset += 1
print()
return res.numpy()
prefill(torch.ones(1, 5, dtype=torch.long))
Out[6]:
tensor([[[119.5000, 110.0000, 113.0000, ..., -28.5000, 14.2500, 18.5000],
[120.0000, 112.0000, 115.0000, ..., -28.6250, 14.5000, 18.8750],
[119.5000, 112.0000, 114.0000, ..., -28.3750, 14.6250, 19.0000],
[120.0000, 110.0000, 116.0000, ..., -28.3750, 14.5000, 19.0000],
[120.0000, 112.0000, 113.0000, ..., -28.3750, 14.2500, 18.7500]]],
device='cuda:0', dtype=torch.bfloat16)
In [7]:
# 构造输入
import jinja2
with open(DIR + "/tokenizer_config.json", "r", encoding="utf-8") as f:
data = json.load(f)
chat_template = data["chat_template"]
tpl = jinja2.Environment().from_string(chat_template)
messages = [{"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "steam现在最火的游戏是什么"}]
rendered = tpl.render(
messages=messages,
add_generation_prompt=True,
enable_thinking=True,
)
tokens = tokenizer.encode(rendered).ids
print(rendered)
print(tokens)
<|im_start|>system You are a helpful assistant.<|im_end|> <|im_start|>user steam现在最火的游戏是什么<|im_end|> <|im_start|>assistant [151644, 8948, 198, 2610, 525, 264, 10950, 17847, 13, 151645, 198, 151644, 872, 198, 46590, 99601, 31235, 79599, 105628, 102021, 151645, 198, 151644, 77091, 198]
In [8]:
# 发起预测
import time
batch_size = 50
start = time.perf_counter()
O = prefill(torch.tensor(tokens, dtype=torch.long).reshape(1, -1).repeat(batch_size, 1))
end = time.perf_counter()
print(f"prefill耗时: {end - start:.2f} 秒")
start = time.perf_counter()
X = lastOutputToken(O)
res = predict(X)
end = time.perf_counter()
print(f"predict: {len(res[0])*batch_size/(end - start):.2f} tokens/s")
print(res)
prefill耗时: 0.11 秒 <think> 好的，用户问的是“Steam现在最火的游戏是什么”。首先，我需要确认用户的问题意图。他们可能想知道最新的热门游戏，或者想了解Steam上最受欢迎的游戏类型。接下来，我得考虑用户可能的背景。可能是刚接触Steam的新手，或者对游戏行业不太熟悉，想了解当前的热门趋势。 然后，我需要回忆一下Steam上最近的热门游戏。根据之前的了解，最近有几个热门游戏，比如《艾尔登法环》、《巫师3》、《最终幻想16》、《原神》和《塞尔达传说：旷野之息》。这些游戏在Steam上都有很高的评分和下载量。不过，用户可能想知道这些游戏的具体信息，比如更新时间、平台支持等。 另外，用户可能没有明确说明他们是在问哪个游戏最火，所以需要给出一个综合的答案，涵盖几个热门游戏，并解释为什么它们受欢迎。同时，要确保信息准确，避免错误。比如，确认《巫师3》的更新时间，以及《原神》的发布日期，这样用户能更清楚地了解。 还要注意用户可能的深层需求，比如他们可能想了解游戏的类型、玩法或者社区氛围，所以回答时可以提到这些方面，帮助用户全面了解。最后，保持回答简洁明了，结构清晰，方便用户快速获取信息。 </think> 目前Steam上最火的游戏包括： 1. **《艾尔登法环》** - 以开放世界冒险和深度剧情著称，2023年发布，全球下载量超10亿。 2. **《巫师3》** - 2023年更新，结合了新世界观和角色深度，评分高达9.5分。 3. **《最终幻想16》** - 2023年发布，融合了新世界观和角色深度，评分9.2分。 4. **《原神》** - 2023年发布，以角色和世界观创新著称，评分9.1分。 5. **《塞尔达传说：旷野之息》** - 2023年发布，结合了新世界观和角色深度，评分9.0分。 这些游戏在Steam的评分和下载量中表现突出，反映了当前游戏行业的热门趋势。如果你有具体偏好（如平台、类型或社区），我可以进一步细化推荐！<|im_end|> predict: 924.85 tokens/s [[151667 198 99692 ... 101914 6313 151645] [151667 198 99692 ... 101914 6313 151645] [151667 198 99692 ... 101914 6313 151645] ... [151667 198 99692 ... 101914 6313 151645] [151667 198 99692 ... 101914 6313 151645] [151667 198 99692 ... 101914 6313 151645]]