Coaching AI Fashions on CPU. Revisiting CPU for ML in an Period of GPU… | by Chaim Rand | Sep, 2024

Revisiting CPU for ML in an Period of GPU Shortage

13 min learn

·

21 hours in the past

The latest successes in AI are sometimes attributed to the emergence and evolutions of the GPU. The GPU’s structure, which generally contains hundreds of multi-processors, high-speed reminiscence, devoted tensor cores, and extra, is especially well-suited to satisfy the intensive calls for of AI/ML workloads. Sadly, the fast development in AI growth has led to a surge within the demand for GPUs, making them troublesome to acquire. Because of this, ML builders are more and more exploring various {hardware} choices for coaching and operating their fashions. In earlier posts, we mentioned the potential for coaching on devoted AI ASICs corresponding to Google Cloud TPU, Haban Gaudi, and AWS Trainium. Whereas these choices supply important cost-saving alternatives, they don’t go well with all ML fashions and might, just like the GPU, additionally undergo from restricted availability. On this submit we return to the great old school CPU and revisit its relevance to ML functions. Though CPUs are usually much less suited to ML workloads in comparison with GPUs, they’re much simpler to accumulate. The power to run (a minimum of a few of) our workloads on CPU might have important implications on growth productiveness.

In earlier posts (e.g., right here) we emphasised the significance of analyzing and optimizing the runtime efficiency of AI/ML workloads as a method of accelerating growth and minimizing prices. Whereas that is essential whatever the compute engine used, the profiling instruments and optimization strategies can fluctuate drastically between platforms. On this submit, we are going to focus on a number of the efficiency optimization choices that pertain to CPU. Our focus can be on Intel® Xeon® CPU processors (with Intel® AVX-512) and on the PyTorch (model 2.4) framework (though comparable strategies could be utilized to different CPUs and frameworks, as properly). Extra particularly, we are going to run our experiments on an Amazon EC2 c7i occasion with an AWS Deep Studying AMI. Please don’t view our alternative of Cloud platform, CPU model, ML framework, or another instrument or library we should always point out, as an endorsement over their options.

Our purpose can be to reveal that though ML growth on CPU will not be our first alternative, there are methods to “soften the blow” and — in some instances — even perhaps make it a viable various.

Disclaimers

Our intention on this submit is to reveal only a few of the ML optimization alternatives obtainable on CPU. Opposite to many of the on-line tutorials on the subject of ML optimization on CPU, we are going to give attention to a coaching workload quite than an inference workload. There are a selection of optimization instruments centered particularly on inference that we’ll not cowl (e.g., see right here and right here).

Please don’t view this submit as a substitute of the official documentation on any of the instruments or strategies that we point out. Understand that given the fast tempo of AI/ML growth, a number of the content material, libraries, and/or directions that we point out could turn into outdated by the point you learn this. Please make sure to consult with the most-up-to-date documentation obtainable.

Importantly, the influence of the optimizations that we focus on on runtime efficiency is prone to fluctuate drastically based mostly on the mannequin and the main points of the surroundings (e.g., see the excessive diploma of variance between fashions on the official PyTorch TouchInductor CPU Inference Efficiency Dashboard). The comparative efficiency numbers we are going to share are particular to the toy mannequin and runtime surroundings that we’ll use. Make sure to reevaluate the entire proposed optimizations by yourself mannequin and runtime surroundings.

Lastly, our focus can be solely on throughput efficiency (as measured in samples per second) — not on coaching convergence. Nonetheless, it must be famous that some optimization strategies (e.g., batch dimension tuning, combined precision, and extra) might have a detrimental impact on the convergence of sure fashions. In some instances, this may be overcome by means of applicable hyperparameter tuning.

We’ll run our experiments on a easy picture classification mannequin with a ResNet-50 spine (from Deep Residual Studying for Picture Recognition). We’ll prepare the mannequin on a faux dataset. The complete coaching script seems within the code block under (loosely based mostly on this instance):

import torch
import torchvision
from torch.utils.information import Dataset, DataLoader
import time
# A dataset with random photographs and labels
class FakeDataset(Dataset):
def __len__(self):
return 1000000
def __getitem__(self, index):
rand_image = torch.randn([3, 224, 224], dtype=torch.float32)
label = torch.tensor(information=index % 10, dtype=torch.uint8)
return rand_image, label
train_set = FakeDataset()
batch_size=128
num_workers=0
train_loader = DataLoader(
dataset=train_set,
batch_size=batch_size,
num_workers=num_workers
)
mannequin = torchvision.fashions.resnet50()
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(mannequin.parameters())
mannequin.prepare()
t0 = time.perf_counter()
summ = 0
depend = 0
for idx, (information, goal) in enumerate(train_loader):
optimizer.zero_grad()
output = mannequin(information)
loss = criterion(output, goal)
loss.backward()
optimizer.step()
batch_time = time.perf_counter() - t0
if idx > 10:  # skip first steps
summ += batch_time
depend += 1
t0 = time.perf_counter()
if idx > 100:
break
print(f'common step time: {summ/depend}')
print(f'throughput: {depend*batch_size/summ}')

Working this script on a c7i.2xlarge (with 8 vCPUs) and the CPU model of PyTorch 2.4, ends in a throughput of 9.12 samples per second. For the sake of comparability, we be aware that the throughput of the identical (unoptimized script) on an Amazon EC2 g5.2xlarge occasion (with 1 GPU and eight vCPUs) is 340 samples per second. Bearing in mind the comparative prices of those two occasion sorts ($0.357 per hour for a c7i.2xlarge and $1.212 for a g5.2xlarge, as of the time of this writing), we discover that coaching on the GPU occasion to provide roughly eleven(!!) occasions higher worth efficiency. Primarily based on these outcomes, the desire for utilizing GPUs to coach ML fashions could be very properly based. Let’s assess a number of the potentialities for decreasing this hole.

On this part we are going to discover some fundamental strategies for growing the runtime efficiency of our coaching workload. Though it’s possible you’ll acknowledge a few of these from our submit on GPU optimization, it is very important spotlight a major distinction between coaching optimization on CPU and GPU platforms. On GPU platforms a lot of our effort was devoted to maximizing the parallelization between (the coaching information preprocessing on) the CPU and (the mannequin coaching on) the GPU. On CPU platforms the entire processing happens on the CPU and our purpose can be to allocate its assets most successfully.

Batch Measurement

Growing the coaching batch dimension can doubtlessly enhance efficiency by decreasing the frequency of the mannequin parameter updates. (On GPUs it has the additional advantage of decreasing the overhead of CPU-GPU transactions corresponding to kernel loading). Nonetheless, whereas on GPU we aimed for a batch dimension that will maximize the utilization of the GPU reminiscence, the identical technique may harm efficiency on CPU. For causes past the scope of this submit, CPU reminiscence is extra difficult and the very best strategy for locating essentially the most optimum batch dimension could also be by means of trial and error. Understand that altering the batch dimension might have an effect on coaching convergence.

The desk under summarizes the throughput of our coaching workload for a number of (arbitrary) selections of batch dimension:

Opposite to our findings on GPU, on the c7i.2xlarge occasion sort our mannequin seems to desire decrease batch sizes.

Multi-process Information Loading

A standard method on GPUs is to assign a number of processes to the info loader in order to cut back the chance of hunger of the GPU. On GPU platforms, a basic rule of thumb is to set the variety of employees in response to the variety of CPU cores. Nonetheless, on CPU platforms, the place the mannequin coaching makes use of the identical assets as the info loader, this strategy might backfire. As soon as once more, the very best strategy for selecting the optimum variety of employees could also be trial and error. The desk under exhibits the typical throughput for various selections of num_workers:

Combined Precision

One other common method is to make use of decrease precision floating level datatypes corresponding to torch.float16 or torch.bfloat16 with the dynamic vary of torch.bfloat16 usually thought of to be extra amiable to ML coaching. Naturally, decreasing the datatype precision can have adversarial results on convergence and must be accomplished fastidiously. PyTorch comes with torch.amp, an computerized combined precision bundle for optimizing the usage of these datatypes. Intel® AVX-512 contains help for the bfloat16 datatype. The modified coaching step seems under:

for idx, (information, goal) in enumerate(train_loader):
optimizer.zero_grad()
with torch.amp.autocast('cpu',dtype=torch.bfloat16):
output = mannequin(information)
loss = criterion(output, goal)
loss.backward()
optimizer.step()

The throughput following this optimization is 24.34 samples per second, a rise of 86%!!

Channels Final Reminiscence Format

Channels final reminiscence format is a beta-level optimization (on the time of this writing), pertaining primarily to imaginative and prescient fashions, that helps storing 4 dimensional (NCHW) tensors in reminiscence such that the channels are the final dimension. This ends in the entire information of every pixel being saved collectively. This optimization pertains primarily to imaginative and prescient fashions. Thought-about to be extra “pleasant to Intel platforms”, this reminiscence format is reported enhance the efficiency of a ResNet-50 on an Intel® Xeon® CPU. The adjusted coaching step seems under:

for idx, (information, goal) in enumerate(train_loader):
information = information.to(memory_format=torch.channels_last)
optimizer.zero_grad()
with torch.amp.autocast('cpu',dtype=torch.bfloat16):
output = mannequin(information)
loss = criterion(output, goal)
loss.backward()
optimizer.step()

The ensuing throughput is 37.93 samples per second — a further 56% enchancment and a complete of 415% in comparison with our baseline experiment. We’re on a job!!

Torch Compilation

In a earlier submit we coated the virtues of PyTorch’s help for graph compilation and its potential influence on runtime efficiency. Opposite to the default keen execution mode through which every operation is run independently (a.ok.a., “eagerly”), the compile API converts the mannequin into an intermediate computation graph which is then JIT-compiled into low-level machine code in a way that’s optimum for the underlying coaching engine. The API helps compilation by way of completely different backend libraries and with a number of configuration choices. Right here we are going to restrict our analysis to the default (TorchInductor) backend and the ipex backend from the Intel® Extension for PyTorch, a library with devoted optimizations for Intel {hardware}. Please see the documentation for applicable set up and utilization directions. The up to date mannequin definition seems under:

import intel_extension_for_pytorch as ipex
mannequin = torchvision.fashions.resnet50()
backend='inductor' # optionally change to 'ipex'
mannequin = torch.compile(mannequin, backend=backend)

Within the case of our toy mannequin, the influence of torch compilation is just obvious when the “channels final” optimization is disabled (and enhance of ~27% for every of the backends). When “channels final” is utilized, the efficiency truly drops. Because of this, we drop this optimization from our subsequent experiments.

There are a selection of alternatives for optimizing the usage of the underlying CPU assets. These embrace optimizing reminiscence administration and thread allocation to the construction of the underlying CPU {hardware}. Reminiscence administration could be improved by means of the usage of superior reminiscence allocators (corresponding to Jemalloc and TCMalloc) and/or decreasing reminiscence accesses which are slower (i.e., throughout NUMA nodes). Threading allocation could be improved by means of applicable configuration of the OpenMP threading library and/or use of Intel’s Open MP library.

Usually talking, these sorts of optimizations require a deep stage understanding of the CPU structure and the options of its supporting SW stack. To simplify issues, PyTorch provides the torch.backends.xeon.run_cpu script for mechanically configuring the reminiscence and threading libraries in order to optimize runtime efficiency. The command under will lead to the usage of the devoted reminiscence and threading libraries. We’ll return to the subject of NUMA nodes once we focus on the choice of distributed coaching.

We confirm applicable set up of TCMalloc (conda set up conda-forge::gperftools) and Intel’s Open MP library (pip set up intel-openmp), and run the next command.

python -m torch.backends.xeon.run_cpu prepare.py

Using the run_cpu script additional boosts our runtime efficiency to 39.05 samples per second. Be aware that the run_cpu script contains many controls for additional tuning efficiency. Make sure to try the documentation to be able to maximize its use.

The Intel® Extension for PyTorch contains further alternatives for coaching optimization by way of its ipex.optimize operate. Right here we reveal its default use. Please see the documentation to be taught of its full capabilities.

 mannequin = torchvision.fashions.resnet50()
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(mannequin.parameters())
mannequin.prepare()
mannequin, optimizer = ipex.optimize(
mannequin, 
optimizer=optimizer,
dtype=torch.bfloat16
)

Mixed with the reminiscence and thread optimizations mentioned above, the resultant throughput is 40.73 samples per second. (Be aware {that a} comparable result’s reached when disabling the “channels final” configuration.)

Intel® Xeon® processors are designed with Non-Uniform Reminiscence Entry (NUMA) through which the CPU reminiscence is split into teams, a.ok.a., NUMA nodes, and every of the CPU cores is assigned to at least one node. Though any CPU core can entry the reminiscence of any NUMA node, the entry to its personal node (i.e., its native reminiscence) is far quicker. This offers rise to the notion of distributing coaching throughout NUMA nodes, the place the CPU cores assigned to every NUMA node act as a single course of in a distributed course of group and information distribution throughout nodes is managed by Intel® oneCCL, Intel’s devoted collective communications library.

We are able to run information distributed coaching throughout NUMA nodes simply utilizing the ipexrun utility. Within the following code block (loosely based mostly on this instance) we adapt our script to run information distributed coaching (in response to utilization detailed right here):

import os, time
import torch
from torch.utils.information import Dataset, DataLoader
from torch.utils.information.distributed import DistributedSampler
import torch.distributed as dist
import torchvision
import oneccl_bindings_for_pytorch as torch_ccl
import intel_extension_for_pytorch as ipex
os.environ["MASTER_ADDR"] = "127.0.0.1"
os.environ["MASTER_PORT"] = "29500"
os.environ["RANK"] = os.environ.get("PMI_RANK", "0")
os.environ["WORLD_SIZE"] = os.environ.get("PMI_SIZE", "1")
dist.init_process_group(backend="ccl", init_method="env://")
rank = os.environ["RANK"]
world_size = os.environ["WORLD_SIZE"]
batch_size = 128
num_workers = 0
# outline dataset and dataloader
class FakeDataset(Dataset):
def __len__(self):
return 1000000
def __getitem__(self, index):
rand_image = torch.randn([3, 224, 224], dtype=torch.float32)
label = torch.tensor(information=index % 10, dtype=torch.uint8)
return rand_image, label
train_dataset = FakeDataset()
dist_sampler = DistributedSampler(train_dataset)
train_loader = DataLoader(
dataset=train_dataset, 
batch_size=batch_size,
num_workers=num_workers,
sampler=dist_sampler
)
# outline mannequin artifacts
mannequin = torchvision.fashions.resnet50()
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(mannequin.parameters())
mannequin.prepare()
mannequin, optimizer = ipex.optimize(
mannequin, 
optimizer=optimizer,
dtype=torch.bfloat16
)
# configure DDP
mannequin = torch.nn.parallel.DistributedDataParallel(mannequin)
# run coaching loop
# destroy the method group
dist.destroy_process_group()

Sadly, as of the time of this writing, the Amazon EC2 c7i occasion household doesn’t embrace a multi-NUMA occasion sort. To check our distributed coaching script, we revert again to a Amazon EC2 c6i.32xlarge occasion with 64 vCPUs and a pair of NUMA nodes. We confirm the set up of Intel® oneCCL Bindings for PyTorch and run the next command (as documented right here):

source $(python -c "import oneccl_bindings_for_pytorch as torch_ccl;print(torch_ccl.cwd)")/env/setvars.sh
# This instance command would make the most of all of the numa sockets of the processor, taking every socket as a rank.
ipexrun cpu --nnodes 1 --omp_runtime intel prepare.py 

The next desk compares the efficiency outcomes on the c6i.32xlarge occasion with and with out distributed coaching:

In our experiment, information distribution did not enhance the runtime efficiency. Please see ipexrun documentation for extra efficiency tuning choices.

In earlier posts (e.g., right here) we mentioned the PyTorch/XLA library and its use of XLA compilation to allow PyTorch based mostly coaching on XLA gadgets corresponding to TPU, GPU, and CPU. Much like torch compilation, XLA makes use of graph compilation to generate machine code that’s optimized for the goal machine. With the institution of the OpenXLA Challenge, one of many acknowledged targets was to help excessive efficiency throughout all {hardware} backends, together with CPU (see the CPU RFC right here). The code block under demonstrates the changes to our unique (unoptimized) script required to coach utilizing PyTorch/XLA:

import torch
import torchvision
import timeimport torch_xla
import torch_xla.core.xla_model as xm
machine = xm.xla_device()
mannequin = torchvision.fashions.resnet50().to(machine)
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(mannequin.parameters())
mannequin.prepare()
for idx, (information, goal) in enumerate(train_loader):
information = information.to(machine)
goal = goal.to(machine)
optimizer.zero_grad()
output = mannequin(information)
loss = criterion(output, goal)
loss.backward()
optimizer.step()
xm.mark_step()

Sadly, (as of the time of this writing) the XLA outcomes on our toy mannequin appear far inferior to the (unoptimized) outcomes we noticed above (— by as a lot as 7X). We count on this to enhance as PyTorch/XLA’s CPU help matures.

We summarize the outcomes of a subset of our experiments within the desk under. For the sake of comparability, we add the throughput of coaching our mannequin on Amazon EC2 g5.2xlarge GPU occasion following the optimization steps mentioned in this submit. The samples per greenback was calculated based mostly on the Amazon EC2 On-demand pricing web page ($0.357 per hour for a c7i.2xlarge and $1.212 for a g5.2xlarge, as of the time of this writing).

Though we succeeded in boosting the coaching efficiency of our toy mannequin on the CPU occasion by a substantial margin (446%), it stays inferior to the (optimized) efficiency on the GPU occasion. Primarily based on our outcomes, coaching on GPU could be ~6.7 occasions cheaper. It’s doubtless that with further efficiency tuning and/or making use of further optimizations methods, we might additional shut the hole. As soon as once more, we emphasize that the comparative efficiency outcomes we’ve got reached are distinctive to this mannequin and runtime surroundings.

Amazon EC2 Spot Situations Reductions

The elevated availability of cloud-based CPU occasion sorts (in comparison with GPU occasion sorts) could indicate better alternative for acquiring compute energy at discounted charges, e.g., by means of Spot Occasion utilization. Amazon EC2 Spot Situations are situations from surplus cloud service capability which are supplied for a reduction of as a lot as 90% off the On-Demand pricing. In trade for the discounted worth, AWS maintains the precise to preempt the occasion with little to no warning. Given the excessive demand for GPUs, it’s possible you’ll discover CPU spot situations simpler to get ahold of than their GPU counterparts. On the time of this writing, c7i.2xlarge Spot Occasion worth is $0.1291 which might enhance our samples per greenback outcome to 1135.76 and additional reduces the hole between the optimized GPU and CPU worth performances (to 2.43X).

Whereas the runtime efficiency outcomes of the optimized CPU coaching of our toy mannequin (and our chosen surroundings) had been decrease than the GPU outcomes, it’s doubtless that the identical optimization steps utilized to different mannequin architectures (e.g., ones that embrace parts that aren’t supported by GPU) could outcome within the CPU efficiency matching or beating that of the GPU. And even in instances the place the efficiency hole is just not bridged, there could very properly be instances the place the scarcity of GPU compute capability would justify operating a few of our ML workloads on CPU.

Given the ubiquity of the CPU, the flexibility to make use of them successfully for coaching and/or operating ML workloads might have large implications on growth productiveness and on end-product deployment technique. Whereas the character of the CPU structure is much less amiable to many ML functions when in comparison with the GPU, there are numerous instruments and strategies obtainable for enhancing its efficiency — a choose few of which we’ve got mentioned and demonstrated on this submit.

On this submit we centered optimizing coaching on CPU. Please make sure to try our many different posts on medium protecting all kinds of subjects pertaining to efficiency evaluation and optimization of machine studying workloads.