GlobalPowerLimitOptimizer works well for single node data parallel training, but in case of distributed data parallel, GPUs in different nodes should make the same final GPU power limit choice. Assuming homogeneous GPUs this is still very likely to happen, but we should make it more robust just in case.
[ ] Try running GlobalPowerLimitOptimizer in a distributed training setting
[ ] Create collective communication helpers that internally call torch.distributed.all_reduce (for AllReduce) if PyTorch is available and something else for JAX. (The time and energy measurements from each ZeusMonitor instance in each node should gather results, and node rank 0 will make global decisions.)
[ ] Update the ImageNet training example to explicitly mention distributed data parallel training.
GlobalPowerLimitOptimizerworks well for single node data parallel training, but in case of distributed data parallel, GPUs in different nodes should make the same final GPU power limit choice. Assuming homogeneous GPUs this is still very likely to happen, but we should make it more robust just in case.GlobalPowerLimitOptimizerin a distributed training settingtorch.distributed.all_reduce(for AllReduce) if PyTorch is available and something else for JAX. (The time and energy measurements from eachZeusMonitorinstance in each node should gather results, and node rank 0 will make global decisions.)