DISTRIBUTED TRAINING OF MACHINE LEARNING MODELS

Smart overview of the Invention

The patent application describes a method for distributed training of large machine learning models using multiple servers and computing devices. These devices are organized into partition groups, each capable of storing a replica of the model's state information. The training process involves a multi-stage gathering of state information, utilizing different communication channels at various stages, including both inter-server and intra-server communications. Once trained, the model is stored for further use.

Background

The application builds on previous work in distributed machine learning training, addressing challenges posed by large models with extensive parameters and datasets. Machine learning models are increasingly applied across various sectors, benefiting from advancements in data processing capabilities and connectivity. However, single devices often lack the capacity to handle the state information required for training such large-scale models, necessitating a distributed approach.

Communication Strategies

The disclosed methods focus on optimizing communication costs during training by considering network heterogeneity. Differences in bandwidth and latency between GPUs on the same server versus those on different servers are significant. The application proposes strategies like scale-aware model partitioning, hierarchical communication, and multi-level gradient synchronization to mitigate these costs, enhancing scalability and efficiency.

Scale-Aware Model Partitioning

In scale-aware model partitioning, training computing devices (TCDs) are divided into partition groups (PGs), each holding a complete replica of the model state. This strategy minimizes the number of servers per PG, reducing inter-server communication and associated costs. PGs operate on different data subsets in parallel, expediting the training process and optimizing resource utilization.

Replication and Efficiency

Within partition groups, TCDs are also organized into replication groups (RGs), where each device maintains the same subset of model state information. This configuration supports efficient data handling and reduces redundant communication, contributing to overall training speed and resource efficiency. The proposed methods aim to maximize the use of available resources while minimizing communication bottlenecks.