Using DFTK on compute clusters

This chapter summarises a few tips and tricks for running DFTK on compute clusters. It assumes you have already installed Julia on the machine in question (see Julia downloads and Julia installation instructions). For a general presentation of using Julia on HPC systems from both the user and the admin perspective, see JuliaOnHPCClusters.

In this presentation we use EPFL's scitas clusters as the example to explain the basic principles.

Julia depot path

By default on Linux-based systems Julia puts all installed packages, including the binary packages into the path $HOME/.julia, which can easily take a few tens of GB. On many compute clusters the /home partition is on a shared filesystem (thus access is slower) and has a tight space quota. Usually it is therefore more advantageous to put Julia packages on a less persistent, faster filesystem. On many systems (such as EPFL scitas) this is the /scratch partition. In your ~/.bashrc or otherwise you should thus redirect the JULIA_DEPOT_PATH to be a subdirectory of /scratch.

EPFL scitas. On scitas the right thing to do is to insert

export JULIA_DEPOT_PATH="/scratch/$USER/.julia"

into your ~/.bashrc.

Installing DFTK into a local Julia environment

When employing a compute cluster, it is often desirable to integrate with the matching cluster-specific libraries, such as vendor-specific versions of BLAS, LAPACK etc. This is easiest achieved by installing DFTK not into the global Julia environment, but instead bundle the installation and the cluster-specific configuration in a local, cluster-specific environment.

On top of the discussion of the main Installation instructions, this requires one additional step, namely the use of a custom Julia package environment.

Setting up a package environment. In the Julia REPL, create a new environment (essentially a new Julia project) using the following commands:

import Pkg
Pkg.activate("path/to/new/environment")

Replace path/to/new/environment with the directory where you wish to create the new environment. This will generate a folder containing Project.toml and Manifest.toml files. Both together provide a reproducible image of the packages you installed in your project. Once the activate call has been issued, you can than install packages as usual. E.g. use Pkg.add("DFTK") to install DFTK. The difference is that instead of tracking this in your global environment, the installations will be tracked in the local Project.toml and Manifest.toml files of the path/to/new/environment folder.

To start a Julia shell directly with such this environment activated, run it as julia --project=path/to/new/environment.

Updating an environment. Start Julia and activate the environment. Then run Pkg.update(), i.e.

import Pkg
Pkg.activate("path/to/new/environment")
Pkg.update()

For more information on Julia environments, see the respective documentation on Code loading and on Managing Environments.

Setting up local preferences

On cluster machines often highly optimised versions of BLAS, LAPACK, FFTW, MPI or other basic packages are provided by the cluster vendor or operator. These can be used with DFTK by configuring an appropriate LocalPreferences.toml file, which tells Julia where the cluster-specific libraries are located. The following sections explain how to generate such a LocalPreferences.toml specific for your cluster. Once this file has been generated and sits next to a Project.toml in a Julia environment, it ensures that the cluster-specific libraries are used instead of the default ones. Therefore this setup only needs to be done once per project.

A useful way to check whether the setup has been successful and DFTK indeed employs the desired cluster-specific libraries provides the

DFTK.versioninfo()

command. It produces an output such as

DFTK Version      0.6.16
Julia Version     1.10.0
FFTW.jl provider  fftw v3.3.10

BLAS.get_config()
  LinearAlgebra.BLAS.LBTConfig
  Libraries:
  └ [ILP64] libopenblas64_.so

MPI.versioninfo()
  MPIPreferences:
    binary:  MPICH_jll
    abi:     MPICH

  Package versions
    MPI.jl:             0.20.19
    MPIPreferences.jl:  0.1.10
    MPICH_jll:          4.1.2+1

  Library information:
    libmpi:  /home/mfh/.julia/artifacts/0ed4137b58af5c5e3797cb0c400e60ed7c308bae/lib/libmpi.so
    libmpi dlpath:  /home/mfh/.julia/artifacts/0ed4137b58af5c5e3797cb0c400e60ed7c308bae/lib/libmpi.so

[...]

which thus specifies in one overview the details of the employed BLAS, LAPACK, FFTW, MPI, etc. libraries.

Switching to MKL for BLAS and FFT

The MKL Julia package provides the Intel MKL library as a BLAS backend in Julia. To use fully use it you need to do two things:

Add a using MKL to your scripts.
Configure a LocalPreferences.jl to employ the MKL also for FFT operations: to do so run the following Julia script:
```
using MKL
using FFTW
FFTW.set_provider!("mkl")
```

Switching to the system-provided MPI library

To use a system-provided MPI library, load the required modules. On scitas that is

module load gcc
module load openmpi

Afterwards follow the MPI system binary instructions and execute

using MPIPreferences
MPIPreferences.use_system_binary()

Using the `--heap-size-hint` flag

Julia uses dynamic memory management, which means that unused memory may not be directly released back to the operating system. Much rather a garbage collection takes care of doing such cleanup periodically. If you are in a memory-bound situation, it can thus be helpful to employ the --heap-size-hint flag, which provides Julia with a hint of the maximal memory Julia may use. For example the call

julia --heap-size-hint 40G

tells Julia to optimise the garbage collection, such that the overall memory consumption remains below 40G. Note, however that this is just a hint, i.e. no hard limits are enforced. Furthermore using too small a heap size hint can have a negative impact on performance.

Estimating and monitoring memory usage

DFTK provides some routines for estimating the memory consumption for storing density, Bloch waves as well as the peak memory consumption during an SCF run. Based on a model this information can be obtained before constructing a basis as

estimate = estimate_memory_usage(model::Model; Ecut=15, kgrid=(1, 1, 1), other_basis_kwargs...)

which takes exactly the same kind of arguments and keyword arguments as a PlaneWaveBasis constructor. The estimate object provides the estimated memory consumption, which in a script can be pretty-printed to stdout as follows

show(stdout, "text/plain", estimate)

Estimated memory usage (per MPI process):
    nonlocal projectors  :   6.7 GiB
    single ψ             : 915.2 MiB
    single ρ             :  25.0 MiB
    peak during SCF      :  14.1 GiB

Notice, that in particular the peak memory during SCF provides an orientation for requesting resources from queuing systems (but at this stage this estimate is still very rough). The same information is also available in the basis printout (i.e. show(stdout, "text/plain", basis)).

During an SCF run the current memory consumption (within the main RAM as well as a potential GPU device) can be monitored by passing the callback ScfDefaultCallback(show_memory=true), see Monitoring self-consistent field calculations.

Running slurm jobs

This example shows how to run a DFTK calculation on a slurm-based system such as scitas. We use the MKL for FFTW and BLAS and the system-provided MPI. This setup will create five files LocalPreferences.toml, Project.toml, dftk.jl, silicon.extxyz and job.sh.

At the time of writing (Dec 2023) following the setup indicated above leads to this LocalPreferences.toml file:

[FFTW]
provider = "mkl"

[MPIPreferences]
__clear__ = ["preloads_env_switch"]
_format = "1.0"
abi = "OpenMPI"
binary = "system"
cclibs = []
libmpi = "libmpi"
mpiexec = "mpiexec"
preloads = []

We place this into a folder next to a Project.toml to define our project:

[deps]
AtomsIO = "1692102d-eeb4-4df9-807b-c9517f998d44"
DFTK = "acf6eb54-70d9-11e9-0013-234b7a5f5337"
FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
MKL = "33e6dc65-8f57-5167-99aa-e5a354878fb2"
MPIPreferences = "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267"
PseudoPotentialData = "5751a51d-ac76-4487-a056-413ecf6fbe19"

We additionally create a small file dftk.jl to run an MPI-parallelised calculation from a passed structure

using MKL
using DFTK
using AtomsIO
using PseudoPotentialData

disable_threading()  # Threading and MPI not compatible

function main(structure; pseudofamily, Ecut, kspacing)
    if mpi_master()
        println("DEPOT_PATH=$DEPOT_PATH")
        println(DFTK.versioninfo())
        println()
    end

    model  = model_DFT(load_system(structure);
                       pseudopotentials=PseudoFamily(pseudofamily),
                       functionals=PBE(),
                       temperature=1e-3,
                       smearing=Smearing.MarzariVanderbilt())
    kgrid = kgrid_from_minimal_spacing(model, kspacing)
    basis = PlaneWaveBasis(model; Ecut, kgrid)

    if mpi_master()
        println()
        show(stdout, MIME("text/plain"), basis)
        println()
        flush(stdout)
    end

    DFTK.reset_timer!(DFTK.timer)
    scfres = self_consistent_field(basis)
    println(DFTK.timer)

    if mpi_master()
        show(stdout, MIME("text/plain"), scfres.energies)
    end
end

and we dump the structure file silicon.extxyz with content

2
pbc=[T, T, T] Lattice="0.00000000 2.71467909 2.71467909 2.71467909 0.00000000 2.71467909 2.71467909 2.71467909 0.00000000" Properties=species:S:1:pos:R:3
Si         0.67866977       0.67866977       0.67866977
Si        -0.67866977      -0.67866977      -0.67866977

Finally the jobscript job.sh for a slurm job looks like this:

#!/bin/bash
# This is the block interpreted by slurm and used as a Job submission script.
# The actual julia payload is in dftk.jl
# In this case it sets the parameters for running with 2 MPI processes using a maximal
# wall time of 2 hours.

#SBATCH --time 2:00:00
#SBATCH --nodes 1
#SBATCH --ntasks 2
#SBATCH --cpus-per-task 1

# Now comes the setup of the environment on the cluster node (the "jobscript")

# IMPORTANT: Set Julia's depot path to be a local scratch
# direction (not your home !).
export JULIA_DEPOT_PATH="$JULIA_DEPOT_PATH:/scratch/$USER/.julia"

# Load modules to setup julia (this is specific to EPFL scitas systems)
module purge
module load gcc
module load openmpi
module load julia

# Run the actual payload of Julia using 1 thread. Use the name of this
# file as an include to make everything self-contained.
srun julia -t 1 --project -e '
    include("dftk.jl")
    main("silicon.extxyz";
         pseudofamily="pd_nc_sr_pbe_standard_0.4.1_upf",
         Ecut=10,
         kspacing=0.3)
'

Running DFTK with MPI and GPUs

Note that DFTK is implemented for GPU-aware MPI. If your cluster does not support such an MPI implementation, keep to CPU execution. This tutorial is illustrated with examples from a NVIDIA cluster using CUDA.jl, but adaptation to AMD GPUs is straightforward. We recommend using a manual download of Julia for finer control.

The easiest way forward is the creation of a Julia environment, with a LocalPreferences.toml file that defines where the GPU-aware MPI library is found on the system, as well as CUDA details. On AMD clusters, only the MPIPreferences are necessary.

# This preference can be removed on AMD GPU clusters
[CUDA_Runtime_jll]
__clear__ = ["local"]
version = "13.1" # CUDA version number should match that of the cluster
cuda_local = false # Use CUDA from a JLL download (see why below)

[MPIPreferences]
__clear__ = ["preloads_env_switch"]
_format = "1.0"
abi = "MPICH"
binary = "system"
cclibs = []
# It is recommended to use an absolute path here
libmpi = "/user-environment/linux-neoverse_v2/cray-mpich-9.1.1-tms3eeelzg36j7u2rqjsxz7pnz5if5bh/lib/libmpi"
mpiexec = ["srun"] # srun for clusters using Slurm
preloads = []

Depending on the cluster, various modules/uenv might need to be loaded to access GPU-aware MPI. For minimal calculations, the following Project.toml is sufficient:

[deps]
#AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" # Replace with AMDGPU if needed
DFTK = "acf6eb54-70d9-11e9-0013-234b7a5f5337"
MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
MPIPreferences = "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267"
PseudoPotentialData = "5751a51d-ac76-4487-a056-413ecf6fbe19"

[extras]
CUDA_Runtime_jll = "76a88914-d11a-5bdc-97e0-2f5a05c973a2" # Not necessary on AMD

Note that before launching a MPI parallel calculation, it is very important to pre-compile DFTK. When pre-compilation takes place on multiple ranks at once, race conditions occur and compilation stalls. Make sure that pre-compilation happens under the same conditions as the subsequent calculation (same module loaded, same Julia environment, etc.), so as to avoid invalidations down the line. This can be done with a single MPI rank job, or in an interactive session. On NVIDIA clusters, it is also important to set a persistent cache for CUDA, e.g. 'export CUDACACHEPATH="/someexistingdirectoryonscratch"' (see below for the reason).

After the code has been pre-compiled, it can be launched normally. Below is an illustrative Slurm submission script:

#!/bin/bash -l

#SBATCH --job-name=DFTK_calc
#SBATCH --time=00:30:00
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=4
#SBATCH --account=<account_id>
#SBATCH --hint=nomultithread
#SBATCH --hint=exclusive

### On this cluster, GPU-aware MPI is obtained via a uenv. Relevent 
### modules might need loading elsewhere.
#SBATCH --uenv=prgenv-gnu/26.3:v1

# Specific cray-mpich env variables
export MPICH_GPU_SUPPORT_ENABLED=1
export MPICH_GPU_IPC_ENABLED=0

# Add Julia exec to the PATH, and set JULIA_DEPOT_PATH to a place on $SCRATCH
export PATH="/path_to_julia_installation_dir/julia-1.12.6/bin:$PATH"
export JULIA_DEPOT_PATH="/path_to_julia_depot_on_scratch/.julia"

# Directory for cached JIT compiled Libxc GPU code (only for NVIDIA)
export CUDA_CACHE_PATH="/path_to_existing_folder_on_scratch"

my_project="/absolute_path_to_julia_project"

srun julia --project=$my_project --compiled-modules=strict example.jl > example.out

Special notes on NVIDIA clusters

Running on NVIDIA GPU clusters is trickier than on AMD because of how Libxc.jl works. With AMD GPUs, the data is always transferred to the host before functionals are evaluated, and the results are copied to the GPU once the operation has finished. With CUDA, the data stays on the device, and functionals are evaluated on the GPU via a binary shipped by LibxcGPUjll. Due to the inner workings of CUDA.jl and BinaryBuilder.jl, the CUDA compatible Libxc JLL is only available when the CUDA toolkit is downloaded through Julia (hence cuda_local = false in the preferences). However, the cluster provided GPU-aware MPI library would have been built with a different CUDA installation. As a result, MPI and Julia might link against different CUDA libraries, which will generate a bunch of warnings. These can be safely ignored if both CUDA installation used the same major version, hence the importance of properly setting version for [CUDA_Runtime_jll] in the preferences.

Defining the CUDA_CACHE_PATH environment variable is very important to reduce calculation starting times. The binaries distributed by LibxcGPUjll are so-called fat binaries, made of generic PTX code, not specialized for specific device architecture (i.e. sm70, sm80, etc.). When the first calculation is underway, the Libxc GPU code is JIT compiled and cached, an operation that can take up to 5 minutes. Because this JIT compilation is not performed by Julia, it is cached in a different location, typically in a temporary folder local to the compute node. As a result, different calculations running on different nodes might not have access to the cache, and the lengthy JIT compilation keeps happening. Setting a CUDA_CACHE_PATH to a persistent, existing, directory that all nodes can access solves that issue.

Using DFTK via the Aiida workflow engine

A preliminary integration of DFTK with the Aiida high-throughput workflow engine is available via the aiida-dftk plugin. This can be a useful alternative if many similar DFTK calculations should be run.