AMP: Automatically Finding Model Parallel Strategies with Heterogeneity Awareness

We present AMP, an automatic algorithm that can find model parallel training strategies with highsystem throughput. AMP is equipped with an expert-design cost model that considers cluster andmodel configurations. We show that when heterogeneity exists in cluster or model setup, currentheuristics are sub-optimal. Our automatic algorithm can find strategies with 1.51× and 1.76×speedup within similar budgets.