Repository: oumi-ai/oumi
File: scripts/frontier/frontier_init.sh (Shell)

Code before:
FI_MR_CACHE_MONITOR=kdreg2     # Required to avoid a deadlock.
FI_CXI_DEFAULT_CQ_SIZE=131072  # Ask the network stack to allocate additional space to process message completions.
FI_CXI_DEFAULT_TX_SIZE=2048    # Ask the network stack to allocate additional space to hold pending outgoing messages.
FI_CXI_RX_MATCH_MODE=hybrid    # Allow the network stack to transition to software mode if necessary.

NCCL_NET_GDR_LEVEL=3           # Typically improves performance, but remove this setting if you encounter a hang/crash.
NCCL_CROSS_NIC=1               # On large systems, this NCCL setting has been found to improve performance
NCCL_SOCKET_IFNAME=hsn0        # NCCL/RCCL will use the high speed network to coordinate startup.
export NCCL_DEBUG=WARN # INFO
## export NCCL_DEBUG_SUBSYS=ALL

# Polaris has 32 "physical" CPU cores, and 64 "logical" cores per node
# (Hyper-threading makes 1 physical core appear as 2 logical cores)
# Physical cores: 0..31. Additional "logical" cores: 32..63.
# https://docs.alcf.anl.gov/polaris/hardware-overview/machine-overview/#polaris-device-affinity-information
NRANKS=1  # Number of MPI ranks to spawn per node (1 worker per node)
NDEPTH=64 # Number of hardware threads per rank (Frontier has 64 CPU physical cores per node, 56 available, 2X logical ones)
CPU_BIND="depth"

# Setup the environment variables.
export all_proxy=socks://proxy.ccs.ornl.gov:3128/
export ftp_proxy=ftp://proxy.ccs.ornl.gov:3128/
export http_proxy=http://proxy.ccs.ornl.gov:3128/
export https_proxy=http://proxy.ccs.ornl.gov:3128/
export no_proxy='localhost,127.0.0.0/8,*.ccs.ornl.gov'
export HF_HOME="/lustre/orion/lrn081/scratch/$USER/.cache/huggingface/slurm_job_${SLURM_JOBID}"
export HF_HUB_CACHE="$HF_HOME/hub"
export HF_ASSETS_CACHE="$HF_HOME/assets"

echo "Loading Frontier modules..."

# Set up default modules.
module load PrgEnv-gnu/8.6.0
module load miniforge3/23.11.0-0
module load rocm/6.2.4
module load craype-accel-amd-gfx90a

# Activate the Oumi Conda environment.
set +x
source activate "/lustre/orion/lrn081/scratch/$USER/miniconda3/envs/oumi"
echo "Conda path: ${CONDA_PREFIX}"
set -x

echo "frontier_init.sh is done!"


Review comment (performance):
Do we need to have a separate HF_HOME per job? Could we remove the /slurm_job_${SLURM_JOBID} and have a single cache across all jobs? Since it just holds content downloaded from HF_HUB, I can't think of a reason they need to be split

Code after:
FI_MR_CACHE_MONITOR=kdreg2     # Required to avoid a deadlock.
FI_CXI_DEFAULT_CQ_SIZE=131072  # Ask the network stack to allocate additional space to process message completions.
FI_CXI_DEFAULT_TX_SIZE=2048    # Ask the network stack to allocate additional space to hold pending outgoing messages.
FI_CXI_RX_MATCH_MODE=hybrid    # Allow the network stack to transition to software mode if necessary.
NCCL_NET_GDR_LEVEL=3           # Typically improves performance, but remove this setting if you encounter a hang/crash.
NCCL_CROSS_NIC=1               # On large systems, this NCCL setting has been found to improve performance
NCCL_SOCKET_IFNAME=hsn0        # NCCL/RCCL will use the high speed network to coordinate startup.
export NCCL_DEBUG=WARN # INFO
## export NCCL_DEBUG_SUBSYS=ALL

# Frontier has 64 "physical" CPU cores, and 128 "logical" cores per node
# (Hyper-threading makes 1 physical core appear as 2 logical cores)
# Physical cores: 0..63. Additional "logical" cores: 64..127.
# Note that 8 out of 64 cores are reserved for the system in the "low noise" mode,
# so there are only 56 allocatable cores by default:
# https://docs.olcf.ornl.gov/systems/frontier_user_guide.html#low-noise-mode-layout
NRANKS=1  # Number of MPI ranks to spawn per node (1 worker per node)
NDEPTH=64 # Number of hardware threads per rank (Frontier has 64 CPU physical cores per node, 56 available, 2X logical ones)
CPU_BIND="depth"

# Setup the environment variables.
export all_proxy=socks://proxy.ccs.ornl.gov:3128/
export ftp_proxy=ftp://proxy.ccs.ornl.gov:3128/
export http_proxy=http://proxy.ccs.ornl.gov:3128/
export https_proxy=http://proxy.ccs.ornl.gov:3128/
export no_proxy='localhost,127.0.0.0/8,*.ccs.ornl.gov'
# TODO: Append "/slurm_job_${SLURM_JOBID}" suffix in case of cache file lock contention.
export HF_HOME="/lustre/orion/lrn081/scratch/$USER/.cache/huggingface"
export HF_HUB_CACHE="$HF_HOME/hub"
export HF_ASSETS_CACHE="$HF_HOME/assets"

echo "Loading Frontier modules..."

# Set up default modules.
module load PrgEnv-gnu/8.6.0
module load miniforge3/23.11.0-0
module load rocm/6.2.4
module load craype-accel-amd-gfx90a

# Activate the Oumi Conda environment.
set +x
source activate "/lustre/orion/lrn081/scratch/$USER/miniconda3/envs/oumi"
echo "Conda path: ${CONDA_PREFIX}"
set -x

echo "frontier_init.sh is done!"