pytorch for gh200

config.yaml

@@ -178,6 +178,12 @@ uenvs:
       deploy:
         eiger: [zen2]
         todi: [gh200]
+  pytorch:
+    "24.7":
+      recipes:
+        gh200: '24.7/gh200'
+      deploy:
+        todi: [gh200]
   quantumespresso:
     "v7.1":
       recipes:

recipes/pytorch/24.7/gh200/compilers.yaml

@@ -0,0 +1,5 @@
+bootstrap:
+  spec: gcc@12
+gcc:
+  specs:
+  - [email protected]

recipes/pytorch/24.7/gh200/config.yaml

@@ -0,0 +1,7 @@
+name: pytorch
+store: /user-environment
+description: PyTorch for GH200
+spack:
+  repo: https://github.com/spack/spack.git
+  commit: releases/v0.22
+modules: true

recipes/pytorch/24.7/gh200/environments.yaml

@@ -0,0 +1,33 @@
+pytorch:
+  compiler:
+  - toolchain: gcc
+    spec: [email protected]
+  mpi:
+    spec: [email protected]
+    gpu: cuda
+  unify: true
+  specs:
+  - aws-ofi-nccl@master
+  - cmake
+  - cuda@12
+  - cray-mpich
+  - nccl
+  - ninja
+  - [email protected]
+  - [email protected]
+  - py-numpy
+  - py-tqdm
+  - py-lightning
+  - py-typing-extensions
+  - [email protected] +distributed +cuda cuda_arch=90 +cudnn +nccl +mpi
+  variants:
+  - +mpi
+  - +cuda
+  - cuda_arch=90
+  views:
+    pytorch:
+      link: roots
+      uenv:
+        add_compilers: true
+        prefix_paths:
+          LD_LIBRARY_PATH: [lib, lib64]

recipes/pytorch/24.7/gh200/modules.yaml

@@ -0,0 +1,23 @@
+modules:
+  # Paths to check when creating modules for all module sets
+  prefix_inspections:
+    bin:
+      - PATH
+    lib:
+      - LD_LIBRARY_PATH
+    lib64:
+      - LD_LIBRARY_PATH
+
+  default:
+    arch_folder: false
+    # Where to install modules
+    roots:
+      tcl: /user-environment/modules
+    tcl:
+      all:
+        autoload: none
+      hash_length: 0
+      exclude_implicits: true
+      exclude: ['%[email protected]', 'gcc %[email protected]']
+      projections:
+        all: '{name}/{version}'