ollama source for Momentry Core verification
This commit is contained in:
Accusys
444
model/models/laguna/model.go
Normal file
444
model/models/laguna/model.go
Normal file
@@ -0,0 +1,444 @@
|
||||
package laguna
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"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/ml/nn/rope"
|
||||
"github.com/ollama/ollama/model"
|
||||
"github.com/ollama/ollama/model/input"
|
||||
"github.com/ollama/ollama/tokenizer"
|
||||
)
|
||||
|
||||
const (
|
||||
cacheTypeSWA = iota
|
||||
cacheTypeCausal
|
||||
)
|
||||
|
||||
type Options struct {
|
||||
hiddenSize int
|
||||
headDim int
|
||||
|
||||
numHeads []int
|
||||
numKVHeads int
|
||||
|
||||
eps float32
|
||||
|
||||
slidingWindow int
|
||||
slidingWindowPattern []bool
|
||||
|
||||
fullRopeDim int
|
||||
fullRopeBase, fullRopeScale float32
|
||||
fullRopeOriginalContextLength int
|
||||
fullRopeAttentionFactor float32
|
||||
fullRopeBetaFast float32
|
||||
fullRopeBetaSlow float32
|
||||
|
||||
swaRopeDim int
|
||||
swaRopeBase, swaRopeScale float32
|
||||
numExperts, numExpertsUsed int
|
||||
normTopKProb bool
|
||||
routedScalingFactor float32
|
||||
decoderSparseStep int
|
||||
denseLayers map[int]bool
|
||||
}
|
||||
|
||||
func (o *Options) numHeadsForLayer(layer int) int {
|
||||
if layer < len(o.numHeads) && o.numHeads[layer] > 0 {
|
||||
return o.numHeads[layer]
|
||||
}
|
||||
if len(o.numHeads) > 0 && o.numHeads[0] > 0 {
|
||||
return o.numHeads[0]
|
||||
}
|
||||
return 1
|
||||
}
|
||||
|
||||
func (o *Options) layerIsSliding(layer int) bool {
|
||||
return layer < len(o.slidingWindowPattern) && o.slidingWindowPattern[layer]
|
||||
}
|
||||
|
||||
func (o *Options) layerUsesMoE(layer int) bool {
|
||||
if o.numExperts == 0 || o.denseLayers[layer] {
|
||||
return false
|
||||
}
|
||||
step := o.decoderSparseStep
|
||||
if step <= 0 {
|
||||
step = 1
|
||||
}
|
||||
return (layer+1)%step == 0
|
||||
}
|
||||
|
||||
func (o *Options) applyRotaryPositionEmbeddings(ctx ml.Context, layer int, states, positions ml.Tensor) ml.Tensor {
|
||||
opts := []func(*rope.Options){rope.WithTypeNeoX()}
|
||||
if o.layerIsSliding(layer) {
|
||||
return nn.RoPE(ctx, states, positions, o.swaRopeDim, o.swaRopeBase, 1./o.swaRopeScale, opts...)
|
||||
}
|
||||
|
||||
opts = append(opts,
|
||||
rope.WithOriginalContextLength(o.fullRopeOriginalContextLength),
|
||||
rope.WithExtrapolationFactor(1),
|
||||
rope.WithAttentionFactor(o.fullRopeAttentionFactor),
|
||||
rope.WithBetaFast(o.fullRopeBetaFast),
|
||||
rope.WithBetaSlow(o.fullRopeBetaSlow),
|
||||
)
|
||||
return nn.RoPE(ctx, states, positions, o.fullRopeDim, o.fullRopeBase, 1./o.fullRopeScale, opts...)
|
||||
}
|
||||
|
||||
type Attention struct {
|
||||
Query *nn.Linear `gguf:"attn_q"`
|
||||
QueryNorm *nn.RMSNorm `gguf:"attn_q_norm"`
|
||||
Key *nn.Linear `gguf:"attn_k"`
|
||||
KeyNorm *nn.RMSNorm `gguf:"attn_k_norm"`
|
||||
Value *nn.Linear `gguf:"attn_v"`
|
||||
Gate *nn.Linear `gguf:"attn_g"`
|
||||
Output *nn.Linear `gguf:"attn_output"`
|
||||
}
|
||||
|
||||
func (sa *Attention) Forward(ctx ml.Context, layer int, hiddenStates, positions ml.Tensor, cache kvcache.Cache, opts *Options) ml.Tensor {
|
||||
batchSize := hiddenStates.Dim(1)
|
||||
numHeads := opts.numHeadsForLayer(layer)
|
||||
|
||||
query := sa.Query.Forward(ctx, hiddenStates)
|
||||
key := sa.Key.Forward(ctx, hiddenStates)
|
||||
value := sa.Value.Forward(ctx, hiddenStates)
|
||||
gate := sa.Gate.Forward(ctx, hiddenStates)
|
||||
|
||||
query = query.Reshape(ctx, opts.headDim, numHeads, batchSize)
|
||||
key = key.Reshape(ctx, opts.headDim, opts.numKVHeads, batchSize)
|
||||
value = value.Reshape(ctx, opts.headDim, opts.numKVHeads, batchSize)
|
||||
|
||||
query = sa.QueryNorm.Forward(ctx, query, opts.eps)
|
||||
key = sa.KeyNorm.Forward(ctx, key, opts.eps)
|
||||
|
||||
query = opts.applyRotaryPositionEmbeddings(ctx, layer, query, positions)
|
||||
key = opts.applyRotaryPositionEmbeddings(ctx, layer, key, positions)
|
||||
|
||||
attention := nn.Attention(ctx, query, key, value, 1./math.Sqrt(float64(opts.headDim)), cache)
|
||||
// Laguna applies the per-head gate softplus in float32, then casts back.
|
||||
gate = gate.Cast(ctx, ml.DTypeF32).Softplus(ctx).Cast(ctx, attention.DType())
|
||||
attention = attention.Mul(ctx, gate.Reshape(ctx, 1, numHeads, batchSize))
|
||||
attention = attention.Reshape(ctx, opts.headDim*numHeads, batchSize)
|
||||
return sa.Output.Forward(ctx, attention)
|
||||
}
|
||||
|
||||
type MLP interface {
|
||||
Forward(ml.Context, ml.Tensor, *Options) ml.Tensor
|
||||
}
|
||||
|
||||
type dense struct {
|
||||
Gate *nn.Linear `gguf:"ffn_gate"`
|
||||
Up *nn.Linear `gguf:"ffn_up"`
|
||||
Down *nn.Linear `gguf:"ffn_down"`
|
||||
}
|
||||
|
||||
func (mlp *dense) Forward(ctx ml.Context, hiddenStates ml.Tensor, _ *Options) ml.Tensor {
|
||||
hiddenStates = mlp.Gate.Forward(ctx, hiddenStates).SILU(ctx, mlp.Up.Forward(ctx, hiddenStates))
|
||||
return mlp.Down.Forward(ctx, hiddenStates)
|
||||
}
|
||||
|
||||
type sparse struct {
|
||||
Router *nn.Linear `gguf:"ffn_gate_inp"`
|
||||
Gate *nn.LinearBatch `gguf:"ffn_gate_exps"`
|
||||
Up *nn.LinearBatch `gguf:"ffn_up_exps"`
|
||||
Down *nn.LinearBatch `gguf:"ffn_down_exps"`
|
||||
SharedExpert *dense `gguf:",suf:_shexp"`
|
||||
ExpProbsBias ml.Tensor `gguf:"exp_probs_b.bias,alt:exp_probs_b"`
|
||||
}
|
||||
|
||||
func (moe *sparse) topKIndices(ctx ml.Context, scores ml.Tensor, opts *Options) ml.Tensor {
|
||||
if moe.ExpProbsBias != nil {
|
||||
scores = scores.Add(ctx, moe.ExpProbsBias)
|
||||
}
|
||||
return scores.TopK(ctx, opts.numExpertsUsed)
|
||||
}
|
||||
|
||||
func (moe *sparse) Forward(ctx ml.Context, hiddenStates ml.Tensor, opts *Options) ml.Tensor {
|
||||
residual := hiddenStates
|
||||
|
||||
scores := moe.Router.Forward(ctx, hiddenStates).Cast(ctx, ml.DTypeF32).Sigmoid(ctx)
|
||||
selectedExperts := moe.topKIndices(ctx, scores, opts)
|
||||
routingWeights := scores.Reshape(ctx, 1, opts.numExperts, hiddenStates.Dim(1)).Rows(ctx, selectedExperts)
|
||||
if opts.normTopKProb {
|
||||
routingWeights = routingWeights.Reshape(ctx, opts.numExpertsUsed, hiddenStates.Dim(1))
|
||||
routingWeights = routingWeights.Div(ctx, routingWeights.SumRows(ctx))
|
||||
routingWeights = routingWeights.Reshape(ctx, 1, opts.numExpertsUsed, hiddenStates.Dim(1))
|
||||
}
|
||||
routingWeights = routingWeights.Scale(ctx, float64(opts.routedScalingFactor))
|
||||
|
||||
hiddenStates = hiddenStates.Reshape(ctx, hiddenStates.Dim(0), 1, hiddenStates.Dim(1))
|
||||
upStates := moe.Up.Forward(ctx, hiddenStates, selectedExperts)
|
||||
hiddenStates = moe.Gate.Forward(ctx, hiddenStates, selectedExperts).SILU(ctx, upStates)
|
||||
|
||||
experts := moe.Down.Forward(ctx, hiddenStates, selectedExperts)
|
||||
experts = experts.Mul(ctx, routingWeights)
|
||||
|
||||
nextStates := experts.View(ctx, 0, experts.Dim(0), experts.Stride(2), experts.Dim(2))
|
||||
for i := 1; i < opts.numExpertsUsed; i++ {
|
||||
nextStates = nextStates.Add(ctx, experts.View(ctx, i*experts.Stride(1), experts.Dim(0), experts.Stride(2), experts.Dim(2)))
|
||||
}
|
||||
|
||||
return nextStates.Add(ctx, moe.SharedExpert.Forward(ctx, residual, opts))
|
||||
}
|
||||
|
||||
type Layer struct {
|
||||
AttentionNorm *nn.RMSNorm `gguf:"attn_norm"`
|
||||
*Attention
|
||||
|
||||
MLPNorm *nn.RMSNorm `gguf:"ffn_norm"`
|
||||
MLP
|
||||
}
|
||||
|
||||
func (l *Layer) Forward(ctx ml.Context, layer int, hiddenStates, positions, outputs ml.Tensor, cache kvcache.Cache, opts *Options) ml.Tensor {
|
||||
residual := hiddenStates
|
||||
hiddenStates = l.AttentionNorm.Forward(ctx, hiddenStates, opts.eps)
|
||||
hiddenStates = l.Attention.Forward(ctx, layer, hiddenStates, positions, cache, opts)
|
||||
|
||||
if outputs != nil {
|
||||
hiddenStates = hiddenStates.Rows(ctx, outputs)
|
||||
residual = residual.Rows(ctx, outputs)
|
||||
}
|
||||
|
||||
hiddenStates = hiddenStates.Add(ctx, residual)
|
||||
residual = hiddenStates
|
||||
|
||||
hiddenStates = l.MLPNorm.Forward(ctx, hiddenStates, opts.eps)
|
||||
hiddenStates = l.MLP.Forward(ctx, hiddenStates, opts)
|
||||
return hiddenStates.Add(ctx, residual)
|
||||
}
|
||||
|
||||
type Model struct {
|
||||
model.Base
|
||||
tokenizer.Tokenizer
|
||||
|
||||
TokenEmbedding *nn.Embedding `gguf:"token_embd"`
|
||||
Layers []Layer `gguf:"blk"`
|
||||
OutputNorm *nn.RMSNorm `gguf:"output_norm"`
|
||||
Output *nn.Linear `gguf:"output,alt:token_embd"`
|
||||
|
||||
*Options
|
||||
}
|
||||
|
||||
func New(c fs.Config) (model.Model, error) {
|
||||
if c.Bool("attention.sink_enabled") {
|
||||
return nil, fmt.Errorf("laguna: SWA attention sinks are not supported")
|
||||
}
|
||||
if c.Uint("attention.gating_type") != 1 {
|
||||
return nil, fmt.Errorf("laguna: unsupported attention gating type %d", c.Uint("attention.gating_type"))
|
||||
}
|
||||
if !c.Bool("attention.qk_norm") {
|
||||
return nil, fmt.Errorf("laguna: Q/K RMSNorm is required")
|
||||
}
|
||||
if gating := c.Uint("expert_gating_func"); gating != 2 {
|
||||
return nil, fmt.Errorf("laguna: unsupported expert gating function %d", gating)
|
||||
}
|
||||
|
||||
numLayers := int(c.Uint("block_count"))
|
||||
opts := newOptions(c, numLayers)
|
||||
layers := make([]Layer, numLayers)
|
||||
for i := range layers {
|
||||
if opts.layerUsesMoE(i) {
|
||||
layers[i].MLP = &sparse{}
|
||||
} else {
|
||||
layers[i].MLP = &dense{}
|
||||
}
|
||||
}
|
||||
|
||||
var pre []string
|
||||
switch c.String("tokenizer.ggml.pre") {
|
||||
case "laguna":
|
||||
pre = []string{
|
||||
`(?:\r?\n)+(?!\r?\n)`,
|
||||
`(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+`,
|
||||
}
|
||||
default:
|
||||
return nil, model.ErrUnsupportedTokenizer
|
||||
}
|
||||
|
||||
m := Model{
|
||||
Tokenizer: tokenizer.NewBytePairEncoding(
|
||||
&tokenizer.Vocabulary{
|
||||
Values: c.Strings("tokenizer.ggml.tokens"),
|
||||
Types: c.Ints("tokenizer.ggml.token_type"),
|
||||
Merges: c.Strings("tokenizer.ggml.merges"),
|
||||
AddBOS: c.Bool("tokenizer.ggml.add_bos_token", true),
|
||||
BOS: []int32{int32(c.Uint("tokenizer.ggml.bos_token_id"))},
|
||||
AddEOS: c.Bool("tokenizer.ggml.add_eos_token", false),
|
||||
EOS: append(
|
||||
[]int32{int32(c.Uint("tokenizer.ggml.eos_token_id"))},
|
||||
c.Ints("tokenizer.ggml.eos_token_ids")...,
|
||||
),
|
||||
},
|
||||
pre...,
|
||||
),
|
||||
Layers: layers,
|
||||
Options: opts,
|
||||
}
|
||||
|
||||
m.Cache = kvcache.NewWrapperCache(
|
||||
kvcache.NewSWACache(int32(opts.slidingWindow), m.Shift),
|
||||
kvcache.NewCausalCache(m.Shift),
|
||||
)
|
||||
return &m, nil
|
||||
}
|
||||
|
||||
func newOptions(c fs.Config, numLayers int) *Options {
|
||||
denseLayers := make(map[int]bool)
|
||||
for _, layer := range configUints(c, "dense_layers") {
|
||||
denseLayers[int(layer)] = true
|
||||
}
|
||||
for i := range c.Uint("leading_dense_block_count") {
|
||||
denseLayers[int(i)] = true
|
||||
}
|
||||
|
||||
fullRopeScale := c.Float("rope.scaling.factor", 1)
|
||||
if fullRopeScale == 0 {
|
||||
fullRopeScale = 1
|
||||
}
|
||||
swaRopeScale := c.Float("rope.swa.scaling.factor", 1)
|
||||
if swaRopeScale == 0 {
|
||||
swaRopeScale = 1
|
||||
}
|
||||
fullRopeType := c.String("rope.scaling.type")
|
||||
fullRopeAttentionFactor := lagunaAttentionFactor(fullRopeType, fullRopeScale, c.Float("rope.scaling.attn_factor"))
|
||||
|
||||
return &Options{
|
||||
hiddenSize: int(c.Uint("embedding_length")),
|
||||
headDim: int(c.Uint("attention.key_length")),
|
||||
numHeads: expandIntArray(configUints(c, "attention.head_count"), numLayers, c.Uint("attention.head_count", 1)),
|
||||
numKVHeads: int(c.Uint("attention.head_count_kv")),
|
||||
eps: c.Float("attention.layer_norm_rms_epsilon", 1e-6),
|
||||
slidingWindow: int(c.Uint("attention.sliding_window", 512)),
|
||||
slidingWindowPattern: slidingWindowPattern(c, numLayers),
|
||||
fullRopeDim: int(c.Uint("rope.dimension_count", c.Uint("attention.key_length"))),
|
||||
fullRopeBase: c.Float("rope.freq_base", 500000),
|
||||
fullRopeScale: fullRopeScale,
|
||||
fullRopeOriginalContextLength: int(c.Uint("rope.scaling.original_context_length", 4096)),
|
||||
fullRopeAttentionFactor: fullRopeAttentionFactor,
|
||||
fullRopeBetaFast: c.Float("rope.scaling.beta_fast", 64),
|
||||
fullRopeBetaSlow: c.Float("rope.scaling.beta_slow", 1),
|
||||
swaRopeDim: int(c.Uint("rope.swa.dimension_count", c.Uint("attention.key_length"))),
|
||||
swaRopeBase: c.Float("rope.swa.freq_base", 10000),
|
||||
swaRopeScale: swaRopeScale,
|
||||
numExperts: int(c.Uint("expert_count")),
|
||||
numExpertsUsed: int(c.Uint("expert_used_count")),
|
||||
normTopKProb: c.Bool("expert_weights_norm", true),
|
||||
routedScalingFactor: c.Float("expert_weights_scale", 1),
|
||||
decoderSparseStep: int(c.Uint("decoder_sparse_step", 1)),
|
||||
denseLayers: denseLayers,
|
||||
}
|
||||
}
|
||||
|
||||
func lagunaAttentionFactor(ropeType string, scaleFactor, attentionFactor float32) float32 {
|
||||
if attentionFactor != 0 {
|
||||
return attentionFactor
|
||||
}
|
||||
if ropeType == "yarn" && scaleFactor > 1 {
|
||||
return float32(0.1*math.Log(float64(scaleFactor)) + 1)
|
||||
}
|
||||
return 1
|
||||
}
|
||||
|
||||
func slidingWindowPattern(c fs.Config, numLayers int) []bool {
|
||||
pattern := c.Bools("attention.sliding_window_pattern")
|
||||
if len(pattern) == numLayers {
|
||||
return pattern
|
||||
}
|
||||
|
||||
layerTypes := configUints(c, "attention.layer_types")
|
||||
if len(layerTypes) == numLayers {
|
||||
pattern = make([]bool, numLayers)
|
||||
for i, layerType := range layerTypes {
|
||||
pattern[i] = layerType == 1
|
||||
}
|
||||
return pattern
|
||||
}
|
||||
|
||||
return make([]bool, numLayers)
|
||||
}
|
||||
|
||||
func configUints(c fs.Config, key string) []uint32 {
|
||||
keyExists := c.Value(c.Architecture()+"."+key) != nil || c.Value(key) != nil
|
||||
if cc, ok := c.(interface {
|
||||
Uints(string, ...[]uint32) []uint32
|
||||
}); ok {
|
||||
if values := cc.Uints(key); len(values) > 0 && (keyExists || !(len(values) == 1 && values[0] == 0)) {
|
||||
return values
|
||||
}
|
||||
}
|
||||
|
||||
ints := c.Ints(key)
|
||||
if len(ints) > 0 && (keyExists || !(len(ints) == 1 && ints[0] == 0)) {
|
||||
values := make([]uint32, len(ints))
|
||||
for i, v := range ints {
|
||||
values[i] = uint32(v)
|
||||
}
|
||||
return values
|
||||
}
|
||||
|
||||
if scalar := c.Uint(key); scalar != 0 {
|
||||
return []uint32{scalar}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func expandIntArray(values []uint32, n int, fallback uint32) []int {
|
||||
if len(values) == 0 {
|
||||
values = []uint32{fallback}
|
||||
}
|
||||
defaultValue := values[0]
|
||||
if len(values) == 1 {
|
||||
defaultValue = values[0]
|
||||
}
|
||||
|
||||
out := make([]int, n)
|
||||
for i := range out {
|
||||
if i < len(values) {
|
||||
out[i] = int(values[i])
|
||||
} else {
|
||||
out[i] = int(defaultValue)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (m *Model) Shift(ctx ml.Context, layer int, key, shift ml.Tensor) (ml.Tensor, error) {
|
||||
return m.Options.applyRotaryPositionEmbeddings(ctx, layer, key, shift), nil
|
||||
}
|
||||
|
||||
func (m *Model) Forward(ctx ml.Context, batch input.Batch) (ml.Tensor, error) {
|
||||
positions := ctx.Input().FromInts(batch.Positions, len(batch.Positions))
|
||||
hiddenStates := m.TokenEmbedding.Forward(ctx, batch.Inputs)
|
||||
|
||||
for i, layer := range m.Layers {
|
||||
if m.Cache != nil {
|
||||
m.Cache.SetLayer(i)
|
||||
if wrapper, ok := m.Cache.(*kvcache.WrapperCache); ok {
|
||||
cacheType := cacheTypeCausal
|
||||
if m.Options.layerIsSliding(i) {
|
||||
cacheType = cacheTypeSWA
|
||||
}
|
||||
wrapper.SetLayerType(cacheType)
|
||||
}
|
||||
}
|
||||
|
||||
var outputs ml.Tensor
|
||||
if i == len(m.Layers)-1 {
|
||||
outputs = batch.Outputs
|
||||
}
|
||||
|
||||
hiddenStates = layer.Forward(ctx, i, hiddenStates, positions, outputs, m.Cache, m.Options)
|
||||
}
|
||||
|
||||
hiddenStates = m.OutputNorm.Forward(ctx, hiddenStates, m.eps)
|
||||
return m.Output.Forward(ctx, hiddenStates), nil
|
||||
}
|
||||
|
||||
func init() {
|
||||
model.Register("laguna", New)
|
||||
}
|
||||
|
||||
var _ model.Model = (*Model)(nil)
|
||||
Reference in New Issue
Block a user