35 lines
1.1 KiB
Markdown
35 lines
1.1 KiB
Markdown
|
|
|
|
|
|
## data flow
|
|
|
|
input_ids = tokenizer.build_chat_input(query, history=history, role=role)
|
|
|
|
for
|
|
input_ids -> [1, 6] 1:batch_num 6:sequence_length
|
|
inputs_embeds -> [6, 1, 4096] 4096:hidden_size
|
|
rotary_pos_emb -> [6, 1, 32, 2] 32:pos的编码维度 2:cos+sin
|
|
|
|
hidden_states = inputs_embeds
|
|
for layers : GLMBlock(hidden_states, rotary_pos_emb)
|
|
hidden_states = RMSNorm(hidden_states)
|
|
hidden_states = hidden_states[-1:] 截取最后一个sequence
|
|
lm_logits = self.output_layer(hidden_states)
|
|
lm_logits = lm_logits.transpose(0, 1).contiguous() -> [1, 1, 65024]
|
|
|
|
probs = softmax(lm_logits) -> [1, 65024]
|
|
next_tokens = torch.multinomial(probs, num_samples=1) 采样 -> [1] 1:batch_num
|
|
|
|
if next_tokens == eos_token_id 推理结束退出循环
|
|
|
|
input_ids = torch.cat([input_ids, next_tokens) -> [1, 7] 1:batch_num
|
|
|
|
response = tokenizer.decode(outputs)
|
|
|
|
## RMSNorm
|
|
|
|
hidden_states -> [6, 1, 4096] 4096:hidden_size
|
|
variance = hidden_states.pow(2).mean(-1, keepdim=True) -> [6, 1, 1]
|
|
hidden_states = hidden_states * torch.rsqrt(variance + self.eps) 平方根倒数
|
|
self.weight -> [4096]
|
|
return (self.weight * hidden_states) -> [6, 1, 4096] |