deepspeedai--deepspeed
wehub-resource-sync
158 行
5.5 KiB
Python
可执行文件
158 行
5.5 KiB
Python
可执行文件
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import hjson
|
|
|
|
from ..constants import AUTOTUNING, AUTOTUNING_METRIC_PATH
|
|
from .base_tuner import BaseTuner
|
|
from .cost_model import XGBoostCostModel
|
|
from .utils import *
|
|
from ..utils import *
|
|
import numbers
|
|
from ..constants import AUTOTUNING_METRIC_LATENCY
|
|
|
|
INIT_NUM = 2
|
|
|
|
|
|
class ModelBasedTuner(BaseTuner):
|
|
"""Exploring the search space with a cost model"""
|
|
|
|
def __init__(self, exps: list, resource_manager, metric, tuning_space):
|
|
super().__init__(exps, resource_manager, metric)
|
|
self.tuning_space = tuning_space
|
|
self.best_iter = 0
|
|
|
|
self.all_configs = [e['ds_config'] for e in exps]
|
|
self.num_all_configs = len(self.all_configs)
|
|
|
|
self.dims = dict_to_dims(self.tuning_space)
|
|
|
|
logger.info(f"Create config dim: {self.dims}, all configs: {self.num_all_configs}")
|
|
|
|
self.visited = set([])
|
|
|
|
self.trials = []
|
|
self.trial_pt = 0
|
|
|
|
init_num = min(INIT_NUM, self.num_all_configs)
|
|
|
|
for _ in range(init_num):
|
|
exp_feature = np.random.randint(self.num_all_configs)
|
|
exp_feature = 0
|
|
while exp_feature in self.visited:
|
|
exp_feature = np.random.randint(self.num_all_configs)
|
|
self.trials.append(exp_feature)
|
|
self.visited.add(exp_feature)
|
|
|
|
self.cost_model = XGBoostCostModel("rank")
|
|
|
|
self.evaluated_configs = []
|
|
self.evaluated_perf = []
|
|
|
|
self.train_ct = 0
|
|
|
|
self.random_exploration_ratio = 0.2 # do random exploration
|
|
|
|
def find_estimated_top_configs(self):
|
|
"""Use the cost model to predict the estimated performance of configurations and find the top ones for the next round of evaluation"""
|
|
|
|
configs = []
|
|
|
|
for c in self.all_configs:
|
|
flattened_ds_config = flatten(c)
|
|
feature_val = []
|
|
for k, v in flattened_ds_config.items():
|
|
if isinstance(v, numbers.Number):
|
|
feature_val.append(v)
|
|
configs.append(feature_val)
|
|
# print(configs)
|
|
# TODO the current implementation requires that all configs have the same shape.
|
|
configs = np.array(configs, dtype=np.float32)
|
|
estimates = self.cost_model.predict(configs)
|
|
|
|
n = len(estimates)
|
|
top_idx = np.argsort(estimates)
|
|
top_idx_ret = top_idx if self.metric == AUTOTUNING_METRIC_LATENCY else top_idx[::-1][:n]
|
|
|
|
# top_configs = [self.all_configs[i] for i in top_idx]
|
|
|
|
return top_idx_ret
|
|
|
|
def next_batch(self, sample_size):
|
|
sampled_batch = []
|
|
|
|
counter = 0
|
|
while counter < sample_size:
|
|
|
|
if len(self.visited) >= self.num_all_configs:
|
|
break
|
|
|
|
while self.trial_pt < len(self.trials):
|
|
logger.debug(f"trials: {self.trials}")
|
|
# Select top promising trials
|
|
index = self.trials[self.trial_pt]
|
|
if index not in self.visited:
|
|
break
|
|
self.trial_pt += 1
|
|
|
|
# To avoid over-exploitation, randomly select one that has not been explored.
|
|
rand = np.random.rand()
|
|
if rand < self.random_exploration_ratio:
|
|
# Do normal selection
|
|
feature = np.random.choice(self.trials)
|
|
while index in self.visited:
|
|
index = np.random.randint(self.num_all_configs)
|
|
|
|
# Need to track both the sampled configs and indices
|
|
|
|
sampled_batch.append(self.all_exps[index])
|
|
self.visited.add(index)
|
|
counter += 1
|
|
|
|
return sampled_batch
|
|
|
|
def has_next(self):
|
|
return len(self.visited) < self.num_all_configs
|
|
|
|
def update(self):
|
|
for exp_id, (exp, err) in self.rm.finished_experiments.items():
|
|
feature_val = []
|
|
if err:
|
|
logger.info(
|
|
f"Skipping exp_id = {exp_id}, exp_name = {exp['name']}, the experiment did not run successfully with error = {err}, thus a metrics.txt does not exist for it. Please check the stderr.log in {exp['result_dir']}"
|
|
)
|
|
ds_config = exp["ds_config"]
|
|
flattened_ds_config = flatten(ds_config)
|
|
for k, v in flattened_ds_config.items():
|
|
if isinstance(v, numbers.Number):
|
|
feature_val.append(v)
|
|
self.evaluated_configs.append(feature_val)
|
|
self.evaluated_perf.append(0.0)
|
|
continue
|
|
|
|
p = exp["ds_config"][AUTOTUNING][AUTOTUNING_METRIC_PATH]
|
|
with open(p, 'r') as f:
|
|
results = hjson.load(f)
|
|
curr_iter = results[self.metric]
|
|
logger.debug(f"parsing the results for {exp_id}， Result is {curr_iter}")
|
|
|
|
ds_config = exp["ds_config"]
|
|
flattened_ds_config = flatten(ds_config)
|
|
for k, v in flattened_ds_config.items():
|
|
if isinstance(v, numbers.Number):
|
|
feature_val.append(v)
|
|
self.evaluated_configs.append(feature_val)
|
|
self.evaluated_perf.append(curr_iter)
|
|
|
|
logger.debug(f"**Evaluated configs: {len(self.evaluated_configs)}, evaluated perf: {self.evaluated_perf}")
|
|
|
|
self.cost_model.fit(self.evaluated_configs, self.evaluated_perf)
|
|
|
|
estimated_top_configs = self.find_estimated_top_configs()
|
|
|
|
self.trials = estimated_top_configs
|
|
self.trial_pt = 0
|
|
self.train_ct += 1
|