Efficient Load Balancing with Ray on Amazon SageMaker | by Chaim Rand | Sep, 2023

A technique for rising DNN coaching effectivity and lowering coaching prices

In earlier posts (e.g., here) we expanded on the significance of profiling and optimizing the efficiency of your DNN coaching workloads. Coaching deep studying fashions — particularly massive ones — could be an costly endeavor. Your means to maximize the utilization of your coaching assets in a way that each accelerates your mannequin convergence and minimizes coaching prices, generally is a decisive issue within the success of your challenge. Efficiency optimization is an iterative course of by which we establish and handle the efficiency bottlenecks in our software, i.e., the parts in our software which might be stopping us from rising useful resource utilization and/or accelerating the run time.

This put up is the third in a collection of posts that target one of many extra frequent efficiency bottlenecks that we encounter when coaching deep studying fashions, the data pre-processing bottleneck. A knowledge pre-processing bottleneck happens when our GPU (or different accelerator) — sometimes the most costly useful resource in our coaching setup — finds itself idle whereas it waits for information enter from overly tasked CPU assets.

In our first post on the subject we mentioned and demonstrated alternative ways of addressing any such bottleneck, together with:

1. Selecting a coaching occasion with a CPU to GPU compute ratio that’s extra suited to your workload,

2. Bettering the workload stability between the CPU and GPU by shifting a few of the CPU operations to the GPU, and

3. Offloading a few of the CPU computation to auxiliary CPU-worker gadgets.

We demonstrated the third possibility utilizing the TensorFlow Data Service API, an answer particular to TensorFlow, by which a portion of the enter information processing could be offloaded onto different gadgets utilizing gRPC because the underlying communication protocol.

In our second post, we proposed a extra general-purpose gRPC-based resolution for utilizing auxiliary CPU staff and demonstrated it on a toy PyTorch mannequin. Though it required a bit extra handbook coding and tuning than the TensorFlow Data Service API, the answer supplied a lot better robustness and allowed for a similar optimization in coaching efficiency.

Load Balancing with Ray

On this put up we are going to reveal a further methodology for utilizing auxiliary CPU staff that goals to mix the robustness of the general-purpose resolution with the simplicity and ease-of-use of the TensorFlow-specific API. The tactic we are going to reveal will use Ray Datasets from the Ray Data library. By leveraging the complete energy of Ray’s resource management and distributed scheduling programs, Ray Knowledge is ready to run our coaching information enter pipeline in method that’s each scalable and distributed. Particularly, we are going to configure our Ray Dataset in such a means that the library will routinely detect and make the most of all the accessible CPU assets for pre-processing the coaching information. We are going to additional wrap our mannequin coaching loop with a Ray AIR Trainer in order to allow seamless scaling to a multi-GPU setting.

Deploying a Ray Cluster on Amazon SageMaker

A prerequisite for utilizing the Ray framework and the utilities it gives in a multi-node surroundings is the deployment of a Ray cluster. Usually, designing, deploying, managing, and sustaining such a compute cluster generally is a daunting process and infrequently requires a devoted devops engineer (or crew of engineers). This may pose an insurmountable impediment for some improvement groups. On this put up we are going to reveal how you can overcome this impediment utilizing AWS’s managed coaching service, Amazon SageMaker. Particularly, we are going to create a SageMaker heterogenous cluster with each GPU situations and CPU situations and use it to deploy a Ray cluster at startup. We are going to then run the Ray AIR coaching software on this Ray cluster whereas counting on Ray’s backend to carry out efficient load balancing throughout all the assets within the cluster. When the coaching software is accomplished, the Ray cluster shall be torn down routinely. Utilizing SageMaker on this method, allows us to deploy and use a Ray cluster with out the overhead that’s generally related to cluster administration.

Ray is a strong framework that permits a variety of machine studying workloads. On this put up we are going to reveal just some of its capabilities and APIs utilizing Ray model 2.6.1. This put up shouldn’t be used as a alternative for the Ray documentation. You’ll want to take a look at the official documentation for probably the most acceptable and up-to-date use of the Ray utilities.

Earlier than we get began, particular because of Boruch Chalk for introducing me to the Ray Knowledge library and its distinctive capabilities.

To facilitate our dialogue, we are going to outline and practice a easy PyTorch (2.0) Vision Transformer-based classification mannequin that we’ll practice on an artificial dataset comprised of random photos and labels. The Ray AIR documentation consists of all kinds of examples that reveal how you can construct several types of coaching workloads utilizing Ray AIR. The script we create right here loosely follows the steps described within the PyTorch image classifier example.

Defining the Ray Dataset and Preprocessor

The Ray AIR Trainer API distinguishes between the uncooked dataset and the preprocessing pipeline that’s utilized to the weather of the dataset earlier than feeding them into the coaching loop. For our uncooked Ray dataset we create a easy range of integers of measurement num_records. Subsequent, we outline the Preprocessor that we want to apply to our dataset. Our Ray Preprocesser comprises two parts: The primary is a BatchMapper that maps the uncooked integers to random image-label pairs. The second is a TorchVisionPreprocessor that performs a torchvision transform on our random batches which converts them to PyTorch tensors and applies a collection of GaussianBlur operations. The GaussianBlur operations are meant to simulate a comparatively heavy information pre-processing pipeline. The 2 Preprocessors are mixed utilizing a Chain Preprocessor. The creation of the Ray dataset and Preprocessor is demonstrated within the code block beneath:

import ray
from typing import Dict, Tuple
import numpy as np
import torchvision.transforms as transforms
from ray.information.preprocessors import Chain, BatchMapper, TorchVisionPreprocessor
def get_ds(batch_size, num_records):
# create a uncooked Ray tabular dataset
ds = ray.information.vary(num_records)
# map an integer to a random image-label pair
def synthetic_ds(batch: Tuple[int]) -> Dict[str, np.ndarray]:
labels = batch['id']
batch_size = len(labels)
photos = np.random.randn(batch_size, 224, 224, 3).astype(np.float32)
labels = np.array([label % 1000 for label in labels]).astype(
dtype=np.int64)
return {"picture": photos, "label": labels}
# step one of the prepocessor maps batches of ints to
# random image-label pairs
synthetic_data = BatchMapper(synthetic_ds, 
batch_size=batch_size, 
batch_format="numpy")
# we outline a torchvision remodel that converts the numpy pairs to 
# tensors after which applies a collection of gaussian blurs to simulate
# heavy preprocessing   
remodel = transforms.Compose(
[transforms.ToTensor()] + [transforms.GaussianBlur(11)]*10
)
# the second step of the prepocessor appplies the torchvision tranform
vision_preprocessor = TorchVisionPreprocessor(columns=["image"], 
remodel=remodel)
# mix the preprocessing steps
preprocessor = Chain(synthetic_data, vision_preprocessor)
return ds, preprocessor

Notice that the Ray information pipeline will routinely use all the CPUs which might be accessible within the Ray cluster. This consists of the CPU assets which might be on the GPU occasion in addition to the CPU assets of any further auxiliary situations within the cluster.

Defining the Coaching Loop

The following step is to outline the coaching sequence that may run on every of the coaching staff (e.g., GPUs). First we outline the mannequin utilizing the favored timm (0.6.13) Python package deal and wrap it utilizing the train.torch.prepare_model API. Subsequent, we extract the suitable shard from the dataset and outline an iterator that yields information batches with the requested batch measurement and copies them to the coaching machine. Then comes the coaching loop itself which is comprised of normal PyTorch code. After we exit the loop, we report again the resultant loss metric. The per-worker coaching sequence is demonstrated within the code block beneath:

import time
from ray import practice
from ray.air import session
import torch.nn as nn
import torch.optim as optim
from timm.fashions.vision_transformer import VisionTransformer
# construct a ViT mannequin utilizing timm
def build_model():
return VisionTransformer()
# outline the coaching loop per employee
def train_loop_per_worker(config):
# wrap the PyTorch mannequin with a Ray object
mannequin = practice.torch.prepare_model(build_model())
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(mannequin.parameters(), lr=0.001, momentum=0.9)
# get the suitable dataset shard
train_dataset_shard = session.get_dataset_shard("practice")
# create an iterator that returns batches from the dataset
train_dataset_batches = train_dataset_shard.iter_torch_batches(
batch_size=config["batch_size"],
prefetch_batches=config["prefetch_batches"],
machine=practice.torch.get_device()
)
t0 = time.perf_counter()
for i, batch in enumerate(train_dataset_batches):
# get the inputs and labels
inputs, labels = batch["image"], batch["label"]
# zero the parameter gradients
optimizer.zero_grad()
# ahead + backward + optimize
outputs = mannequin(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
# print statistics
if i % 100 == 99:  # print each 100 mini-batches
avg_time = (time.perf_counter()-t0)/100
print(f"Iteration {i+1}: avg time per step {avg_time:.3f}")
t0 = time.perf_counter()
metrics = dict(running_loss=loss.merchandise())
session.report(metrics)

Defining the Ray Torch Coach

As soon as we’ve outlined our information pipeline and coaching loop, we are able to transfer on to establishing the Ray TorchTrainer. We configure the Coach in a way that takes under consideration the accessible assets within the cluster. Particularly, we set the variety of coaching staff in line with the variety of GPUs and we set the batch measurement in line with the reminiscence accessible on our goal GPU. We construct our dataset with the variety of information required to coach for exactly 1000 steps.

from ray.practice.torch import TorchTrainer
from ray.air.config import ScalingConfig
def train_model():
# we are going to configure the variety of staff, the dimensions of our
# dataset, and the dimensions of the information storage in line with the
# accessible assets 
num_gpus = int(ray.available_resources().get("GPU", 0))
# set the variety of coaching staff in line with the variety of GPUs
num_workers = num_gpus if num_gpus > 0 else 1
# we set the batch measurement based mostly on the GPU reminiscence capability of the
# Amazon EC2 g5 occasion household
batch_size = 64
# create an artificial dataset with sufficient information to coach for 1000 steps
num_records = batch_size * 1000 * num_workers
ds, preprocessor = get_ds(batch_size, num_records)
ds = preprocessor(ds) 
coach = TorchTrainer(
train_loop_per_worker=train_loop_per_worker,
train_loop_config={"batch_size": batch_size},
datasets={"practice": ds},
scaling_config=ScalingConfig(num_workers=num_workers, 
use_gpu=num_gpus > 0),
)
coach.match()

Deploy a Ray Cluster and Run the Coaching Sequence

We now outline the entry level of our coaching script. It’s right here that we setup the Ray cluster and provoke the coaching sequence on the top node. We use the Environment class from the sagemaker-training library to find the situations within the heterogenous SageMaker cluster as described in this tutorial. We outline the primary node of the GPU occasion group as our Ray cluster head node and run the suitable command on all the different nodes to attach them to the cluster. (See the Ray documentation for extra particulars on creating clusters.) We program the top node to attend till all of the nodes have linked after which begin the coaching sequence. This ensures that Ray will make the most of all the accessible assets when defining and distributing the underlying Ray duties.

import time
import subprocess
from sagemaker_training import surroundings
if __name__ == "__main__":
# use the Atmosphere() class to auto-discover the SageMaker cluster
env = surroundings.Atmosphere()
if env.current_instance_group == 'gpu' and 
env.current_instance_group_hosts.index(env.current_host) == 0:
# the top node begins a ray cluster
p = subprocess.Popen('ray begin --head --port=6379',
shell=True).wait()
ray.init()
# calculate the whole variety of nodes within the cluster
teams = env.instance_groups_dict.values()
cluster_size = sum(len(v['hosts']) for v in record(teams))
# wait till all SageMaker nodes have linked to the Ray cluster
connected_nodes = 1
whereas connected_nodes < cluster_size:
time.sleep(1)
assets = ray.available_resources().keys()
connected_nodes = sum(1 for s in record(assets) if 'node' in s)
# name the coaching sequence
train_model()
# tear down the ray cluster
p = subprocess.Popen("ray down", shell=True).wait()
else:
# employee nodes connect to the top node
head = env.instance_groups_dict['gpu']['hosts'][0]
p = subprocess.Popen(
f"ray begin --address='{head}:6379'",
shell=True).wait()
# utility for checking if the cluster continues to be alive
def is_alive():
from subprocess import Popen
p = Popen('ray standing', shell=True)
p.talk()[0]
return p.returncode
# preserve node alive till the method on head node completes
whereas is_alive() == 0:
time.sleep(10)

Coaching on an Amazon SageMaker Heterogenous Cluster

With our coaching script full, we are actually tasked with deploying it to an Amazon SageMaker Heterogenous Cluster. To do that we observe the steps described in this tutorial. We begin by making a source_dir listing into which we place the our practice.py script and a necessities.txt file containing the 2 pip packages our script relies on, timm and ray[air]. These are routinely put in on every of the nodes within the SageMaker cluster. We outline two SageMaker Instance Groups, the primary with a single ml.g5.xlarge occasion (containing 1 GPU and 4 vCPUs), and the second with a single ml.c5.4xlarge occasion (containing 16 vCPUs). We then use the SageMaker PyTorch estimator to outline and deploy our coaching job to the cloud.

from sagemaker.pytorch import PyTorch
from sagemaker.instance_group import InstanceGroup
cpu_group = InstanceGroup("cpu", "ml.c5.4xlarge", 1)
gpu_group = InstanceGroup("gpu", "ml.g5.xlarge", 1)
estimator = PyTorch(
entry_point='practice.py',
source_dir='./source_dir',
framework_version='2.0.0',
position='<arn position>',
py_version='py310',
job_name='hetero-cluster',
instance_groups=[gpu_group, cpu_group]
)
estimator.match()

Within the desk beneath we evaluate the runtime outcomes of operating our coaching script in two totally different settings: a single ml.g5.xlarge GPU occasion and a heterogenous cluster containing an ml.g5.xlarge occasion and an ml.c5.4xlarge. We consider the system useful resource utilization utilizing Amazon CloudWatch and estimate the coaching price utilizing the Amazon SageMaker pricing accessible as of the time of this writing ($0.816 per hour for the ml.c5.4xlarge occasion and $1.408 for the ml.g5.xlarge).

The comparatively excessive CPU utilization mixed with the low GPU utilization of the only occasion experiment signifies a efficiency bottleneck within the information pre-processing pipeline. These are clearly addressed when shifting to the heterogenous cluster. Not solely does the GPU utilization enhance, however so does the coaching pace. Total, the value effectivity of coaching will increase by 23%.

We must always emphasize that these toy experiments had been created purely for the aim of demonstrating the automated load balancing options enabled by the Ray ecosystem. It’s potential that tuning of the management parameters might have led to improved efficiency. Additionally it is possible that selecting a distinct resolution for addressing the CPU bottleneck (akin to selecting an occasion from the EC2 g5 household with extra CPUs) might have resulted in higher price efficiency.

On this put up we’ve got demonstrated how Ray datasets can be utilized to stability the load of a heavy information pre-processing pipeline throughout all the accessible CPU staff within the cluster. This permits us to simply handle CPU bottlenecks by merely including auxiliary CPU situations to the coaching surroundings. Amazon SageMaker’s heterogenous cluster assist is a compelling method to run a Ray coaching job within the cloud because it handles all sides of the cluster administration avoiding the necessity for devoted devops assist.

Take into account that the answer offered right here is only one of many various methods of addressing CPU bottlenecks. The most effective resolution for you’ll extremely depend upon the main points of your challenge.

As traditional, please be at liberty to achieve out with feedback, corrections, and questions.

The post Efficient Load Balancing with Ray on Amazon SageMaker | by Chaim Rand | Sep, 2023 appeared first on AIPressRoom.