Chapter 5,7-9: Update CPU bindings when using torchrun

[marlon-tobaben]

Quite clear that the CPU bindings are sensible, but would be good to check if it is smart to set them twice with srun. This was mentioned by @mitjasai in #89

• in the runscript
  

  LUMI-AI-Guide/5-multi-gpu-and-node/run_ds_srun.sh

  Lines 35 to 40 in 6ca88fc

  	CPU_BIND_MASKS="0x00fe000000000000,0xfe00000000000000,0x0000000000fe0000,0x00000000fe000000,0x00000000000000fe,0x000000000000fe00,0x000000fe00000000,0x0000fe0000000000"
  	export SINGULARITYENV_PREPEND_PATH=/user-software/bin # gives access to packages inside the container
  	srun --cpu-bind=v,mask_cpu=$CPU_BIND_MASKS singularity run -B ../resources/ai-guide-env.sqsh:/user-software:image-src=/ \
  	$SIF bash -c 'export RANK=$SLURM_PROCID; export LOCAL_RANK=$SLURM_LOCALID; python ds_visiontransformer.py --deepspeed --deepspeed_config ds_config.json'

• in the python script
  

  LUMI-AI-Guide/5-multi-gpu-and-node/ds_visiontransformer.py

  Lines 18 to 35 in 6ca88fc

  	def set_cpu_affinity(local_rank):
  	LUMI_GPU_CPU_map = {
  	# A mapping from GCD to the closest CPU cores in a LUMI-G node
  	# Note that CPU cores 0, 8, 16, 24, 32, 40, 48, 56 are reserved for the
  	# system and not available for the user
  	# See https://docs.lumi-supercomputer.eu/hardware/lumig/
  	0: [49, 50, 51, 52, 53, 54, 55],
  	1: [57, 58, 59, 60, 61, 62, 63],
  	2: [17, 18, 19, 20, 21, 22, 23],
  	3: [25, 26, 27, 28, 29, 30, 31],
  	4: [1, 2, 3, 4, 5, 6, 7],
  	5: [9, 10, 11, 12, 13, 14, 15],
  	6: [33, 34, 35, 36, 37, 38, 39],
  	7: [41, 42, 43, 44, 45, 46, 47],
  	}
  	cpu_list = LUMI_GPU_CPU_map[local_rank]
  	print(f"Rank {rank} (local {local_rank}) binding to cpus: {cpu_list}")
  	psutil.Process().cpu_affinity(cpu_list)

We think it does not hurt but it might not be pedagogical sensible to claim that with srun the PyTorch script is portable.

[mitjasai]

Something like this might make sense:

diff --git a/5-multi-gpu-and-node/ddp_visiontransformer.py b/5-multi-gpu-and-node/ddp_visiontransformer.py
index c757787..ddcaf0b 100644
--- a/5-multi-gpu-and-node/ddp_visiontransformer.py
+++ b/5-multi-gpu-and-node/ddp_visiontransformer.py
@@ -84,7 +84,9 @@ if __name__ == "__main__":
     local_rank = int(os.environ["LOCAL_RANK"])
     torch.cuda.set_device(local_rank)
     rank = int(os.environ["RANK"])
-    set_cpu_affinity(local_rank)
+
+    if os.getenv("TORCHELASTIC_RUN_ID"):
+        set_cpu_affinity(local_rank)
 
     # Define transformations
     transform = transforms.Compose(

[mitjasai]

I looked into this a bit, and Elastic Launch supports NUMA binding since last year. Passing --numa-binding=exclusive when using Elastic Launch applies the same bindings as the set_cpu_affinity function in the script. Specifically:

rank=6, local_rank=6, os.sched_getaffinity(0)={33, 34, 35, 36, 37, 38, 39}
rank=1, local_rank=1, os.sched_getaffinity(0)={57, 58, 59, 60, 61, 62, 63}
rank=3, local_rank=3, os.sched_getaffinity(0)={25, 26, 27, 28, 29, 30, 31}
rank=0, local_rank=0, os.sched_getaffinity(0)={49, 50, 51, 52, 53, 54, 55}
rank=4, local_rank=4, os.sched_getaffinity(0)={1, 2, 3, 4, 5, 6, 7}
rank=2, local_rank=2, os.sched_getaffinity(0)={17, 18, 19, 20, 21, 22, 23}
rank=7, local_rank=7, os.sched_getaffinity(0)={41, 42, 43, 44, 45, 46, 47}
rank=5, local_rank=5, os.sched_getaffinity(0)={9, 10, 11, 12, 13, 14, 15}