DYNAMIX: RL-based Adaptive Batch Size Optimization in Distributed Machine Learning Systems

Fuente: arXiv
Salvato in:
Dettagli Bibliografici
Autori principali: Dai, Yuanjun, He, Keqiang, Wang, An
Natura: Preprint
Pubblicazione: 2025
Soggetti:
Accesso online:
Tags: Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
_version_ 1866918157637124096
author Dai, Yuanjun
He, Keqiang
Wang, An
author_facet Dai, Yuanjun
He, Keqiang
Wang, An
contents Existing batch size selection approaches in distributed machine learning rely on static allocation or simplistic heuristics that fail to adapt to heterogeneous, dynamic computing environments. We present DYNAMIX, a reinforcement learning framework that formulates batch size optimization as a sequential decision-making problem using Proximal Policy Optimization (PPO). Our approach employs a multi-dimensional state representation encompassing network-level metrics, system-level resource utilization, and training statistical efficiency indicators to enable informed decision-making across diverse computational resources. Our approach eliminates the need for explicit system modeling while integrating seamlessly with existing distributed training frameworks. Through evaluations across diverse workloads, hardware configurations, and network conditions, DYNAMIX achieves up to 6.3% improvement in the final model accuracy and 46% reduction in the total training time. Our scalability experiments demonstrate that DYNAMIX maintains the best performance as cluster size increases to 32 nodes, while policy transfer experiments show that learned policies generalize effectively across related model architectures.
format Preprint
id arxiv_https___arxiv_org_abs_2510_08522
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle DYNAMIX: RL-based Adaptive Batch Size Optimization in Distributed Machine Learning Systems
Dai, Yuanjun
He, Keqiang
Wang, An
Machine Learning
Distributed, Parallel, and Cluster Computing
Existing batch size selection approaches in distributed machine learning rely on static allocation or simplistic heuristics that fail to adapt to heterogeneous, dynamic computing environments. We present DYNAMIX, a reinforcement learning framework that formulates batch size optimization as a sequential decision-making problem using Proximal Policy Optimization (PPO). Our approach employs a multi-dimensional state representation encompassing network-level metrics, system-level resource utilization, and training statistical efficiency indicators to enable informed decision-making across diverse computational resources. Our approach eliminates the need for explicit system modeling while integrating seamlessly with existing distributed training frameworks. Through evaluations across diverse workloads, hardware configurations, and network conditions, DYNAMIX achieves up to 6.3% improvement in the final model accuracy and 46% reduction in the total training time. Our scalability experiments demonstrate that DYNAMIX maintains the best performance as cluster size increases to 32 nodes, while policy transfer experiments show that learned policies generalize effectively across related model architectures.
title DYNAMIX: RL-based Adaptive Batch Size Optimization in Distributed Machine Learning Systems
topic Machine Learning
Distributed, Parallel, and Cluster Computing
url https://arxiv.org/abs/2510.08522