zeus
zeus copied to clipboard
`GlobalPowerLimitOptimizer` for distributed data parallel training
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
GlobalPowerLimitOptimizerin 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 eachZeusMonitorinstance 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.