mirror of https://gitee.com/namelin2022/ollama
Michael Yang
1 year ago
committed by
Michael Yang
Michael Yang
13 changed files with 833 additions and 15 deletions
@ -0,0 +1,167 @@ |
|||||
|
package convert |
||||
|
|
||||
|
import ( |
||||
|
"slices" |
||||
|
"strings" |
||||
|
|
||||
|
"github.com/pdevine/tensor" |
||||
|
"github.com/pdevine/tensor/native" |
||||
|
|
||||
|
"github.com/ollama/ollama/fs/ggml" |
||||
|
) |
||||
|
|
||||
|
type llama4Model struct { |
||||
|
ModelParameters |
||||
|
TextModel struct { |
||||
|
llamaModel |
||||
|
NumExpertsPerToken uint32 `json:"num_experts_per_tok"` |
||||
|
NumLocalExperts uint32 `json:"num_local_experts"` |
||||
|
InterleaveMOELayerStep uint32 `json:"interleave_moe_layer_step"` |
||||
|
UseQKNorm bool `json:"use_qk_norm"` |
||||
|
IntermediateSizeMLP uint32 `json:"intermediate_size_mlp"` |
||||
|
} `json:"text_config"` |
||||
|
VisionModel struct { |
||||
|
NumHiddenLayers uint32 `json:"num_hidden_layers"` |
||||
|
HiddenSize uint32 `json:"hidden_size"` |
||||
|
IntermediateSize uint32 `json:"intermediate_size"` |
||||
|
NumAttentionHeads uint32 `json:"num_attention_heads"` |
||||
|
ImageSize uint32 `json:"image_size"` |
||||
|
PatchSize uint32 `json:"patch_size"` |
||||
|
RopeTheta float32 `json:"rope_theta"` |
||||
|
NormEpsilon float32 `json:"norm_eps"` |
||||
|
PixelShuffleRatio float32 `json:"pixel_shuffle_ratio"` |
||||
|
} `json:"vision_config"` |
||||
|
} |
||||
|
|
||||
|
// KV implements ModelConverter.
|
||||
|
func (p *llama4Model) KV(t *Tokenizer) ggml.KV { |
||||
|
kv := p.ModelParameters.KV(t) |
||||
|
kv["general.architecture"] = "llama4" |
||||
|
|
||||
|
for k, v := range p.TextModel.KV(t) { |
||||
|
if strings.HasPrefix(k, "llama.") { |
||||
|
kv[strings.ReplaceAll(k, "llama.", "llama4.")] = v |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
kv["llama4.intermediate_size"] = p.TextModel.IntermediateSizeMLP |
||||
|
kv["llama4.intermediate_size_moe"] = p.TextModel.IntermediateSize |
||||
|
|
||||
|
kv["llama4.expert_count"] = p.TextModel.NumLocalExperts |
||||
|
kv["llama4.expert_used_count"] = p.TextModel.NumExpertsPerToken |
||||
|
kv["llama4.interleave_moe_layer_step"] = p.TextModel.InterleaveMOELayerStep |
||||
|
kv["llama4.use_qk_norm"] = p.TextModel.UseQKNorm |
||||
|
|
||||
|
kv["llama4.vision.block_count"] = p.VisionModel.NumHiddenLayers |
||||
|
kv["llama4.vision.embedding_length"] = p.VisionModel.HiddenSize |
||||
|
kv["llama4.vision.feed_forward_length"] = p.VisionModel.IntermediateSize |
||||
|
kv["llama4.vision.attention.head_count"] = p.VisionModel.NumAttentionHeads |
||||
|
kv["llama4.vision.image_size"] = p.VisionModel.ImageSize |
||||
|
kv["llama4.vision.patch_size"] = p.VisionModel.PatchSize |
||||
|
kv["llama4.vision.rope.freq_base"] = p.VisionModel.RopeTheta |
||||
|
kv["llama4.vision.layer_norm_epsilon"] = p.VisionModel.NormEpsilon |
||||
|
kv["llama4.vision.pixel_shuffle_ratio"] = p.VisionModel.PixelShuffleRatio |
||||
|
return kv |
||||
|
} |
||||
|
|
||||
|
// Replacements implements ModelConverter.
|
||||
|
func (p *llama4Model) Replacements() []string { |
||||
|
return append( |
||||
|
p.TextModel.Replacements(), |
||||
|
"language_model.", "", |
||||
|
"vision_model", "v", |
||||
|
"multi_modal_projector", "mm", |
||||
|
"feed_forward.down_proj", "ffn_down", |
||||
|
"feed_forward.up_proj", "ffn_up", |
||||
|
"feed_forward.gate_proj", "ffn_gate", |
||||
|
"feed_forward.", "ffn_", |
||||
|
"shared_expert.down_proj", "down_shexp", |
||||
|
"shared_expert.gate_proj", "gate_shexp", |
||||
|
"shared_expert.up_proj", "up_shexp", |
||||
|
"experts.down_proj", "down_exps.weight", |
||||
|
"experts.gate_up_proj", "gate_up_exps.weight", |
||||
|
"router", "gate_inp", |
||||
|
"patch_embedding.linear", "patch_embedding", |
||||
|
) |
||||
|
} |
||||
|
|
||||
|
// Tensors implements ModelConverter.
|
||||
|
func (p *llama4Model) Tensors(ts []Tensor) []ggml.Tensor { |
||||
|
var out []ggml.Tensor |
||||
|
|
||||
|
var textTensors []Tensor |
||||
|
for _, t := range ts { |
||||
|
if strings.HasPrefix(t.Name(), "v.") || strings.HasPrefix(t.Name(), "mm.") { |
||||
|
out = append(out, ggml.Tensor{ |
||||
|
Name: t.Name(), |
||||
|
Kind: t.Kind(), |
||||
|
Shape: t.Shape(), |
||||
|
WriterTo: t, |
||||
|
}) |
||||
|
} else if strings.Contains(t.Name(), "ffn_gate_up_exps") { |
||||
|
// gate and up projectors are fused
|
||||
|
// dims[1], dims[2] must be swapped
|
||||
|
// [experts, hidden_size, intermediate_size * 2] --> [experts, intermediate_size, hidden_size]
|
||||
|
halfDim := int(t.Shape()[2]) / 2 |
||||
|
|
||||
|
newShape := slices.Clone(t.Shape()) |
||||
|
newShape[1], newShape[2] = newShape[2]/2, newShape[1] |
||||
|
for i, name := range []string{"ffn_gate_exps", "ffn_up_exps"} { |
||||
|
// clone tensor since we need separate repackers
|
||||
|
tt := t.Clone() |
||||
|
tt.SetRepacker(p.repack(nil, nil, tensor.S(i*halfDim, (i+1)*halfDim))) |
||||
|
out = append(out, ggml.Tensor{ |
||||
|
Name: strings.ReplaceAll(tt.Name(), "ffn_gate_up_exps", name), |
||||
|
Kind: tt.Kind(), |
||||
|
Shape: newShape, |
||||
|
WriterTo: tt, |
||||
|
}) |
||||
|
} |
||||
|
} else if strings.Contains(t.Name(), "ffn_down_exps") { |
||||
|
// dims[1], dims[2] must be swapped
|
||||
|
// [experts, intermediate_size, hidden_size] --> [experts, hidden_size, intermediate_size]
|
||||
|
t.SetRepacker(p.repack()) |
||||
|
newShape := slices.Clone(t.Shape()) |
||||
|
newShape[1], newShape[2] = newShape[2], newShape[1] |
||||
|
out = append(out, ggml.Tensor{ |
||||
|
Name: t.Name(), |
||||
|
Kind: t.Kind(), |
||||
|
Shape: newShape, |
||||
|
WriterTo: t, |
||||
|
}) |
||||
|
} else { |
||||
|
textTensors = append(textTensors, t) |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
p.TextModel.skipRepack = true |
||||
|
out = append(out, p.TextModel.Tensors(textTensors)...) |
||||
|
return out |
||||
|
} |
||||
|
|
||||
|
func (p *llama4Model) repack(slice ...tensor.Slice) Repacker { |
||||
|
return func(name string, data []float32, shape []uint64) ([]float32, error) { |
||||
|
dims := make([]int, len(shape)) |
||||
|
for i, dim := range shape { |
||||
|
dims[i] = int(dim) |
||||
|
} |
||||
|
|
||||
|
var t tensor.Tensor = tensor.New(tensor.WithShape(dims...), tensor.WithBacking(data)) |
||||
|
t, err := t.Slice(slice...) |
||||
|
if err != nil { |
||||
|
return nil, err |
||||
|
} |
||||
|
|
||||
|
if err := t.T(0, 2, 1); err != nil { |
||||
|
return nil, err |
||||
|
} |
||||
|
|
||||
|
t = tensor.Materialize(t) |
||||
|
// flatten tensor so it can be return as a vector
|
||||
|
if err := t.Reshape(t.Shape().TotalSize()); err != nil { |
||||
|
return nil, err |
||||
|
} |
||||
|
|
||||
|
return native.VectorF32(t.(*tensor.Dense)) |
||||
|
} |
||||
|
} |
||||
@ -0,0 +1,100 @@ |
|||||
|
package llama4 |
||||
|
|
||||
|
import ( |
||||
|
"bytes" |
||||
|
"image" |
||||
|
|
||||
|
"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" |
||||
|
"github.com/ollama/ollama/model/input" |
||||
|
) |
||||
|
|
||||
|
type Model struct { |
||||
|
model.Base |
||||
|
model.BytePairEncoding |
||||
|
|
||||
|
*VisionModel `gguf:"v,vision"` |
||||
|
*Projector `gguf:"mm"` |
||||
|
*TextModel |
||||
|
} |
||||
|
|
||||
|
type Projector struct { |
||||
|
Linear1 *nn.Linear `gguf:"linear_1"` |
||||
|
} |
||||
|
|
||||
|
func (p *Projector) Forward(ctx ml.Context, visionOutputs ml.Tensor) ml.Tensor { |
||||
|
return p.Linear1.Forward(ctx, visionOutputs) |
||||
|
} |
||||
|
|
||||
|
func New(c fs.Config) (model.Model, error) { |
||||
|
m := Model{ |
||||
|
BytePairEncoding: model.NewBytePairEncoding( |
||||
|
c.String("tokenizer.ggml.pretokenizer", `(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+`), |
||||
|
&model.Vocabulary{ |
||||
|
Values: c.Strings("tokenizer.ggml.tokens"), |
||||
|
Types: c.Uints("tokenizer.ggml.token_type"), |
||||
|
Merges: c.Strings("tokenizer.ggml.merges"), |
||||
|
BOS: int32(c.Uint("tokenizer.ggml.bos_token_id")), |
||||
|
AddBOS: c.Bool("tokenizer.ggml.add_bos_token", true), |
||||
|
EOS: int32(c.Uint("tokenizer.ggml.eos_token_id")), |
||||
|
AddEOS: c.Bool("tokenizer.ggml.add_eos_token", false), |
||||
|
}, |
||||
|
), |
||||
|
VisionModel: newVisionModel(c), |
||||
|
TextModel: newTextModel(c), |
||||
|
} |
||||
|
|
||||
|
m.Cache = kvcache.NewWrapperCache( |
||||
|
// TODO: pretend this is chunked attention for now
|
||||
|
kvcache.NewSWACache(8192, m.Shift), |
||||
|
kvcache.NewCausalCache(m.Shift), |
||||
|
) |
||||
|
|
||||
|
return &m, nil |
||||
|
} |
||||
|
|
||||
|
func (m *Model) EncodeMultimodal(ctx ml.Context, multimodalData []byte) (any, error) { |
||||
|
if len(m.VisionModel.Layers) < 1 { |
||||
|
return nil, model.ErrNoVisionModel |
||||
|
} |
||||
|
|
||||
|
img, _, err := image.Decode(bytes.NewReader(multimodalData)) |
||||
|
if err != nil { |
||||
|
return nil, err |
||||
|
} |
||||
|
|
||||
|
f32s, aspectRatio, err := m.ProcessImage(ctx, img) |
||||
|
if err != nil { |
||||
|
return nil, err |
||||
|
} |
||||
|
|
||||
|
pixelValues, err := ctx.Input().FromFloatSlice(f32s, len(f32s)) |
||||
|
if err != nil { |
||||
|
return nil, err |
||||
|
} |
||||
|
|
||||
|
visionOutputs := m.VisionModel.Forward(ctx, pixelValues) |
||||
|
visionOutputs = visionOutputs.Reshape(ctx, visionOutputs.Dim(0), visionOutputs.Dim(1)*visionOutputs.Dim(2)*visionOutputs.Dim(3)) |
||||
|
return m.Projector.Forward(ctx, visionOutputs), nil |
||||
|
} |
||||
|
|
||||
|
func (m *Model) Forward(ctx ml.Context, batch input.Batch) (ml.Tensor, error) { |
||||
|
positions, err := ctx.Input().FromIntSlice(batch.Positions, len(batch.Positions)) |
||||
|
if err != nil { |
||||
|
return nil, err |
||||
|
} |
||||
|
|
||||
|
outputs, err := ctx.Input().FromIntSlice(batch.Outputs, len(batch.Outputs)) |
||||
|
if err != nil { |
||||
|
return nil, err |
||||
|
} |
||||
|
|
||||
|
return m.TextModel.Forward(ctx, batch.Inputs, positions, outputs, batch, m.Cache), nil |
||||
|
} |
||||
|
|
||||
|
func init() { |
||||
|
model.Register("llama4", New) |
||||
|
} |
||||
@ -0,0 +1,223 @@ |
|||||
|
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 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) |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
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 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, 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, 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 |
||||
|
useQKNorm bool |
||||
|
} |
||||
|
|
||||
|
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)), |
||||
|
useQKNorm: c.Bool("use_qk_norm", true), |
||||
|
}, |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
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) |
||||
|
|
||||
|
for i, layer := range m.Layers { |
||||
|
cache.SetLayer(i) |
||||
|
wc := cache.(*kvcache.WrapperCache) |
||||
|
wc.SetLayerType(1) |
||||
|
useChunkedAttention := (i+1)%4 != 0 |
||||
|
if useChunkedAttention { |
||||
|
wc.SetLayerType(0) |
||||
|
} |
||||
|
|
||||
|
var lastLayerOutputs ml.Tensor |
||||
|
if i == len(m.Layers)-1 { |
||||
|
lastLayerOutputs = outputs |
||||
|
} |
||||
|
|
||||
|
hiddenStates = layer.Forward(ctx, hiddenStates, positions, 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 |
||||
|
} |
||||
@ -0,0 +1,256 @@ |
|||||
|
package llama4 |
||||
|
|
||||
|
import ( |
||||
|
"math" |
||||
|
|
||||
|
"github.com/ollama/ollama/fs" |
||||
|
"github.com/ollama/ollama/ml" |
||||
|
"github.com/ollama/ollama/ml/nn" |
||||
|
) |
||||
|
|
||||
|
type VisionAttention 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"` |
||||
|
} |
||||
|
|
||||
|
// applyVisionRotaryEmbedding applies 2D rotary embedding to the input tensor.
|
||||
|
// This is equivalent to the Pytorch implmentation using half rotations:
|
||||
|
//
|
||||
|
// cos, sin = torch.cos(freqs), torch.sin(freqs)
|
||||
|
// cos = cos.unsqueeze(-1)
|
||||
|
// sin = sin.unsqueeze(-1)
|
||||
|
// t = t.reshape(*t.shape[:-1], -1, 2)
|
||||
|
// t_out = (t * cos) + (_rotate_half(t) * sin)
|
||||
|
// t_out = t_out.flatten(3)
|
||||
|
//
|
||||
|
// Which is equivalent to the Pytorch implementation using complex numbers:
|
||||
|
//
|
||||
|
// t_ = torch.view_as_complex(t.float().reshape(*t.shape[:-1], -1, 2))
|
||||
|
// freqs_ci = reshape_for_broadcast(freqs_ci=freq_cis, t=t_) # freqs_ci[:,:,None,:]
|
||||
|
// freqs_ci = freqs_ci.to(t_.device)
|
||||
|
// t_out = torch.view_as_real(t_ * freqs_ci).flatten(3)
|
||||
|
//
|
||||
|
// Due to the 1) the dimensional and 2) the datatype limitations of current backends,
|
||||
|
// we need to use a different approach to achieve the same result.
|
||||
|
func applyVisionRotaryEmbedding(ctx ml.Context, t, cos, sin ml.Tensor) ml.Tensor { |
||||
|
width, height, channels, tiles := t.Dim(0), t.Dim(1), t.Dim(2), t.Dim(3) |
||||
|
|
||||
|
t = t.Reshape(ctx, 2, t.Dim(0)/2, t.Dim(1)*t.Dim(2)*t.Dim(3)) |
||||
|
|
||||
|
// t1 = t[..., 0::2]
|
||||
|
t1 := t.View(ctx, 0, 1, t.Stride(1), t.Dim(1), t.Stride(2), t.Dim(2)).Contiguous(ctx) |
||||
|
t1 = t1.Reshape(ctx, width/2, height, channels, tiles) |
||||
|
|
||||
|
// t2 = t[..., 1::2]
|
||||
|
t2 := t.View(ctx, t.Stride(0), 1, t.Stride(1), t.Dim(1), t.Stride(2), t.Dim(2)).Contiguous(ctx) |
||||
|
t2 = t2.Reshape(ctx, width/2, height, channels, tiles) |
||||
|
|
||||
|
// cos_out = torch.stack((t1 * cos, t2 * cos), dim=-1)
|
||||
|
cosOut := t1.Mul(ctx, cos).Concat(ctx, t2.Mul(ctx, cos), 0) |
||||
|
cosOut = cosOut.Reshape(ctx, cosOut.Dim(0)/2, 2, cosOut.Dim(1)*cosOut.Dim(2)*cosOut.Dim(3)) |
||||
|
cosOut = cosOut.Permute(ctx, 1, 0, 2, 3).Contiguous(ctx) |
||||
|
cosOut = cosOut.Reshape(ctx, width, height, channels, tiles) |
||||
|
|
||||
|
// sin_out = torch.stack((-t2 * sin, t1 * sin), dim=-1)
|
||||
|
sinOut := t2.Neg(ctx).Mul(ctx, sin).Concat(ctx, t1.Mul(ctx, sin), 0) |
||||
|
sinOut = sinOut.Reshape(ctx, sinOut.Dim(0)/2, 2, sinOut.Dim(1)*sinOut.Dim(2)*sinOut.Dim(3)) |
||||
|
sinOut = sinOut.Permute(ctx, 1, 0, 2, 3).Contiguous(ctx) |
||||
|
sinOut = sinOut.Reshape(ctx, width, height, channels, tiles) |
||||
|
|
||||
|
return cosOut.Add(ctx, sinOut) |
||||
|
} |
||||
|
|
||||
|
func (sa *VisionAttention) Forward(ctx ml.Context, hiddenState, cos, sin ml.Tensor, opts *VisionOptions) ml.Tensor { |
||||
|
headDim := opts.hiddenSize / opts.numHeads |
||||
|
|
||||
|
query := sa.Query.Forward(ctx, hiddenState) |
||||
|
key := sa.Key.Forward(ctx, hiddenState) |
||||
|
value := sa.Value.Forward(ctx, hiddenState) |
||||
|
|
||||
|
query = query.Reshape(ctx, headDim, opts.numHeads, query.Dim(1), query.Dim(2)) |
||||
|
key = key.Reshape(ctx, headDim, opts.numHeads, key.Dim(1), key.Dim(2)) |
||||
|
value = value.Reshape(ctx, headDim, opts.numHeads, value.Dim(1), value.Dim(2)) |
||||
|
|
||||
|
query = applyVisionRotaryEmbedding(ctx, query, cos, sin) |
||||
|
key = applyVisionRotaryEmbedding(ctx, key, cos, sin) |
||||
|
|
||||
|
attention := nn.Attention(ctx, query, key, value, 1./math.Sqrt(float64(headDim)), nil) |
||||
|
attention = attention.Reshape(ctx, opts.hiddenSize, attention.Dim(2), attention.Dim(3)) |
||||
|
return sa.Output.Forward(ctx, attention) |
||||
|
} |
||||
|
|
||||
|
type VisionMLP struct { |
||||
|
FC1 *nn.Linear `gguf:"fc1"` |
||||
|
FC2 *nn.Linear `gguf:"fc2"` |
||||
|
} |
||||
|
|
||||
|
func (mlp *VisionMLP) Forward(ctx ml.Context, hiddenStates ml.Tensor, opts *VisionOptions) ml.Tensor { |
||||
|
hiddenStates = mlp.FC1.Forward(ctx, hiddenStates).GELU(ctx) |
||||
|
hiddenStates = mlp.FC2.Forward(ctx, hiddenStates) |
||||
|
return hiddenStates |
||||
|
} |
||||
|
|
||||
|
type VisionLayer struct { |
||||
|
InputLayerNorm *nn.LayerNorm `gguf:"attn_norm"` |
||||
|
*VisionAttention |
||||
|
|
||||
|
PostAttentionNorm *nn.LayerNorm `gguf:"ffn_norm"` |
||||
|
*VisionMLP `gguf:"mlp"` |
||||
|
} |
||||
|
|
||||
|
func (e *VisionLayer) Forward(ctx ml.Context, hiddenStates, cos, sin ml.Tensor, opts *VisionOptions) ml.Tensor { |
||||
|
residual := hiddenStates |
||||
|
|
||||
|
// self attention
|
||||
|
hiddenStates = e.InputLayerNorm.Forward(ctx, hiddenStates, opts.eps) |
||||
|
hiddenStates = e.VisionAttention.Forward(ctx, hiddenStates, cos, sin, opts) |
||||
|
hiddenStates = hiddenStates.Add(ctx, residual) |
||||
|
|
||||
|
// MLP
|
||||
|
residual = hiddenStates |
||||
|
hiddenStates = e.PostAttentionNorm.Forward(ctx, hiddenStates, opts.eps) |
||||
|
hiddenStates = e.VisionMLP.Forward(ctx, hiddenStates, opts) |
||||
|
hiddenStates = hiddenStates.Add(ctx, residual) |
||||
|
|
||||
|
return hiddenStates |
||||
|
} |
||||
|
|
||||
|
type VisionAdapter struct { |
||||
|
FC1 *nn.Linear `gguf:"mlp.fc1"` |
||||
|
FC2 *nn.Linear `gguf:"mlp.fc2"` |
||||
|
} |
||||
|
|
||||
|
func (a *VisionAdapter) Forward(ctx ml.Context, hiddenStates ml.Tensor, opts *VisionOptions) ml.Tensor { |
||||
|
patches := hiddenStates.Dim(1) |
||||
|
patchSize := int(math.Sqrt(float64(patches))) |
||||
|
|
||||
|
hiddenStates = hiddenStates.Reshape(ctx, hiddenStates.Dim(0), patchSize, patchSize, hiddenStates.Dim(2)) |
||||
|
|
||||
|
channels, width, height, tiles := hiddenStates.Dim(0), hiddenStates.Dim(1), hiddenStates.Dim(2), hiddenStates.Dim(3) |
||||
|
|
||||
|
channels, width = int(float32(channels)/opts.pixelShuffleRatio), int(float32(width)*opts.pixelShuffleRatio) |
||||
|
hiddenStates = hiddenStates.Reshape(ctx, channels, width, height, tiles) |
||||
|
hiddenStates = hiddenStates.Permute(ctx, 0, 2, 1, 3).Contiguous(ctx) |
||||
|
|
||||
|
channels, height = int(float32(channels)/opts.pixelShuffleRatio), int(float32(height)*opts.pixelShuffleRatio) |
||||
|
hiddenStates = hiddenStates.Reshape(ctx, channels, width, height, tiles) |
||||
|
hiddenStates = hiddenStates.Permute(ctx, 0, 2, 1, 3).Contiguous(ctx) |
||||
|
|
||||
|
hiddenStates = hiddenStates.Reshape(ctx, channels, width*height, tiles) |
||||
|
|
||||
|
hiddenStates = a.FC1.Forward(ctx, hiddenStates).GELU(ctx) |
||||
|
hiddenStates = a.FC2.Forward(ctx, hiddenStates).GELU(ctx) |
||||
|
return hiddenStates |
||||
|
} |
||||
|
|
||||
|
type VisionOptions struct { |
||||
|
hiddenSize, numHeads int |
||||
|
imageSize, patchSize int |
||||
|
|
||||
|
ropeTheta float32 |
||||
|
eps float32 |
||||
|
pixelShuffleRatio float32 |
||||
|
} |
||||
|
|
||||
|
type PatchEmbedding struct { |
||||
|
*nn.Linear |
||||
|
} |
||||
|
|
||||
|
func (p *PatchEmbedding) Forward(ctx ml.Context, hiddenStates ml.Tensor, opts *VisionOptions) ml.Tensor { |
||||
|
kernel := ctx.Input().Empty(ml.DTypeF32, opts.patchSize, opts.patchSize, hiddenStates.Dim(2)) |
||||
|
hiddenStates = kernel.IM2Col(ctx, hiddenStates, opts.patchSize, opts.patchSize, 0, 0, 1, 1) |
||||
|
hiddenStates = hiddenStates.Reshape(ctx, hiddenStates.Dim(0), hiddenStates.Dim(1)*hiddenStates.Dim(2), hiddenStates.Dim(3)) |
||||
|
return p.Linear.Forward(ctx, hiddenStates) |
||||
|
} |
||||
|
|
||||
|
type VisionModel struct { |
||||
|
Layers []VisionLayer `gguf:"blk"` |
||||
|
|
||||
|
*PatchEmbedding `gguf:"patch_embedding"` |
||||
|
ClassEmbedding ml.Tensor `gguf:"class_embedding"` |
||||
|
PositionalEmbedding ml.Tensor `gguf:"positional_embedding_vlm"` |
||||
|
|
||||
|
LayerNormPre *nn.LayerNorm `gguf:"layernorm_pre"` |
||||
|
LayerNormPost *nn.LayerNorm `gguf:"layernorm_post"` |
||||
|
|
||||
|
*VisionAdapter `gguf:"vision_adapter"` |
||||
|
|
||||
|
*VisionOptions |
||||
|
} |
||||
|
|
||||
|
func newVisionModel(c fs.Config) *VisionModel { |
||||
|
return &VisionModel{ |
||||
|
Layers: make([]VisionLayer, c.Uint("vision.block_count")), |
||||
|
VisionOptions: &VisionOptions{ |
||||
|
hiddenSize: int(c.Uint("vision.embedding_length")), |
||||
|
numHeads: int(c.Uint("vision.attention.head_count")), |
||||
|
imageSize: int(c.Uint("vision.image_size")), |
||||
|
patchSize: int(c.Uint("vision.patch_size")), |
||||
|
ropeTheta: float32(c.Float("vision.rope.freq_base")), |
||||
|
eps: c.Float("vision.layer_norm_epsilon"), |
||||
|
pixelShuffleRatio: float32(c.Float("vision.pixel_shuffle_ratio")), |
||||
|
}, |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
func (m *VisionModel) Forward(ctx ml.Context, pixelValues ml.Tensor) ml.Tensor { |
||||
|
hiddenStates := m.PatchEmbedding.Forward(ctx, pixelValues, m.VisionOptions) |
||||
|
hiddenStates = hiddenStates.Concat(ctx, m.ClassEmbedding.Repeat(ctx, 2, hiddenStates.Dim(2)), 1) |
||||
|
|
||||
|
hiddenStates = hiddenStates.Add(ctx, m.PositionalEmbedding) |
||||
|
hiddenStates = m.LayerNormPre.Forward(ctx, hiddenStates, m.eps) |
||||
|
|
||||
|
cos, sin := m.rotaryEmbedding(ctx) |
||||
|
for _, layer := range m.Layers { |
||||
|
hiddenStates = layer.Forward(ctx, hiddenStates, cos, sin, m.VisionOptions) |
||||
|
} |
||||
|
|
||||
|
hiddenStates = m.LayerNormPost.Forward(ctx, hiddenStates, m.eps) |
||||
|
hiddenStates = hiddenStates.Unpad(ctx, 0, 1, 0, 0) |
||||
|
hiddenStates = m.VisionAdapter.Forward(ctx, hiddenStates, m.VisionOptions) |
||||
|
return hiddenStates |
||||
|
} |
||||
|
|
||||
|
// floorDiv is a helper function to perform floor division. This mimics PyTorch's div(round_mode='floor') function
|
||||
|
// which in turn mimics Python's // operator.
|
||||
|
func floorDiv[T int | int16 | int32 | int64 | uint | uint16 | uint32 | uint64](a, b T) T { |
||||
|
if b == 0 { |
||||
|
panic("division by zero") |
||||
|
} |
||||
|
|
||||
|
if (a >= 0 && b > 0) || (a <= 0 && b < 0) || a%b == 0 { |
||||
|
return a / b |
||||
|
} |
||||
|
|
||||
|
return a/b - 1 |
||||
|
} |
||||
|
|
||||
|
func (m *VisionModel) rotaryEmbedding(ctx ml.Context) (ml.Tensor, ml.Tensor) { |
||||
|
patchesPerSide := m.imageSize / m.patchSize |
||||
|
numPatches := patchesPerSide*patchesPerSide + 1 |
||||
|
|
||||
|
headDim := m.hiddenSize / m.numHeads |
||||
|
freqDim := headDim / 2 |
||||
|
|
||||
|
freqs := make([]float32, numPatches*freqDim) |
||||
|
for i := range numPatches - 1 { |
||||
|
for j := 0; j < freqDim; j += 2 { |
||||
|
positionX := i*freqDim/2 + j/2 |
||||
|
positionY := (i+numPatches)*freqDim/2 + j/2 |
||||
|
ropeFreq := math.Pow(float64(m.ropeTheta), float64(j)*2/float64(headDim)) |
||||
|
freqs[positionX] = float32(float64(1+i-floorDiv(i, patchesPerSide)*patchesPerSide) / ropeFreq) |
||||
|
freqs[positionY] = float32(float64(1+floorDiv(i, patchesPerSide)) / ropeFreq) |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
ropeFreqs, err := ctx.Input().FromFloatSlice(freqs, freqDim/2, numPatches, 2) |
||||
|
if err != nil { |
||||
|
panic(err) |
||||
|
} |
||||
|
|
||||
|
ropeFreqs = ropeFreqs.Permute(ctx, 0, 2, 1, 3).Contiguous(ctx) |
||||
|
ropeFreqs = ropeFreqs.Reshape(ctx, freqDim, 1, numPatches) |
||||
|
return ropeFreqs.Cos(ctx), ropeFreqs.Sin(ctx) |
||||
|
} |
||||
Loading…
Reference in new issue