Source-linked AI summary
Tune: A Research Platform for Distributed Model Selection and Training
Richard Liaw, Eric Liang, Robert Nishihara, Philipp Moritz, Joseph E. Gonzalez, Ion Stoica
TL;DR
Distributed model selection is costly and existing search methods are often ad hoc or infrastructure-heavy. Tune provides a unified narrow-waist framework separating training scripts from schedulers, and supports integration of many state-of-the-art algorithms. The framework is presented as extensible, scalable to distributed workloads, and easier for model developers and algorithm researchers to use.
Problem
Distributed model selection is computationally costly, while existing techniques are often tied to frameworks, difficult to reproduce, or require substantial infrastructure.
Method
Tune separates user-facing training and researcher-facing scheduling APIs and implements them on Ray for distributed execution and resource management.
Results
Tune’s APIs enable reproduction and integration of a wide variety of state-of-the-art hyperparameter search algorithms.
Takeaways & Limitations
Tune provides an extensible framework that model developers can incorporate into model-design processes while exposing distributed search algorithms to users.
Takeaways & Limitations
Tune keeps active-trial metadata in memory and relies on checkpoints for fault tolerance, though schedulers may use external storage when necessary.
Abstract
from arXiv · showhide
Modern machine learning algorithms are increasingly computationally demanding, requiring specialized hardware and distributed computation to achieve high performance in a reasonable time frame. Many hyperparameter search algorithms have been proposed for improving the efficiency of model selection, however their adaptation to the distributed compute environment is often ad-hoc. We propose Tune, a unified framework for model selection and training that provides a narrow-waist interface between training scripts and search algorithms. We show that this interface meets the requirements for a broad range of hyperparameter search algorithms, allows straightforward scaling of search to large clusters, and simplifies algorithm implementation. We demonstrate the implementation of several state-of-the-art hyperparameter search algorithms in Tune. Tune is available at http://ray.readthedocs.io/en/latest/tune.html.
1. Introduction
Tune addresses the computational cost and infrastructure burden of distributed model selection with an open-source framework and narrow interfaces between training scripts and search algorithms.
- Motivation: Model selection can consume weeks when training large models sequentially, motivating distributed training and parallel model selection.ResNet-101 training is described as taking around 24 hours on one GPU.
- Problem: Existing model-selection techniques are often tied to specific frameworks, closed source, unreproducible, or costly to implement infrastructurally.
- Contribution: Tune is introduced as an open-source framework for distributed model selection.
- Contribution: Tune’s APIs support reproducing and integrating a wide variety of state-of-the-art hyperparameter search algorithms.
2. Related work
Prior model-selection systems provide useful search or automation capabilities but commonly couple algorithms to infrastructure, frameworks, or execution models.
- Distributed model selection: HyperOpt, Spearmint, and HPOLib manage distributed search and evaluation but couple systems to search structures and require manual cluster resource management.
- Distributed model selection: Google Vizier provides parallel trials, optimization algorithms, and performance analysis, but is tied to closed-source infrastructure.
- Model debugging: Mistique emphasizes model debugging and memory-footprint minimization rather than executing the model-selection process.
- AutoML systems: Auto-SKLearn and Auto-WEKA automate model selection through meta-learning and ensembling but support only Scikit-Learn and WEKA, respectively.
3. Requirements for API generality
A general model-search platform must coordinate sequential decisions with parallel, irregular training while supporting arbitrary resources, intermediate results, and usable monitoring.
- Core requirements: Model search evaluates many trials in parallel while search algorithms examine results sequentially and affect ongoing computation.
- Core requirements: Trials may have irregular durations and resource usage, and user code may require arbitrary CPU, GPU, and parallel resources.
- Core requirements: Schedulers need intermediate trial results to support early stopping and mid-training parameter cloning or mutation.
- User experience: A usable platform should monitor and visualize trial progress and outcomes while keeping experiment specification simple.
- Proposed solution: Tune proposes user-facing and scheduling APIs on Ray, whose task and actor abstractions support irregular computation and intermediate-result decisions.
4. Tune API
Tune separates user training interfaces from scheduler interfaces, allowing schedulers to control distributed trials through cooperative or direct execution APIs and event-based decisions.
- API design: Tune separates a user API for model training from a scheduling API for researchers developing model-search algorithms.This division gives users a choice of search algorithms and researchers access to diverse workloads.
- User API: Cooperative control exposes intermediate results, checkpoints, and mid-training hyperparameter changes with minimal training-code modifications.
- User API: Tune also offers a class-based interface in which schedulers directly invoke training, checkpoint, and restore methods.This mode has debuggability advantages over cooperative control.
- Scheduler API: Supported scheduling actions include early stopping, hyperparameter adjustment, cloning promising trials, querying shared results, and prioritizing scarce resources.
- Scheduler API: Schedulers choose trials when resources are available and can stop, checkpoint, restart, or reconfigure trials as results arrive.
- Scheduler API: Tune keeps active-trial metadata in memory and uses checkpoints for fault tolerance, while allowing schedulers to use external storage when needed.
- Scaling computation: Tune integrates model-selection algorithms, including original and asynchronous HyperBand, with the asynchronous variant simpler to implement in distributed settings.
5. Implementation
Tune uses Ray to execute distributed hyperparameter-search trials, including nested parallel computations, while avoiding a central scheduling bottleneck.
- Tune implements distributed model-selection algorithms using Ray’s actor abstraction to run trials.Ray provides a more flexible programming model than Spark or MPI, enabling Tune’s trial schedulers.
- Ray’s two-level distributed scheduler supports nested computations such as hyperparameter optimization.Scheduling occurs locally when possible and spills over to other cluster machines when local resources are exhausted.
- Local scheduling avoids a central bottleneck when trials themselves use parallel computations.
6. Conclusion and Future Work
Tune provides a general API and system for extensible distributed hyperparameter search that model developers can incorporate into their workflows. Future development targets tuning analysis and debugging alongside new functionality.
- Tune supports extensible distributed hyperparameter-search algorithms through a general API and system.
- The system is designed to be incorporated into end-user model-development processes.
- Ongoing development aims to add functionality for analyzing and debugging intermediate tuning results.