mirror of
https://github.com/dogkeeper886/ollama37.git
synced 2025-12-12 00:37:04 +00:00
3c14461d5d
For some multimodal models (such as gemma3), we create a single graph that generates the image embedding and then use this in the text model. The embedding tensor is completely opaque to the runner. However, this doesn't work if we need to use the embedding in multiple batches. This can arise if the embedding is larger than the batch size. In these cases (as with llama4), we would like to create views that are more appropriately sized. However, if we do this then the original source tensor is used in multiple graphs, which isn't allowed. To avoid that problem, models with this pattern compute the embedding tensor on first use and recreate the individual views. There is no longer a single vision and text graph. This codifies the pattern of separating vision and text graphs. The logic of computing tensors on demand is moved to the runner, so models no longer have to worry about this. It also gives the runner visibility into the multimodal tensors, which is important for memory management.
255 lines
9.0 KiB
Go
255 lines
9.0 KiB
Go
package llama4
|
|
|
|
import (
|
|
"cmp"
|
|
"math"
|
|
|
|
"github.com/ollama/ollama/fs"
|
|
"github.com/ollama/ollama/kvcache"
|
|
"github.com/ollama/ollama/ml"
|
|
"github.com/ollama/ollama/ml/nn"
|
|
"github.com/ollama/ollama/model/input"
|
|
)
|
|
|
|
type TextAttention struct {
|
|
Query *nn.Linear `gguf:"attn_q"`
|
|
Key *nn.Linear `gguf:"attn_k"`
|
|
Value *nn.Linear `gguf:"attn_v"`
|
|
Output *nn.Linear `gguf:"attn_output"`
|
|
RopeFactors ml.Tensor `gguf:"rope_factors"`
|
|
}
|
|
|
|
func (sa *TextAttention) Forward(ctx ml.Context, hiddenStates, positions, attentionScales ml.Tensor, cache kvcache.Cache, useRope bool, opts *TextOptions) ml.Tensor {
|
|
batchSize, headDim := hiddenStates.Dim(1), cmp.Or(opts.headDim, opts.hiddenSize/opts.numHeads)
|
|
|
|
query := sa.Query.Forward(ctx, hiddenStates)
|
|
key := sa.Key.Forward(ctx, hiddenStates)
|
|
value := sa.Value.Forward(ctx, hiddenStates)
|
|
|
|
query = query.Reshape(ctx, headDim, opts.numHeads, batchSize)
|
|
key = key.Reshape(ctx, headDim, opts.numKVHeads, batchSize)
|
|
value = value.Reshape(ctx, headDim, opts.numKVHeads, batchSize)
|
|
|
|
if useRope {
|
|
query = query.RoPE(ctx, positions, sa.RopeFactors, uint32(opts.ropeDim), uint32(0), opts.ropeBase, opts.ropeScale)
|
|
key = key.RoPE(ctx, positions, sa.RopeFactors, uint32(opts.ropeDim), uint32(0), opts.ropeBase, opts.ropeScale)
|
|
}
|
|
|
|
if opts.useQKNorm {
|
|
query = query.RMSNorm(ctx, nil, opts.eps)
|
|
key = key.RMSNorm(ctx, nil, opts.eps)
|
|
}
|
|
|
|
if attentionScales != nil && !useRope {
|
|
query = query.Mul(ctx, attentionScales)
|
|
}
|
|
|
|
attention := nn.Attention(ctx, query, key, value, 1./math.Sqrt(float64(headDim)), cache)
|
|
attention = attention.Reshape(ctx, opts.hiddenSize, batchSize)
|
|
return sa.Output.Forward(ctx, attention)
|
|
}
|
|
|
|
type TextMLP struct {
|
|
Gate *nn.Linear `gguf:"ffn_gate"`
|
|
Up *nn.Linear `gguf:"ffn_up"`
|
|
Down *nn.Linear `gguf:"ffn_down"`
|
|
}
|
|
|
|
func (mlp *TextMLP) Forward(ctx ml.Context, hiddenStates ml.Tensor, opts *TextOptions) ml.Tensor {
|
|
hiddenStates = mlp.Gate.Forward(ctx, hiddenStates).SILU(ctx).Mul(ctx, mlp.Up.Forward(ctx, hiddenStates))
|
|
return mlp.Down.Forward(ctx, hiddenStates)
|
|
}
|
|
|
|
type TextExperts struct {
|
|
Gate ml.Tensor `gguf:"ffn_gate_exps.weight"`
|
|
Up ml.Tensor `gguf:"ffn_up_exps.weight"`
|
|
Down ml.Tensor `gguf:"ffn_down_exps.weight"`
|
|
}
|
|
|
|
func (e *TextExperts) Forward(ctx ml.Context, hiddenStates, routerLogits ml.Tensor, opts *TextOptions) ml.Tensor {
|
|
experts := routerLogits.TopK(ctx, opts.numExpertsUsed)
|
|
scores := routerLogits.Sigmoid(ctx).Reshape(ctx, 1, opts.numExperts, hiddenStates.Dim(1)).Rows(ctx, experts)
|
|
|
|
hiddenStates = hiddenStates.Reshape(ctx, hiddenStates.Dim(0), 1, hiddenStates.Dim(1))
|
|
hiddenStates = hiddenStates.Repeat(ctx, 1, opts.numExpertsUsed)
|
|
hiddenStates = hiddenStates.Mul(ctx, scores)
|
|
|
|
upStates := e.Up.MulmatID(ctx, hiddenStates, experts)
|
|
gateStates := e.Gate.MulmatID(ctx, hiddenStates, experts)
|
|
downStates := e.Down.MulmatID(ctx, upStates.Mul(ctx, gateStates.SILU(ctx)), experts)
|
|
|
|
nextStates := downStates.View(ctx, 0, hiddenStates.Dim(0), downStates.Stride(2), hiddenStates.Dim(2))
|
|
for i := 1; i < opts.numExpertsUsed; i++ {
|
|
nextStates.Add(ctx, downStates.View(ctx, i*downStates.Stride(1), hiddenStates.Dim(0), downStates.Stride(2), hiddenStates.Dim(2)))
|
|
}
|
|
|
|
return nextStates
|
|
}
|
|
|
|
// TextSharedExpert is TextMLP with different tensor names
|
|
type TextSharedExpert struct {
|
|
Gate *nn.Linear `gguf:"ffn_gate_shexp"`
|
|
Up *nn.Linear `gguf:"ffn_up_shexp"`
|
|
Down *nn.Linear `gguf:"ffn_down_shexp"`
|
|
}
|
|
|
|
func (mlp *TextSharedExpert) Forward(ctx ml.Context, hiddenStates ml.Tensor, opts *TextOptions) ml.Tensor {
|
|
hiddenStates = mlp.Gate.Forward(ctx, hiddenStates).SILU(ctx).Mul(ctx, mlp.Up.Forward(ctx, hiddenStates))
|
|
return mlp.Down.Forward(ctx, hiddenStates)
|
|
}
|
|
|
|
type TextMOE struct {
|
|
Router *nn.Linear `gguf:"ffn_gate_inp"`
|
|
Experts *TextExperts
|
|
SharedExpert *TextSharedExpert
|
|
}
|
|
|
|
func (moe *TextMOE) Forward(ctx ml.Context, hiddenStates ml.Tensor, opts *TextOptions) ml.Tensor {
|
|
hiddenDim, sequenceLength, batchSize := hiddenStates.Dim(0), hiddenStates.Dim(1), hiddenStates.Dim(2)
|
|
hiddenStates = hiddenStates.Reshape(ctx, hiddenDim, sequenceLength*batchSize)
|
|
routerLogits := moe.Router.Forward(ctx, hiddenStates)
|
|
|
|
sharedStates := moe.SharedExpert.Forward(ctx, hiddenStates, opts)
|
|
routedStates := moe.Experts.Forward(ctx, hiddenStates, routerLogits, opts)
|
|
return sharedStates.Add(ctx, routedStates)
|
|
}
|
|
|
|
type TextFeedForward interface {
|
|
Forward(ctx ml.Context, hiddenStates ml.Tensor, opts *TextOptions) ml.Tensor
|
|
}
|
|
|
|
type TextLayer struct {
|
|
AttentionNorm *nn.LayerNorm `gguf:"attn_norm"`
|
|
Attention *TextAttention
|
|
|
|
FFNNorm *nn.LayerNorm `gguf:"ffn_norm"`
|
|
FeedForward TextFeedForward
|
|
}
|
|
|
|
func (d *TextLayer) Forward(ctx ml.Context, hiddenStates, positions, attentionScales, outputs ml.Tensor, cache kvcache.Cache, useRope bool, opts *TextOptions) ml.Tensor {
|
|
residual := hiddenStates
|
|
|
|
// self attention
|
|
hiddenStates = d.AttentionNorm.Forward(ctx, hiddenStates, opts.eps)
|
|
hiddenStates = d.Attention.Forward(ctx, hiddenStates, positions, attentionScales, cache, useRope, opts)
|
|
|
|
if outputs != nil {
|
|
hiddenStates = hiddenStates.Rows(ctx, outputs)
|
|
residual = residual.Rows(ctx, outputs)
|
|
}
|
|
|
|
hiddenStates = hiddenStates.Add(ctx, residual)
|
|
residual = hiddenStates
|
|
|
|
hiddenStates = d.FFNNorm.Forward(ctx, hiddenStates, opts.eps)
|
|
hiddenStates = d.FeedForward.Forward(ctx, hiddenStates, opts)
|
|
|
|
return residual.Add(ctx, hiddenStates)
|
|
}
|
|
|
|
type TextOptions struct {
|
|
hiddenSize int
|
|
numHeads, numKVHeads, headDim int
|
|
numExperts, numExpertsUsed int
|
|
ropeDim int
|
|
ropeBase, ropeScale float32
|
|
eps float32
|
|
interleaveLayerStep int
|
|
noRopeInterval int
|
|
useQKNorm bool
|
|
attentionTemperatureTuning bool
|
|
attentionScale float64
|
|
attentionFloorScale float64
|
|
}
|
|
|
|
type TextModel struct {
|
|
Layers []TextLayer `gguf:"blk"`
|
|
|
|
TokenEmbedding *nn.Embedding `gguf:"token_embd"`
|
|
OutputNorm *nn.LayerNorm `gguf:"output_norm"`
|
|
Output *nn.Linear `gguf:"output,alt:token_embd"`
|
|
|
|
*TextOptions
|
|
}
|
|
|
|
func newTextModel(c fs.Config) *TextModel {
|
|
layers := make([]TextLayer, c.Uint("block_count"))
|
|
interleaveLayerStep := c.Uint("interleave_moe_layer_step", 1)
|
|
for i := range layers {
|
|
if (i+1)%int(interleaveLayerStep) == 0 {
|
|
layers[i] = TextLayer{FeedForward: &TextMOE{}}
|
|
} else {
|
|
layers[i] = TextLayer{FeedForward: &TextMLP{}}
|
|
}
|
|
}
|
|
|
|
return &TextModel{
|
|
Layers: layers,
|
|
TextOptions: &TextOptions{
|
|
hiddenSize: int(c.Uint("embedding_length")),
|
|
numHeads: int(c.Uint("attention.head_count")),
|
|
numKVHeads: int(c.Uint("attention.head_count_kv")),
|
|
headDim: int(c.Uint("attention.head_dim", 128)),
|
|
numExperts: int(c.Uint("expert_count")),
|
|
numExpertsUsed: int(c.Uint("expert_used_count")),
|
|
ropeDim: int(c.Uint("rope.dimension_count")),
|
|
ropeBase: c.Float("rope.freq_base"),
|
|
ropeScale: c.Float("rope.freq_scale", 1),
|
|
eps: c.Float("attention.layer_norm_rms_epsilon"),
|
|
interleaveLayerStep: int(c.Uint("interleave_moe_layer_step", 1)),
|
|
noRopeInterval: int(c.Uint("no_rope_interval", 4)),
|
|
useQKNorm: c.Bool("use_qk_norm", true),
|
|
attentionTemperatureTuning: c.Bool("attention.temperature_tuning", true),
|
|
attentionScale: float64(c.Float("attention.scale", 0.1)),
|
|
attentionFloorScale: float64(c.Float("attention.floor_scale", 8192)),
|
|
},
|
|
}
|
|
}
|
|
|
|
func (m *TextModel) Forward(ctx ml.Context, inputs, positions, outputs ml.Tensor, batch input.Batch, cache kvcache.Cache) ml.Tensor {
|
|
hiddenStates := m.TokenEmbedding.Forward(ctx, inputs).Duplicate(ctx)
|
|
|
|
for _, mi := range batch.Multimodal {
|
|
img := mi.Multimodal[0].Tensor
|
|
ctx.Forward(img.Copy(ctx, hiddenStates.View(ctx, mi.Index*hiddenStates.Stride(1), img.Dim(0)*img.Dim(1))))
|
|
}
|
|
|
|
var attentionScales ml.Tensor
|
|
if m.attentionTemperatureTuning {
|
|
scales := make([]float32, len(batch.Positions))
|
|
for i, p := range batch.Positions {
|
|
scales[i] = float32(math.Log(math.Floor(((float64(p)+1.0)/float64(m.attentionFloorScale))+1.0))*m.attentionScale + 1.0)
|
|
}
|
|
|
|
var err error
|
|
attentionScales, err = ctx.Input().FromFloatSlice(scales, 1, 1, len(scales))
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
}
|
|
|
|
for i, layer := range m.Layers {
|
|
cache.SetLayer(i)
|
|
wc := cache.(*kvcache.WrapperCache)
|
|
wc.SetLayerType(1)
|
|
useChunkedAttention := (i+1)%m.noRopeInterval != 0
|
|
if useChunkedAttention {
|
|
wc.SetLayerType(0)
|
|
}
|
|
|
|
var lastLayerOutputs ml.Tensor
|
|
if i == len(m.Layers)-1 {
|
|
lastLayerOutputs = outputs
|
|
}
|
|
|
|
hiddenStates = layer.Forward(ctx, hiddenStates, positions, attentionScales, lastLayerOutputs, cache, useChunkedAttention, m.TextOptions)
|
|
}
|
|
|
|
hiddenStates = m.OutputNorm.Forward(ctx, hiddenStates, m.eps)
|
|
return m.Output.Forward(ctx, hiddenStates)
|
|
}
|
|
|
|
func (m *TextModel) Shift(ctx ml.Context, layer int, key, shift ml.Tensor) (ml.Tensor, error) {
|
|
return key.RoPE(ctx, shift, m.Layers[layer].Attention.RopeFactors, uint32(0), uint32(m.ropeDim), m.ropeBase, m.ropeScale), nil
|
|
}
|