Repository: ollama/ollama
File: llm/server.go (Go)

Code before:

func (s *llamaServer) Load(ctx context.Context, systemInfo ml.SystemInfo, systemGPUs []ml.DeviceInfo, requireFull bool) ([]ml.DeviceID, error) {
	slog.Info("loading model", "model layers", s.totalLayers, "requested", s.options.NumGPU)

	gpus := append(make([]ml.DeviceInfo, 0, len(systemGPUs)), systemGPUs...)

	// Synthesize memory allocation information based on our estimates
	s.mem = &ml.BackendMemory{CPU: ml.DeviceMemory{
		Name:    "CPU",
		Weights: make([]uint64, s.totalLayers),
		Cache:   make([]uint64, s.totalLayers),
	}, GPUs: make([]ml.DeviceMemory, len(gpus))}

	for i := range s.mem.GPUs {
		s.mem.GPUs[i].Name = gpus[i].Name
		s.mem.GPUs[i].DeviceID = gpus[i].DeviceID
		s.mem.GPUs[i].Weights = make([]uint64, s.totalLayers)
		s.mem.GPUs[i].Cache = make([]uint64, s.totalLayers)
	}

	// Check if embedding model and adjust batch size accordingly
	batchSize := s.loadRequest.BatchSize
	_, isEmbedding := s.ggml.KV()[fmt.Sprintf("%s.pooling_type", s.ggml.KV().Architecture())]
	if isEmbedding {
		batchSize = s.options.NumCtx
		s.loadRequest.BatchSize = batchSize
		slog.Info("embedding model detected, setting batch size to context length", "batch_size", batchSize)
	}

	kv, graphPartialOffload, graphFullOffload := s.ggml.GraphSize(uint64(s.options.NumCtx), uint64(batchSize),
		s.loadRequest.Parallel, s.loadRequest.KvCacheType, s.loadRequest.FlashAttention)

	// Use the size of one layer as a buffer
	layers := s.ggml.Tensors().GroupLayers()
	if blk0, ok := layers["blk.0"]; ok {
		for i := range gpus {
			gpus[i].FreeMemory -= blk0.Size() + kv[0]
		}
	} else {
		slog.Warn("model missing blk.0 layer size")
	}

	// Assign all the layers to the CPU for now, they will get reassigned later
	for i := range s.ggml.KV().BlockCount() {
		if blk, ok := layers[fmt.Sprintf("blk.%d", i)]; ok {
			s.mem.CPU.Weights[i] = blk.Size()
			s.mem.CPU.Cache[i] += kv[i]
		}
	}

	// We historically haven't included InputWeights in the model size

Review comment (performance):
Good catch that we weren't actually setting `loadRequest`. It probably would be better to just operate on that rather than introducing a new variable to avoid potential issues like this.

You could also check that the batch size is less than the context to avoid logging when the batch size is already sufficient, similar to the Ollama engine.

Code after:

func (s *llamaServer) Load(ctx context.Context, systemInfo ml.SystemInfo, systemGPUs []ml.DeviceInfo, requireFull bool) ([]ml.DeviceID, error) {
	slog.Info("loading model", "model layers", s.totalLayers, "requested", s.options.NumGPU)

	gpus := append(make([]ml.DeviceInfo, 0, len(systemGPUs)), systemGPUs...)

	// Synthesize memory allocation information based on our estimates
	s.mem = &ml.BackendMemory{CPU: ml.DeviceMemory{
		Name:    "CPU",
		Weights: make([]uint64, s.totalLayers),
		Cache:   make([]uint64, s.totalLayers),
	}, GPUs: make([]ml.DeviceMemory, len(gpus))}

	for i := range s.mem.GPUs {
		s.mem.GPUs[i].Name = gpus[i].Name
		s.mem.GPUs[i].DeviceID = gpus[i].DeviceID
		s.mem.GPUs[i].Weights = make([]uint64, s.totalLayers)
		s.mem.GPUs[i].Cache = make([]uint64, s.totalLayers)
	}

	// Check if embedding model and adjust batch size accordingly
	_, isEmbedding := s.ggml.KV()[fmt.Sprintf("%s.pooling_type", s.ggml.KV().Architecture())]
	if isEmbedding && s.loadRequest.BatchSize < s.options.NumCtx {
		s.loadRequest.BatchSize = s.options.NumCtx
		slog.Info("embedding model detected, setting batch size to context length", "batch_size", s.loadRequest.BatchSize)
	}

	kv, graphPartialOffload, graphFullOffload := s.ggml.GraphSize(uint64(s.options.NumCtx), uint64(s.loadRequest.BatchSize),
		s.loadRequest.Parallel, s.loadRequest.KvCacheType, s.loadRequest.FlashAttention)

	// Use the size of one layer as a buffer
	layers := s.ggml.Tensors().GroupLayers()
	if blk0, ok := layers["blk.0"]; ok {
		for i := range gpus {
			gpus[i].FreeMemory -= blk0.Size() + kv[0]
		}
	} else {
		slog.Warn("model missing blk.0 layer size")
	}

	// Assign all the layers to the CPU for now, they will get reassigned later
	for i := range s.ggml.KV().BlockCount() {
		if blk, ok := layers[fmt.Sprintf("blk.%d", i)]; ok {
			s.mem.CPU.Weights[i] = blk.Size()
			s.mem.CPU.Cache[i] += kv[i]
		}
	}

	// We historically haven't included InputWeights in the model size
	var outputWeights uint64
	if layer, ok := layers["output_norm"]; ok {