162 lines
5.3 KiB
Markdown
162 lines
5.3 KiB
Markdown
|
|
|
|
|
|
## data flow
|
|
|
|
```
|
|
|
|
query
|
|
|
|
|
tokenizer -> input_ids
|
|
|
|
|
rotary_pos_emb embedding -> [1, 6, 4096]
|
|
\ /
|
|
GLMBlock x 28 -> [6, 1, 4096] <━━━┓
|
|
┃
|
|
RMSNorm -> [6, 1, 4096] ┃
|
|
┃
|
|
[-1:] -> [1, 1, 4096] ┃
|
|
┃
|
|
Linear -> [1, 1, 65024] ┃
|
|
┃
|
|
softmax -> [1, 65024] ┃
|
|
┃
|
|
multinomial -> [1] ┃
|
|
┃
|
|
cat([input_ids, next_tokens]) ━━━┛
|
|
|
|
|
|
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) # final_layernorm -> [6, 1, 4096]
|
|
hidden_states = hidden_states[-1:] 截取最后一个sequence -> [1, 1, 4096]
|
|
lm_logits = Linear(hidden_states) -> [1, 1, 65024]
|
|
lm_logits = lm_logits.transpose(0, 1).contiguous() -> [1, 1, 65024]
|
|
|
|
probs = softmax(lm_logits) -> [1, 65024] {Softmax}(x_{i}) = \frac{\exp(x_i)}{\sum_j \exp(x_j)}
|
|
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]
|
|
/ \
|
|
| pow(2) -> [6, 1, 4096]
|
|
| |
|
|
| mean -> [6, 1, 1]
|
|
| ↓
|
|
| rsqrt( + eps) -> [6, 1, 1]
|
|
\ /
|
|
mul -> [6, 1, 4096]
|
|
\ weight -> [4096]
|
|
\ /
|
|
mul -> [6, 1, 4096]
|
|
|
|
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]
|
|
```
|
|
|
|
## MLP
|
|
|
|
```
|
|
hidden_states -> [6, 1, 4096]
|
|
Linear -> [6, 1, 27392]
|
|
/ \
|
|
chunk1 chunk0 -> [6, 1, 13696]
|
|
| | \
|
|
| | sigmoid
|
|
| | /
|
|
| mul
|
|
\ /
|
|
mul -> [6, 1, 13696]
|
|
Linear -> [6, 1, 4096]
|
|
|
|
Linear(hidden_states) no bias -> [6, 1, 27392]
|
|
silu (x) = [6, 1, 13696] * sigmoid([6, 1, 13696])
|
|
Linear(intermediate_parallel) no bias -> [6, 1, 4096]
|
|
```
|
|
|
|
## self_attention
|
|
```
|
|
x -> [6, 1, 4096]
|
|
|
|
|
Linear -> [6, 1, 4608]
|
|
/ | \
|
|
[6, 1, 32, 128] <- q k v
|
|
/ | \
|
|
pos_emb pos_emb \
|
|
| | |
|
|
| expand expand -> [6, 1, 32, 128]
|
|
\ / |
|
|
┏---- dot |
|
|
┃ softmax /
|
|
attention┃ \ /
|
|
┗---- dot -> [1, 32, 6, 128] -> [6, 1, 4096]
|
|
Linear -> [6, 1, 4096]
|
|
|
|
hidden_states: [s, b, h]
|
|
mixed_x_layer = Linear(hidden_states) -> [6, 1, 4608] 4608:4096+256+256
|
|
|
|
(query_layer, key_layer, value_layer) = mixed_x_layer.split -> [6, 1, 4096], [6, 1, 256], [6, 1, 256]
|
|
query_layer = query_layer.view -> [6, 1, 32, 128]
|
|
key_layer = key_layer.view -> [6, 1, 2, 128]
|
|
value_layer = value_layer.view -> [6, 1, 2, 128]
|
|
|
|
query_layer = self.apply_rotary_pos_emb(query_layer, rotary_pos_emb)
|
|
key_layer = self.apply_rotary_pos_emb(key_layer, rotary_pos_emb)
|
|
|
|
key_layer = key_layer.unsqueeze(-2) -> [6, 1, 2, 1, 128]
|
|
key_layer = key_layer.expand -> [6, 1, 2, 16, 128]
|
|
key_layer = key_layer.contiguous().view -> [6, 1, 32, 128]
|
|
|
|
value_layer = value_layer.unsqueeze(-2) -> [6, 1, 2, 1, 128]
|
|
value_layer = value_layer.expand -> [6, 1, 2, 16, 128]
|
|
value_layer = value_layer.contiguous().view -> [6, 1, 32, 128]
|
|
|
|
query_layer permute(1, 2, 0, 3) -> [1, 32, 6, 128]
|
|
key_layer permute(1, 2, 0, 3) -> [1, 32, 6, 128]
|
|
value_layer permute(1, 2, 0, 3) -> [1, 32, 6, 128]
|
|
context_layer = scaled_dot_product_attention(query_layer, key_layer, value_layer) -> [1, 32, 6, 128]
|
|
att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
|
|
att = F.softmax(att, dim=-1)
|
|
y = att @ v -> (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
|
|
context_layer = context_layer.permute(2, 0, 1, 3).reshape() -> [6, 1, 4096]
|
|
|
|
return Linear(context_layer) -> [6, 1, 4096]
|
|
```
|
|
|
|
## GLMBlock
|
|
```
|
|
input
|
|
| \
|
|
| RMSNorm
|
|
| self_attention
|
|
| dropout
|
|
| /
|
|
Add
|
|
| \
|
|
| RMSNorm
|
|
| MLP
|
|
| dropout
|
|
| /
|
|
Add
|
|
```
|
|
所有的输出shape都是[6, 1, 4096], 6:sequence_length 1:batch_num 4096:hidden_size |