python3Packages.pytorch: migrate to cudaPackages

[SomeoneSerge]

Use the redistributable cuda packages, instead of runfile-based
cudatoolkit

Build log: https://gist.github.com/SomeoneSerge/4e5b5e9d9ee82c410eebad60d0cdc8f7

All went awry, it used cudatoolkit through cudnn. Reworking

EDIT later on 2022-04-15: Now it actually builds with cuda-redist! Rebased on the cudnn PR with removed propagatedBuildInputs, adjusted, and pushed

CC @NixOS/cuda-maintainers

Description of changes

Things done

• Built on platform(s)
  • x86_64-linux
  • aarch64-linux
  • x86_64-darwin
  • aarch64-darwin
• For non-Linux: Is sandbox = true set in nix.conf? (See Nix manual)
• Tested, as applicable:
  • NixOS test(s) (look inside nixos/tests)
  • and/or package tests
  • or, for functions and "core" functionality, tests in lib/tests or pkgs/test
  • made sure NixOS tests are linked to the relevant packages
• Tested compilation of all packages that depend on this change using nix-shell -p nixpkgs-review --run "nixpkgs-review rev HEAD". Note: all changes have to be committed, also see nixpkgs-review usage
• Tested basic functionality of all binary files (usually in ./result/bin/)
• 22.05 Release Notes (or backporting 21.11 Release notes)
  • (Package updates) Added a release notes entry if the change is major or breaking
  • (Module updates) Added a release notes entry if the change is significant
  • (Module addition) Added a release notes entry if adding a new NixOS module
  • (Release notes changes) Ran nixos/doc/manual/md-to-db.sh to update generated release notes
• Fits CONTRIBUTING.md.

[samuela]

hell yeah! this looks dope. does it build and everything?

[FRidh]

cudnn using libcublas instead of cudatoolkit if possible #168755

[SomeoneSerge]

Hooray!

[SomeoneSerge]

@FRidh same point arises: this assumes redist packages are present unconditionally, i.e. it only works for cuda>=11. I think we don't have to support cuda10 in downstream derivations

[SomeoneSerge]

I'm not sure I'm ready for this to be merged:

• I want to ensure that everything that is meant to support GPU still supports it. There are e.g. some suspicious new lines in the build log that probably ought be investigated: https://gist.github.com/SomeoneSerge/3b5d7e2e3af5318f4ef6a380f3ef5c35#file-pytorch-redist-log-diff-L160
• I maybe should rework magma too, otherwise between the PRs our pytorch will be including both distributions of cuda in the closure, which is not nice...

[SomeoneSerge]

Now magma uses redist packages too, reducing the reported closure size of pytorch from 12.75GiB to 10.15GiB (which is still too much, imo)
Kineto/cupti issue resolved just with adding libcupti to the closure

[SomeoneSerge]

...note that this change doesn't support building magma with cuda10 (pre-cuda11.4)

We could account for those as well, but I thought it's ok to ignore because magma is outside cudaPackages namespace

[SomeoneSerge]

Interesting that magma has platforms = unix, but neither cudatoolkit nor redist cuda packages are marked with darwin support

[samuela]

Interesting that magma has platforms = unix, but neither cudatoolkit nor redist cuda packages are marked with darwin support

Although I think it may technically be possible to hook up NVIDIA hardware to a macOS system, I've never heard of anyone actually doing this... except maybe for kicks?

[SomeoneSerge]

RE: magma.cudaPackages == cudaPackages

I've actually just noticed a bug with the original expression.
Evaluating passthru attributes of pytorchWithCuda on master¹ triggers evaluation of cudatoolkit, resulting in an error:

❯ nix eval --impure --expr '((builtins.getFlake github:NixOS/nixpkgs/master).legacyPackages.x86_64-linux.python3Packages.pytorch).cudaSupport'
false
❯ nix eval --impure --expr '((builtins.getFlake github:NixOS/nixpkgs/master).legacyPackages.x86_64-linux.python3Packages.pytorchWithCuda).cudaSupport'
error: Package ‘cudatoolkit-11.6.1’ ... has an unfree license (‘unfree’), refusing to evaluate

This happens because of the following line in master:

assert !cudaSupport || magma.cudatoolkit == cudatoolkit;

So, just to verify, in this PR we don't seem to be repeating that same error, although the whole thing kind of smells

❯ nix eval --impure --expr '((builtins.getFlake (toString ./.)).legacyPackages.x86_64-linux.python3Packages.pytorch).cudaSupport'
false
❯ nix eval --impure --expr '((builtins.getFlake (toString ./.)).legacyPackages.x86_64-linux.python3Packages.pytorchWithCuda).cudaSupport'
true

Footnotes

1. master is at 97eb40e0f8a9ec4e7306645dfc446f3cb4b9eed4 at the moment ↩

[SomeoneSerge]

I suspect we might still run into similar evaluation issues if we make an overlay resulting in lack of pointer-equality 🤔

[SomeoneSerge]

pytorch = callPackage ../development/python-modules/pytorch {
    # ...
    cudaPackages = pkgs.cudaPackages.overrideScope' (final: prev: { });
}

Then pytorchWithCuda.cudaSupport results in

error: assertion '(cudaSupport -> ((magma).cudaPackages == cudaPackages))' failed

...and pytorch.cudaSupport evaluates into false

[FRidh]

Might want to write a function that takes a list of sets and a list of packages and tests for every package whether it is equal in every set.

[SomeoneSerge]

Just one thing I don't like about this is that the user cannot disable those asserts downstream, but I can live with that

[SomeoneSerge]

@FRidh

Might want to write a function that takes a list of sets and a list of packages and tests for every package whether it is equal in every set.

1. I have a feeling that referential equality is somehow more appropriate? Except we're relying on nix's internals
2. If we compare the packages, we'd be at risk of getting that refusing to evaluate error again 🤔

[SomeoneSerge]

We definitely should write our own equality function just for isolation, regardless of whether it falls back to == or tests only packages. This is out of scope of the PR, however

[SomeoneSerge]

^

[ConnorBaker]

@SomeoneSerge if you don't mind, I'd like to take this up and continue where you left off.

[SomeoneSerge]

@ConnorBaker this would be just great, please do!