89 lines
3.6 KiB
Markdown
89 lines
3.6 KiB
Markdown
|
## data flow
|
|||
|
|
|
|||
|
|
```
|
|||
|
|
query -> "你好"
|
|||
|
|
┃
|
|||
|
|
tokenizer -> input_ids [6]
|
|||
|
|
┃
|
|||
|
|
rotary_pos_emb embedding -> [1, 6, 4096]
|
|||
|
|
╲ ╱
|
|||
|
|
GLMBlock x 28 -> [6, 1, 4096] <━━━┓
|
|||
|
|
RMSNorm -> [6, 1, 4096] ┃ final_layernorm
|
|||
|
|
[-1:] -> [1, 1, 4096] ┃
|
|||
|
|
Linear -> [1, 1, 65024] ┃ output_layer 4096->65024
|
|||
|
|
softmax -> [1, 65024] ┃
|
|||
|
|
multinomial -> [1] ┃
|
|||
|
|
cat([input_ids, next_tokens]) ━━━┛
|
|||
|
|
↓
|
|||
|
|
tokenizer.decode( )
|
|||
|
|
|
|||
|
|
# GLMBlock
|
|||
|
|
|
|||
|
|
input
|
|||
|
|
╱ ╲
|
|||
|
|
╱ 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]
|
|||
|
|
┃ ┋ ╲ ╱
|
|||
|
|
┃ RMSNorm mul -> [6, 1, 4096]
|
|||
|
|
┃ ╲
|
|||
|
|
┃ SelfAttention x -> [6, 1, 4096]
|
|||
|
|
┃ ┋ ┃
|
|||
|
|
┃ ┋ Linear -> [6, 1, 4608] 4096->4608
|
|||
|
|
┃ ┋ ╱ ┃ ╲
|
|||
|
|
┃ ┋ q k v [6, 1, 32, 128] [6, 1, 2, 128] [6, 1, 2, 128]
|
|||
|
|
┃ ┋ ╱ ┃ ╲
|
|||
|
|
┃ ┋ pos_emb pos_emb ╲ -> cat( x0*y0-x1*y1, x1*y0-x0*y1, x, y)
|
|||
|
|
┃ ┋ ┃ ┃ ┃
|
|||
|
|
┃ ┋ ┃ expand expand -> [6, 1, 32, 128] [6, 1, 32, 128]
|
|||
|
|
┃ ┋ permute permute permute -> [1, 32, 6, 128] [1, 32, 6, 128] [1, 32, 6, 128]
|
|||
|
|
┃ ┋ ╲ ╱ ┃
|
|||
|
|
┃ ┋ ┏---- matmul ┃ -> [1, 32, 6, 128] [1, 32, 128, 6] -> [1, 32, 6, 6]
|
|||
|
|
┃ ┋ ┃ add(mask) ╱ -> [1, 32, 6, 6]
|
|||
|
|
┃ ┋ attention┃ softmax ╱ -> [1, 32, 6, 6] dim:-1
|
|||
|
|
┃ ┋ ┃ ╲ ╱
|
|||
|
|
┃ ┋ ┗---- matmul -> [1, 32, 6, 6] [1, 32, 6, 128] -> [1, 32, 6, 128] -> [6, 1, 4096]
|
|||
|
|
┃ SelfAttention Linear -> [6, 1, 4096] 4096->4096
|
|||
|
|
┃ ╱
|
|||
|
|
┃ dropout
|
|||
|
|
╲ ╱
|
|||
|
|
Add
|
|||
|
|
╱ ╲
|
|||
|
|
┃ 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]
|
|||
|
|
┃ ┋ ╲ ╱
|
|||
|
|
┃ RMSNorm mul -> [6, 1, 4096]
|
|||
|
|
┃ ╱
|
|||
|
|
┃ mlp ╱
|
|||
|
|
┃ ┋ Linear -> [6, 1, 27392] 4096->27392
|
|||
|
|
┃ ┋ ╱ ╲
|
|||
|
|
┃ ┋ chunk1 chunk0 -> [6, 1, 13696]
|
|||
|
|
┃ ┋ ┃ ┃ ╲
|
|||
|
|
┃ ┋ ┃ ┃ sigmoid
|
|||
|
|
┃ ┋ ┃ ┃ ╱
|
|||
|
|
┃ ┋ ┃ mul
|
|||
|
|
┃ ┋ ╲ ╱
|
|||
|
|
┃ ┋ mul -> [6, 1, 13696]
|
|||
|
|
┃ mlp Linear -> [6, 1, 4096] 13696->4096
|
|||
|
|
┃ ╱
|
|||
|
|
┃ dropout
|
|||
|
|
┃ ╱
|
|||
|
|
Add
|
|||
|
|
|
|||
|
|
```
|