world of deep studying coaching, the position of the ML developer could 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} parts — 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 one in all these parts can severely hamper the general coaching course of.
On this complicated panorama, it’s of important significance that you’ve got an intimate understanding of your system’s underlying topology and that you know the way to use it towards optimum runtime efficiency. In a earlier submit, we explored the important 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 answerable for making ready and pre-processing knowledge, for loading GPU kernels, and for processing output, whereas the GPU is answerable for the mannequin execution. This cooperation isn’t merely a hand-off — it’s a continuing, high-speed alternate of knowledge and instructions, in what could 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 consider 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 trendy multi-GPU Excessive-Efficiency Computing (HPC) methods encompass 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 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, well-liked AI/ML frameworks — most notably PyTorch — don’t account for NUMA structure by default. Nonetheless, as we’ll display on this submit, you’ll be able to 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 are going to then display tips on how 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 one in all its implications: its impression on deep studying. For extra complete particulars on the subject, please consult with different authoritative sources.
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 some 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 might be operating on. On this part, we’ll display tips on how to discover the NUMA format of an Amazon EC2 p4d.96xlarge occasion utilizing generally out there Linux command-line instruments.
CPU NUMA Node Discovery
The lscpu command gives details about the CPU structure of a Linux system, together with a bit describing the NUMA format. 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 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 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 studies 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 section of its command-line output which studies 4 PCI IDs on NUMA node 0. These are marked with “(3D)” tags—widespread identifiers of 3D accelerators, in any other case generally known as GPUs.
Machine (1122GB whole)
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:
out there: 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 measurement: 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 measurement: 574411 MB
node 1 free: 572420 MB
node distances:
node 0 1
0: 10 21
1: 21 10This gives helpful info akin to reminiscence sizes and CPU assignments per NUMA node, in addition to inter-node reminiscence entry prices (increased numbers = better latency).
NUMA Topology Abstract
To summarize the topology we’ve found, here’s a Python illustration of the CPU and GPU format:
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 are going to use this later to implement NUMA-aware coaching.
The Impression of NUMA Placement on Knowledge 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 parts (e.g., sequential algorithms akin to non-maximum suppression) on the CPU. On this submit, we’ll give attention to the switch of enter knowledge from the CPU to the GPU — a important 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 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 principle 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 principle coaching course of for a given GPU — and all of its related dataloader employee processes — will probably be situated 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 lead to efficiency bottlenecks.
Let’s think about a very unhealthy setup:
- GPU i is situated 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 ends in the next inefficient sequence:
- Particular person samples are created on NUMA node 0.
- The samples are transmitted by means of the interconnect to the principle 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 situation could also be uncommon, it illustrates how the default Linux scheduler — if left unmanaged — can lead to inefficient placement and redundant visitors over the NUMA interconnect. And with the excessive price of GPU coaching, counting on the “luck of the scheduler” isn’t beneficial.
When NUMA Placement Issues Most
The efficiency impression of poor NUMA placement relies upon closely on the workload traits. Particularly, coaching steps that encompass a lot of massive knowledge transactions will endure greater than coaching steps with few transactions and small knowledge sizes.
In relation to dataloading, the impression of inefficient NUMA Placement may also 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 would possibly 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 impression could be vital.
Benchmark Impression of NUMA Pinning
As a result of the impact of NUMA-aware pinning can fluctuate broadly, it’s important to benchmark its impression on a per-workload foundation.
In some conditions, NUMA pinning may even harm efficiency. As an example, in methods the place CPUs on one NUMA node are designated for different duties, or methods 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 display the impression of NUMA consciousness on runtime efficiency, we design a toy distributed coaching experiment. Our baseline implementation merely studies 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 akin to lscpu and nvidia-smi.
The next code block incorporates 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():
end result = subprocess.run(['numactl', '--show'],
verify=True,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
textual content=True)
output = end result.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 widespread 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 massive reminiscence transactions. On the GPU, pictures are downscaled to 224×224 earlier than being handed to the mannequin. This setup ends in a bottleneck within the enter knowledge pipeline. The bottleneck could 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 studies its NUMA task utilizing the utilities we outlined above. For the dataloader staff that is executed 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 vital to notice that when making use of NUMA-binding utilizing numa_bind inside a course of, the CPU-binding isn’t all the time inherited by subprocesses. It’s due to this fact 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 pictures 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 isn't None:
numa_bind(bind_to_node)
print(f'GPU {rank} employee {worker_id} NUMA properties:n'
f'{get_numa_placement(rank)}')
# normal coaching loop
def prepare(
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)
machine = torch.cuda.current_device()
mannequin = DDP(resnet18().to(machine), [local_rank])
remodel = 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(machine, non_blocking=True)
targets = goal.to(machine, 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 / 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"])
prepare(local_rank, world_size, bind2gpu)
else:
world_size = torch.cuda.device_count()
torch.multiprocessing.spawn(
fn=prepare,
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 had 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 occasions. 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, decreasing 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
An alternate 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 quite than on entry. This avoids the opportunity of early reminiscence allocations on the incorrect 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 might fluctuate between methods, so it is best to affirm 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 determined by these utilities, you could want to copy a few of their logic manually. Moreover, some high-level frameworks (e.g., Lightning) might not expose management over course of initialization, stopping using 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 prepare.py &
executed
waitOutcomes:
The desk beneath 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 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 important.
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 personal AI/ML initiatives.
For extra suggestions, methods, and strategies for optimizing PyTorch mannequin growth, you’ll want to try the opposite posts in this collection.
