Integrated hardware architecture and distribution strategy optimization for deep learning models

A method of training optimization for a deep learning model, the method comprising:

• receiving input data which includes at least one training script for the deep learning model and defines an area constraint for accelerator optimization;
• extracting operator graphs pertaining to each layer in the deep learning model from the at least one training script using a graph extractor component;
• estimating a latency for each operator in each of the operator graphs using a runtime estimator component;
• generating a set of feasible architecture configurations that satisfy the area constraint using an architecture generator component;
• performing an evaluation process that includes:
• selecting a current feasible architecture configuration from the set of feasible architecture configurations having a largest area that has not been evaluated;
• providing the operator graphs, the latency for each operator in each of the operator graphs, and the current feasible architecture configuration to a solver component;
• using the solver component to solve an integer linear program (ILP) for each layer in the deep learning model based on the operator graphs and the latency of each operator in each of the operator graphs to generate an optimal schedule of operators to execute on a single accelerator architecture;
• using the solver component to solve a dynamic programming (DP) algorithm with reference to the optimal schedule of operators and the current feasible architecture configuration to determine a distribution strategy that identifies an optimal combination of pipeline parallel depth, data parallel width, and tensor model parallel width for executing the layers across multiple accelerators;
• providing feedback to the architecture generator component indicating a throughput for the current feasible architecture configuration;
• and
• repeating the evaluation process until a convergence is reached, the convergence being indicated by a characteristic change in performance for the current feasible architecture configuration;
• once the convergence has been reached, selecting an optimal architecture configuration from the feasible architecture configurations that has been evaluated and the distribution strategy determined for the optimal architecture configuration to use as a basis for an accelerator design for training the deep learning model.

A non-transitory computer readable medium on which are stored instructions that, when executed, cause a programmable device to perform functions of:

• receiving input data which includes at least one training script for a deep learning model and defines an area constraint for accelerator optimization;
• extracting operator graphs pertaining to each layer in the deep learning model from the at least one training script using a graph extractor component;
• estimating a latency for each operator in each of the operator graphs using a runtime estimator component;
• generating a set of feasible architecture configurations that satisfy the area constraint using an architecture generator component;
• performing an evaluation process that includes:
• selecting a current feasible architecture configuration from the set of feasible architecture configurations having a largest area that has not been evaluated;
• providing the operator graphs, the latency for each operator in each of the operator graphs, and the current feasible architecture configuration to a solver component;
• using the solver component to solve an integer linear program (ILP) for each layer in the deep learning model based on the operator graphs and the latency of each operator in each of the operator graphs to generate an optimal schedule of operators to execute on a single accelerator architecture;
• using the solver component to solve a dynamic programming (DP) algorithm with reference to the optimal schedule of operators and the current feasible architecture configuration to determine a distribution strategy that identifies an optimal combination of pipeline parallel depth, data parallel width, and tensor model parallel width for executing the layers across multiple accelerators;
• providing feedback to the architecture generator component indicating a throughput for the current feasible architecture configuration;
• and
• repeating the evaluation process until a convergence is reached, the convergence being indicated by a characteristic change in performance for the current feasible architecture configuration;
• once the convergence has been reached, selecting an optimal architecture configuration from the feasible architecture configurations that has been evaluated and the distribution strategy determined for the optimal architecture configuration to use as a basis for an accelerator design for training the deep learning model.