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 enterprise. Your capacity to maximize the utilization of your coaching sources in a way that each accelerates your mannequin convergence and minimizes coaching prices, generally is a decisive issue within the success of your mission. Efficiency optimization is an iterative course of wherein we determine and deal with the efficiency bottlenecks in our utility, i.e., the parts in our utility which can be stopping us from rising useful resource utilization and/or accelerating the run time.

This submit is the third in a sequence of posts that concentrate on one of many extra widespread 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 various accelerator) — usually the most costly useful resource in our coaching setup — finds itself idle whereas it waits for information enter from overly tasked CPU sources.

In our first post on the subject we mentioned and demonstrated other ways of addressing this sort of bottleneck, together with:

1. Selecting a coaching occasion with a CPU to GPU compute ratio that’s extra suited to your workload,
2. Enhancing the workload steadiness between the CPU and GPU by transferring among the CPU operations to the GPU, and
3. Offloading among the CPU computation to auxiliary CPU-worker gadgets.

We demonstrated the third choice utilizing the TensorFlow Data Service API, an answer particular to TensorFlow, wherein 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 answer 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 offered a lot higher robustness and allowed for a similar optimization in coaching efficiency.

Load Balancing with Ray

On this submit we’ll show an extra methodology for utilizing auxiliary CPU staff that goals to mix the robustness of the general-purpose answer with the simplicity and ease-of-use of the TensorFlow-specific API. The strategy we’ll show will use Ray Datasets from the Ray Data library. By leveraging the complete energy of Ray’s resource management and distributed scheduling techniques, Ray Knowledge is ready to run our coaching information enter pipeline in method that’s each scalable and distributed. Specifically, we’ll configure our Ray Dataset in such a means that the library will robotically detect and make the most of the entire out there CPU sources for pre-processing the coaching information. We’ll 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 affords in a multi-node atmosphere is the deployment of a Ray cluster. On the whole, designing, deploying, managing, and sustaining such a compute cluster generally is a daunting job and infrequently requires a devoted devops engineer (or group of engineers). This could pose an insurmountable impediment for some growth groups. On this submit we’ll show overcome this impediment utilizing AWS’s managed coaching service, Amazon SageMaker. Specifically, we’ll create a SageMaker heterogenous cluster with each GPU situations and CPU situations and use it to deploy a Ray cluster at startup. We’ll then run the Ray AIR coaching utility on this Ray cluster whereas counting on Ray’s backend to carry out efficient load balancing throughout the entire sources within the cluster. When the coaching utility is accomplished, the Ray cluster will probably be torn down robotically. 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 robust framework that allows a variety of machine studying workloads. On this submit we’ll show only a few of its capabilities and APIs utilizing Ray model 2.6.1. This submit shouldn’t be used as a alternative for the Ray documentation. Make sure you try the official documentation for probably the most applicable and up-to-date use of the Ray utilities.

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

To facilitate our dialogue, we’ll 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 contains all kinds of examples that show 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 sequence 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 under:

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 sequence 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

Word that the Ray information pipeline will robotically use the entire CPUs which can be out there within the Ray cluster. This contains the CPU sources which can be on the GPU occasion in addition to the CPU sources of any extra 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 gadget. Then comes the coaching loop itself which is comprised of ordinary PyTorch code. Once we exit the loop, we report again the resultant loss metric. The per-worker coaching sequence is demonstrated within the code block under:

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"],
gadget=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 into consideration the out there sources within the cluster. Particularly, we set the variety of coaching staff in accordance with the variety of GPUs and we set the batch measurement in accordance with the reminiscence out there 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'll configure the variety of staff, the scale of our
# dataset, and the scale of the info storage in accordance with the
# out there sources 
num_gpus = int(ray.available_resources().get("GPU", 0))
# set the variety of coaching staff in accordance 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 pinnacle 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 the entire different nodes to attach them to the cluster. (See the Ray documentation for extra particulars on creating clusters.) We program the pinnacle 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 the entire out there sources when defining and distributing the underlying Ray duties.

import time
import subprocess
from sagemaker_training import atmosphere
if __name__ == "__main__":
# use the Surroundings() class to auto-discover the SageMaker cluster
env = atmosphere.Surroundings()
if env.current_instance_group == 'gpu' and 
env.current_instance_group_hosts.index(env.current_host) == 0:
# the pinnacle node begins a ray cluster
p = subprocess.Popen('ray begin --head --port=6379',
shell=True).wait()
ray.init()
# calculate the overall variety of nodes within the cluster
teams = env.instance_groups_dict.values()
cluster_size = sum(len(v['hosts']) for v in checklist(teams))
# wait till all SageMaker nodes have linked to the Ray cluster
connected_nodes = 1
whereas connected_nodes < cluster_size:
time.sleep(1)
sources = ray.available_resources().keys()
connected_nodes = sum(1 for s in checklist(sources) 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 pinnacle 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 remains 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 at the moment are tasked with deploying it to an Amazon SageMaker Heterogenous Cluster. To do that we comply with 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 will depend on, timm and ray[air]. These are robotically 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 under 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 value utilizing the Amazon SageMaker pricing out there 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 one occasion experiment signifies a efficiency bottleneck within the information pre-processing pipeline. These are clearly addressed when transferring to the heterogenous cluster. Not solely does the GPU utilization improve, however so does the coaching velocity. General, 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 could have led to improved efficiency. Additionally it is doubtless that selecting a special answer for addressing the CPU bottleneck (comparable to selecting an occasion from the EC2 g5 household with extra CPUs) could have resulted in higher value efficiency.

On this submit we’ve demonstrated how Ray datasets can be utilized to steadiness the load of a heavy information pre-processing pipeline throughout the entire out there CPU staff within the cluster. This allows us to simply deal with CPU bottlenecks by merely including auxiliary CPU situations to the coaching atmosphere. Amazon SageMaker’s heterogenous cluster help is a compelling solution to run a Ray coaching job within the cloud because it handles all sides of the cluster administration avoiding the necessity for devoted devops help.

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

As ordinary, please be happy to achieve out with feedback, corrections, and questions.