3.6 KiB
3.6 KiB
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