Ollama: Using lib.mk and different GPUs

nPrevail · June 3, 2024, 8:16pm

I understand that Ollama can use either AMD or NVIDIA gpus, and that it’s normally defined like so:

# enable Ollama
services.ollama = {
	enable = true;
	acceleration = "cuda"; # "cuda" or "rocm"
};

I also found this on the GitHub page, and I’m trying to make some sense out of it as well:

github.com

NixOS/nixpkgs/blob/d6f6eb2a984f2ba9a366c31e4d36d65465683450/nixos/modules/services/misc/ollama.nix

{ config, lib, pkgs, ... }:
let
  inherit (lib) types;

  cfg = config.services.ollama;
  ollamaPackage = cfg.package.override {
    inherit (cfg) acceleration;
    linuxPackages = config.boot.kernelPackages // {
      nvidia_x11 = config.hardware.nvidia.package;
    };
  };
in
{
  options = {
    services.ollama = {
      enable = lib.mkEnableOption "ollama server for local large language models";
      package = lib.mkPackageOption pkgs "ollama" { };
      home = lib.mkOption {
        type = types.str;
        default = "%S/ollama";

This file has been truncated. show original

acceleration = lib.mkOption {
        type = types.nullOr (types.enum [ false "rocm" "cuda" ]);
        default = null;
        example = "rocm";
        description = ''
          What interface to use for hardware acceleration.

          - `null`: default behavior
            if `nixpkgs.config.rocmSupport` is enabled, uses `"rocm"`
            if `nixpkgs.config.cudaSupport` is enabled, uses `"cuda"`
            otherwise defaults to `false`
          - `false`: disable GPU, only use CPU
          - `"rocm"`: supported by most modern AMD GPUs
          - `"cuda"`: supported by most modern NVIDIA GPUs
        '';
      };

But If I were using this drive on other devices with different GPUs (NVIDIA, or AMD, or vice versa), how can I define GPU “acceleration”? So for instance, could I use lib.mkForce as eitiher "true or false" to indicate which GPU to use?

How can I write this in a config.nix file?