Login| Sign Up| Help| Contact|

Patent Searching and Data


Title:
A-SLOWMO: MEMORY-OPTIMIZED MODEL PARALLEL ASGD SOLUTION
Document Type and Number:
WIPO Patent Application WO/2024/032873
Kind Code:
A1
Abstract:
A system and a method for distributed training of a machine learning model are disclosed. Staleness of ASGD algorithms leads to degradation of their convergence rate, generalization and scalability. The disclosure is of distributed training of a machine learning model, using model parallel (MP) asynchronous stochastic gradient descent (ASGD), while mitigating gradient staleness using a parameter prediction method, wherein the momentum is recovered from the parameters of a neighbor virtual worker, thus reducing memory requirements. Some implementation also apply gradient normalization.

Inventors:
TALYANSKY ROMAN (DE)
MELAMED ZACH (DE)
KATZ MICHAEL (DE)
KISILEV PAVEL (DE)
Application Number:
PCT/EP2022/072283
Publication Date:
February 15, 2024
Filing Date:
August 08, 2022
Export Citation:
Click for automatic bibliography generation   Help
Assignee:
HUAWEI TECH CO LTD (CN)
TALYANSKY ROMAN (DE)
International Classes:
G06N3/063; G06N3/084; G06N3/098; G06N3/044; G06N3/045; G06N3/0475
Other References:
ELIAD SAAR ET AL: "Fine-tuning giant neural networks on commodity hardware with automatic pipeline model parallelism Fine-tuning giant neural networks on commodity hardware with automatic pipeline model parallelism", THIS PAPER IS INCLUDED IN THE PROCEEDINGS OF THE 2021 USENIX ANNUAL TECHNICAL CONFERENCE, 16 July 2021 (2021-07-16), XP093026805, Retrieved from the Internet [retrieved on 20230224]
IDO HAKIMI ET AL: "Taming Momentum in a Distributed Asynchronous Environment", ARXIV.ORG, CORNELL UNIVERSITY LIBRARY, 201 OLIN LIBRARY CORNELL UNIVERSITY ITHACA, NY 14853, 13 October 2020 (2020-10-13), XP081782707
Attorney, Agent or Firm:
KREUZ, Georg M. (DE)
Download PDF:
Claims:
WHAT IS CLAIMED IS:

1. A system for training machine learning based models, using a plurality of computing devices, the plurality of computing devices are configured for model parallel pipeline training, wherein at least one computing device is configured for: designating a plurality of virtual workers arranged in a round robin form; designating a plurality of worker units to each virtual worker of the plurality of virtual workers; assigning to each worker unit from the plurality of worker units associated with at least one of the virtual workers, a buffer for storing a parameter vector for a machine learning based model part; and in at least one iteration: calculating a reconstructed momentum part using the parameter vector part of at least one worker unit, and the parameter vector part of an additional worker unit, selected by a proximity measure in a circular form; calculating an updated parameter vector part using the parameter vector part of the additional worker unit, a gradient vector, and the reconstructed momentum.

2. The system of claim 1, wherein the proximity measure selects the parameter vector from additional worker unit associated with the least recent virtual worker, and the updated parameter vector storage location is updated circularly from most recent former iteration.

3. The system of claim 1, wherein the at least one iteration comprises a plurality of subiterations.

4. The system of claim 1, further comprising a second moment buffer, and the at least one computing device is further configured for updating the second moment buffer using the momentum.

5. The system of claim 4, further comprising gradient normalization.

6. The system of claim 4, wherein at least one computing device is further configured for in at least one iteration: calculating a second moment by elementwise multiplication of the momentum by the momentum; applying smoothing and bias correction to the second moment; calculating a normalization vector using the second moment and a former normalization momentum; applying smoothing and bias correction on the normalization vector; updating a normalization momentum using the gradient; and applying bias correction to the normalization momentum.

7. The system of claim 6, wherein the applying bias correction to the normalization momentum comprises division of the normalization momentum by a factor based on a normalization factor.

8. A computer-implemented method for training machine learning based models, using a plurality of computing devices, the plurality of computing devices are configured for model parallel pipeline training, the method comprising: designating a plurality of virtual workers arranged in a round robin form; designating a plurality of worker units to each virtual worker of the plurality of virtual workers; assigning to each worker unit from the plurality of worker units associated with at least one of the virtual workers, a buffer for storing a parameter vector for a machine learning based model part; and at least one iteration of: calculating a reconstructed momentum part using the parameter vector part of at least one worker unit, and the parameter vector part of an additional worker unit, selected by a proximity measure in a circular form; calculating an updated parameter vector part using the parameter vector part of the additional worker unit, a gradient vector, and the reconstructed momentum.

9. The computer-implemented method of claim 8, wherein the proximity measure selects the parameter vector from additional worker unit associated with the least recent virtual worker, and the updated parameter vector storage location is updated circularly from most recent former iteration.

10. The computer-implemented method of claim 8, wherein the at least one iteration comprises a plurality of sub-iterations.

11. The computer-implemented method of claim 8, further comprising a second moment buffer, and updating the second moment buffer using the momentum.

12. The computer-implemented method of claim 11, further comprising gradient normalization.

13. The computer-implemented method of claim 11, wherein at least one computing device is further configured for in at least one iteration: calculating a second moment by elementwise multiplication of the momentum by the momentum; applying smoothing and bias correction to the second moment; calculating a normalization vector using the second moment and a former normalization momentum; applying smoothing and bias correction on the normalization vector; updating a normalization momentum using the gradient; and applying bias correction to the normalization momentum.

14. The computer-implemented method of claim 13, wherein the applying bias correction to the normalization momentum comprises division of the normalization momentum by a factor based on a normalization factor.

15. A computer program product comprising instructions for training machine learning based models, using a plurality of computing devices, the plurality of computing devices are configured for model parallel pipeline training, wherein execution of the instructions by one or more processors of a computing system is to cause a computing system to: designate a plurality of virtual workers arranged in a round robin form; designate a plurality of worker units to each virtual worker of the plurality of virtual workers; assign to each worker unit from the plurality of worker units associated with at least one of the virtual workers, a buffer for storing a parameter vector for a machine learning based model part; and in at least one iteration: calculate a reconstructed momentum part using the parameter vector part of at least one worker unit, and the parameter vector part of an additional worker unit, selected by a proximity measure in a circular form; calculate an updated parameter vector part using the parameter vector part of the additional worker unit, a gradient vector, and the reconstructed momentum.

Description:
A-SLOWMO: MEMORY-OPTIMIZED MODEL PARALLEL ASGD SOLUTION

BACKGROUND

Some embodiments described in the present disclosure relate to a machine learning model training and, more specifically, but not exclusively, to a model parallel (MP) asynchronous stochastic gradient descent (ASGD) using parameters distributed training, using a reconstructed momentum.

Using asynchronous stochastic gradient descent (SGD) on a system comprising a computing device is a ubiquitous method of training machine learning modes, and particularly neural networks. ASGD is a distributed algorithm for training very large scale deep learning models.

The training of very large machine learning models that may not fit into the main memory of a single computing device, an acceleration device, or a graphic processing unit (GPU), may be more effective when the model is partitioned into several parts and these parts are placed over a set of physical acceleration devices.

The algorithm may be applied using a plurality of computing devices, which may designate a plurality of virtual workers, each virtual worker may comprise a plurality of worker units assigned each to a model part. Additionally, a process, peripheral, and/or computing device may also be referred to as a master or parameter storage.

However, ASGD is characterized by a disadvantage of gradient staleness. The problem of gradient staleness arises, since the gradient, computed at iteration k, is merged into model at a latency of T iteration, or at [fc+r]: [fc+r+1] = [fc+r] - ([fc]). Since the gradient computed at [fc+r] may be significantly different from the gradient, computed at [fc], the precision may be lower, and the convergence may be slower, and difficult to reach.

Model Parallel (MP) pipeline execution model is designed for training very large DL models, which may not fit into a single device. Available devices may be partitioned into stages, wherein a stage is associated with a virtial worker. As the model may be partitioned and each part may be assigned to a corresponding stage. A virtial worker may compute forward and backward passes for the mode, using worker units, each may correspond to model part. A minibatch may be partitioned into micro-batches. Gradients for micro-batches may be computed separately and then aggregated into a gradient for the entire mini-batch. To compute a gradient, a micro-batch may undergo a forward pass through stages and then a backward pass. D. Narayanan, Amar Phanishayee, Kaiyu Shi, X. Chen, and M. Zaharia, propose in “Memory-Efficient Pipeline-Parallel DNN Training”, on arXiv preprint arXiv:2006.09503 (2020), a MP pipeline execution model, formulated for the case, where #micro_batches > #stages. In this solution, the pipeline manages 2 model replica.

Each virtual worker may receive parameters, which may also be referred to a as weights, at a given time t, namely w[t] from the master, compute a gradient, namely Vf(w[t — 1]) and apply an update rule:

In cases of two virtual workers, where the gradient is computed at time t-1 and is merged into model of time t, resulting in staleness of 1.

The concept of momentum, is ubiquitously applied both in SGD and ASGD. An ASGD worker may start computing a gradient at time t from the current version of the central model parameters, which may also be referred to as the starting parameters (SP). When the worker finishes computing the gradient, other workers might have updated the central model, so that the worker merges its gradient into the updated version of the central model parameters, which may also be referred to as the final parameters (FP). The number of updates that other workers applied to the central model between SP and FP is stalenes, which may be denoted r[t] .

Staleness of ASGD algorithms leads to degradation of their convergence rate, generalization and scalability. Below we review three approaches to mitigate staleness.

Delay compensation methods, as suggested by Shuxin Zheng, Qi Meng, Taifeng Wang, Wei Chen, Nenghai Yu, Zhiming Ma, and Tie-Yan Liu, in “Asynchronous Stochastic Gradient Descent with Delay Compensation”, at ICML 2017, wherein a worker computes the gradient on starting parameters (SP), then it computes a first order approximation of the gradient at FP and, finally, merges the approximated gradient to final parameters (FP).

However, the approximation quality may drop as the distance between SP and FP grows or, in turn, as the number of workers grows, leading to poor scalability and generalization.

Parameter prediction methods, such as suggested by Saar Barkai, Ido Hakimi, and A. Schuster, in “Gap Aware Mitigation of Gradient Staleness”, at ICLR 2020, wherein a worker reads SP and predicts FP were also introduced. Followingly, the worker computes the gradient on the predicted parameters and merges it FP. However, the high memory requirement of 2 memory buffers per each worker limits the applicability. One buffer is required for the parameters and the other for the momentum buffer.

Gradient penalizing methods may use either staleness or gap for staleness mitigation. Staleness-aware methods may penalize gradients using staleness as a penalizing factor for example by dividing it by staleness. These methods suffer from over-penalization, since while staleness may be large, the actual distance between the starting and final parameters in the parameter space may be small.

Some parameter prediction methods may be more accurate, however at the cost of managing two buffers per each worker: for parameters and momentum, leading to high memory demand. The use of gap for gradient penalization, which is computed as a ratio of the actual distance between the starting and final parameters and average optimization step size. Since gap is defined over the actual distance between the starting and the final parameters, it does not suffer from over-penalization, as staleness-aware methods, yielding better results than staleness-aware methods. When gap-aware methods are combined with parameter-prediction methods, they achieve good generalization.

SUMMARY

It is an object of the present disclosure to describe a system and a method for distributed training of a machine learning model, using model parallel (MP) asynchronous stochastic gradient descent (ASGD), while mitigating gradient staleness using a parameter prediction method, wherein the momentum is recovered from the parameters of a neighbor virtual worker, thus reducing memory requirements.

The foregoing and other objects are achieved by the features of the independent claims. Further implementation forms are apparent from the dependent claims, the description and the figures.

According to an aspect of some embodiments of the present invention there is provided a system for training machine learning based models, using a plurality of computing devices, the plurality of computing devices are configured for model parallel pipeline training, wherein at least one computing device is configured for: designating a plurality of virtual workers arranged in a round robin form; designating a plurality of worker units to each virtual worker of the plurality of virtual workers; assigning to each worker unit from the plurality of worker units associated with at least one of the virtual workers, a buffer for storing a parameter vector for a machine learning based model part; and in at least one iteration: calculating a reconstructed momentum part using the parameter vector part of at least one worker unit, and the parameter vector part of an additional worker unit, selected by a proximity measure in a circular form; calculating an updated parameter vector part using the parameter vector part of the additional worker unit, a gradient vector, and the reconstructed momentum.

According to an aspect of some embodiments of the present invention there is provided a computer-implemented method for training machine learning based models, using a plurality of computing devices, the plurality of computing devices are configured for model parallel pipeline training, the method comprising: designating a plurality of virtual workers arranged in a round robin form; designating a plurality of worker units to each virtual worker of the plurality of virtual workers; assigning to each worker unit from the plurality of worker units associated with at least one of the virtual workers, a buffer for storing a parameter vector for a machine learning based model part; and at least one iteration of: calculating a reconstructed momentum part using the parameter vector part of at least one worker unit, and the parameter vector part of an additional worker unit, selected by a proximity measure in a circular form; calculating an updated parameter vector part using the parameter vector part of the additional worker unit, a gradient vector, and the reconstructed momentum.

According to an aspect of some embodiments of the present invention there is provided a computer program product comprising instructions for training machine learning based models, using a plurality of computing devices, the plurality of computing devices are configured for model parallel pipeline training, wherein execution of the instructions by one or more processors of a computing system is to cause a computing system to: designate a plurality of virtual workers arranged in a round robin form; designate a plurality of worker units to each virtual worker of the plurality of virtual workers; assign to each worker unit from the plurality of worker units associated with at least one of the virtual workers, a buffer for storing a parameter vector for a machine learning based model part; and in at least one iteration: calculate a reconstructed momentum part using the parameter vector part of at least one worker unit, and the parameter vector part of an additional worker unit, selected by a proximity measure in a circular form; calculate an updated parameter vector part using the parameter vector part of the additional worker unit, a gradient vector, and the reconstructed momentum.

Optionally, the proximity measure selects the parameter vector from additional worker unit associated with the least recent virtual worker, and the updated parameter vector storage location is updated circularly from most recent former iteration.

Optionally, the at least one iteration comprises a plurality of sub-iterations.

Optionally, the, further comprising a second moment buffer, and the at least one computing device is further configured for updating the second moment buffer using the momentum.

Optionally, the further comprising gradient normalization.

Optionally, the at least one computing device is further configured for in at least one iteration: calculating a second moment by elementwise multiplication of the momentum by the momentum; applying smoothing and bias correction to the second moment; calculating a normalization vector using the second moment and a former normalization momentum; applying smoothing and bias correction on the normalization vector; updating a normalization momentum using the gradient; and applying bias correction to the normalization momentum.

Optionally, the applying bias correction to the normalization momentum comprises division of the normalization momentum by a factor based on a normalization factor.

Other systems, methods, features, and advantages of the present disclosure will be or become apparent to one with skill in the art upon examination of the following drawings and detailed description. It is intended that all such additional systems, methods, features, and advantages be included within this description, be within the scope of the present disclosure, and be protected by the accompanying claims.

Unless otherwise defined, all technical and/or scientific terms used herein have the same meaning as commonly understood by one of ordinary skill in the art to which embodiments. Although methods and materials similar or equivalent to those described herein can be used in the practice or testing of embodiments, exemplary methods and/or materials are described below. In case of conflict, the patent specification, including definitions, will control. In addition, the materials, methods, and examples are illustrative only and are not intended to be necessarily limiting.

BRIEF DESCRIPTION OF THE SEVERAL VIEWS OF THE DRAWINGS )

Some embodiments are herein described, by way of example only, with reference to the accompanying drawings. With specific reference now to the drawings in detail, it is stressed that the particulars shown are by way of example and for purposes of illustrative discussion of embodiments. In this regard, the description taken with the drawings makes apparent to those skilled in the art how embodiments may be practiced.

In the drawings:

FIG. 1 is a schematic illustration of an exemplary system for training of a neural network, according to some embodiments of the present disclosure;

FIG. 2A is a flowchart schematically representing an optional flow of operations for distributed step of training of a machine learning model, using recovered momentum, according to some embodiments of the present disclosure;

FIG. 2B is a flowchart schematically representing an optional flow of operations for calculate a second moment, according to some embodiments of the present disclosure;

FIG. 3A is a schematic diagram of an exemplary process of forward and backward passes of MP ASGD training, according to some embodiments of the present disclosure;

FIG. 3B is a schematic diagram of an exemplary worker hierarchy for MP ASGD training, according to some embodiments of the present disclosure;

FIG. 4 is a schematic diagram of a snapshot of an exemplary pipeline MP of ASGD execution implementation, according to some embodiments of the present disclosure;

FIG. 5A is a schematic graph depicting training using an exemplary experiment, according to some embodiments of the present disclosure;

FIG. 5B is a schematic graph depicting training using an exemplary experiment, according to some embodiments of the present disclosure;

FIG. 6 is an exemplary procedure, comprising gradient normalization, according to some embodiments of the present disclosure; and

FIG. 7 is an additional schematic diagram of an exemplary process of forward and backward passes of MP ASGD training, according to some embodiments of the present disclosure. DETAILED DESCRIPTION

Some embodiments described in the present disclosure relate to a machine learning model training and, more specifically, but not exclusively, to a model parallel (MP) asynchronous stochastic gradient descent (ASGD) using parameters distributed training, using a reconstructed momentum.

Training a complex machine learning model, such as a deep neural network, using large dataset, is time, memory, and energy consuming, and may be slow or even impractical on a single computing device. Therefore methods for distributed training were developed.

Distributed training between a plurality of computing devices, also referred to as worker nodes, may improve training speed, and enable training of large machine learning models.

The gradient descent acceleration may be data parallel, i.e. based on distribution of the data from the training dataset between different computing devices, model parallel pipeline, i.e. based on partitioning the model and distributing it over different devices, or a combination thereof, for example by defining a two-dimensional array of computing devices, wherein each computing device is assigned to a pipeline stage according to an associated row location, and to the data part from the training set according to an associated column location.

Training may alternatively be done by genetic algorithms, however gradient descent is the ubiquitous choice, particularly stochastic gradient descent, which may be accelerated by parallel computing used for the distributed training using synchronous stochastic gradient descent (SSGD), or ASGD. SSGD requires waiting for the slowest, or most remote, computing device, from which the updated gradients arrive last, and therefore may be slower than ASGD.

Model Parallel (MP) pipeline execution model may lead to an ASGD training, with round-robin order of model updates, requiring ASGD optimizers. State of the art (SotA) ASGD optimizers lead to poor convergence rate, generalization and scalability, as in delay compensation methods, gradient penalizing methods or gradient-aware methods. DANA-GA is proposes good convergence rate, generalization and scalability properties. However, DANA- GA has a high memory demand, since it requires 2 buffers per each worker: for weights and for the first moment.

The disclosure comprises a method for reduceing these memory requirement almost by factor of 2, as compared to SotA solutions, while preserving convergence rate and generalization of SotA solutions. DANA-GA memory requirement: 2 N + 3 buffers, some implementations of the disclosure, for example some proposed improvements of SlowMo requires managing a single buffer per worker for both model parameters and for momentum. amd additional buffers for storing the value of previous parameters of the previous worker, computing second moment, and computing normalization vector. In total, these implementations for N workers: N + 3 buffers, as compared to 2N + 3 buffers in DANA-GA.

Note that for staleness T, a gradient G is merged into master model after other workers merged r gradients into it. If during these r updates gradient dynamics changed drastically, gradient G may be irrelevant and merging it into the master model may destabilize optimization. To avoid such a destabilizing behavior, we propose a gradient normalization algorithm.

For better understanding of the disclosed method, recall the notation r[t] for the staleness of a gradient that a computing device, or a worker, starts computing at time t, i.e. that the computing device may merge this gradient to the central model at time t+r[t], after other workers have updated the central model r[t] times. Note that for staleness T, a gradient G is merged into master model after other workers merged r gradients into it. If during these r updates gradient dynamics changed drastically, gradient G may be irrelevant and merging it into the master model may destabilize optimization.

Staleness of ASGD algorithms leads to degradation of their convergence rate, generalization and scalability. Known methods to mitigate the gradient staleness are either limited in scale or require extensive additional memory. The disclosure reduces almost by factor of 2 memory requirement, as compared to known SotA methods, while preserving convergence rate and generalization of SotA solutions.

Some implementations of the disclosure replace explicit storage of per-worker first moment with an on-the-fly computation. This may be done by reconstruction of per-worker first moment on-the-fly allows us to avoid storing it permanently at the worker main memory, leading to memory reduction of the algorithm almost by factor of 2. Some implementations of the disclosure accordingly modifying a master update rule to allow on-the-fly re-computation of per-worker first moment.

Some implementations of the disclosure comprise a gradient-normalization procedure, in accordance with a ratio of the fast and slow dynamics of ASGD optimization process. The gradient-normalization procedure may detect iteration spans with high rate of gradient change and normalizes the gradients accordingly, avoiding optimization instability due to very outdated gradients.

Some implementations of the disclosure apply discretization of the second moment in our MP implementation of AdamW for reducing resource consumption.

Benefits of the disclosure comprise reducing the memory requirements of SotA Model Parallel pipeline algorithms, formulating the entire memory-optimized ASGD solution to fit the model-parallel pipeline execution model, extending the synchronous SGD SlowMo framework to an asynchronous framework a-SLOWMO that fits the above, and extending the asynchronous framework a-SlowMo to SGD with momentum and AdamW optimizer.

As used herein, the terms computing devices, worker nodes or workers, refer to computing environments, computers, GPU cards, field programmable gate arrays (FPGA), application specific integrated circuits (ASIC), workstations, servers, digital signal processing (DSP) modules, a combination thereof, and/or similar devices, which are apt to execute instructions for training complex machine learning models, and store the parameters of the machine learning models. The computing devices may communicate using a backbone bus, cables arranged at specific topologies, communication protocols such as Ethernet, the internet, and/or the like.

Before explaining at least one embodiment in detail, it is to be understood that embodiments are not necessarily limited in its application to the details of construction and the arrangement of the components and/or methods set forth in the following description and/or illustrated in the drawings and/or the Examples. Implementations described herein are capable of other embodiments or of being practiced or carried out in various ways.

Embodiments may be a system, a method, and/or a computer program product. The computer program product may include a computer readable storage device medium (or media) having computer readable program instructions thereon for causing a processor to carry out aspects of the embodiments.

The flowchart and block diagrams in the Figures illustrate the architecture, functionality, and operation of possible implementations of systems, methods, and computer program products according to various embodiments. In this regard, each block in the flowchart or block diagrams may represent a module, segment, or portion of instructions, which comprises one or more executable instructions for implementing the specified logical function(s). In some alternative implementations, the functions noted in the block may occur out of the order noted in the figures. For example, two blocks shown in succession may, in fact, be executed substantially concurrently, or the blocks may sometimes be executed in the reverse order, depending upon the functionality involved. It will also be noted that each block of the block diagrams and/or flowchart illustration, and combinations of blocks in the block diagrams and/or flowchart illustration, can be implemented by special purpose hardware-based systems that perform the specified functions or acts or carry out combinations of special purpose hardware and computer instructions. The disclosure comprises an exemplary algorithm referred to as slow momentum (SlowMo) for SGD with momentum, has input of a base optimizer LR y[t]; fast momentum factor p; slow momentum factor P; number of workers N; initial point x[0,0] and initial slow buffer u=0, and proceeds as follows:

1. For t in {0,1,. . ,,T-1 { at worker i do

2. For k in {0} do

3. Compute gradient: g[i,t] of worker i

4. g [i, t] =normalize_gradi ent(g [i ,t] )

5. Update momentum: m[i,t,k]=pm[i,t-N,k]+g[i,t]

6. End

7. Update master: x[t+l]=x[t]- y[t]m[i,t,k]-(y[t- l]-y[t]P)u[t]

8. Update slow momentum: u[t+l]= Pu[t]+m[i,t,k]

9. Update outer iterates of worker i: x[i,t+l ]=x[t+l ]- y[t]u[t+l ]

10. End

Therefore, storing model parameters is sufficient to reconstruct the first moment of a worker numbered i.

Thus, we do not need to manage explicitly master and slow momentum.

This gives rise to memory optimized pipeline version of a-SLOWMO.

Note that in the algorithm as described includes an additional term (y[t- 1 ]- y[t]P)u[t] to update a master rule. Adding this term enables reconstructing momentum m[i,t-N,k] of worker i on-the-fly and, thus, avoid the need to store it.

The disclosure is based on some properties, such as of the momentum: m[i,t-N,k]=-(l/(2y[t]))(x[i,t+l]-x[i-l,t]). Therefore a pipeline version of the disclosed algoritm may be equivalent, without explicit management of master parameters and slow momentum.

As shown in line 9 of the following procedure, the computation of x[i,t+l] may be replaced by the following equivalent recursive expression: x[i,t+l]= x[i- 1 ,t]+ p(y[t]/y[t-N])(x[i,t-N+l]- x[i- 1 ,t-N]])-2y[t]g[i,t]

Therefore, an exemplary procedure, referred to as pipeline a-SlowMo for SGD with a momentum may bases on an optimizer LRy[t]; fast momentum factor p; slow momentum factor P; number of workers N; initial point x[0,0], and proceed as follows:

1. For t in {0,1,. . ,,T-1 } at worker i do

2. For k in {0} do

3. Compute gradient: g[i,t] of worker i

4. g [i, t] =normalize_gradi ent(g [i ,t] )

5. Compute parameters difference: A=-(x[i,t-N+l] “Xprev)

6. Store current weights of worker i . Xprev X [i,t-N+l]

7. End

8. Update outer iterates of worker i:

9. x[i,t+l]=x[i-l,t]-p(y[t]/y[t-N])A-2y[t]g[i,t]

10. End

A comparison between DANA-GA and the disclosure for N workers. DANA-GA requires 2N+3 buffers, while the disclosure requires N+3 buffers, a reduction almost by a factor of 2. For ADAMW formulation in A-SLOWMO framework present our non-MP-pipeline procedure, bases on an optimizer LR y[t]; fast momentum factor p; slow momentum factor P; number of workers N; initial master point x[0] and per-worker momentum m[i,t]=0 and initial slow buffer u=0; constants P2=0.999 and s=10' 9 , second moment estimate y[0]=0, its previous discrete value y’=l; weight decay factor :

1. For t in {1,...,T} at worker i do

2. For k in {0} do

3. Compute gradient: g[i,t] of worker i

4. g[i,t], V’ ’ [t]=normalize_gradient(g[i,t]) 5. Update first moment estimate: m[i,t]= pm[i,t-N]+g[i,t]

6. Update second moment estimate: y[t]= P2y [t- 1 ]+g[i,t] A 2

7. Compute bias-corrected first moment estimate: m’[i,t]= pm[i,t](l-p)/(l- p A ceil(t/N))

8. If (t-1) mod N==0 then

9. Compute and discretize bias-corrected second moment estimate: y”=y[t](i-P2)/(i-P2 A t)

10. Compute ratio: r[t]=(l/V”[t])m’[i,t]/(sqrt(y”)+s)+kx[i-l,t]

11. End

12. Update master: x[t+l ]=x[t]-y[t]r[t]-(y[t- l]-y[t]P)u[t]

13. Update slow momentum: u[t+l ]=Pu[t]+r[t]

14. Update outer iterates of worker i: x[i,t+l]=x[t+l]-y[t]u[t+l]

15. End

Now, elaborations on the additions to the procedures, that do not occur in plain AdamW are elaborated. The disclosed procedure reconstruct momentum of a worker using linear algebraic operations. However, Adam update rule involves non-linear element-wise division of the first moment by the second moment, breaking the basis of the proposed momentum reconstruction. Since for estimating the second moment, Adam algorithm uses a big constant, e.g. 0.999, leading to a very slow change in second moment estimation. This enables picking up a value of the second moment estimation in line 9 and use it along several subsequent iterations. This, in turn, allows us to fall back to linear operations to reconstruct the first moment of a worker.

Next, note that in long iteration spans, where at each iteration gradient normalization in line 4 considerably changes the gradient, this normalization affects both first and second moment estimations. Thus, division of the first moment by second moment in line 10, leads to cancellation of that gradient normalization. To avoid this cancellation, we use in line 10 estimation of the norm of normalization factor V’ ’ [t] that we receive in line 4 to normalize the ratio of moments.

As in the former procdure, the factor (y[t- 1 ]- y[t]P)u[t] at line 12 may allow on the fly first moment reconstruction.

That gives rise an equivalent MP pipeline version and get rid of explicit management of master and slow momentum. For this purpose, delta is defined as: Delta=-(x[i,t-N+l]-x[i- l,t-N]). Followingly, as shown in line 14, an implementation may replace computation of x[i,t+ 1 ] with the following equivalent expression: x[i,t+l]=x[i-l,t]+ pa[t](sqrt(y”[t-N])+e)/(sqrt(y”[t])+e)(-b[t]Delta+c[t]x[ i-l,t-N])- d[t]g[i,t]/(sqrt(y”[t])+e)-c[t]x[i-l,t], where a[t]=(V’ ’ [t-N]/V’ ’ [t])-(l-p A ceil((t-N)/N))/(l-p A ceil(t/N)), b[t]=y[t]/y[t-N], c[t]=2y[t]Z. and d[t]=(2y[t]/V”[t])-(l-p)/(l-p A ceil(t/N)) are scalars.

This may give rise ti an equivalent MP pipeline form which implements a-SlowMo for MP pipeline AdamW, based on an optimizer LR y[t]; fast momentum factor p; slow momentum factor P; number of workers N; initial master point and per-worker points x[0]; constants 32=0.999 and s=10' 9 , second moment estimate y[0]=0, its previous discrete value y’_prev=l; weight decay factor :

1. For t in {1,...,T} at worker i do

2. For k in {0} do

3. Compute gradient: g[i,t] of worker i

4. If (t-1) mod N==0 then

5. Store the previous discrete second moment estimate: y”_prev=y”

6. Compute the difference over the parameters: Delta=-(x[i,t-N+l]-x_prev).

7. Store the current weights of worker i: x_prev=x[i,t-N+l]

8. g[i,t], V”[t]=normalize_gradient(g[i,t])

9. Update second moment estimate: y[t]= P2 y[t- 1 ]+g[i,t] A 2

10. If (t- 1 ) mod N==0 then

11. Compute and discretize bias-corrected second moment estimate: y”=y[t](i-P2)/(i-P2 A t) 13. Compute scalars: a[t], b[t], c[t] and d[t]

14. Update outer iterates of worker i:

15. x[i,t+ 1 ]=x[i- 1 ,t]+pa[t](sqrt(y’ ’_prev)+s)/(sqrt(y ’ ’)+e)(-b [t]D elta+c [t]x [i- 1 ,t-N])- d[t]g[i,t]/(sqrt(y”)+e)-c[t]x[i-l,t],

16. End

The disclosure comprises reconstruction of per-worker first moment on-the-fly allows us to avoid storing it permanently at the worker main memory, leading to memory reduction of the algorithm almost by factor of 2. Gradient-normalization procedure detects iteration spans with high rate of gradient change and normalizes the gradients accordingly, avoiding optimization instability due to very out-dated gradients is also comprised in the disclosure. The disclosure also comprises a discretization of the second moment in our MP implementation of AdamW algorithm makes this implementation feasible.

Referring now to the drawings, FIG. 1 is a schematic illustration of an exemplary system for training of a neural network, according to some embodiments of the present disclosure. An exemplary client computer system 100 may be used for executing execute processes such as 200 for training a network using recovered momentum. Further details about these exemplary processes follow as FIG. 2A is described.

Various aspects of the present disclosure are described by narrative text, flowcharts, block diagrams of computer systems and/or block diagrams of the machine logic included in computer program product (CPP) embodiments. With respect to any flowcharts, depending upon the technology involved, the operations may be performed in a different order than what is shown in a given flowchart. For example, again depending upon the technology involved, two operations shown in successive flowchart blocks may be performed in reverse order, as a single integrated step, concurrently, or in a manner at least partially overlapping in time.

A computer program product embodiment ("CPP embodiment" or “CPP”) is a term used in the present disclosure to describe any set of one, or more, storage media (also called "mediums") collectively included in a set of one, or more, storage devices that collectively include machine readable code corresponding to instructions and/or data for performing computer operations specified in a given CPP claim. A "storage device" is any tangible device that can retain and store instructions for use by a computer processor. Without limitation, the computer readable storage medium may be an electronic storage medium, a magnetic storage medium, an optical storage medium, an electromagnetic storage medium, a semiconductor storage medium, a mechanical storage medium, or any suitable combination of the foregoing. Some known types of storage devices that include these mediums include: diskette, hard disk, random access memory (RAM), read-only memory (ROM), erasable programmable read-only memory (EPROM or Flash memory), static random access memory (SRAM), compact disc read-only memory (CD-ROM), digital versatile disk (DVD), memory stick, floppy disk, mechanically encoded device (such as punch cards or pits / lands formed in a major surface of a disc) or any suitable combination of the foregoing. A computer readable storage medium, as that term is used in the present disclosure, is not to be construed as storage in the form of transitory signals per se, such as radio waves or other freely propagating electromagnetic waves, electromagnetic waves propagating through a waveguide, light pulses passing through a fiber optic cable, electrical signals communicated through a wire, and/or other transmission media. As will be understood by those of skill in the art, data is typically moved at some occasional points in time during normal operations of a storage device, such as during access, defragmentation or garbage collection, but this does not render the storage device as transitory because the data is not transitory while it is stored.

Computing environment 100 contains an example of an environment for the execution of at least some of the computer code involved in performing the inventive methods, such as momentum recovering ASGD neural network training 180. In addition to block 180, computing environment 100 includes, for example, computer 102, wide area network (WAN) 108, end user device (EUD) 132, remote server 104, public cloud 150, and private cloud 106. In this embodiment, computer 102 includes processor set 110 (including processing circuitry 120 and cache 134), communication fabric 160, volatile memory 112, persistent storage 116 (including operating system 122 and block 180, as identified above), peripheral device set 114 (including user interface (UI), device set 126, storage 124, and Internet of Things (loT) sensor set 128), and network module 118. Remote server 104 includes remote database 130. Public cloud 150 includes gateway 140, cloud orchestration module 146, host physical machine set 142, virtual machine set 148, and container set 144.

COMPUTER 102 may take the form of a desktop computer, laptop computer, tablet computer, smart phone, smart watch or other wearable computer, mainframe computer, quantum computer or any other form of computer or mobile device now known or to be developed in the future that is capable of running a program, accessing a network or querying a database, such as remote database 130. As is well understood in the art of computer technology, and depending upon the technology, performance of a computer-implemented method may be distributed among multiple computers and/or between multiple locations. On the other hand, in this presentation of computing environment 100, detailed discussion is focused on a single computer, specifically computer 102, to keep the presentation as simple as possible. Computer 102 may be located in a cloud, even though it is not shown in a cloud in Figure 1. On the other hand, computer 102 is not required to be in a cloud except to any extent as may be affirmatively indicated.

PROCESSOR SET 110 includes one, or more, computer processors of any type now known or to be developed in the future. For example, a processor set may include one or more of a central processing unit (CPU), a microcontroller, a parallel processor, supporting multiple data such as a digital signal processing (DSP) unit, a graphical processing unit (GPU) module, and the like, as well as optical processors, quantum processors, and processing units based on technologies that may be developed in the future. Processing circuitry 120 may be distributed over multiple packages, for example, multiple, coordinated integrated circuit chips. Processing circuitry 120 may implement multiple processor threads and/or multiple processor cores. Cache 134 is memory that is located in the processor chip package(s) and is typically used for data or code that should be available for rapid access by the threads or cores running on processor set 110. Cache memories are typically organized into multiple levels depending upon relative proximity to the processing circuitry. Alternatively, some, or all, of the cache for the processor set may be located “off chip.” In some computing environments, processor set 110 may be designed for working with qubits and performing quantum computing.

Computer readable program instructions are typically loaded onto computer 102 to cause a series of operational steps to be performed by processor set 110 of computer 102 and thereby effect a computer-implemented method, such that the instructions thus executed will instantiate the methods specified in flowcharts and/or narrative descriptions of computer-implemented methods included in this document (collectively referred to as “the inventive methods”). These computer readable program instructions are stored in various types of computer readable storage media, such as cache 134 and the other storage media discussed below. The program instructions, and associated data, are accessed by processor set 110 to control and direct performance of the inventive methods. In computing environment 100, at least some of the instructions for performing the inventive methods may be stored in block 180 in persistent storage 116.

COMMUNICATION FABRIC 160 is the signal conduction paths that allow the various components of computer 102 to communicate with each other. Typically, this fabric is made of switches and electrically conductive paths, such as the switches and electrically conductive paths that make up busses, bridges, physical input / output ports and the like. Other types of signal communication paths may be used, such as fiber optic communication paths and/or wireless communication paths.

VOLATILE MEMORY 112 is any type of volatile memory now known or to be developed in the future. Examples include dynamic type random access memory (RAM) or static type RAM. Typically, the volatile memory is characterized by random access, but this is not required unless affirmatively indicated. In computer 102, the volatile memory 112 is located in a single package and is internal to computer 102, but, alternatively or additionally, the volatile memory may be distributed over multiple packages and/or located externally with respect to computer 102.

PERSISTENT STORAGE 116 is any form of non-volatile storage for computers that is now known or to be developed in the future. The non- volatility of this storage means that the stored data is maintained regardless of whether power is being supplied to computer 102 and/or directly to persistent storage 116. Persistent storage 116 may be a read only memory (ROM), but typically at least a portion of the persistent storage allows writing of data, deletion of data and re-writing of data. Some familiar forms of persistent storage include magnetic disks and solid state storage devices. Operating system 122 may take several forms, such as various known proprietary operating systems or open source Portable Operating System Interface type operating systems that employ a kernel. The code included in block 180 typically includes at least some of the computer code involved in performing the inventive methods.

PERIPHERAL DEVICE SET 114 includes the set of peripheral devices of computer 102. Data communication connections between the peripheral devices and the other components of computer 102 may be implemented in various ways, such as Bluetooth connections, Near-Field Communication (NFC) connections, connections made by cables (such as universal serial bus (USB) type cables), insertion type connections (for example, secure digital (SD) card), connections made though local area communication networks and even connections made through wide area networks such as the internet. In various embodiments, UI device set 126 may include components such as a display screen, speaker, microphone, wearable devices (such as goggles and smart watches), keyboard, mouse, printer, touchpad, game controllers, and haptic devices. Storage 124 is external storage, such as an external hard drive, or insertable storage, such as an SD card. Storage 124 may be persistent and/or volatile. In some embodiments, storage 124 may take the form of a quantum computing storage device for storing data in the form of qubits. In embodiments where computer 102 is required to have a large amount of storage (for example, where computer 102 locally stores and manages a large database) then this storage may be provided by peripheral storage devices designed for storing very large amounts of data, such as a storage area network (SAN) that is shared by multiple, geographically distributed computers. loT sensor set 128 is made up of sensors that can be used in Internet of Things applications. For example, one sensor may be a thermometer and another sensor may be a motion detector.

NETWORK MODULE 118 is the collection of computer software, hardware, and firmware that allows computer 102 to communicate with other computers through WAN 108. Network module 118 may include hardware, such as modems or Wi-Fi signal transceivers, software for packetizing and/or de-packetizing data for communication network transmission, and/or web browser software for communicating data over the internet. In some embodiments, network control functions and network forwarding functions of network module 118 are performed on the same physical hardware device. In other embodiments (for example, embodiments that utilize software-defined networking (SDN)), the control functions and the forwarding functions of network module 118 are performed on physically separate devices, such that the control functions manage several different network hardware devices. Computer readable program instructions for performing the inventive methods can typically be downloaded to computer 102 from an external computer or external storage device through a network adapter card or network interface included in network module 118.

WAN 108 is any wide area network (for example, the internet) capable of communicating computer data over non-local distances by any technology for communicating computer data, now known or to be developed in the future. In some embodiments, the WAN may be replaced and/or supplemented by local area networks (LANs) designed to communicate data between devices located in a local area, such as a Wi-Fi network. The WAN and/or LANs typically include computer hardware such as copper transmission cables, optical transmission fibers, wireless transmission, routers, firewalls, switches, gateway computers and edge servers.

END USER DEVICE (EUD) 132 is any computer system that is used and controlled by an end user (for example, a customer of an enterprise that operates computer 102), and may take any of the forms discussed above in connection with computer 102. EUD 132 typically receives helpful and useful data from the operations of computer 102. For example, in a hypothetical case where computer 102 is designed to provide a recommendation to an end user, this recommendation would typically be communicated from network module 118 of computer 102 through WAN 108 to EUD 132. In this way, EUD 132 can display, or otherwise present, the recommendation to an end user. In some embodiments, EUD 132 may be a client device, such as thin client, heavy client, mainframe computer, desktop computer and so on. REMOTE SERVER 104 is any computer system that serves at least some data and/or functionality to computer 102. Remote server 104 may be controlled and used by the same entity that operates computer 102. Remote server 104 represents the machine(s) that collect and store helpful and useful data for use by other computers, such as computer 102. For example, in a hypothetical case where computer 102 is designed and programmed to provide a recommendation based on historical data, then this historical data may be provided to computer 102 from remote database 130 of remote server 104.

PUBLIC CLOUD 150 is any computer system available for use by multiple entities that provides on-demand availability of computer system resources and/or other computer capabilities, especially data storage (cloud storage) and computing power, without direct active management by the user. Cloud computing typically leverages sharing of resources to achieve coherence and economies of scale. The direct and active management of the computing resources of public cloud 150 is performed by the computer hardware and/or software of cloud orchestration module 146. The computing resources provided by public cloud 150 are typically implemented by virtual computing environments that run on various computers making up the computers of host physical machine set 142, which is the universe of physical computers in and/or available to public cloud 150. The virtual computing environments (VCEs) typically take the form of virtual machines from virtual machine set 148 and/or containers from container set 144. It is understood that these VCEs may be stored as images and may be transferred among and between the various physical machine hosts, either as images or after instantiation of the VCE. Cloud orchestration module 146 manages the transfer and storage of images, deploys new instantiations of VCEs and manages active instantiations of VCE deployments. Gateway 140 is the collection of computer software, hardware, and firmware that allows public cloud 150 to communicate through WAN 108.

Some further explanation of virtualized computing environments (VCEs) will now be provided. VCEs can be stored as “images.” A new active instance of the VCE can be instantiated from the image. Two familiar types of VCEs are virtual machines and containers. A container is a VCE that uses operating-system-level virtualization. This refers to an operating system feature in which the kernel allows the existence of multiple isolated user-space instances, called containers. These isolated user-space instances typically behave as real computers from the point of view of programs running in them. A computer program running on an ordinary operating system can utilize all resources of that computer, such as connected devices, files and folders, network shares, CPU power, and quantifiable hardware capabilities. However, programs running inside a container can only use the contents of the container and devices assigned to the container, a feature which is known as containerization.

PRIVATE CLOUD 106 is similar to public cloud 150, except that the computing resources are only available for use by a single enterprise. While private cloud 106 is depicted as being in communication with WAN 108, in other embodiments a private cloud may be disconnected from the internet entirely and only accessible through a local/private network. A hybrid cloud is a composition of multiple clouds of different types (for example, private, community or public cloud types), often respectively implemented by different vendors. Each of the multiple clouds remains a separate and discrete entity, but the larger hybrid cloud architecture is bound together by standardized or proprietary technology that enables orchestration, management, and/or data/appli cation portability between the multiple constituent clouds. In this embodiment, public cloud 150 and private cloud 106 are both part of a larger hybrid cloud.

Referring now to FIG. 2A which is a flowchart schematically representing an optional flow of operations for distributed step of training of a machine learning model, using recovered momentum, according to some embodiments of the present disclosure.

The processing circuitry 120 may execute the exemplary process 200 for a variety of purposes involving training machine learning based models, using a plurality of computing devices, configured for model parallel pipeline training, for language models, computer vision, inference from video, multimodality, analytics and/or the like. Alternatively, the process 200 or parts thereof may be executing using a remote system, an auxiliary system, and/or the like.

The exemplary process 200 starts, as shown in 202, with designating a plurality of virtual workers arranged in a round robin form.

Work may be distributed between virtual workers according to various aspects of the training, Exemplary examples of aspects comprise handling different parts of the models, different time snapshots of the model, different input data, different contents of memory inside a recurrent network, or other network comprising memory, and/or the like.

In some implementations of the disclosure, at least one iteration may comprise a plurality of sub-iterations. A virtual worker may be assigned to stages, phases of forward and backward passes, or time snapshots of a model, or parts thereof.

The communication pattern between virtual workers may be based on a proximity relation, wherein a virtual worker may communicate with adjacent virtual workers, arranged in a round robin form, which gives rise to proximity measures in a circular form. The exemplary process 200 continues, as shown in 204, with designating a plurality of worker units to each virtual worker of the plurality of virtual workers.

The system may comprise several modules, such as GPU, which may be installed on one or several computing devices. Parts of the training tasks may be distributed according to variety of known methods and methods that may be developed in the future. One or more virtual processes may be associated with one or more modules. Some virtual processes may be designated as worker units may be assigned to one of the virtual workers of the plurality of virtual workers.

Each module, may have different memory modules available for faster access, and thus model parallel training may further increase speed, and enable training larger models.

The exemplary process 200 continues, as shown in 206, with assigning to each worker unit from the plurality of worker units associated with at least one of the virtual workers, a buffer for storing a parameter vector for a machine learning based model part.

The buffer may be located in the volatile memory 112, and store parameters such as weights, biases, routing, for example for deformable networks, and/or the like. The buffer may store parameters as integers, floating point, or other encodings.

The exemplary process 200 continues, as shown in 210, with at least one iteration of calculating a reconstructed momentum part using the parameter vector part of at least one worker unit, and the parameter vector part of an additional worker unit, selected by a proximity measure in a circular form.

The reconstruction may be based on the formula:

Wherein the number in parenthesis enumerates the virtual worker, and t is the time.

And subsequently, as shown in 212, the process 200 may continue the at least one iteration by using the processing circuitry 120, for calculating an updated parameter vector part using the parameter vector part of the additional worker unit, a gradient vector, and the reconstructed momentum.

The updated parameter vector may be calculated using the formula: Wherein g is the gradient. This evades the need to manage explicitly master and slow momentum, and gives rise to a memory optimized pipeline version of a-SLOWMO.

Optionally, the processing circuitry or system executing the process 200 may continue, as shown in 220, with updating the second moment buffer using the momentum.

Maintaining a second moment buffer the second moment buffer and updating it, using the momentum may be used for gradient normalization. The second moment buffer may be initialized to a given value, and each update may be based on a weighted averaging of the current value, and the second moment of the recent gradient. The motivation behind gradient normalization comprises that in ASGD settings, short-term bursts of gradient activity may lead to instability of optimization. Therefore, the disclosure comprises the following gradient normalization. In an exemplary setup, where ? x » P, e.g. ? x = 0.999 and P = 0.9, z t , captures long-range dynamics of momentum, while v t approximates short range dynamics. The value of V k may capture the extent of the change in short-term momentum dynamics, and division by V k penalizes gradients, may be merged after sharp changes in momentum dynamics.

Reference is also made to FIG. 2B which is a flowchart schematically representing an optional flow of operations for calculate a second moment, according to some embodiments of the present disclosure.

The exemplary process 250 may be executed for. The process 250 may be executed by the processing circuitry 120, for improving the stability, speed, and robustness of processes such as 200, and may benefit from hardware such as GPU.

The process 250 may start, as shown in 252 by calculating a second moment by elementwise multiplication of the momentum by the momentum m, generating the value m[i,t- N,k]) A 2.

The exemplary process 250 continues, as shown in 254, with applying smoothing and bias correction to the second moment

The second moment z[t+l] may be updated by a smoothing additive weighting factor i, between the former value z[t] and the squared momentum. Thereby assigning a value such as pi z[t]+(l-pi)(m[i,t-N,k]) A 2 to the second moment. Other smoothing methods may also be applied. The bias corrected second moment may be z’[t+l]=z[t+l]/(l- i A (t+l))

The exemplary process 250 continues, as shown in 256, with calculating a normalization vector using the second moment and a former normalization momentum. The normalization vector V[t] may be used for normalizing the gradient g[i,t], The normalization vector may be affected by a learning rate y ratio and may be calculated according to the formula: V[t]=(y[t]/y[O])(|v’[t]|/(sqrt(z’[t+l])+e))+l-d

And the normalized gradient may be g[i,t]=g[i,t]/V[t]

The exemplary process 250 continues, as shown in 258, with applying smoothing and bias correction on the norm of the normalization vector.

The smoothing may be a moving average based normalization scalar, based on a bias corrected norm of the normalization vector, which may be calculated according to:

V’[t+l]=P’V’[t]+(l-P’)||V[t]||, and/or v’tt ^v’tt+W- P 1 ))

The exemplary process 250 continues, as shown in 260, with updating a normalization momentum using the gradient.

The normalization momentum may be updated using the gradient, for example by: v[t+l]=Pv[t]+g[i,t]

And subsequently, the exemplary process 250 continues, as shown in 262, with applying bias correction to the normalization momentum. Applying the bias correction to the normalization momentum may comprise division of the normalization momentum by a factor based on a normalization factor, by an estimated factor, or by other methods.

The bias corrected normalization momentum may be updated, for example according to: v’ [t+ 1 ]=v[t+ 1 ]/( 1 -P A (t+ 1 ))

Reference is also made to FIG. 3A which is a schematic diagram of an exemplary process of forward and backward passes of MP ASGD training, according to some embodiments of the present disclosure.

The figure describes an exemplary Model Parallel (MP) pipeline execution mode, which may be used for training very large DL models, which may not fit into a single device.

Devices may be partitioned into stages; each stage may be associated with a worker. A model may be partitioned and each part may be assigned to the corresponding stage. A stage may compute forward and backward passes for the corresponding model part. A mini-batch may be associated with an iteration and partitioned into micro-batches, associated with gradients for micro-batches are computed separately and then aggregated into a gradient for the entire mini-batch.

The data is processed in the form as in the figure. Each row represents a worker and time progresses from left to right. The numbers in the darker squares indicate the micro stage, and the brighter squares indicate the backward pass. The patterned background indicates updating the parameters as done in the last sub-iteration of each iteration, where gradients are merged into the most up-to-date model.

In this exemplary setting, the number of stages equals the number of micro-batches, resulting in staleness greater than one. In this example, there are 4 stages and 4 micro-batches per mini-batch.

To compute a gradient, a micro-batch undergoes a forward pass through stages and then a backward pass. Some exemplary pipelines manage 2 model replica, where the gradient is computed at time t — 1 and is merged into model of time t, resulting in staleness of 1.

A new model is produced when all its parts have been updated. Models are produced at times of backward passes. The Pipeline may start a forward pass on micro-batches 5 and 6 on the initial model, and may finishes merging it at time of a second updated model, resulting in staleness of two. The number of models at each time in the pipeline may also be three, or other numbers.

Reference is now made to FIG. 3B which is a schematic diagram of an exemplary worker hierarchy for MP ASGD training, according to some embodiments of the present disclosure.

In ubiquitous implementations of the disclosure, with number of stages #stages and number of micro-batches #micro_batches, the number of models m that pipeline manages: m = #stages — #micro_batches + 1

And the gradient staleness:

T = m — 1 .

In Data parallel (DP) setting, a worker is a device that stores a model replica and computes gradients over the model replica. In MP setting, a worker is associated with a stage that computes forward and backward passes of the stage for the corresponding part of the model.

In the setup of some implementations of the disclosure, a worker, in the sense of a worker unit may be designated to a virtual worker with one of model replica that pipeline manages at a certain time.

This figure helps understand some aspects of MP training. The system may assign worker units to different model parts, as well as different sub iterations within an iteration, and the number of workers may be, for example a multiplication of the number of stages, or subiterations, and the number of model parts in the partition.

Additional workers may be added for control, load management, and other associated tasks.

It should be noted that data parallelism may be also implemented and systems applying combined data and model parallel training, as well as other parallelism methods that may be developed in the future, are within the scope of the claims.

Reference is now made to FIG. 4 which is a schematic diagram of a snapshot of an exemplary pipeline MP of ASGD execution implementation, according to some embodiments of the present disclosure.

In this exemplary snapshot, L is the loss function of the model, x^: model parameters of worker i, assigned at time t.

The disclosure may manipulate worker indexes i modulo number of workers m. Workers may update their parameters in a round-robin manner.

The disclosure may simulate a pipeline execution model in the form of a cyclic buffer of length m. The update rule is of the form in this exemplary snapshot is: where f represents an optimizer, such as momentum, Adam update rule, or the like.

Reference is now made to FIG. 5A which is a schematic graph depicting training using an exemplary experiment, according to some embodiments of the present disclosure.

The experiment is based on a-SlowMo for MP pipeline Adam using a simulation environment allowing simulating executions that follow a strict order, such as round robin order. First, we describe our simulation environment fo MP pipeline execution model. In the exemplary experiment settings, #stages > #micro_batches, resulting in staleness > 1. In this case, MP pipeline Execution Model manages N replica of model weights for the number of stages:

N=#stages-#micro_batches+l .

In this example, a logical worker is designtaed to each model replica, and logical workers may be referred to as workers. Thus, the MP pipeline may be represented as Execution Model (EM) as a cyclic buffer. Each entry in the buffer stores a model replica. Model updates follow the following rule. EM computes a gradient G[x[h]] over the element at the head of the buffer. Then EM loads the element at the tail x[t], t=(h- 1 ) mod N, updates it using: x[h]=x[t]-f(r|G[x[h]])

In this example, the updated model is pushed back to the head of the cyclic buffer. Note that function f in the above expression represents any optimization technique, such as momentum, Adam update rule, etc. In this way, the model version for which the gradient is calculated is N-l iterations older than the model version, to which this gradient is merged, leading to gradient staleness of N-l. Finally, head is advanced to h=(h+l) mod N.

The exemplary experiment test setup comprise MoRe-GPT is based on porting Nvidia’s Megatron GPT code from https://github.com/NVIDIA/Megatron-LM. MoRe-GPT adapted to fpl6 arithmetic. Variants tested: GPT-Small and GPT-Medium, using fp32 and fpl6 arithmetic.

Data was EnWiki-4000 (about 965,000 articles), MoRe ASGD scaling rule were 512 Baseline mini-batch size, 8 MoRe workers with per worker mini-batch size: (64=512 / 8) , 8 MoRe workers with per worker, and lr=lr_baseline / 8.

FIG. 5A shows the loss progress during training for the Generative Pre-trained Transformer (GPT) Small model.

Reference is also made to FIG. 5B which is a schematic graph depicting training using an exemplary experiment, according to some embodiments of the present disclosure.

FIG. 5B shows the loss progress during training for the Generative Pre-trained Transformer (GPT) Medium model, in a setup as described in FIG. 5A. Reference is now made to FIG. 6 which is an exemplary procedure, comprising gradient normalization, according to some embodiments of the present disclosure.

The gradient normalization is incorporated in an iteration, with a base optimizer with learning rate y[t]; slow momentum factor P; momentum pi for computing second moment; P’ momentum factor for computing long-tem normalization scalar; fast momentum factor p; momentum m[i,t-N] of worker i;

Note that in the algorithm, all binary vector operations are entry-wise. The notation comprises x A y: x in power y for scalar x, x A 2: dot product of x with itself for vector x, l A d: vector of length d with value 1 in each entry, |v|: vector with the same entries as in vector v and absolute values in the corresponding entries, and ||v||: 12 norm of vector v.

The input comprises gradient g[i,t] of worker i.and the output comprises normalized gradient g[i,t] of worker i and long-term normalization scalar.

An exemplary procedure is descirbed in the figure. An exemplary variant of the procedure may proceed as follows:

1. Update second moment: z[t+l]= pi z[t]+(l-pi)(m[i,t-N,k]) A 2

2. Compute bias corrected second moment: z’[t+l]=z[t+l]/(l-pi A (t+l))

3. Calculate normalization vector: V[t]=(y[t]/y[0])(|v’ [t]|/(sqrt(z’ [t+l])+s))+l A d

4. Normalize gradient: g[i,t]=g[i,t]/V[t]

5. Compute long-term normalization scalar:

6. V’[t+l]=P’V’[t]+(l-P’)||V[t]||, V”[t+l]=V’[t+l]/(l- P’ (t+l))

7. Update normalization momentum:

8. v[t+l]=Pv[t]+g[i,t]

9. Update bias corrected normalization momentum: v’[t+l]=v[t+l]/(l-P A (t+l))

10. Return normalized gradient and long-term normalization scalar: g[i,t], V”[t+1]

The reasoning behind normalizing the gradient comprises that when the momentum coefficient pi of second moment z[t] is much larger than the momentum coefficient P of normalization momentum v[t], e.g. pi=0.999 and P=0.9, the second moment z[t] smooths a short burst, whilethe normalization momentum v[t] might be amplified thereby. In this case, entry j of normalization vector V may be large and, thus, normalizes the large value at entry j of gradient g[i,t].

Reference is now made to FIG. 7 which is an additional schematic diagram of an exemplary process of forward and backward passes of MP ASGD training, according to some embodiments of the present disclosure. In this an exemplary process, the number of stages is greater than the number of micro-batches, resulting in staleness > 1. Having four stages and two micro-batches per mini-batch.

At the darker cells, the gradients are merged into the most up-to-date model. A new model marked M (for example Mi) is produced when all its parts have been updated Models are produced at times 0 (initial model), 14 and 26.

The pipeline may starts forward pass on micro-batches 5 and 6 on model Mo while it may finish merging it at time t 26 into M2, That means staleness = 2, and the number of different models at each time in pipeline is 3.

The descriptions of the various embodiments have been presented for purposes of illustration, but are not intended to be exhaustive or limited to the embodiments disclosed. Many modifications and variations will be apparent to those of ordinary skill in the art without departing from the scope and spirit of the described embodiments. The terminology used herein was chosen to best explain the principles of the embodiments, the practical application or technical improvement over technologies found in the marketplace, or to enable others of ordinary skill in the art to understand the embodiments disclosed herein.

It is expected that during the life of a patent maturing from this application many relevant machine learning models, neural network variants, and training systems and methods will be developed and the scopes of the terms machine learning model, neural network, ASGD, training step and the like are intended to include all such new technologies a priori.

As used herein the term “about” refers to ± 10 %.

The terms "comprises", "comprising", "includes", "including", “having” and their conjugates mean "including but not limited to". This term encompasses the terms "consisting of' and "consisting essentially of'.

The phrase "consisting essentially of means that the composition or method may include additional ingredients and/or steps, but only if the additional ingredients and/or steps do not materially alter the basic and novel characteristics of the claimed composition or method.

As used herein, the singular form "a", "an" and "the" include plural references unless the context clearly dictates otherwise. For example, the term "a compound" or "at least one compound" may include a plurality of compounds, including mixtures thereof.

The word “exemplary” is used herein to mean “serving as an example, instance or illustration”. Any embodiment described as “exemplary” is not necessarily to be construed as preferred or advantageous over other embodiments and/or to exclude the incorporation of features from other embodiments. The word “optionally” is used herein to mean “is provided in some embodiments and not provided in other embodiments”. Any particular embodiment may include a plurality of “optional” features unless such features conflict.

Throughout this application, various embodiments may be presented in a range format. It should be understood that the description in range format is merely for convenience and brevity and should not be construed as an inflexible limitation on the scope of embodiments. Accordingly, the description of a range should be considered to have specifically disclosed all the possible subranges as well as individual numerical values within that range. For example, description of a range such as from 1 to 6 should be considered to have specifically disclosed subranges such as from 1 to 3, from 1 to 4, from 1 to 5, from 2 to 4, from 2 to 6, from 3 to 6 etc., as well as individual numbers within that range, for example, 1, 2, 3, 4, 5, and 6. This applies regardless of the breadth of the range.

Whenever a numerical range is indicated herein, it is meant to include any cited numeral (fractional or integral) within the indicated range. The phrases “ranging/ranges between” a first indicate number and a second indicate number and “ranging/ranges from” a first indicate number “to” a second indicate number are used herein interchangeably and are meant to include the first and second indicated numbers and all the fractional and integral numerals therebetween.

It is appreciated that certain features of embodiments, which are, for clarity, described in the context of separate embodiments, may also be provided in combination in a single embodiment. Conversely, various features of embodiments, which are, for brevity, described in the context of a single embodiment, may also be provided separately or in any suitable subcombination or as suitable in any other described embodiment. Certain features described in the context of various embodiments are not to be considered essential features of those embodiments, unless the embodiment is inoperative without those elements.

Although embodiments have been described in conjunction with specific embodiments thereof, it is evident that many alternatives, modifications, equivalences, approximations and variations will be apparent to those skilled in the art. Accordingly, it is intended to embrace all such alternatives, modifications and variations that fall within the spirit and broad scope of the appended claims.

It is the intent of the applicant(s) that all publications, patents and patent applications referred to in this specification are to be incorporated in their entirety by reference into the specification, as if each individual publication, patent or patent application was specifically and individually noted when referenced that it is to be incorporated herein by reference. In addition, citation or identification of any reference in this application shall not be construed as an admission that such reference is available as prior art to the present invention. To the extent that section headings are used, they should not be construed as necessarily limiting. In addition, any priority document(s) of this application is/are hereby incorporated herein by reference in its/their entirety.