jeans/ollama
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

model: document qwen2 forward pass

Browse Source
This commit is contained in:
Bruce MacDonald 2025-02-14 14:55:30 -08:00
parent 9dc1fb8a91
commit f93bd92027
1 changed files with 91 additions and 65 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

156
model/models/qwen2/model.go
View File

@ -10,10 +10,15 @@ import (
) )
type Options struct { type Options struct {
RopeFactors ml.Tensor `gguf:"rope_freqs.weight"` RopeFactors ml.Tensor `gguf:"rope_freqs.weight"`
ctxLen, hiddenSize, numHeads, numKVHeads int contextLength int
eps, ropeBase, ropeScale float32 hiddenSize int
ropeDim uint32 numAttnHeads int
numKVHeads int
modelEpsilon float32
ropeBaseFreq float32
ropeFreqScale float32
ropeDimensions uint32
} }
type Model struct { type Model struct {
@ -42,14 +47,14 @@ func New(c ml.Config) (model.Model, error) {
), ),
Layers: make([]Layer, c.Uint("block_count")), Layers: make([]Layer, c.Uint("block_count")),
Options: &Options{ Options: &Options{
hiddenSize: int(c.Uint("embedding_length")), hiddenSize: int(c.Uint("embedding_length")),
numHeads: int(c.Uint("attention.head_count")), numAttnHeads: int(c.Uint("attention.head_count")),
numKVHeads: int(c.Uint("attention.head_count_kv")), numKVHeads: int(c.Uint("attention.head_count_kv")),
eps: c.Float("attention.layer_norm_rms_epsilon"), modelEpsilon: c.Float("attention.layer_norm_rms_epsilon"),
ctxLen: int(c.Uint("context_length")), contextLength: int(c.Uint("context_length")),
ropeBase: c.Float("rope.freq_base"), ropeBaseFreq: c.Float("rope.freq_base"),
ropeScale: c.Float("rope.freq_scale", 1), ropeFreqScale: c.Float("rope.freq_scale", 1),
ropeDim: c.Uint("rope.dimension_count", 64), ropeDimensions: c.Uint("rope.dimension_count", 64),
}, },
} }
@ -58,21 +63,24 @@ func New(c ml.Config) (model.Model, error) {
return m, nil return m, nil
} }
// Shift applies rotary position embeddings to the key tensor for causal attention caching
func (m *Model) Shift(ctx ml.Context, layer int, key, shift ml.Tensor) (ml.Tensor, error) { func (m *Model) Shift(ctx ml.Context, layer int, key, shift ml.Tensor) (ml.Tensor, error) {
return key.RoPE( return key.RoPE(
ctx, ctx,
ml.RopeConfig{ ml.RopeConfig{
PositionIDs: shift, PositionIDs: shift,
RopeFactors: m.Options.RopeFactors, RopeFactors: m.Options.RopeFactors,
RopeDim: m.Options.ropeDim, RopeDim: m.Options.ropeDimensions,
RopeType: ml.RopeTypeNeoX, RopeType: ml.RopeTypeNeoX,
OrigCtxLen: m.Options.ctxLen, OrigCtxLen: m.Options.contextLength,
RopeBase: m.Options.ropeBase, RopeBase: m.Options.ropeBaseFreq,
RopeScale: m.Options.ropeScale, RopeScale: m.Options.ropeFreqScale,
}, },
), nil ), nil
} }
// SelfAttention implements the multi-head self-attention mechanism
// with separate projections for query, key, value and output transformations
type SelfAttention struct { type SelfAttention struct {
Query *nn.Linear `gguf:"attn_q"` Query *nn.Linear `gguf:"attn_q"`
Key *nn.Linear `gguf:"attn_k"` Key *nn.Linear `gguf:"attn_k"`
@ -81,49 +89,59 @@ type SelfAttention struct {
} }
func (sa *SelfAttention) Forward(ctx ml.Context, hiddenState, inputPositions ml.Tensor, cache kvcache.Cache, opts *Options) ml.Tensor { func (sa *SelfAttention) Forward(ctx ml.Context, hiddenState, inputPositions ml.Tensor, cache kvcache.Cache, opts *Options) ml.Tensor {
// Initialize dimensions and configuration
batchSize := hiddenState.Dim(1) batchSize := hiddenState.Dim(1)
headDim := opts.hiddenSize / opts.numHeads headDimension := opts.hiddenSize / opts.numAttnHeads
rc := ml.RopeConfig{ ropeConfig := ml.RopeConfig{
PositionIDs: inputPositions, PositionIDs: inputPositions,
RopeFactors: nil, RopeFactors: nil,
RopeDim: opts.ropeDim, RopeDim: opts.ropeDimensions,
RopeType: ml.RopeTypeNeoX, RopeType: ml.RopeTypeNeoX,
OrigCtxLen: opts.ctxLen, OrigCtxLen: opts.contextLength,
RopeBase: opts.ropeBase, RopeBase: opts.ropeBaseFreq,
RopeScale: opts.ropeScale, RopeScale: opts.ropeFreqScale,
} }
q := sa.Query.Forward(ctx, hiddenState) // Project and reshape query states with rotary embeddings
queryStates := sa.Query.Forward(ctx, hiddenState)
queryStates = queryStates.Reshape(ctx, headDimension, opts.numAttnHeads, batchSize)
queryStates = queryStates.RoPE(ctx, ropeConfig)
q = q.Reshape(ctx, headDim, opts.numHeads, batchSize) // Project and reshape key states with rotary embeddings
q = q.RoPE(ctx, rc) keyStates := sa.Key.Forward(ctx, hiddenState)
keyStates = keyStates.Reshape(ctx, headDimension, opts.numKVHeads, batchSize)
keyStates = keyStates.RoPE(ctx, ropeConfig)
k := sa.Key.Forward(ctx, hiddenState) // Project and reshape value states
k = k.Reshape(ctx, headDim, opts.numKVHeads, batchSize) valueStates := sa.Value.Forward(ctx, hiddenState)
k = k.RoPE(ctx, rc) valueStates = valueStates.Reshape(ctx, headDimension, opts.numKVHeads, batchSize)
v := sa.Value.Forward(ctx, hiddenState) // Update and retrieve from KV cache
v = v.Reshape(ctx, headDim, opts.numKVHeads, batchSize) cache.Put(ctx, keyStates, valueStates)
keyStates, valueStates, attentionMask := cache.Get(ctx)
cache.Put(ctx, k, v) // Prepare tensors for attention computation
k, v, mask := cache.Get(ctx) queryStates = queryStates.Permute(ctx, 0, 2, 1, 3).Contiguous(ctx)
keyStates = keyStates.Permute(ctx, 0, 2, 1, 3).Contiguous(ctx)
valueStates = valueStates.Permute(ctx, 1, 2, 0, 3).Contiguous(ctx)
q = q.Permute(ctx, 0, 2, 1, 3).Contiguous(ctx) // Apply scaling and attention mask to scores
k = k.Permute(ctx, 0, 2, 1, 3).Contiguous(ctx) attentionScores := keyStates.MulmatFullPrec(ctx, queryStates)
v = v.Permute(ctx, 1, 2, 0, 3).Contiguous(ctx) attentionScores = attentionScores.Scale(ctx, 1.0/math.Sqrt(float64(headDimension)))
attentionScores = attentionScores.Add(ctx, attentionMask)
// Compute scaled dot-product attention
attentionProbs := attentionScores.Softmax(ctx)
kq := k.MulmatFullPrec(ctx, q) // Apply attention weights and reshape
kq = kq.Scale(ctx, 1.0/math.Sqrt(float64(headDim))) weightedStates := valueStates.Mulmat(ctx, attentionProbs)
kq = kq.Add(ctx, mask) weightedStates = weightedStates.Permute(ctx, 0, 2, 1, 3).Contiguous(ctx)
kq = kq.Softmax(ctx) weightedStates = weightedStates.Reshape(ctx, opts.hiddenSize, batchSize)
kqv := v.Mulmat(ctx, kq) // Project to output dimension
kqv = kqv.Permute(ctx, 0, 2, 1, 3).Contiguous(ctx) return sa.Output.Forward(ctx, weightedStates)
kqv = kqv.Reshape(ctx, opts.hiddenSize, batchSize)
return sa.Output.Forward(ctx, kqv)
} }
// MLP implements the feed-forward network component with SwiGLU activation
type MLP struct { type MLP struct {
Up *nn.Linear `gguf:"ffn_up"` Up *nn.Linear `gguf:"ffn_up"`
Down *nn.Linear `gguf:"ffn_down"` Down *nn.Linear `gguf:"ffn_down"`
@ -131,10 +149,16 @@ type MLP struct {
} }
func (mlp *MLP) Forward(ctx ml.Context, hiddenState ml.Tensor, opts *Options) ml.Tensor { func (mlp *MLP) Forward(ctx ml.Context, hiddenState ml.Tensor, opts *Options) ml.Tensor {
hiddenState = mlp.Gate.Forward(ctx, hiddenState).SILU(ctx).Mul(ctx, mlp.Up.Forward(ctx, hiddenState)) // Apply SwiGLU activation gating
return mlp.Down.Forward(ctx, hiddenState) gateActivation := mlp.Gate.Forward(ctx, hiddenState).SILU(ctx)
upProjection := mlp.Up.Forward(ctx, hiddenState)
intermediateStates := gateActivation.Mul(ctx, upProjection)
// Project back to hidden dimension
return mlp.Down.Forward(ctx, intermediateStates)
} }
// Layer represents a single transformer layer combining self-attention and feed-forward components
type Layer struct { type Layer struct {
AttentionNorm *nn.RMSNorm `gguf:"attn_norm"` AttentionNorm *nn.RMSNorm `gguf:"attn_norm"`
SelfAttention *SelfAttention SelfAttention *SelfAttention
@ -143,52 +167,54 @@ type Layer struct {
} }
func (l *Layer) Forward(ctx ml.Context, hiddenState, positionIDs ml.Tensor, cache kvcache.Cache, opts *Options) ml.Tensor { func (l *Layer) Forward(ctx ml.Context, hiddenState, positionIDs ml.Tensor, cache kvcache.Cache, opts *Options) ml.Tensor {
// Self-attention branch with residual connection
residual := hiddenState residual := hiddenState
hiddenState = l.AttentionNorm.Forward(ctx, hiddenState, opts.eps) normalizedAttention := l.AttentionNorm.Forward(ctx, hiddenState, opts.modelEpsilon)
attentionOutput := l.SelfAttention.Forward(ctx, normalizedAttention, positionIDs, cache, opts)
hiddenState = attentionOutput.Add(ctx, residual)
hiddenState = l.SelfAttention.Forward(ctx, hiddenState, positionIDs, cache, opts) // Feed-forward branch with residual connection
hiddenState = hiddenState.Add(ctx, residual)
residual = hiddenState residual = hiddenState
normalizedMLP := l.MLPNorm.Forward(ctx, hiddenState, opts.modelEpsilon)
hiddenState = l.MLPNorm.Forward(ctx, hiddenState, opts.eps) mlpOutput := l.MLP.Forward(ctx, normalizedMLP, opts)
output := mlpOutput.Add(ctx, residual)
hiddenState = l.MLP.Forward(ctx, hiddenState, opts)
output := hiddenState.Add(ctx, residual)
return output return output
} }
func (m *Model) Forward(ctx ml.Context, opts model.Options) (ml.Tensor, error) { func (m *Model) Forward(ctx ml.Context, opts model.Options) (ml.Tensor, error) {
inputs, err := ctx.FromIntSlice(opts.Inputs, len(opts.Inputs)) // Convert input tokens and positions to tensors
inputTensor, err := ctx.FromIntSlice(opts.Inputs, len(opts.Inputs))
if err != nil { if err != nil {
return nil, err return nil, err
} }
positions, err := ctx.FromIntSlice(opts.Positions, len(opts.Positions)) positionsTensor, err := ctx.FromIntSlice(opts.Positions, len(opts.Positions))
if err != nil { if err != nil {
return nil, err return nil, err
} }
hiddenState := m.TokenEmbedding.Forward(ctx, inputs) // Initial token embedding
hiddenStates := m.TokenEmbedding.Forward(ctx, inputTensor)
// Process through transformer layers
for i, layer := range m.Layers { for i, layer := range m.Layers {
m.Cache.SetLayer(i) m.Cache.SetLayer(i)
hiddenState = layer.Forward(ctx, hiddenState, positions, m.Cache, m.Options) hiddenStates = layer.Forward(ctx, hiddenStates, positionsTensor, m.Cache, m.Options)
} }
hiddenState = m.OutputNorm.Forward(ctx, hiddenState, m.eps) // Final layer normalization and output projection
normalizedOutput := m.OutputNorm.Forward(ctx, hiddenStates, m.modelEpsilon)
logits := m.Output.Forward(ctx, normalizedOutput)
hiddenState = m.Output.Forward(ctx, hiddenState) // Extract requested output token positions
outputsTensor, err := ctx.FromIntSlice(opts.Outputs, len(opts.Outputs))
outputs, err := ctx.FromIntSlice(opts.Outputs, len(opts.Outputs))
if err != nil { if err != nil {
return nil, err return nil, err
} }
return hiddenState.Rows(ctx, outputs), nil return logits.Rows(ctx, outputsTensor), nil
} }
func init() { func init() {