Add NUMA affinity control for NVIDIA GPUs

[muellerzr]

What does this PR do?

This PR adds the ability to launch using numactl, which acts as an inbetween stage when launching from torchrun etc.

Before:

trampoline.sh

#!/usr/bin/bash

# Query the bus ID for device LOCAL_RANK
BUS_ID=$(nvidia-smi --query-gpu=pci.bus_id -i $LOCAL_RANK --format=csv,noheader)
BUS_ID=${BUS_ID,,}

# Find the numa node for device LOCAL_RANK
NODE=$(cat /sys/bus/pci/devices/${BUS_ID:4}/numa_node)

echo "Starting local rank $RANK on numa node $NODE"
numactl --cpunodebind=$NODE --membind=$NODE "$@"

torchrun --nproc_per_node=8 --monitor-interval=1 --no-python ./trampoline.sh python3 -c "print('hello')"

After:

accelerate launch --num_processes 8 --monitor-interaval 1 --enable-numa-affinity python3 -c "print('hello')"

Or through the config:

compute_environment: LOCAL_MACHINE
deepspeed_config: {}
distributed_type: 'MULTI_GPU'
...
enable_numa_affinity: true

accelerate launch myfile.py

Fixes #2241

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@BenjaminBossan @stas00

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[BenjaminBossan]

Thanks, Zach, for enabling NUMA affinity control to be easily enabled with accelerate. I have no experience whatsoever with numactl, so I skipped those specific parts of the review and focused on the rest.

[muellerzr]

@stas00 I think this should fully enable you now. You can either set enable_cpu_affinity to True in your config.yaml, or set the env variable ACCELERATE_CPU_AFFINITY=1

[BenjaminBossan]

Thanks a lot. The code LGTM. Regarding whether this solves the initial problem, I defer to Stas.

[stas00]

This looks good code-wise, but I won't have a chance to try it out till next week as I'm on the road.

Did you get a chance to test it on the H100 nodes at HF? If so it's probably good to go. If not and you don't mind waiting I will test next week.

[muellerzr]

Despite my best efforts couldn't get away from using pynvml 😭 However it seems to come auto-installed (probably with PyTorch?) so did a simple import check. As we get more libs/backend support we can look at appropriate interfaces rather than reinvent the wheel :)

[muellerzr]

Confirmed with @stas00 the load balancing looks proper, once tests pass will merge ✅