  The convergence of SGD based distributed training algorithms is tied to the
data distribution across workers. Standard partitioning techniques try to
achieve equal-sized partitions with per-class population distribution in
proportion to the total dataset. Partitions having the same overall population
size or even the same number of samples per class may still have Non-IID
distribution in the feature space. In heterogeneous computing environments,
when devices have different computing capabilities, even-sized partitions
across devices can lead to the straggler problem in distributed SGD. We develop
a framework for distributed SGD in heterogeneous environments based on a novel
data partitioning algorithm involving submodular optimization. Our data
partitioning algorithm explicitly accounts for resource heterogeneity across
workers while achieving similar class-level feature distribution and
maintaining class balance. Based on this algorithm, we develop a distributed
SGD framework that can accelerate existing SOTA distributed training algorithms
by up to 32%.
