Pytorch with cuda support

Help
1

Hi everyone,

I’m new to NixOS and recently set up a PC to develop ML models using PyTorch with CUDA support. I’ve been experimenting with a few nix-shell configurations, but I keep running into issues with long build times, especially with compiling CUDA and other dependencies.

Here are the configurations I’ve tried:

Configuration 1:

let
  # Pin to a specific nixpkgs commit for reproducibility.
  pkgs = import (fetchTarball "https://github.com/NixOS/nixpkgs/archive/24bb1b20a9a57175965c0a9fb9533e00e370c88b.tar.gz") {config.allowUnfree = true;};
in pkgs.mkShell {
  buildInputs = [
    pkgs.python311
    pkgs.python311Packages.torch-bin
    pkgs.python311Packages.unidecode
    pkgs.python311Packages.inflect
    pkgs.python311Packages.librosa
    pkgs.python311Packages.pip
  ];

  shellHook = ''
    echo "You are now using a NIX environment"
    export CUDA_PATH=${pkgs.cudatoolkit}
  '';
}

Configuration 2:

nix-shell --arg config '{ allowUnfree = true; }' -p 'python311.withPackages (ps: with ps; [torchWithCuda])' -I nixpkgs=https://github.com/NixOS/nixpkgs/archive/26c6b9f5394b2e64649ff18394d61ce5594874e2.tar.gz’

Configuration 3:


{ pkgs ? import <nixpkgs> {
  config = {
    allowUnfree = true;
    cudaSupport = true;
  };
} }:
  pkgs.mkShell {
    # nativeBuildInputs is usually what you want -- tools you need to run
    nativeBuildInputs = with pkgs.buildPackages; [
      python311
      cudaPackages_12.cudatoolkit
      python311Packages.pytorch-bin
      python311Packages.pip
    ];

    shellHook = ''
      echo "You are now using a NIX environment"
      export CUDA_PATH=${pkgs.cudatoolkit}
    '';
}

Issue: In all cases, I end up with long build times because of the compilation of CUDA, PyTorch, and other dependencies. The builds can take over an hour, which is quite frustrating.

Question: Is there a way to set up a development environment for PyTorch with CUDA on NixOS without having to compile everything from scratch? Are there any pre-built binaries or optimized configurations that I can use to speed up this process? ( I have tried torch-bin, as suggested here, but it still was compiling everything in the same way )

I appreciate any advice or tips you can provide!

Thanks in advance!

2

Some additional info :

nix --version
nix (Nix) 2.18.5

systemStateVersion : 24.05

nvidia-smi output:

Mon Aug 26 13:55:12 2024
±----------------------------------------------------------------------------------------+
| NVIDIA-SMI 550.78 Driver Version: 550.78 CUDA Version: 12.4 |
|-----------------------------------------±-----------------------±---------------------+
| GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|=========================================+========================+======================|
| 0 NVIDIA GeForce RTX 3090 Off | 00000000:01:00.0 On | N/A |
| 0% 37C P8 26W / 370W | 417MiB / 24576MiB | 4% Default |
| | | N/A |
±----------------------------------------±-----------------------±---------------------+

±----------------------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=========================================================================================|
| 0 N/A N/A 1345 G …mzcsvfvcv-xorg-server-21.1.13/bin/X 162MiB |
| 0 N/A N/A 1716 G …irefox-128.0.3/bin/.firefox-wrapped 136MiB |
| 0 N/A N/A 2068 G …erProcess --variations-seed-version 86MiB |
| 0 N/A N/A 16208 G …/profiles/per-user/nazara/bin/kitty 16MiB |
±----------------------------------------------------------------------------------------+

Nvidia driver:


  services.xserver.videoDrivers = ["nvidia"];

  hardware.nvidia = {


    modesetting.enable = true;


    powerManagement.enable = false;


    powerManagement.finegrained = false;

   
    open = false;


    nvidiaSettings = true;


    package = config.boot.kernelPackages.nvidiaPackages.stable;
  };
3

I found a solution for this issue. Here’s how I addressed it:

  1. Cachix Setup:
  • I used Cachix not to compile/build CUDA.
  • Also, I used package pytorch-bin in my nix-shell file.
  1. Nix-Shell Configuration:
  • I did not use cudaSupport = true; in the nix-shell file. Including this option would compile CUDA, which is unnecessary in this case.
  • The same issue occurred with the torchWithCuda package (or similar).
  1. Cachix Installation:
  • When installing, Cachix will generate a new caching.nix file and a caching directory in /etc/nixos. Make sure to move these to your dotfiles.
  • After installing Cachix, you’ll need to import it into your configuration.nix file and do nixos-rebuild.

After all this steps I successfully was able to run nix-shell shell.nix and had cuda available in pytorch

Hope this helps someone, I might also post here my shell.nix later on.

Useful links:
https://nixos.wiki/wiki/CUDA

https://app.cachix.org/cache/cuda-maintainers

1 Like
4

Here is the nix-shell that I used


let
  # Pin to a specific nixpkgs commit for reproducibility.
  pkgs = import (fetchTarball "https://github.com/NixOS/nixpkgs/archive/24bb1b20a9a57175965c0a9fb9533e00e370c88b.tar.gz") {config.allowUnfree = true; };
in pkgs.mkShell {
  nativeBuildInputs = [
    pkgs.python311
    pkgs.python311Packages.torch-bin
    pkgs.python311Packages.torchaudio-bin
    pkgs.python311Packages.torch-audiomentations
    pkgs.python311Packages.librosa
    pkgs.python311Packages.jiwer
    pkgs.python311Packages.datasets
    pkgs.python311Packages.transformers
    pkgs.python311Packages.evaluate
    pkgs.python311Packages.accelerate
    pkgs.python311Packages.pip
    
  ];

  shellHook = ''
    echo "You are now using a NIX environment"
    export CUDA_PATH=${pkgs.cudatoolkit}
    echo $CUDA_PATH
  '';
}