world of deep studying coaching, the position of the ML developer will be likened to that of the conductor of an orchestra. Simply as a conductor should time the entry of every instrument to supply the right concord, so should the ML practitioner orchestrate a mess of {hardware} elements — CPUs and GPUs with their related reminiscence, high-speed storage, community controllers, varied communication buses, and many others. — to work collectively seamlessly to maximise runtime efficiency. Simply as a single off-key be aware can disrupt a complete musical manufacturing, a bottleneck or inefficiency in any certainly one of these elements can severely hamper the general coaching course of.
On this advanced panorama, it’s of crucial significance that you’ve an intimate understanding of your system’s underlying topology and that you understand how to use it towards optimum runtime efficiency. In a earlier put up, we explored the crucial position of topology consciousness in a distributed coaching setting and mentioned the benefit of topology-aware gradient sharing algorithms in minimizing cross-node communication and boosting efficiency.
On this put up, the tenth in our collection on PyTorch mannequin evaluation and optimization, we zoom in on the collaboration between the CPU and GPU in coaching and operating AI/ML fashions. In a typical coaching pipeline, the CPU is liable for getting ready and pre-processing knowledge, for loading GPU kernels, and for processing output, whereas the GPU is liable for the mannequin execution. This cooperation isn’t merely a hand-off — it’s a continuing, high-speed alternate of information and instructions, in what will be likened to an intricate dance — the place precision timing and bodily proximity are essential. For this dance to be carried out optimally, it have to be choreographed in a fashion that accounts for the underlying system topology. Specifically, it should take into consideration the system’s Non-Uniform Reminiscence Entry (NUMA) structure.
NUMA Structure
The NUMA structure is designed to optimize reminiscence transactions by associating native reminiscence banks instantly with particular CPU sockets. Most fashionable multi-GPU Excessive-Efficiency Computing (HPC) programs include two or extra NUMA nodes, the place CPUs and GPUs are divided into disjoint teams, every connected to 1 node. NUMA is best when reminiscence banks are accessed from throughout the similar node. Accessing reminiscence on a distant node requires knowledge traversal over a devoted NUMA interconnect, which is considerably slower than accessing native reminiscence. In memory-intensive functions like AI/ML workloads, cross-NUMA reminiscence accesses can introduce efficiency bottlenecks.
Sadly, in style AI/ML frameworks — most notably PyTorch — don’t account for NUMA structure by default. Nonetheless, as we’ll reveal on this put up, you possibly can introduce NUMA-awareness into your PyTorch script with out a lot issue.
Within the subsequent part, we’ll discover the NUMA structure of the favored Amazon EC2 p4d.96xlarge occasion (containing 8 NVIDIA A100 GPUs and 96 vCPUs) operating a PyTorch (2.6) Deep Studying AMI (DLAMI). We’ll then reveal tips on how to implement a NUMA-aware PyTorch script and consider its affect on runtime efficiency.
Disclaimers
The NUMA structure is a posh and nuanced matter. On this put up, we discover simply certainly one of its implications: its affect on deep studying. For extra complete particulars on the subject, please check with different authoritative assets.
The code we’ll share is meant for demonstrative functions and shouldn’t be relied on for correctness or optimality. Please don’t interpret our selection of platform, framework, or another instrument or library as an endorsement for its use.
NUMA Structure Discovery
There are a number of methods to detect the NUMA structure of the system you’re operating on. On this part, we’ll reveal tips on how to discover the NUMA structure of an Amazon EC2 p4d.96xlarge occasion utilizing generally obtainable Linux command-line instruments.
CPU NUMA Node Discovery
The lscpu command offers details about the CPU structure of a Linux system, together with a piece describing the NUMA structure. By operating the command on an Amazon EC2 p4d.96xlarge occasion we study that it consists of 96 vCPUs divided between two NUMA nodes:
NUMA:
NUMA node(s): 2
NUMA node0 CPU(s): 0-23,48-71
NUMA node1 CPU(s): 24-47,72-95GPU NUMA Node Discovery
To find out which NUMA node every GPU is connected to, we use a two-step course of: First, we determine the PCI ID related to every GPU, after which we glance up the NUMA node related to that PCI ID.
The PCI ID is among the properties of the GPUs reported by the nvidia-smi utility. Within the following snippet, we see the PCI Bus IDs of the primary two out of the eight GPUs on our Amazon EC2 p4d.96xlarge occasion:
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 570.133.20 Driver Model: 570.133.20 CUDA Model: 12.8 |
|-----------------------------------------+------------------------+----------------------+
| GPU Title Persistence-M | Bus-Id Disp.A | Unstable Uncorr. ECC |
| Fan Temp Perf Pwr:Utilization/Cap | Reminiscence-Utilization | GPU-Util Compute M. |
| | | MIG M. |
|=========================================+========================+======================|
| 0 NVIDIA A100-SXM4-40GB On | 00000000:10:1C.0 Off | 0 |
| N/A 48C P0 57W / 400W | 0MiB / 40960MiB | 0% Default |
| | | Disabled |
+-----------------------------------------+------------------------+----------------------+
| 1 NVIDIA A100-SXM4-40GB On | 00000000:10:1D.0 Off | 0 |
| N/A 45C P0 56W / 400W | 0MiB / 40960MiB | 0% Default |
| | | Disabled |
+-----------------------------------------+------------------------+----------------------+Subsequent, we use these PCI IDs to find out the corresponding NUMA node by studying from the /sys/bus/pci/units/ path:
ubuntu@XX:~$ cat /sys/bus/pci/units/0000:10:1c.0/numa_node
0
ubuntu@XX:~$ cat /sys/bus/pci/units/0000:10:1d.0/numa_node
0This means that GPUs 0 and 1 are linked to NUMA node 0.
Further Instruments
An alternate technique for locating the NUMA node task of the PCI IDs is utilizing lstopo — a command-line utility that reviews the topology of a pc system. Although it isn’t included by default within the DLAMI, it may be simply put in by operating:
sudo apt set up hwlocHere’s a small phase of its command-line output which reviews 4 PCI IDs on NUMA node 0. These are marked with “(3D)” tags—frequent identifiers of 3D accelerators, in any other case often known as GPUs.
Machine (1122GB complete)
Package deal L#0
NUMANode L#0 (P#0 561GB)
HostBridge
2 x { PCI 10:1c.0-1d.0 (3D) }
HostBridge
2 x { PCI 20:1c.0-1d.0 (3D) }One other useful gizmo is numactl — a command-line utility in Linux inspecting and managing NUMA insurance policies. To put in numactl, run:
sudo apt set up numactlYou’ll be able to examine the NUMA configuration by operating:
numactl --hardwareOn our Amazon EC2 p4d.96xlarge occasion this produces the next output:
obtainable: 2 nodes (0-1)
node 0 cpus: 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71
node 0 dimension: 574309 MB
node 0 free: 572012 MB
node 1 cpus: 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95
node 1 dimension: 574411 MB
node 1 free: 572420 MB
node distances:
node 0 1
0: 10 21
1: 21 10This offers helpful info corresponding to reminiscence sizes and CPU assignments per NUMA node, in addition to inter-node reminiscence entry prices (greater numbers = better latency).
NUMA Topology Abstract
To summarize the topology we’ve found, here’s a Python illustration of the CPU and GPU structure:
cpus_per_numa_node = [
list(range(0, 24)) + list(range(48, 72)), # NUMA node 0
list(range(24, 48)) + list(range(72, 96)) # NUMA node 1
]
gpus_per_numa_node = [
[0, 1, 2, 3], # NUMA node 0
[4, 5, 6, 7] # NUMA node 1
]We’ll use this later to implement NUMA-aware coaching.
The Affect of NUMA Placement on Information Loading
Reminiscence transactions between CPU and GPU happen at varied levels throughout mannequin execution — for instance, when offloading tensors to CPU reminiscence, or when executing sure mannequin elements (e.g., sequential algorithms corresponding to non-maximum suppression) on the CPU. On this put up, we’ll deal with the switch of enter knowledge from the CPU to the GPU — a crucial a part of each AI/ML workflow.
The CPU Processes in a Typical Distributed Coaching Job
In a typical distributed coaching setting, new CPU processes are created on two events:
- At startup: A separate coaching course of is created for every GPU. These processes deal with mannequin setup and coaching execution on their assigned GPUs. Within the script we’ll introduce later, these are launched by way of torch.multiprocessing.spawn.
- Per dataloader: Every coaching course of creates its personal DataLoader occasion to supply knowledge batches for its GPU. Every dataloader sometimes creates a number of employee processes, which generate particular person coaching samples. These samples are then grouped by the primary course of into batches.
Within the case of our Amazon EC2 p4d.96xlarge occasion, every of those processes is assigned to a CPU, which resides on one of many two NUMA nodes.
Why NUMA Placement Issues
Ideally, the primary coaching course of for a given GPU — and all of its related dataloader employee processes — will probably be positioned on the identical NUMA node because the GPU. In any other case, we could find yourself seeing a substantial quantity of visitors on the NUMA interconnects, which may end in efficiency bottlenecks.
Let’s think about a very unhealthy setup:
- GPU i is positioned on NUMA node 0.
- The principle coaching course of assigned to GPU i is scheduled on a CPU on NUMA node 1.
- The employee processes spawned by the coaching course of are all assigned to CPUs on NUMA node 0.
This leads to the next inefficient sequence:
- Particular person samples are created on NUMA node 0.
- The samples are transmitted by the interconnect to the primary course of on node 1, the place they’re grouped collectively right into a coaching batch.
- The batch is shipped again throughout the interconnect to node 0, the place it’s fed to the GPU.
Sounds horrendous, proper?
Whereas this precise state of affairs could also be uncommon, it illustrates how the default Linux scheduler — if left unmanaged — may end up in inefficient placement and redundant visitors over the NUMA interconnect. And with the excessive value of GPU coaching, counting on the “luck of the scheduler” shouldn’t be really useful.
When NUMA Placement Issues Most
The efficiency affect of poor NUMA placement relies upon closely on the workload traits. Particularly, coaching steps that include a lot of giant knowledge transactions will undergo greater than coaching steps with few transactions and small knowledge sizes.
With regards to dataloading, the affect of inefficient NUMA Placement may also depend upon the dimensions of the mannequin. Recall that AI/ML workloads are designed to run the dataloading on the CPU in parallel with mannequin execution on the GPU. Thus, if the GPU execution takes considerably longer than the dataloading, inefficient NUMA placement may go unnoticed. But when dataloading time is much like or longer than GPU execution time — or if you happen to’re already experiencing GPU hunger — the affect will be important.
Benchmark Affect of NUMA Pinning
As a result of the impact of NUMA-aware pinning can fluctuate broadly, it’s important to benchmark its affect on a per-workload foundation.
In some conditions, NUMA pinning may even damage efficiency. For example, in programs the place CPUs on one NUMA node are designated for different duties, or programs the place one NUMA node accommodates CPUs bit no GPUs, NUMA pinning may restrict entry to CPU energy, finally straining throughput efficiency.
A Toy PyTorch Experiment
To reveal the affect of NUMA consciousness on runtime efficiency, we design a toy distributed coaching experiment. Our baseline implementation merely reviews the NUMA task of every spawned course of. We then apply NUMA-based CPU and reminiscence affinity and measure the affect on throughput.
NUMA Discovery and Pinning Utilities
We start by defining utility capabilities for NUMA node discovery and pinning. The implementation proven right here makes use of the hardcoded NUMA topology we summarized earlier. A extra strong model would dynamically uncover topology by parsing the output of system utilities corresponding to lscpu and nvidia-smi.
The next code block accommodates utilities for wanting up NUMA placement. For every course of we report each the NUMA node of the host CPU and the NUMA node of its allotted reminiscence it’s certain to. We use numactl --show to detect the reminiscence binding of the present course of.
import os, re, psutil, ctypes, subprocess
# Uncover NUMA node of course of
def discover_cpu_numa_placement():
cpu_id = psutil.Course of().cpu_num()
for node in vary(len(cpus_per_numa_node)):
if cpu_id in cpus_per_numa_node[node]:
return node
# Uncover NUMA node of GPU
def discover_gpu_numa_placement(rank):
for node in vary(len(gpus_per_numa_node)):
if rank in gpus_per_numa_node[node]:
return node
# Use numactl to get mememory binding of CPU course of
def get_membinding():
consequence = subprocess.run(['numactl', '--show'],
verify=True,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
textual content=True)
output = consequence.stdout
match = re.search(r"membind:s*([0-9s]+)", output)
nodes = [int(n) for n in match.group(1).split()]
return nodes
# Detect NUMA placement of course of
def get_numa_placement(rank):
cpu_node = discover_cpu_numa_placement()
gpu_node = discover_gpu_numa_placement(rank)
m_bind = get_membinding()
node_match = cpu_node == gpu_node
standing = f"GPU node: {gpu_node}n"
f"CPU node: {cpu_node}n"
f"mem binding {m_bind[0] if len(m_bind)==1 else m_bind}n"
if not node_match:
standing += "GPU and CPU NUMA nodes do NOT matchn"
return standingOne frequent technique for setting CPU affinity in Python is by way of the os.sched_setaffinity operate. Nonetheless, this technique is inadequate for our functions as a result of it solely pins the CPU — it doesn’t bind the reminiscence it makes use of. To bind each CPU and reminiscence binding we use the numa_bind operate from the libnuma library. (Run sudo apt set up libnuma-dev to put in).
# Set course of affinity by NUMA node ID
def set_affinity_by_node(node):
pid = os.getpid()
target_cpus = cpus_per_numa_node[node]
os.sched_setaffinity(pid, target_cpus)
# Bind a course of and reminiscence to given NUMA node
def numa_bind(node):
libnuma = ctypes.CDLL("libnuma.so")
libnuma.numa_allocate_nodemask.restype = ctypes.c_void_p
libnuma.numa_bitmask_clearall.argtypes = [ctypes.c_void_p]
libnuma.numa_bitmask_setbit.argtypes = [ctypes.c_void_p, ctypes.c_uint]
libnuma.numa_bind.argtypes = [ctypes.c_void_p]
nodemask_ptr = libnuma.numa_allocate_nodemask()
libnuma.numa_bitmask_clearall(nodemask_ptr)
libnuma.numa_bitmask_setbit(nodemask_ptr, node)
libnuma.numa_bind(nodemask_ptr)Mannequin Definition
Subsequent, we outline a easy distributed coaching script utilizing a ResNet-18 picture classification mannequin and an artificial dataset. Every artificial pattern is a randomly generated 1024×1024 picture, simulating giant reminiscence transactions. On the GPU, photos are downscaled to 224×224 earlier than being handed to the mannequin. This setup leads to a bottleneck within the enter knowledge pipeline. The bottleneck will be detected by evaluating throughput (in steps per second) throughout regular coaching versus when operating on a cached batch. For extra on figuring out dataloader bottlenecks, see our earlier posts (e.g., right here and right here).
Every time a brand new course of is began, it reviews its NUMA task utilizing the utilities we outlined above. For the dataloader staff that is finished utilizing a customized worker_init_fn operate. We embrace a numa_aware management flag that determines whether or not to use NUMA pinning.
It’s essential to notice that when making use of NUMA-binding utilizing numa_bind inside a course of, the CPU-binding shouldn’t be all the time inherited by subprocesses. It’s subsequently important to reapply NUMA binding explicitly throughout the dataloader staff.
import time
import torch
from functools import partial
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.knowledge import Dataset, DataLoader
from torchvision.fashions import resnet18
from torchvision.transforms import Resize
# An artificial dataset with random photos and labels
class FakeDataset(Dataset):
def __init__(self, n_items):
tremendous().__init__()
self.n_items = n_items
def __len__(self):
return self.n_items
def __getitem__(self, index):
rand_image = torch.randn([3, 1024, 1024], dtype=torch.float32)
label = torch.tensor(knowledge=index % 1000, dtype=torch.int64)
return rand_image, label
# Callback for DataLoader staff to detect their NUMA placement.
def worker_init_fn(worker_id, rank=0, bind_to_node=None):
if bind_to_node shouldn't be None:
numa_bind(bind_to_node)
print(f'GPU {rank} employee {worker_id} NUMA properties:n'
f'{get_numa_placement(rank)}')
# commonplace coaching loop
def practice(
local_rank,
world_size,
numa_aware=False
):
bind_to_node = None
if numa_aware:
bind_to_node = discover_gpu_numa_placement(local_rank)
numa_bind(bind_to_node)
print(f'GPU {local_rank} coaching course of NUMA properties:n'
f'{get_numa_placement(local_rank)}')
torch.cuda.set_device(local_rank)
# DDP setup
os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = str(2222)
dist.init_process_group('nccl', rank=local_rank,
world_size=world_size)
gadget = torch.cuda.current_device()
mannequin = DDP(resnet18().to(gadget), [local_rank])
rework = Resize(224)
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(mannequin.parameters())
# num steps
warmup = 10
lively = 100
total_steps = warmup + lively
# distribute evenly throughout GPUs
num_workers = os.cpu_count() // world_size
batch_size = 128
data_loader = DataLoader(
FakeDataset(total_steps * batch_size),
batch_size=batch_size,
num_workers=num_workers,
pin_memory=True,
worker_init_fn=partial(
worker_init_fn,
rank=local_rank,
bind_to_node=bind_to_node
)
)
for idx, (inputs, goal) in enumerate(data_loader, begin=1):
inputs = inputs.to(gadget, non_blocking=True)
targets = goal.to(gadget, non_blocking=True)
optimizer.zero_grad()
outputs = mannequin(rework(inputs))
loss = criterion(outputs, targets)
loss.backward()
optimizer.step()
if idx == warmup:
torch.cuda.synchronize()
t0 = time.perf_counter()
elif idx == total_steps:
break
if local_rank == 0:
torch.cuda.synchronize()
total_time = time.perf_counter() - t0
print(f'common step time: {total_time / lively}')
print(f'common throughput: {lively / total_time}')
dist.destroy_process_group()
if __name__ == '__main__':
bind2gpu = False
if os.environ.get("LOCAL_RANK", None):
# initialized with torchrun or bash script
local_rank = int(os.environ["LOCAL_RANK"])
world_size = int(os.environ["WORLD_SIZE"])
practice(local_rank, world_size, bind2gpu)
else:
world_size = torch.cuda.device_count()
torch.multiprocessing.spawn(
fn=practice,
args=(world_size, bind2gpu),
nprocs=world_size,
be a part of=True
)Observing NUMA Placement
Here’s a pattern output from operating the script on a single GPU with 4 dataloader staff and no NUMA binding. On this run, all processes have been scheduled on NUMA node 1, whereas the GPU resides on NUMA node 0:
GPU 0 coaching course of NUMA properties:
GPU node: 0
CPU node: 1
mem binding [0, 1]
GPU and CPU NUMA nodes do NOT match
GPU 0 employee 1 NUMA properties:
GPU node: 0
CPU node: 1
mem binding [0, 1]
GPU and CPU NUMA nodes do NOT match
GPU 0 employee 3 NUMA properties:
GPU node: 0
CPU node: 1
mem binding [0, 1]
GPU and CPU NUMA nodes do NOT match
GPU 0 employee 0 NUMA properties:
GPU node: 0
CPU node: 1
mem binding [0, 1]
GPU and CPU NUMA nodes do NOT match
GPU 0 employee 2 NUMA properties:
GPU node: 0
CPU node: 1
mem binding [0, 1]
GPU and CPU NUMA nodes do NOT matchBaseline Outcomes
NUMA placement can fluctuate between runs, so we repeated the baseline experiment ten instances. The resultant common throughput was 1.04 steps per second.
NUMA-Conscious Coaching
To allow NUMA-aware coaching, we set the numa_aware flag to True. This causes every coaching course of to run on a CPU from the identical NUMA node as its assigned GPU and allocate reminiscence on that very same NUMA node. This configuration ensures NUMA-locality throughout CPU, reminiscence, and GPU, lowering the visitors over the NUMA interconnect.
The typical throughput on this setting elevated to 1.24 steps per second — a 19% enchancment over the baseline experiment.
CPU Binding with numactl
An alternate strategy to NUMA pinning is to launch every coaching course of from the command-line by way of the numactl command. The benefit of this technique is that the binding is utilized earlier than the method is began reasonably than on entry. This avoids the opportunity of early reminiscence allocations on the fallacious node earlier than pinning. One other benefit is that the NUMA placement is inherited by subprocesses, making it pointless to re-pin dataloader staff manually. Be aware, that the inheritance conduct could fluctuate between programs, so it’s best to verify it in your particular setup earlier than counting on it.
One draw back of this technique is that it can’t be simply built-in with PyTorch’s launch utilities like torch.multiprocessing.spawn or torchrun. In case your code is dependent upon these utilities, you could want to duplicate a few of their logic manually. Moreover, some high-level frameworks (e.g., Lightning) could not expose management over course of initialization, stopping the usage of binding by way of numactl.
Right here’s a pattern Bash script that wraps our coaching script with NUMA pinning utilizing numactl:
#!/bin/bash
# Outline GPU-to-NUMA mapping
GPU_LIST=(0 1 2 3 4 5 6 7)
GPU_TO_NUMA=(0 0 0 0 1 1 1 1)
NUM_GPUS=${#GPU_LIST[@]}
WORLD_SIZE=$NUM_GPUS
for i in "${!GPU_LIST[@]}"; do
GPU_ID=${GPU_LIST[$i]}
NUMA_NODE=${GPU_TO_NUMA[$i]}
LOCAL_RANK=$i
echo "Launch GPU $LOCAL_RANK on NUMA $NUMA_NODE" >&1
numactl --cpunodebind=$NUMA_NODE --membind=$NUMA_NODE
env
LOCAL_RANK=$LOCAL_RANK
WORLD_SIZE=$WORLD_SIZE
python practice.py &
finished
waitOutcomes:
The desk under summarizes the outcomes of our experiments.
On this toy instance, the advantages of NUMA-aware coaching are clear. Nonetheless, as famous earlier, the precise affect can fluctuate relying in your mannequin structure, knowledge loading traits, and system configuration.
Abstract
In our fixed pursuit of AI/ML workload optimization, topology consciousness — together with NUMA node placement — is crucial.
On this put up, we continued our exploration of PyTorch mannequin profiling and optimization by demonstrating how NUMA pinning can enhance throughput efficiency. We hope you can find this technique helpful in your personal AI/ML initiatives.
For extra suggestions, methods, and strategies for optimizing PyTorch mannequin improvement, make sure to take a look at the opposite posts in this collection.
