How to properly setup an environment with Pytorch?

Help
1

Hi, I only switched to nix recently and I’ve been attempting to setup an environment where I can do development in pytorch. I found some previous threads on the subject here, but none of them seemed to have a full example of a flake that you can nix develop and with my limited skills I haven’t yet been able to patch together something that works. I would strongly prefer to use the cuda-maintainers cachix where possible, as every time I have tried to compile torch it times out. I’d also prefer to do things the nix way wherever possible, but ultimately I just want something that will work. (And by work I mean something that at least lets you run the pytorch beginner example with cuda)

Here is the flake I have come up with so far:

{
  description = "Pytorch development environment";

  nixConfig = {
    extra-substituters = [
      "https://cuda-maintainers.cachix.org"
    ];
    extra-trusted-public-keys = [
      "cuda-maintainers.cachix.org-1:0dq3bujKpuEPMCX6U4WylrUDZ9JyUG0VpVZa7CNfq5E="
    ];
  };

  inputs.nixpkgs.url = "https://flakehub.com/f/NixOS/nixpkgs/0";

  outputs = {
    self,
    nixpkgs,
  }: {
    devShells.x86_64-linux.default = let
      pkgs = import nixpkgs {
        system = "x86_64-linux";
        allowUnfree = true;
        cudaSupport = true;
      };
    in
      pkgs.mkShell {
        packages = with pkgs; [
          (python313.withPackages (p:
            with p; [
              # python packages
              torchWithCuda
              torchvision
            ]))
          # non-python packages
        ];
      };
  };
}

This runs into the issue error: Package ‘cuda_nvcc-12.8.93’ in /nix/store/x06djk7jw02qb6y8fxd0mjgaksw2kd9n-source/pkgs/development/cuda-modules/generic-builders/manifest.nix:323 has an unfree license (‘CUDA EULA’), refusing to evaluate. Which is confusing to me since I already specify allowUnfree = true; in the nixpkgs config. Is the config somehow not visible? Any help would be appreciated.

2

I think it’s config.allowUnfree

import nixpkgs { system = "x86_64-linux"; config.allowUnfree = true; };

see here Allow unfree in flakes

there’s a way via the CLI but you’ll have to pass --impure along with it I think.

3

Thanks for pointing that out, I feel dumb for not noticing that sooner. I think this is almost working now, the only problem is that its not actually using the cached packages. It still is doing hours of compilation when I run it which my machine is really able to handle - it hasn’t successfully finished compilation yet in any of the times I have tried it. Do you know what I could be doing wrong there?

4

CUDA packages are not cached in the main cache.nixos.org due to the licensing situation. I think there’s a community cache available right now, and a new initiative happening to improve things but I’m not too sure on the specifics. There was a recent announcement. Nix, Flox, Nvidia Opening Up CUDA Redistribution on Nix

1 Like
5

Thanks @Lun for linking that info! I was able to get this to work using the flox cache (I think the other caches are being phased out now that an officially sponsored solution exists). Here is the flake I arrived at:

{
  description = "Pytorch development environment";

  nixConfig = {
    extra-substituters = [
      "https://cache.flox.dev"
    ];
    extra-trusted-public-keys = [
      "flox-cache-public-1:7F4OyH7ZCnFhcze3fJdfyXYLQw/aV7GEed86nQ7IsOs="
    ];
  };

  inputs.nixpkgs.url = "github:flox/nixpkgs";

  outputs = {
    self,
    nixpkgs,
  }: {
    devShells.x86_64-linux.default = let
      pkgs = import nixpkgs {
        system = "x86_64-linux";
        config.allowUnfree = true;
        config.cudaSupport = true;
      };
    in
      pkgs.mkShell {
        packages = with pkgs; [
          (python313.withPackages (p:
            with p; [
              # python packages
              torch
              torchvision
            ]))
          # non-python packages
        ];
      };
  };
}