Jump to content

Improve training time of distributed machine learning with NCCL Fast Socket

Recommended Posts

Large Machine Learning (ML) models – such as large language models, generative AI, and vision models – are dramatically increasing the number of trainable parameters and are achieving state-of-the-art results. Increasing the number of parameters results in the model being too large to fit on a single VM instance thus demands distributed compute to spread the model across multiple nodes. Google Kubernetes Engine (GKE) has built-in support for NCCL Fast Socket, to help improve the time to train large ML models with distributed, multi-node clusters.

Enterprises are looking for faster and cheaper performance to train their ML models. With distributed training, communicating gradients across nodes is a performance bottleneck. Optimizing inter-node latency is critical to reduce training time and costs. Distributed training uses collective communication as a transport layer over the network between the multiple hosts. Collective communication primitives such as all-gather, all-reduce, broadcast, reduce, reduce-scatter, and point-to-point send and receive are used in distributed training in Machine Learning. 

The NVIDIA Collective Communication Library (NCCL) is commonly used by popular ML frameworks such as TensorFlow and PyTorch. It is a highly optimized implementation for high bandwidth and low latency between NVIDIA GPUs. Google developed a proprietary version of NCCL called NCCL Fast Socket to optimize performance for deep learning on Google Cloud.

NCCL Fast Socket uses a number of techniques to achieve better and more consistent NCCL performance.

  • Use of multiple network flows to attain maximum throughput. NCCL Fast Socket introduces additional optimizations over NCCL’s built-in multi-stream support, including better overlapping of multiple communication requests.

  • Dynamic load balancing of multiple network flows. NCCL can adapt to changing network and host conditions. With this optimization, straggler network flows will not significantly slow down the entire NCCL collective operation.

  • Integration with Google Cloud’s Andromeda virtual network stack.This increases overall network throughput by avoiding contentions in virtual machines (VMs).

We tested (NVIDIA NCCL tests) the performance of NCCL Fast Socket vs NCCL on various machine shapes with 2 node GKE clusters.

NCCL Fast Socket.jpg

The following chart shows the results. For each machine shape, the NCCL performance without Fast Socket is normalized to 1. In each case, using NCCL Fast Socket demonstrated increased performance in a range of 1.3 to 2.6 times faster internetwork communication speed.

NCCL Fast Socket Blog Benchmarks.jpg

As a built-in feature, GKE users can take advantage of NCCL Fast Socket without changing or recompiling their applications, ML frameworks (such as TensorFlow or PyTorch), or even the NCCL library itself. To start using NCCL Fast Socket, create a node pool that uses the plugin with the --enable-fast-socket and --enable-gvnic flags. You can also update an existing node pool using gcloud container node-pools update.

[StructValue([(u'code', u'gcloud container node-pools create NODEPOOL_NAME \\\r\n --accelerator type=ACCELERATOR_TYPE, count=ACCELERATOR_COUNT \\\r\n --machine-type=MACHINE_TYPE \\\r\n --cluster=CLUSTER_NAME \\\r\n --enable-fast-socket \\\r\n --enable-gvnic'), (u'language', u''), (u'caption', <wagtail.wagtailcore.rich_text.RichText object at 0x3e8fe9903c10>)])]

To achieve better network throughput with NCCL, Google Virtual NICs (gVNICs) must be enabled when creating VM instances. For detailed instructions on how to use gVNICs, please refer to the gVNIC guide

To verify that NCCL Fast Socket has been enabled, view the kube-system pods:

[StructValue([(u'code', u'kubectl get pods -n kube-system'), (u'language', u''), (u'caption', <wagtail.wagtailcore.rich_text.RichText object at 0x3e8fe9de5d90>)])]

And the output should b similar to:

[StructValue([(u'code', u'NAME READY STATUS RESTARTS AGE\r\nfast-socket-installer-qvfdw 2/2 Running 0 10m\r\nfast-socket-installer-rtjs4 2/2 Running 0 10m\r\nfast-socket-installer-tm294 2/2 Running 0 10m'), (u'language', u''), (u'caption', <wagtail.wagtailcore.rich_text.RichText object at 0x3e8fe98f9cd0>)])]

To learn more visit GKE NCCL Fast Socket documentation. We look forward to hearing how NCCL Fast Socket improves your ML Training experience on GKE.

Link to comment
Share on other sites

Join the conversation

You can post now and register later. If you have an account, sign in now to post with your account.

Reply to this topic...

×   Pasted as rich text.   Paste as plain text instead

  Only 75 emoji are allowed.

×   Your link has been automatically embedded.   Display as a link instead

×   Your previous content has been restored.   Clear editor

×   You cannot paste images directly. Upload or insert images from URL.

  • Create New...