Systems, apparatuses, methods, and computer-readable media are provided for load partitioning in distributed machine learning (ML) training using heterogeneous compute nodes in a heterogeneous computing environment, where the heterogeneous compute nodes are connected to a master node via respective wireless links. ML computations are performed by individual heterogeneous compute nodes on respective load partitions. The ML computations are balanced across the heterogeneous compute nodes based on knowledge of respective computational and link parameters of the heterogeneous compute nodes. Other embodiments may be described and/or claimed.