world of deep studying coaching, the function 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 large number of {hardware} elements — CPUs and GPUs with their related reminiscence, high-speed storage, community controllers, numerous communication buses, and many others. — to work collectively seamlessly to maximise runtime efficiency. Simply as a single off-key be aware can disrupt a whole musical manufacturing, a bottleneck or inefficiency in any considered one of these elements can severely hamper the general coaching course of.
On this complicated panorama, it’s of important 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 important function 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 sequence 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 chargeable for getting ready and pre-processing information, for loading GPU kernels, and for processing output, whereas the GPU is chargeable 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 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. Particularly, it should take note of 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) techniques include two or extra NUMA nodes, the place CPUs and GPUs are divided into disjoint teams, every connected to at least one node. NUMA is most effective when reminiscence banks are accessed from throughout the identical 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. Nevertheless, as we are going to show on this put up, you’ll be able to introduce NUMA-awareness into your PyTorch script with out a lot issue.
Within the subsequent part, we are going to 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 show methods to implement a NUMA-aware PyTorch script and consider its influence on runtime efficiency.
Disclaimers
The NUMA structure is a posh and nuanced subject. On this put up, we discover simply considered one of its implications: its influence on deep studying. For extra complete particulars on the subject, please confer with different authoritative sources.
The code we are going to 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 might be operating on. On this part, we are going to show methods 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 supplies 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 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 | Risky 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 related 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 experiences 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 experiences 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 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 possibly can 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 supplies helpful info corresponding to reminiscence sizes and CPU assignments per NUMA node, in addition to inter-node reminiscence entry prices (greater 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 Affect of NUMA Placement on Knowledge Loading
Reminiscence transactions between CPU and GPU happen at numerous 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 information 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 through 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 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 positioned on the identical NUMA node because the GPU. In any other case, we could find yourself seeing a substantial quantity of site 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 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 and grouped into batches by employees on NUMA node 0.
- Every batch is transmitted via the interconnect to the principle course of on node 1.
- The batch is shipped again throughout the interconnect to node 0, the place it’s fed to the GPU.
Sounds horrendous, proper??
Whereas this actual state of affairs could also be uncommon, it illustrates how the default Linux scheduler — if left unmanaged — can lead to inefficient placement and redundant site visitors over the NUMA interconnect. And with the excessive value of GPU coaching, counting on the “luck of the scheduler” isn’t really useful.
When NUMA Placement Issues Most
The efficiency influence of poor NUMA placement relies upon closely on the workload traits. Particularly, coaching steps that include numerous giant information transactions will endure greater than coaching steps with few transactions and small information sizes.
In the case of dataloading, the influence of inefficient NUMA Placement will even rely on 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 just like or longer than GPU execution time — or in the event you’re already experiencing GPU hunger — the influence could be vital.
Benchmark Affect of NUMA Pinning
As a result of the impact of NUMA-aware pinning can range extensively, it’s important to benchmark its influence on a per-workload foundation.
In some conditions, NUMA pinning may even damage efficiency. As an example, in techniques the place CPUs on one NUMA node are designated for different duties, or techniques 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 show the influence of NUMA consciousness on runtime efficiency, we design a toy distributed coaching experiment. Our baseline implementation merely experiences the NUMA task of every spawned course of. We then apply NUMA-based CPU and reminiscence affinity and measure the influence 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 sturdy 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 sure 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 through the os.sched_setaffinity operate. Nevertheless, 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 information 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 experiences its NUMA task utilizing the utilities we outlined above. For the dataloader employees that is executed utilizing a customized worker_init_fn operate. We embody a numa_aware management flag that determines whether or not to use NUMA pinning.
It’s necessary to notice that when making use of NUMA-binding utilizing numa_bind inside a course of, the CPU-binding isn’t at all times 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 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(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 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)
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"])
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 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 range 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, lowering the site 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 strategy to NUMA pinning is to launch every coaching course of from the command-line through 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 potential for 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 employees manually. Notice, that the inheritance habits could range between techniques, so you need 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 relies on 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 through 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 under summarizes the outcomes of our experiments.
On this toy instance, the advantages of NUMA-aware coaching are clear. Nevertheless, as famous earlier, the precise influence can range 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 important.
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 will see this technique helpful in your personal AI/ML initiatives.
For extra suggestions, tips, and methods for optimizing PyTorch mannequin growth, make sure to try the opposite posts in this sequence.
