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 provide 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 observe can disrupt a complete musical manufacturing, a bottleneck or inefficiency in any considered one of these elements can severely hamper the general coaching course of.
On this advanced panorama, it’s of vital significance that you’ve got an intimate understanding of your system’s underlying topology and that you understand how to use it towards optimum runtime efficiency. In a previous post, we explored the vital 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 submit, 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 accountable for getting ready and pre-processing information, for loading GPU kernels, and for processing output, whereas the GPU is accountable for the mannequin execution. This cooperation isn’t merely a hand-off — it’s a continuing, high-speed trade of knowledge 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 way that accounts for the underlying system topology. Particularly, it should consider the system’s Non-Uniform Memory Access (NUMA) structure.
NUMA Structure
The NUMA structure is designed to optimize reminiscence transactions by associating native reminiscence banks straight with particular CPU sockets. Most trendy multi-GPU Excessive-Efficiency Computing (HPC) techniques include two or extra NUMA nodes, the place CPUs and GPUs are divided into disjoint teams, every hooked up to 1 node. NUMA is best when reminiscence banks are accessed from throughout the similar node. Accessing reminiscence on a distant node requires information 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, widespread AI/ML frameworks — most notably PyTorch — don’t account for NUMA structure by default. Nonetheless, as we’ll reveal on this submit, you may introduce NUMA-awareness into your PyTorch script with out a lot problem.
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 Learning AMI (DLAMI). We’ll then reveal the right way to implement a NUMA-aware PyTorch script and consider its impression on runtime efficiency.
Disclaimers
The NUMA structure is a fancy and nuanced subject. On this submit, we discover simply considered one of its implications: its impression 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 alternative of platform, framework, or every other device 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 the right way to discover the NUMA structure of an Amazon EC2 p4d.96xlarge occasion utilizing generally accessible Linux command-line instruments.
CPU NUMA Node Discovery
The lscpu command offers details about the CPU structure of a Linux system, together with a bit describing the NUMA structure. By operating the command on an Amazon EC2 p4d.96xlarge occasion we be taught 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 hooked up to, we use a two-step course of: First, we establish the PCI ID related to every GPU, after which we glance up the NUMA node related to that PCI ID.
The PCI ID is without doubt one of 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 Identify 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/gadgets/ path:
ubuntu@XX:~$ cat /sys/bus/pci/gadgets/0000:10:1c.0/numa_node
0
ubuntu@XX:~$ cat /sys/bus/pci/gadgets/0000:10:1d.0/numa_node
0This means that GPUs 0 and 1 are linked to NUMA node 0.
Extra Instruments
Another technique for locating the NUMA node task of the PCI IDs is utilizing lstopo — a command-line utility that stories 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 stories 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 whole)
Bundle 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 great tool 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:
accessible: 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 data equivalent to reminiscence sizes and CPU assignments per NUMA node, in addition to inter-node reminiscence entry prices (increased numbers = larger 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 Influence of NUMA Placement on Information Loading
Reminiscence transactions between CPU and GPU happen at varied phases throughout mannequin execution — for instance, when offloading tensors to CPU reminiscence, or when executing sure mannequin elements (e.g., sequential algorithms equivalent to non-maximum suppression) on the CPU. On this submit, we’ll give attention to the switch of enter information from the CPU to the GPU — a vital 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 offer information batches for its GPU. Every dataloader usually 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 might find yourself seeing a substantial quantity of visitors on the NUMA interconnects, which may end in efficiency bottlenecks.
Let’s think about a very dangerous setup:
- GPU i is positioned on NUMA node 0.
- The primary 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 price of GPU coaching, counting on the “luck of the scheduler” shouldn’t be advisable.
When NUMA Placement Issues Most
The efficiency impression of poor NUMA placement relies upon closely on the workload traits. Particularly, coaching steps that include a lot of massive information transactions will undergo greater than coaching steps with few transactions and small information sizes.
Relating to dataloading, the impression of inefficient NUMA Placement may even depend upon the scale 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 just like or longer than GPU execution time — or in case you’re already experiencing GPU hunger — the impression will be vital.
Benchmark Influence of NUMA Pinning
As a result of the impact of NUMA-aware pinning can differ extensively, it’s important to benchmark its impression on a per-workload foundation.
In some conditions, NUMA pinning may even harm efficiency. For example, in techniques the place CPUs on one NUMA node are designated for different duties, or techniques the place one NUMA node incorporates CPUs bit no GPUs, NUMA pinning may restrict entry to CPU energy, finally straining throughput efficiency.
A Toy PyTorch Experiment
To reveal the impression of NUMA consciousness on runtime efficiency, we design a toy distributed coaching experiment. Our baseline implementation merely stories the NUMA task of every spawned course of. We then apply NUMA-based CPU and reminiscence affinity and measure the impression on throughput.
NUMA Discovery and Pinning Utilities
We start by defining utility features 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 equivalent to lscpu and nvidia-smi.
The next code block incorporates utilities for trying 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():
outcome = subprocess.run(['numactl', '--show'],
verify=True,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
textual content=True)
output = outcome.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 perform. 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 perform 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 massive reminiscence transactions. On the GPU, photographs are downscaled to 224×224 earlier than being handed to the mannequin. This setup leads to a bottleneck within the enter information 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., here and here).
Every time a brand new course of is began, it stories its NUMA task utilizing the utilities we outlined above. For the dataloader employees that is accomplished utilizing a customized worker_init_fn perform. We embrace a numa_aware management flag that determines whether or not to use NUMA pinning.
It’s vital 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 due to this fact important to reapply NUMA binding explicitly throughout the dataloader employees.
import time
import torch
from functools import partial
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.information import Dataset, DataLoader
from torchvision.fashions import resnet18
from torchvision.transforms import Resize
# An artificial dataset with random photographs 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(information=index % 1000, dtype=torch.int64)
return rand_image, label
# Callback for DataLoader employees 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)}')
# customary 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)
system = torch.cuda.current_device()
mannequin = DDP(resnet18().to(system), [local_rank])
remodel = Resize(224)
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(mannequin.parameters())
# num steps
warmup = 10
energetic = 100
total_steps = warmup + energetic
# 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(system, non_blocking=True)
targets = goal.to(system, non_blocking=True)
optimizer.zero_grad()
outputs = mannequin(remodel(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 / energetic}')
print(f'common throughput: {energetic / 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 part of=True
)Observing NUMA Placement
Here’s a pattern output from operating the script on a single GPU with 4 dataloader employees 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 differ 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 common throughput on this setting elevated to 1.24 steps per second — a 19% enchancment over the baseline experiment.
CPU Binding with numactl
Another method 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 relatively than on entry. This avoids the potential for early reminiscence allocations on the improper node earlier than pinning. One other benefit is that the NUMA placement is inherited by subprocesses, making it pointless to re-pin dataloader employees manually. Be aware, that the inheritance conduct might differ between techniques, so you need 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 depends upon these utilities, it’s possible you’ll want to duplicate a few of their logic manually. Moreover, some high-level frameworks (e.g., Lightning) might 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 &
accomplished
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 impression can differ relying in your mannequin structure, information loading traits, and system configuration.
Abstract
In our fixed pursuit of AI/ML workload optimization, topology consciousness — together with NUMA node placement — is vital.
On this submit, we continued our exploration of PyTorch mannequin profiling and optimization by demonstrating how NUMA pinning can enhance throughput efficiency. We hope you’ll find this technique helpful in your individual AI/ML initiatives.
For extra ideas, tips, and methods for optimizing PyTorch mannequin growth, be sure you take a look at the opposite posts in this series.
