Searching parallel schedules for execution of artificial intelligence workloads

A computing system, comprising:

• memory;
• a processor system coupled to the memory;
• and
• one or more computer readable storage media storing instructions that, when loaded into the memory, cause the processor system to perform operations comprising:
• receiving an internal representation of a transformer model, wherein the internal representation defines one or more repeating blocks, each block comprising a sequence of cells, and each cell comprising a set of tasks of the transformer model;
• and
• searching for a plurality of parallel schedules for partitioning devices included in a device cluster for parallel execution of the transformer model, wherein the searching comprises:
• determining a number of model replicas, wherein each model replica represents a copy of the transformer model, wherein devices included in the device cluster are partitioned into the number of model replicas;
• determining a number of stages that divide the one or more repeating blocks, wherein devices partitioned into each model replica are partitioned into the number of stages;
• determining a number of cell replicas for each cell in a block, wherein each cell replica represents a copy of the corresponding cell, wherein devices partitioned into each stage are partitioned into the number of cell replicas;
• and
• for each cell replica of a cell, generating a task mapping which maps the set of tasks included in the cell to devices partitioned into the cell replica.

One or more computer-readable media having encoded thereon computer-executable instructions causing one or more processors to perform a method, the method comprising:

• receiving an internal representation of a transformer model, wherein the internal representation defines one or more repeating blocks, each block comprising a sequence of cells, and each cell comprising a set of tasks of the transformer model;
• and
• searching for a plurality of parallel schedules for partitioning devices included in a device cluster for parallel execution of the transformer model, wherein the searching comprises:
• determining a number of model replicas, wherein each model replica represents a copy of the transformer model, wherein devices included in the device cluster are partitioned into the number of model replicas;
• determining a number of stages that divide the one or more repeating blocks, wherein devices partitioned into each model replica are partitioned into the number of stages;
• determining a number of cell replicas for each cell in a block, wherein each cell replica represents a copy of the corresponding cell, wherein devices partitioned into each stage are partitioned into the number of cell replicas;
• and
• for each cell replica of a cell, generating a task mapping which maps the set of tasks included in the cell to devices partitioned into the cell replica.