[Paper Review] Tune: A Research Platform for Distributed Model Selection and Training

Richard Liaw, Eric Liang+4 more|arXiv (Cornell University)|Jul 13, 2018

Machine Learning and Data Classification11 references615 citations

TL;DR

Tune provides a unified, open-source API and scheduling framework for distributed model selection and training, enabling easy integration of various hyperparameter search algorithms on top of Ray. It decouples training scripts from search logic to scale experiments across clusters.

ABSTRACT

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.

Motivation & Objective

• Motivate the need for scalable, reproducible distributed model selection and training.
• Introduce Tune as a narrow-waist API between training scripts and hyperparameter search algorithms.
• Demonstrate that Tune enables a broad range of search strategies and easy integration across frameworks.

Proposed method

• Propose a two-API design: a user API for training scripts and a scheduling API for search algorithms.
• Implement a cooperative control model or direct-trail control to interact with Tune during training.
• Provide a TrialScheduler interface with on_result and choose_trial_to_run to manage parallel trials.
• Build on the Ray framework to handle distributed execution, resource management, and data handling across trials.
• Implement and integrate multiple state-of-the-art hyperparameter search algorithms (e.g., HyperBand variants, Median Stopping Rule, HyperOpt, Population-Based Training).
• Offer a minimal example and a domain-specific DSL for defining initial trial configurations.

Experimental results

Research questions

• RQ1 Can Tune support a broad range of hyperparameter optimization algorithms with a single, general API?
• RQ2 Does the Ray-based implementation enable scalable, distributed execution of many concurrent trials?
• RQ3 Can intermediate trial results be effectively used to make dynamic scheduling decisions and hyperparameter adaptations?
• RQ4 Is the user experience simple enough to enable easy integration into existing training scripts while preserving reproducibility?
• RQ5 How does Tune facilitate reproduction, visualization, and comparison of AutoML experiments across different schedulers?

Key findings

• Tune provides narrow-waist user and scheduling APIs enabling easy integration of various hyperparameter search algorithms.
• The framework supports irregular and heterogeneous trial workloads and intermediate-result-driven scheduling decisions.
• Multiple algorithms are implemented or integrated in Tune, including asynchronous HyperBand, HyperBand, Median Stopping Rule, HyperOpt, and Population-Based Training.
• Trials run as Ray tasks/actors with resource management and data handling via Ray, enabling nested distributed computations.
• Tune offers a minimal example and a DSL for grid/search configuration, and logs progress via console and TensorBoard integrations.