Optimizing Distributed Machine Learning Training on Google Vertex AI: Fast Socket and Reduction Server

Google Vertex AI is a managed cloud platform that integrates AutoML and AI Platform to support the full lifecycle of machine learning, from data preprocessing to model serving. The presentation focuses on performance optimizations for large‑scale distributed training.

Background and Memory Wall – As model sizes and data volumes grow, training on a single GPU becomes insufficient due to limited compute and memory bandwidth. Distributed training is now the mainstream solution, but the memory‑bandwidth gap (memory bandwidth improving ~30× versus compute ~90,000× over 20 years) creates a bottleneck, especially during gradient aggregation (all‑reduce), which can consume up to two‑thirds of training time.

Optimization Path – Optimizations are categorized into three layers: framework‑level (e.g., GSPMD, GPipe, DeepSpeed ZeRO), precision‑level (low‑precision compression), and communication‑level. The talk concentrates on the communication layer, which is framework‑agnostic and can benefit all users.

Fast Socket: High‑Performance NCCL Network Stack

Large‑message throughput is limited by many TCP connections and uneven bandwidth across sockets, causing straggler effects.

Fast Socket introduces fine‑grained data slicing, dynamic load‑balancing based on per‑socket progress, and zero‑copy sends, reducing CPU/GPU overhead and improving large‑message throughput.

For small messages, Fast Socket eliminates extra thread hops by letting the proxy thread send directly, inlines control messages, and uses kernel busy‑polling to lower latency and jitter.

End‑to‑end tests on 100 GbE show 60 %+ bandwidth improvements for all‑reduce across message sizes and 30 %+ step‑per‑second speed‑up for BERT‑Large fine‑tuning, without any user‑side code changes.

Reduction Server: Accelerating Gradient Aggregation

Inspired by parameter‑server architecture, a lightweight CPU‑based reduction server aggregates gradients from workers and returns the reduced result, halving the data transferred compared to traditional ring all‑reduce.

The design reduces all‑reduce latency from O(N) to O(1) for small messages and improves algorithmic efficiency, effectively doubling the algorithmic bandwidth.

Implementation uses NCCL’s communication layer on the worker side and a high‑performance Fiber‑based network stack with SIMD‑optimized reduction kernels on the server side.

Performance tests demonstrate up to 75 % training speed‑up and a corresponding reduction in total cost of ownership, even after accounting for the modest CPU node overhead.

Conclusion and Outlook – Memory‑wall issues will persist, so platform‑level, framework‑agnostic optimizations like Fast Socket and Reduction Server are essential. Future work will explore additional parallelism strategies and deeper integration with AI frameworks, while maintaining the goal of transparent, cost‑effective acceleration for large‑scale AI training.