refactor memory check

This commit is contained in:
Michael Yang 2023-10-12 09:47:17 -07:00
parent 92189a5855
commit bd6e38fb1a

View file

@ -58,38 +58,27 @@ func New(workDir, model string, adapters []string, opts api.Options) (LLM, error
} }
} }
totalResidentMemory := memory.TotalMemory() var requiredMemory int64
var f16Multiplier int64 = 2
totalResidentMemory := int64(memory.TotalMemory())
switch ggml.ModelType() { switch ggml.ModelType() {
case "3B", "7B": case "3B", "7B":
if ggml.FileType() == "F16" && totalResidentMemory < 16*format.GigaByte { requiredMemory = 8 * format.GigaByte
return nil, fmt.Errorf("F16 model requires at least 16 GB of memory")
} else if totalResidentMemory < 8*format.GigaByte {
return nil, fmt.Errorf("model requires at least 8 GB of memory")
}
case "13B": case "13B":
if ggml.FileType() == "F16" && totalResidentMemory < 32*format.GigaByte { requiredMemory = 16 * format.GigaByte
return nil, fmt.Errorf("F16 model requires at least 32 GB of memory")
} else if totalResidentMemory < 16*format.GigaByte {
return nil, fmt.Errorf("model requires at least 16 GB of memory")
}
case "30B", "34B", "40B": case "30B", "34B", "40B":
if ggml.FileType() == "F16" && totalResidentMemory < 64*format.GigaByte { requiredMemory = 32 * format.GigaByte
return nil, fmt.Errorf("F16 model requires at least 64 GB of memory")
} else if totalResidentMemory < 32*format.GigaByte {
return nil, fmt.Errorf("model requires at least 32 GB of memory")
}
case "65B", "70B": case "65B", "70B":
if ggml.FileType() == "F16" && totalResidentMemory < 128*format.GigaByte { requiredMemory = 64 * format.GigaByte
return nil, fmt.Errorf("F16 model requires at least 128 GB of memory")
} else if totalResidentMemory < 64*format.GigaByte {
return nil, fmt.Errorf("model requires at least 64 GB of memory")
}
case "180B": case "180B":
if ggml.FileType() == "F16" && totalResidentMemory < 512*format.GigaByte { requiredMemory = 128 * format.GigaByte
return nil, fmt.Errorf("F16 model requires at least 512GB of memory") f16Multiplier = 4
} else if totalResidentMemory < 128*format.GigaByte { }
return nil, fmt.Errorf("model requires at least 128GB of memory")
} if ggml.FileType() == "F16" && requiredMemory*f16Multiplier > totalResidentMemory {
return nil, fmt.Errorf("F16 model requires at least %s of memory", format.HumanBytes(requiredMemory))
} else if requiredMemory > totalResidentMemory {
return nil, fmt.Errorf("model requires at least %s of memory", format.HumanBytes(requiredMemory))
} }
switch ggml.Name() { switch ggml.Name() {