Skip to content

Update gpu affinity on pm-gpu #7818

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: master
Choose a base branch
from
Open

Update gpu affinity on pm-gpu #7818

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6e6daee
Add mpi_task to gpu affinity script
amametjanov Oct 21, 2025
162842f
Add S/M/L pe-layouts for ne256-wcyclxx on pm-gpu
amametjanov Oct 21, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
76 changes: 76 additions & 0 deletions cime_config/allactive/config_pesall.xml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2548,4 +2548,80 @@
</pes>
</mach>
</grid>
<grid name="a%ne256np4.pg2.*">
<mach name="pm-gpu">
<pes compset=".*SCREAM.*ELM.*MPASSI.*MPASO.*" pesize="S">
<comment>pm-gpu ne256 for fully coupled cases with ATM on GPU, MPAS on CPU -- WCYCLXX2010 64 nodes, 64x1</comment>
<MAX_MPITASKS_PER_NODE>64</MAX_MPITASKS_PER_NODE>
<ntasks>
<ntasks_atm>256</ntasks_atm>
<ntasks_cpl>-64</ntasks_cpl>
<ntasks_lnd>-64</ntasks_lnd>
<ntasks_rof>-64</ntasks_rof>
<ntasks_ice>-64</ntasks_ice>
<ntasks_ocn>-64</ntasks_ocn>
</ntasks>
<rootpe>
<rootpe_glc>1</rootpe_glc>
<rootpe_wav>1</rootpe_wav>
<rootpe_esp>1</rootpe_esp>
<rootpe_iac>1</rootpe_iac>
</rootpe>
<pstrid>
<pstrid_atm>16</pstrid_atm>
</pstrid>
<excl_stride>
<excl_stride_atm>16</excl_stride_atm>
</excl_stride>
</pes>
<pes compset=".*SCREAM.*ELM.*MPASSI.*MPASO.*" pesize="M">
<comment>pm-gpu ne256 for fully coupled cases with ATM on GPU, MPAS on CPU -- WCYCLXX2010 128 nodes, 64x1</comment>
<MAX_MPITASKS_PER_NODE>64</MAX_MPITASKS_PER_NODE>
<ntasks>
<ntasks_atm>512</ntasks_atm>
<ntasks_cpl>-128</ntasks_cpl>
<ntasks_lnd>-128</ntasks_lnd>
<ntasks_rof>-128</ntasks_rof>
<ntasks_ice>-128</ntasks_ice>
<ntasks_ocn>-128</ntasks_ocn>
</ntasks>
<rootpe>
<rootpe_glc>1</rootpe_glc>
<rootpe_wav>1</rootpe_wav>
<rootpe_esp>1</rootpe_esp>
<rootpe_iac>1</rootpe_iac>
</rootpe>
<pstrid>
<pstrid_atm>16</pstrid_atm>
</pstrid>
<excl_stride>
<excl_stride_atm>16</excl_stride_atm>
</excl_stride>
</pes>
<pes compset=".*SCREAM.*ELM.*MPASSI.*MPASO.*" pesize="M">
<comment>pm-gpu ne256 for fully coupled cases with ATM on GPU, MPAS on CPU -- WCYCLXX2010 128 nodes, 64x1</comment>
<MAX_MPITASKS_PER_NODE>64</MAX_MPITASKS_PER_NODE>
<ntasks>
<ntasks_atm>1024</ntasks_atm>
<ntasks_cpl>-256</ntasks_cpl>
<ntasks_lnd>-256</ntasks_lnd>
<ntasks_rof>-256</ntasks_rof>
<ntasks_ice>-256</ntasks_ice>
<ntasks_ocn>-256</ntasks_ocn>
</ntasks>
<rootpe>
<rootpe_glc>1</rootpe_glc>
<rootpe_wav>1</rootpe_wav>
<rootpe_esp>1</rootpe_esp>
<rootpe_iac>1</rootpe_iac>
</rootpe>
<pstrid>
<pstrid_atm>16</pstrid_atm>
</pstrid>
<excl_stride>
<excl_stride_atm>16</excl_stride_atm>
</excl_stride>
</pes>
</mach>
</grid>
</config_pes>
10 changes: 1 addition & 9 deletions cime_config/machines/config_batch.xml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -411,15 +411,7 @@
<directives>
<directive> --constraint=gpu</directive>
</directives>
<directives COMPSET="!.*MMF.*" compiler="gnugpu">
<directive> --gpus-per-node=4</directive>
<directive> --gpu-bind=none</directive>
</directives>
<directives COMPSET=".*MMF.*" compiler="gnugpu">
<directive> --gpus-per-task=1</directive>
<directive> --gpu-bind=map_gpu:0,1,2,3</directive>
</directives>
<directives compiler="nvidiagpu">
<directives compiler=".*gpu">
<directive> --gpus-per-node=4</directive>
<directive> --gpu-bind=none</directive>
</directives>
Expand Down
15 changes: 9 additions & 6 deletions cime_config/machines/config_machines.xml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -350,8 +350,8 @@
<BATCH_SYSTEM>nersc_slurm</BATCH_SYSTEM>
<SUPPORTED_BY>e3sm</SUPPORTED_BY>
<MAX_TASKS_PER_NODE>128</MAX_TASKS_PER_NODE>
<MAX_TASKS_PER_NODE compiler="gnu">256</MAX_TASKS_PER_NODE>
<MAX_TASKS_PER_NODE compiler="nvidia">256</MAX_TASKS_PER_NODE>
<MAX_TASKS_PER_NODE compiler="gnu">128</MAX_TASKS_PER_NODE>
<MAX_TASKS_PER_NODE compiler="nvidia">128</MAX_TASKS_PER_NODE>
<MAX_MPITASKS_PER_NODE>4</MAX_MPITASKS_PER_NODE>
<MAX_MPITASKS_PER_NODE compiler="gnu">64</MAX_MPITASKS_PER_NODE>
<MAX_MPITASKS_PER_NODE compiler="nvidia">64</MAX_MPITASKS_PER_NODE>
Expand All @@ -362,8 +362,9 @@
<arg name="label"> --label</arg>
<arg name="num_tasks"> -n {{ total_tasks }} -N {{ num_nodes }}</arg>
<arg name="thread_count">-c $SHELL{echo 128/`./xmlquery --value MAX_MPITASKS_PER_NODE`|bc}</arg>
<arg name="binding"> $SHELL{if [ 64 -ge `./xmlquery --value MAX_MPITASKS_PER_NODE` ]; then echo "--cpu_bind=cores"; else echo "--cpu_bind=threads";fi;} </arg>
<arg name="binding"> $SHELL{if [ 64 -ge `./xmlquery --value MAX_MPITASKS_PER_NODE` ]; then echo "--cpu-bind=cores"; else echo "--cpu-bind=threads";fi;} </arg>
<arg name="placement"> -m plane=$SHELL{echo `./xmlquery --value MAX_MPITASKS_PER_NODE`}</arg>
<arg name="gpu-bind"> /global/cfs/cdirs/e3sm/tools/set_affinity_npergpu.sh $SHELL{echo `./xmlquery --value MAX_MPITASKS_PER_NODE`}</arg>
Copy link
Member

@rljacob rljacob Oct 21, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What does this set_affinity_npergu.sh file do?

</arguments>
</mpirun>
<module_system type="module" allow_error="true">
Expand Down Expand Up @@ -453,16 +454,18 @@
<env name="MPICH_ENV_DISPLAY">1</env>
<env name="MPICH_VERSION_DISPLAY">1</env>
<env name="MPICH_MPIIO_DVS_MAXNODES">1</env>
<env name="OMP_STACKSIZE">128M</env>
<env name="OMP_PROC_BIND">spread</env>
<env name="OMP_PLACES">threads</env>
<env name="HDF5_USE_FILE_LOCKING">FALSE</env>
<env name="PERL5LIB">/global/cfs/cdirs/e3sm/perl/lib/perl5-only-switch</env>
<env name="FI_MR_CACHE_MONITOR">kdreg2</env>
<env name="MPICH_COLL_SYNC">MPI_Bcast</env>
<env name="NETCDF_PATH">$ENV{CRAY_NETCDF_HDF5PARALLEL_PREFIX}</env>
<env name="PNETCDF_PATH">$ENV{CRAY_PARALLEL_NETCDF_PREFIX}</env>
</environment_variables>
<environment_variables BUILD_THREADED="TRUE">
<env name="OMP_STACKSIZE">128M</env>
<env name="OMP_PROC_BIND">spread</env>
<env name="OMP_PLACES">threads</env>
</environment_variables>
<environment_variables compiler="gnugpu">
<env name="MPICH_GPU_SUPPORT_ENABLED">1</env>
</environment_variables>
Expand Down