mirror of
https://github.com/ollama/ollama.git
synced 2026-03-08 23:04:13 -05:00
386 lines
11 KiB
Go
386 lines
11 KiB
Go
package convert
|
|
|
|
import (
|
|
"cmp"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io/fs"
|
|
"math"
|
|
"slices"
|
|
"strings"
|
|
|
|
"github.com/ollama/ollama/fs/ggml"
|
|
)
|
|
|
|
type hybridPattern string
|
|
|
|
func (p *hybridPattern) UnmarshalJSON(data []byte) error {
|
|
if string(data) == "null" {
|
|
*p = ""
|
|
return nil
|
|
}
|
|
|
|
var single string
|
|
if err := json.Unmarshal(data, &single); err == nil {
|
|
*p = hybridPattern(strings.TrimSpace(single))
|
|
return nil
|
|
}
|
|
|
|
var parts []string
|
|
if err := json.Unmarshal(data, &parts); err == nil {
|
|
*p = hybridPattern(strings.Join(parts, ""))
|
|
return nil
|
|
}
|
|
|
|
return fmt.Errorf("hybrid_override_pattern must be a string or string array")
|
|
}
|
|
|
|
type nemotronHModel struct {
|
|
ModelParameters
|
|
MaxPositionEmbeddings uint32 `json:"max_position_embeddings"`
|
|
HiddenSize uint32 `json:"hidden_size"`
|
|
NumHiddenLayers uint32 `json:"num_hidden_layers"`
|
|
NumAttentionHeads uint32 `json:"num_attention_heads"`
|
|
NumKeyValueHeads uint32 `json:"num_key_value_heads"`
|
|
HeadDim uint32 `json:"head_dim"`
|
|
LayerNormEpsilon float32 `json:"layer_norm_epsilon"`
|
|
NormEpsilon float32 `json:"norm_eps"`
|
|
RopeTheta float32 `json:"rope_theta"`
|
|
PartialRotaryFactor float32 `json:"partial_rotary_factor"`
|
|
ConvKernel uint32 `json:"conv_kernel"`
|
|
SSMStateSize uint32 `json:"ssm_state_size"`
|
|
MambaNumHeads uint32 `json:"mamba_num_heads"`
|
|
MambaHeadDim uint32 `json:"mamba_head_dim"`
|
|
NGroups uint32 `json:"n_groups"`
|
|
IntermediateSize uint32 `json:"intermediate_size"`
|
|
HybridOverridePattern hybridPattern `json:"hybrid_override_pattern"`
|
|
|
|
// MoE
|
|
NumExperts uint32 `json:"num_experts"`
|
|
NumSharedExperts uint32 `json:"num_shared_experts"`
|
|
NRoutedExperts uint32 `json:"n_routed_experts"`
|
|
NSharedExperts uint32 `json:"n_shared_experts"`
|
|
NumExpertsPerTok uint32 `json:"num_experts_per_tok"`
|
|
MoEIntermediateSize uint32 `json:"moe_intermediate_size"`
|
|
MoESharedExpertIntermediate uint32 `json:"moe_shared_expert_intermediate_size"`
|
|
NormTopKProb bool `json:"norm_topk_prob"`
|
|
RoutedScalingFactor float32 `json:"routed_scaling_factor"`
|
|
ExpertGroupCount uint32 `json:"n_group"`
|
|
ExpertGroupUsedCount uint32 `json:"topk_group"`
|
|
}
|
|
|
|
var _ ModelConverter = (*nemotronHModel)(nil)
|
|
|
|
func (n *nemotronHModel) parseMore(_ fs.FS) error {
|
|
if n.NumHiddenLayers == 0 {
|
|
return fmt.Errorf("nemotron_h: num_hidden_layers must be set")
|
|
}
|
|
if n.HiddenSize == 0 {
|
|
return fmt.Errorf("nemotron_h: hidden_size must be set")
|
|
}
|
|
if n.NumAttentionHeads == 0 {
|
|
return fmt.Errorf("nemotron_h: num_attention_heads must be set")
|
|
}
|
|
if n.HeadDim == 0 {
|
|
if n.HiddenSize%n.NumAttentionHeads != 0 {
|
|
return fmt.Errorf("nemotron_h: hidden_size (%d) must be divisible by num_attention_heads (%d)", n.HiddenSize, n.NumAttentionHeads)
|
|
}
|
|
n.HeadDim = n.HiddenSize / n.NumAttentionHeads
|
|
}
|
|
if n.NumKeyValueHeads == 0 {
|
|
n.NumKeyValueHeads = n.NumAttentionHeads
|
|
}
|
|
if n.ConvKernel == 0 {
|
|
return fmt.Errorf("nemotron_h: conv_kernel must be set")
|
|
}
|
|
if n.SSMStateSize == 0 {
|
|
return fmt.Errorf("nemotron_h: ssm_state_size must be set")
|
|
}
|
|
if n.ssmHeadCount() == 0 {
|
|
return fmt.Errorf("nemotron_h: mamba_num_heads must be set")
|
|
}
|
|
if n.MambaHeadDim == 0 {
|
|
return fmt.Errorf("nemotron_h: mamba_head_dim must be set")
|
|
}
|
|
if n.NGroups == 0 {
|
|
n.NGroups = 1
|
|
}
|
|
|
|
if _, _, err := n.layerArrays(); err != nil {
|
|
return err
|
|
}
|
|
|
|
if n.isMoE() {
|
|
if n.routedExpertCount() == 0 {
|
|
return fmt.Errorf("nemotron_h: routed expert count must be set for MoE models")
|
|
}
|
|
if n.NumExpertsPerTok == 0 {
|
|
return fmt.Errorf("nemotron_h: num_experts_per_tok must be set for MoE models")
|
|
}
|
|
if n.NumExpertsPerTok > n.routedExpertCount() {
|
|
return fmt.Errorf("nemotron_h: num_experts_per_tok (%d) cannot exceed expert_count (%d)", n.NumExpertsPerTok, n.routedExpertCount())
|
|
}
|
|
if n.moeIntermediateSize() == 0 {
|
|
return fmt.Errorf("nemotron_h: moe_intermediate_size must be set for MoE models")
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (n *nemotronHModel) isMoE() bool {
|
|
return cmp.Or(n.routedExpertCount(), n.NumExpertsPerTok, n.MoEIntermediateSize) > 0
|
|
}
|
|
|
|
func (n *nemotronHModel) routedExpertCount() uint32 {
|
|
return cmp.Or(n.NRoutedExperts, n.NumExperts)
|
|
}
|
|
|
|
func (n *nemotronHModel) sharedExpertCount() uint32 {
|
|
return cmp.Or(n.NSharedExperts, n.NumSharedExperts)
|
|
}
|
|
|
|
func (n *nemotronHModel) ssmHeadCount() uint32 {
|
|
return n.MambaNumHeads
|
|
}
|
|
|
|
func (n *nemotronHModel) ssmInnerSize() uint32 {
|
|
return n.MambaHeadDim * n.ssmHeadCount()
|
|
}
|
|
|
|
func (n *nemotronHModel) epsilon() float32 {
|
|
return cmp.Or(n.NormEpsilon, n.LayerNormEpsilon, float32(1e-5))
|
|
}
|
|
|
|
func (n *nemotronHModel) moeIntermediateSize() uint32 {
|
|
return cmp.Or(n.MoEIntermediateSize, n.IntermediateSize)
|
|
}
|
|
|
|
func (n *nemotronHModel) denseIntermediateSize() uint32 {
|
|
return cmp.Or(n.IntermediateSize, n.MoEIntermediateSize)
|
|
}
|
|
|
|
func (n *nemotronHModel) layerArrays() (headCountKV []uint32, ffnLengths []uint32, err error) {
|
|
pattern := strings.TrimSpace(string(n.HybridOverridePattern))
|
|
if pattern == "" {
|
|
return nil, nil, fmt.Errorf("nemotron_h: hybrid_override_pattern must be set")
|
|
}
|
|
|
|
runes := []rune(pattern)
|
|
if len(runes) != int(n.NumHiddenLayers) {
|
|
return nil, nil, fmt.Errorf("nemotron_h: hybrid_override_pattern length (%d) must match num_hidden_layers (%d)", len(runes), n.NumHiddenLayers)
|
|
}
|
|
|
|
headCountKV = make([]uint32, n.NumHiddenLayers)
|
|
ffnLengths = make([]uint32, n.NumHiddenLayers)
|
|
|
|
attnKVHeads := cmp.Or(n.NumKeyValueHeads, n.NumAttentionHeads)
|
|
moeFFN := n.moeIntermediateSize()
|
|
denseFFN := n.denseIntermediateSize()
|
|
|
|
for i, layerType := range runes {
|
|
switch layerType {
|
|
case 'M':
|
|
// Recurrent layer: no KV heads and no FFN.
|
|
case '*', 'A':
|
|
// Attention-only layer.
|
|
headCountKV[i] = attnKVHeads
|
|
case 'E':
|
|
// MoE layer.
|
|
if moeFFN == 0 {
|
|
return nil, nil, fmt.Errorf("nemotron_h: moe layer at index %d but moe_intermediate_size is zero", i)
|
|
}
|
|
ffnLengths[i] = moeFFN
|
|
case '-':
|
|
// Dense FFN layer.
|
|
if denseFFN == 0 {
|
|
return nil, nil, fmt.Errorf("nemotron_h: dense FFN layer at index %d but intermediate_size is zero", i)
|
|
}
|
|
ffnLengths[i] = denseFFN
|
|
default:
|
|
return nil, nil, fmt.Errorf("nemotron_h: unsupported layer type %q in hybrid_override_pattern at index %d", layerType, i)
|
|
}
|
|
}
|
|
|
|
return headCountKV, ffnLengths, nil
|
|
}
|
|
|
|
func (n *nemotronHModel) KV(t *Tokenizer) KV {
|
|
kv := n.ModelParameters.KV(t)
|
|
|
|
arch := "nemotron_h"
|
|
if n.isMoE() {
|
|
arch = "nemotron_h_moe"
|
|
}
|
|
kv["general.architecture"] = arch
|
|
kv["block_count"] = n.NumHiddenLayers
|
|
kv["context_length"] = n.MaxPositionEmbeddings
|
|
kv["embedding_length"] = n.HiddenSize
|
|
kv["attention.head_count"] = n.NumAttentionHeads
|
|
kv["attention.key_length"] = n.HeadDim
|
|
kv["attention.value_length"] = n.HeadDim
|
|
kv["attention.layer_norm_epsilon"] = n.epsilon()
|
|
kv["attention.layer_norm_rms_epsilon"] = n.epsilon()
|
|
kv["rope.freq_base"] = cmp.Or(n.RopeTheta, float32(10000))
|
|
if n.PartialRotaryFactor > 0 && n.PartialRotaryFactor <= 1 {
|
|
kv["rope.dimension_count"] = uint32(float32(n.HeadDim) * n.PartialRotaryFactor)
|
|
}
|
|
|
|
if headCountKV, ffnLengths, err := n.layerArrays(); err == nil {
|
|
kv["attention.head_count_kv"] = headCountKV
|
|
kv["feed_forward_length"] = ffnLengths
|
|
}
|
|
|
|
kv["ssm.conv_kernel"] = n.ConvKernel
|
|
kv["ssm.inner_size"] = n.ssmInnerSize()
|
|
kv["ssm.state_size"] = n.SSMStateSize
|
|
kv["ssm.group_count"] = n.NGroups
|
|
kv["ssm.time_step_rank"] = n.ssmHeadCount()
|
|
|
|
if n.isMoE() {
|
|
kv["expert_count"] = n.routedExpertCount()
|
|
kv["expert_used_count"] = n.NumExpertsPerTok
|
|
kv["expert_feed_forward_length"] = n.moeIntermediateSize()
|
|
if n.sharedExpertCount() > 0 {
|
|
kv["expert_shared_count"] = n.sharedExpertCount()
|
|
}
|
|
if n.MoESharedExpertIntermediate > 0 {
|
|
kv["expert_shared_feed_forward_length"] = n.MoESharedExpertIntermediate
|
|
}
|
|
kv["expert_weights_norm"] = n.NormTopKProb
|
|
kv["expert_weights_scale"] = n.RoutedScalingFactor
|
|
if n.ExpertGroupCount > 0 {
|
|
kv["expert_group_count"] = n.ExpertGroupCount
|
|
}
|
|
if n.ExpertGroupUsedCount > 0 {
|
|
kv["expert_group_used_count"] = n.ExpertGroupUsedCount
|
|
}
|
|
}
|
|
|
|
return kv
|
|
}
|
|
|
|
func normalizeVectorShapeToColumn(shape []uint64) []uint64 {
|
|
switch len(shape) {
|
|
case 1:
|
|
return []uint64{shape[0], 1}
|
|
case 2:
|
|
if shape[0] == 1 && shape[1] > 1 {
|
|
return []uint64{shape[1], 1}
|
|
}
|
|
if shape[1] == 1 && shape[0] > 1 {
|
|
return []uint64{shape[0], 1}
|
|
}
|
|
}
|
|
|
|
return slices.Clone(shape)
|
|
}
|
|
|
|
func (n *nemotronHModel) Tensors(ts []Tensor) []*ggml.Tensor {
|
|
var out []*ggml.Tensor
|
|
|
|
remaining := ts
|
|
if n.isMoE() {
|
|
merges := make([]merge, 0, n.NumHiddenLayers*2)
|
|
for i := range n.NumHiddenLayers {
|
|
merges = append(merges, merge{
|
|
fmt.Sprintf("blk.%d.mixer.experts.*.up_proj.weight", i),
|
|
fmt.Sprintf("blk.%d.ffn_up_exps.weight", i),
|
|
}, merge{
|
|
fmt.Sprintf("blk.%d.mixer.experts.*.down_proj.weight", i),
|
|
fmt.Sprintf("blk.%d.ffn_down_exps.weight", i),
|
|
})
|
|
}
|
|
|
|
merged, rest := mergeTensors(ts, merges...)
|
|
out = append(out, merged...)
|
|
remaining = rest
|
|
}
|
|
|
|
nGroups := uint64(cmp.Or(n.NGroups, uint32(1)))
|
|
for _, t := range remaining {
|
|
name := t.Name()
|
|
shape := slices.Clone(t.Shape())
|
|
|
|
switch {
|
|
case strings.HasSuffix(name, ".ssm_a"):
|
|
shape = normalizeVectorShapeToColumn(shape)
|
|
t.SetRepacker(func(_ string, data []float32, _ []uint64) ([]float32, error) {
|
|
out := make([]float32, len(data))
|
|
for i, v := range data {
|
|
out[i] = -float32(math.Exp(float64(v)))
|
|
}
|
|
return out, nil
|
|
})
|
|
case strings.HasSuffix(name, ".ssm_d"):
|
|
shape = normalizeVectorShapeToColumn(shape)
|
|
case strings.HasSuffix(name, ".ssm_norm.weight"):
|
|
switch len(shape) {
|
|
case 1:
|
|
if nGroups > 0 && shape[0]%nGroups == 0 {
|
|
shape = []uint64{nGroups, shape[0] / nGroups}
|
|
}
|
|
case 2:
|
|
if shape[0] == 1 && nGroups > 0 && shape[1]%nGroups == 0 {
|
|
shape = []uint64{nGroups, shape[1] / nGroups}
|
|
}
|
|
}
|
|
case strings.HasSuffix(name, ".ssm_conv1d.weight"):
|
|
if len(shape) == 3 {
|
|
if shape[0] == 1 {
|
|
shape = []uint64{shape[1], shape[2]}
|
|
} else if shape[1] == 1 {
|
|
shape = []uint64{shape[0], shape[2]}
|
|
}
|
|
}
|
|
}
|
|
|
|
out = append(out, &ggml.Tensor{
|
|
Name: name,
|
|
Kind: t.Kind(),
|
|
Shape: shape,
|
|
WriterTo: t,
|
|
})
|
|
}
|
|
|
|
return out
|
|
}
|
|
|
|
func (n *nemotronHModel) Replacements() []string {
|
|
return []string{
|
|
// Embedding and output
|
|
"lm_head", "output",
|
|
"backbone.embeddings", "token_embd",
|
|
"backbone.norm_f", "output_norm",
|
|
"backbone.layers", "blk",
|
|
|
|
// Recurrent (Mamba2) tensors
|
|
"mixer.in_proj", "ssm_in",
|
|
"mixer.out_proj", "ssm_out",
|
|
"mixer.dt_bias", "ssm_dt.bias",
|
|
"mixer.A_log", "ssm_a",
|
|
"mixer.D", "ssm_d",
|
|
"mixer.conv1d", "ssm_conv1d",
|
|
"mixer.norm.weight", "ssm_norm.weight",
|
|
|
|
// Attention tensors
|
|
"mixer.q_proj", "attn_q",
|
|
"mixer.k_proj", "attn_k",
|
|
"mixer.v_proj", "attn_v",
|
|
"mixer.o_proj", "attn_output",
|
|
|
|
// FFN / MoE tensors
|
|
"mixer.gate.e_score_correction_bias", "exp_probs_b.bias",
|
|
"mixer.gate", "ffn_gate_inp",
|
|
"mixer.fc1_latent_proj", "ffn_latent_in",
|
|
"mixer.fc2_latent_proj", "ffn_latent_out",
|
|
"mixer.shared_experts.up_proj", "ffn_up_shexp",
|
|
"mixer.shared_experts.down_proj", "ffn_down_shexp",
|
|
"mixer.up_proj", "ffn_up",
|
|
"mixer.down_proj", "ffn_down",
|
|
|
|
// Per-layer pre-norm
|
|
".norm.weight", ".attn_norm.weight",
|
|
}
|
|
}
|