Mixture-of-experts layer with switchable parallel modes

1 . A computing system comprising:

a plurality of processing devices configured to execute a Mixture-of-Experts (MoE) layer included in an MoE model, wherein:

the MoE layer includes a plurality of expert sub-models that each have a respective plurality of parameter values; and

the MoE layer is configured to be switchable between a data parallel mode and an expert-data-model parallel mode without conveying the respective parameter values of the expert sub-models among the plurality of processing devices, wherein:

in the data parallel mode, the plurality of processing devices are configured to:

receive a first plurality of input tokens;

at a gating function, select one or more first destination expert sub-models based at least in part on the first plurality of input tokens;

at a plurality of first expert sub-model partitions of the one or more first destination expert sub-models, compute a plurality of output tokens based at least in part on the first plurality of input tokens;

perform an all-gather operation on the plurality of output tokens; and

output the plurality of output tokens subsequently to performing the all-gather operation as a first MoE layer output; and

in the expert-data-model parallel mode, the plurality of processing devices are configured to:

receive a second plurality of input tokens;

at the gating function, select one or more second destination expert sub-models based at least in part on the second plurality of input tokens;

within each of a plurality of processing device groups, compute a respective set of output tokens at a respective plurality of second expert sub-model partitions of the one or more second destination expert sub-models based at least in part on respective copies of the second plurality of input tokens;

perform respective all-gather operations on each of the respective sets of output tokens;

perform a respective plurality of local reduce operations on the sets of output tokens; and

output results of the plurality of local reduce operations as a second MoE layer output.

2 . The computing system of claim 1 , wherein the first plurality of input tokens and the second plurality of input tokens are each partitioned into respective sets of input tokens that are input into the gating function at respective processing devices of the plurality of processing devices.

3 . The computing system of claim 1 , wherein the plurality of processing devices are further configured to:

via an MoE layer application-programming interface (API), receive, for each of a plurality of iterations, respective processing device group counts that indicate respective numbers of processing device groups at the plurality of iterations; and

at each of the plurality of iterations, select between the data parallel mode and the expert-data-model parallel mode based at least in part on the processing device group count of that iteration.

4 . The computing system of claim 1 , wherein, in the expert-data-model parallel mode, the plurality of processing devices are further configured to perform a first collective communication of the respective copies of the second plurality of input tokens between the plurality of processing devices subsequently to executing the gating function.

5 . The computing system of claim 4 , wherein, in the expert-data-model parallel mode, the plurality of processing devices are further configured to perform a second collective communication of the sets of output tokens between the plurality of processing devices.

6 . The computing system of claim 4 , wherein, in the expert-data-model parallel mode:

the plurality of processing devices are further configured to perform a respective plurality of local repeat operations on the second plurality of input tokens prior to the first collective communication; and

the local reduce operations are local sum operations.

7 . The computing system of claim 1 , wherein the plurality of processing devices are further configured to:

compute a first communication complexity of computing an MoE layer output of the MoE layer via the data parallel mode;

compute a second communication complexity of computing the MoE layer output via the expert-data-model parallel mode; and

compute the MoE layer output using whichever of the data parallel mode and the expert-data-model parallel mode has a lower respective communication complexity.

8 . The computing system of claim 1 , wherein:

the plurality of processing devices are further configured to pipeline input tokens processed at the MoE layer into a plurality of streams when computing an MoE layer output of the MoE layer;

the plurality of streams each include a respective plurality of partitioned input tensors into which a plurality of input tensors are divided; and

the plurality of streams are processed in parallel at each of a first collective communication stage, the plurality of expert sub-models, and a second collective communication stage of the MoE layer, respectively.

9 . The computing system of claim 8 , wherein the plurality of processing devices are further configured to determine a number of the streams based at least in part on a capacity factor of the plurality of expert sub-models.

10 . The computing system of claim 9 , wherein the plurality of processing devices are configured to compute the number of the streams in each of a plurality of iterations at least in part by:

iteratively computing, over the plurality of iterations:

a plurality of capacity factor bins that include respective ranges of capacity factor values; and

respective numbers of the streams associated with the plurality of capacity factor bins; and

determining the number of streams based at least in part on a capacity factor bin of the plurality of capacity factor bins within which the capacity factor is located.

11 . A method for use with a computing system, the method comprising:

executing a Mixture-of-Experts (MoE) layer included in an MoE model at a plurality of processing devices, wherein:

the MoE layer includes a plurality of expert sub-models that each have a respective plurality of parameter values; and

executing the MoE layer includes switching the MoE layer between a data parallel mode and an expert-data-model parallel mode without conveying the respective parameter values of the expert sub-models among the plurality of processing devices, wherein:

in the data parallel mode, executing the MoE layer includes:

receiving a first plurality of input tokens;

at a gating function, selecting one or more first destination expert sub-models based at least in part on the first plurality of input tokens;

at a plurality of first expert sub-model partitions of the one or more first destination expert sub-models, computing a plurality of output tokens based at least in part on the first plurality of input tokens;

performing an all-gather operation on the plurality of output tokens; and

outputting the plurality of output tokens subsequently to performing the all-gather operation as a first MoE layer output; and

in the expert-data-model parallel mode, executing the MoE layer includes:

receiving a second plurality of input tokens;

at the gating function, selecting one or more second destination expert sub-models based at least in part on the second plurality of input tokens;

within each of a plurality of processing device groups, computing a respective set of output tokens at a respective plurality of second expert sub-model partitions of the one or more second expert destination sub-models based at least in part on respective copies of the second plurality of input tokens;

performing respective all-gather operations on each of the respective sets of output tokens;

performing a respective plurality of local reduce operations on the sets of output tokens; and

outputting results of the plurality of local reduce operations as a second MoE layer output.

12 . The method of claim 11 , further comprising partitioning each of the first plurality of input tokens and the second plurality of input tokens into respective sets of input tokens that are input into the gating function at respective processing devices of the plurality of processing devices.

13 . The method of claim 11 , further comprising, in the expert-data-model parallel mode:

performing a first collective communication of the respective copies of the second plurality of input tokens between the plurality of processing devices subsequently to executing the gating function; and

performing a second collective communication of the sets of output tokens between the plurality of processing devices.

14 . The method of claim 11 , further comprising:

computing a first communication complexity of computing an MoE layer output of the MoE layer via the data parallel mode;

computing a second communication complexity of computing the MoE layer output via the expert-data-model parallel mode; and

computing the MoE layer output using whichever of the data parallel mode and the expert-data-model parallel mode has a lower respective communication complexity.

15 . The method of claim 11 , further comprising:

pipelining input tokens processed at the MoE layer into a plurality of streams when computing an MoE layer output of the MoE layer, wherein the plurality of streams each include a respective plurality of partitioned input tensors into which a plurality of input tensors are divided; and

processing the plurality of streams in parallel at each of a first collective communication stage, the plurality of expert sub-models, and a second collective communication stage of the MoE layer, respectively.

16 . A computing system comprising:

a plurality of processing devices configured to execute a Mixture-of-Experts (MoE) layer included in an MoE model, wherein:

the MoE layer is configured to be switchable between a first parallel mode and a second parallel mode, in which a first MoE layer output and a second MoE layer output are respectively computed in a first iteration and a second iteration of a plurality of iterations;

in the first parallel mode, the plurality of processing devices are configured to:

receive a first plurality of input tokens;

at a gating function, select one or more first destination expert sub-models based at least in part on the first plurality of input tokens;

at a plurality of first expert sub-model partitions of the one or more first destination expert sub-models, compute a plurality of output tokens based at least in part on the first plurality of input tokens;

perform an all-gather operation on the plurality of output tokens; and

output the plurality of output tokens subsequently to performing the all-gather operation as the first MoE layer output; and

in the second parallel mode, the plurality of processing devices are configured to:

receive a second plurality of input tokens;

at the gating function, select one or more second destination expert sub-models based at least in part on the second plurality of input tokens;

within each of a plurality of processing device groups, compute a respective set of output tokens at a respective plurality of second expert sub-model partitions of the one or more second destination expert sub-models based at least in part on respective copies of the second plurality of input tokens;

perform respective all-gather operations on each of the respective sets of output tokens;

perform a respective plurality of local reduce operations on the sets of output tokens; and

output results of the plurality of local reduce operations as the second MoE layer output.