add ML specific flags for running on pm-gpu

E3SM-Project · Jul 23, 2024 · 68fb14a · 68fb14a
1 parent d65ee59
commit 68fb14a
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/cime_config/machines/config_machines.xml b/cime_config/machines/config_machines.xml
@@ -391,6 +391,7 @@
       <modules compiler="gnugpu">
         <command name="load">cudatoolkit/12.2</command>
         <command name="load">craype-accel-nvidia80</command>
+        <command name="load">cudnn/8.9.3_cuda12</command>
       </modules>
 
       <modules compiler="nvidiagpu">
@@ -435,6 +436,8 @@
       <env name="NETCDF_PATH">$ENV{CRAY_NETCDF_HDF5PARALLEL_PREFIX}</env>
       <env name="PNETCDF_PATH">$ENV{CRAY_PARALLEL_NETCDF_PREFIX}</env>
       <env name="TF_FORCE_GPU_ALLOW_GROWTH">true</env>
+      <env name="XLA_FLAGS">--xla_gpu_cuda_data_dir=/opt/nvidia/hpc_sdk/Linux_x86_64/23.9/cuda/12.2</env>
+      <env name="TF_USE_LEGACY_KERAS">1</env>
     </environment_variables>
     <environment_variables compiler="gnugpu">
       <env name="MPICH_GPU_SUPPORT_ENABLED">1</env>

diff --git a/components/eamxx/cmake/machine-files/pm-gpu.cmake b/components/eamxx/cmake/machine-files/pm-gpu.cmake
@@ -18,10 +18,8 @@ endif()
 
 include (${EKAT_MACH_FILES_PATH}/mpi/srun.cmake)
 
-set(PYBIND11_PYTHON_VERSION 3.8 CACHE STRING "")
-set(Python_EXECUTABLE /pscratch/sd/a/andrep/conda_venvs/fv3net-local/bin/python CACHE STRING "" FORCE)
-set(PYTHON_EXECUTABLE /pscratch/sd/a/andrep/conda_venvs/fv3net-local/bin/python CACHE STRING "" FORCE)
-set(Python_INCLUDE_DIRS /pscratch/sd/a/andrep/conda_venvs/fv3net-local/include/python3.9 CACHE STRING "" FORCE)
+set(PYBIND11_PYTHON_VERSION 3.9 CACHE STRING "")
+execute_process(COMMAND conda activate /global/common/software/m4492/fv3net-shared-py39)
 option (SCREAM_ENABLE_ML_CORRECTION "Whether to enable ML correction parametrization" ON)
 set(EKAT_MPI_EXTRA_ARGS "${EKAT_MPI_EXTRA_ARGS} --gpus-per-task=1" CACHE STRING "" FORCE)