Files
paddlepaddle--paddle/python/paddle/distributed/auto_parallel/static/tuner/optimization_tuner.py
T
2026-07-13 12:40:42 +08:00

644 lines
23 KiB
Python

# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import copy
import json
import logging
# import yaml
import os
import pathlib
import pickle
import shlex
import shutil
import subprocess
import sys
import time
import paddle
from paddle.distributed.auto_parallel.static.completion import Completer
from paddle.distributed.auto_parallel.static.dist_context import (
DistributedContext,
)
from paddle.distributed.auto_parallel.static.partitioner import Partitioner
from paddle.distributed.auto_parallel.static.process_group import (
clear_all_process_groups,
get_all_process_groups,
new_process_group,
)
from paddle.distributed.auto_parallel.static.reshard import Resharder
from paddle.distributed.auto_parallel.static.utils import debug_program
from paddle.distributed.passes import PassContext, new_pass
from paddle.static import append_backward, program_guard
from ..utils import get_logger
from .algorithms import new_algorithm
from .config import TuningConfig
from .trial import TrialStatus
def _get_new_params_grads(target_program, ref_program, ref_params_grads):
ref_block = ref_program.global_block()
target_block = target_program.global_block()
target_params_grads = []
for p, g in ref_params_grads:
# NOTE grad var might not be generated
assert ref_block.has_var(p.name)
assert target_block.has_var(p.name)
new_p = target_block.var(p.name)
if g:
new_g = target_block.var(g.name)
else:
new_g = None
target_params_grads.append((new_p, new_g))
return target_params_grads
def _get_new_loss(target_program, ref_program, loss):
ref_block = ref_program.global_block()
target_block = target_program.global_block()
assert ref_block.has_var(loss.name)
return target_block.var(loss.name)
def parse_process_groups():
group_map = {}
all_process_groups = get_all_process_groups()
for process_group in all_process_groups:
group_map[process_group.id] = process_group.ranks
return group_map
def get_metric(results):
assert isinstance(results, dict), (
f"results should be type of dictionary, but got {type(results)}."
)
if 'Throughput' in results and isinstance(results['Throughput'], float):
return float(results['Throughput'])
else:
return -1.0
def parse_results(results):
if results['Throughput'] > 0:
return "Throughput: {} step / s.".format(results['Throughput'])
et = results.get("ErrorType", None)
if et == "ResourceExhaustedError":
return "Fail with OOM"
else:
return "Fail with UNKNOWN ERROR"
# TODO only dependent on dist context
# all env need to be start a new pass are member of dist context
def _copy_context(ref_dist_context):
# clear all process groups and recover the world process group
clear_all_process_groups()
ranks = []
for process_mesh in ref_dist_context._process_meshes:
ranks.extend(process_mesh.process_ids)
new_process_group(list(set(ranks)))
new_dist_context = DistributedContext()
new_dist_context._serial_main_program = (
ref_dist_context.serial_main_program.clone(for_test=False)
)
new_dist_context._serial_startup_program = (
ref_dist_context.serial_startup_program.clone(for_test=False)
)
# mapping variable into new dist context
if getattr(ref_dist_context, '_params_grads', None):
new_dist_context._params_grads = _get_new_params_grads(
new_dist_context.serial_main_program,
ref_dist_context.serial_main_program,
ref_dist_context._params_grads,
)
new_dist_context._serial_loss = _get_new_loss(
new_dist_context.serial_main_program,
ref_dist_context.serial_main_program,
ref_dist_context.serial_loss,
)
for key, var_list in ref_dist_context._serial_feed_vars.items():
new_var_list = []
for var in var_list:
block_idx = var.block.idx
var_name = var.name
var = new_dist_context._serial_main_program.blocks[
block_idx
]._var_recursive(var_name)
new_var_list.append(var)
new_dist_context._serial_feed_vars[key] = new_var_list
for key, var_list in ref_dist_context._serial_fetch_vars.items():
new_var_list = []
# metrics is a list of list
if key == "metrics":
for inner_var_list in var_list:
new_inner_var_list = []
for var in inner_var_list:
block_idx = var.block.idx
var_name = var.name
var = new_dist_context._serial_main_program.blocks[
block_idx
]._var_recursive(var_name)
new_inner_var_list.append(var)
new_var_list.append(new_inner_var_list)
else:
for var in var_list:
block_idx = var.block.idx
var_name = var.name
var = new_dist_context._serial_main_program.blocks[
block_idx
]._var_recursive(var_name)
new_var_list.append(var)
new_dist_context._serial_fetch_vars[key] = new_var_list
# copy information in forward and backward
new_dist_context._serial_optimizer = copy.deepcopy(
ref_dist_context.serial_optimizer
)
new_dist_context._dist_tensors_for_program = copy.deepcopy(
ref_dist_context._dist_tensors_for_program
)
new_dist_context._dist_ops_for_program = copy.deepcopy(
ref_dist_context._dist_ops_for_program
)
for pm in ref_dist_context.process_meshes:
new_dist_context.add_process_mesh(pm)
new_dist_context._dist_op_context = copy.deepcopy(
ref_dist_context._dist_op_context
)
new_dist_context._block_state = copy.deepcopy(ref_dist_context.block_state)
return new_dist_context
class OptimizationTuner:
"""
OptimizationTuner is used to manage the tuning procedure of hyper-parameters (configs)
of Optimization Pass in AutoParallel.
"""
def __init__(
self,
dist_context,
dataset,
inputs_spec,
labels_spec,
batch_size,
rank,
):
self._config = TuningConfig(dist_context.strategy)
# should not modify dist context from calling function
self._baseline_dist_context = _copy_context(dist_context)
self._baseline_completer = Completer(self._baseline_dist_context)
self._rank = rank
self._inputs_spec = inputs_spec
self._labels_spec = labels_spec
self._dataset = dataset
self._batch_size = batch_size
self._finished_trials = []
self._best_metric = None
self._best_iter = float("-inf")
self._logger = get_logger(logging.INFO)
self._build_programs_without_optimization()
self._select_tuning_algorithm()
@property
def project_dir(self):
dirname = self._config.project_dir
if not os.path.exists(dirname):
if self.rank == 0:
pathlib.Path(dirname).mkdir(parents=True, exist_ok=True)
return dirname
@property
def rank(self):
return self._rank
@property
def device_id(self):
return paddle.distributed.ParallelEnv().device_id
# TODO Generate complete program with all parts like forward, backward, update
# as well as parallelism transformation.
def _build_programs_without_optimization(self):
serial_main_program = self._baseline_dist_context.serial_main_program
serial_startup_program = (
self._baseline_dist_context.serial_startup_program
)
serial_loss = self._baseline_dist_context.serial_loss
with program_guard(serial_main_program, serial_startup_program):
params_grads = append_backward(
serial_loss,
distop_context=self._baseline_dist_context.dist_op_context,
)
self._baseline_completer.complete_backward_annotation(
serial_main_program
)
self._baseline_dist_context.block_state.parse_backward_blocks(
serial_main_program
)
self._baseline_dist_context._params_grads = params_grads
if self._config.debug:
baseline_dir = os.path.join(self.project_dir, "baseline")
if not os.path.exists(baseline_dir):
pathlib.Path(baseline_dir).mkdir(parents=True, exist_ok=True)
debug_program(
self._baseline_dist_context._serial_main_program,
baseline_dir,
"main",
)
debug_program(
self._baseline_dist_context._serial_startup_program,
baseline_dir,
"startup",
)
def _select_tuning_algorithm(self):
selected_passes_set = self._config.tuning_passes_name
algorithm_name = "_".join(sorted(selected_passes_set))
self._algorithm = new_algorithm(algorithm_name, self._config)
def _apply_optimization(self, trial):
new_strategy = trial.space
dist_context = _copy_context(self._baseline_dist_context)
pass_context = PassContext()
completer = Completer(dist_context)
main_program = dist_context.serial_main_program
startup_program = dist_context.serial_startup_program
# applying optimization pass
if new_strategy.amp.enable:
config = copy.deepcopy(new_strategy.amp.to_dict())
config["dist_context"] = dist_context
config["params_grads"] = dist_context._params_grads
# TODO AMP Pass should not use loss var
config["loss"] = dist_context.serial_loss
config["input_data"] = (
self._baseline_dist_context.serial_feed_vars["inputs"]
+ self._baseline_dist_context.serial_feed_vars["labels"]
)
if config["dtype"] == "float16" and config["level"] == "o2":
config["base_opt"] = dist_context.serial_optimizer
auto_parallel_fp16_pass = new_pass("auto_parallel_fp16", config)
auto_parallel_fp16_pass.apply(
[main_program], [startup_program], pass_context
)
dist_context._serial_loss = auto_parallel_fp16_pass.get_loss()
else:
auto_parallel_amp_pass = new_pass("auto_parallel_amp", config)
auto_parallel_amp_pass.apply(
[main_program], [startup_program], pass_context
)
dist_context._serial_loss = auto_parallel_amp_pass.get_loss()
if new_strategy.recompute.enable:
config = copy.deepcopy(new_strategy.recompute.to_dict())
config["dist_context"] = dist_context
config["no_grad_set"] = None
config["loss"] = dist_context.serial_loss
auto_parallel_recompute_pass = new_pass(
"auto_parallel_recompute", config
)
auto_parallel_recompute_pass.apply(
[main_program], [startup_program], pass_context
)
# Do logical partition
partitioner = Partitioner(dist_context, self.rank)
(
dist_main_prog,
dist_startup_prog,
dist_params_grads,
) = partitioner.partition(
main_program, startup_program, dist_context._params_grads
)
# Generate optimizer
# FIXME should be remove from apply pass after pass support optimizers
with (
program_guard(dist_main_prog, dist_startup_prog),
dist_main_prog.switch_name_generator_guard("opt_"),
):
optimizer_ops = dist_context.serial_optimizer.apply_gradients(
dist_params_grads
)
completer.complete_update_annotation(dist_main_prog)
resharder = Resharder(
dist_main_prog,
dist_startup_prog,
self.rank,
dist_context,
dist_params_grads,
)
resharder.reshard()
config = {}
config["dist_context"] = dist_context
config["global_rank"] = self.rank
config["use_sharding"] = new_strategy.sharding.enable
dp_pass = new_pass("auto_parallel_data_parallel_optimization", config)
dp_pass.apply([dist_main_prog], [dist_startup_prog], pass_context)
if new_strategy.sharding.enable:
config = copy.deepcopy(new_strategy.sharding.to_dict())
config["dist_context"] = dist_context
config["params_grads"] = dist_params_grads
config["global_rank"] = self.rank
auto_parallel_sharding_pass = new_pass(
"auto_parallel_sharding", config
)
auto_parallel_sharding_pass.apply(
[dist_main_prog], [dist_startup_prog], pass_context
)
dist_params_grads = pass_context.get_attr("params_grads")
# gradient clip
config = copy.deepcopy(new_strategy.sharding.to_dict())
config["dist_context"] = dist_context
config["params_grads"] = dist_params_grads
config["rank_id"] = self.rank
auto_parallel_clip_pass = new_pass("auto_parallel_grad_clip", config)
auto_parallel_clip_pass.apply(
[dist_main_prog], [dist_startup_prog], pass_context
)
if new_strategy.gradient_merge.enable:
config = copy.deepcopy(new_strategy.gradient_merge.to_dict())
config["dist_context"] = dist_context
config["params_grads"] = dist_params_grads
auto_parallel_gradient_merge_pass = new_pass(
"auto_parallel_gradient_merge_pass", config
)
auto_parallel_gradient_merge_pass.apply(
[dist_main_prog], [dist_startup_prog], pass_context
)
trial.main_program, trial.startup_program = (
dist_main_prog,
dist_startup_prog,
)
return trial
def _get_profile_context(self, trial, result_path):
profile_ctx = {}
profile_ctx['distributed_env'] = copy.deepcopy(
paddle.distributed.ParallelEnv()
)
profile_ctx['group_map'] = parse_process_groups()
profile_ctx["loss_var_name"] = (
self._baseline_dist_context.serial_loss.name
)
profile_ctx["main_program_decs"] = (
trial.main_program.desc.serialize_to_string()
)
profile_ctx["startup_program_decs"] = (
trial.startup_program.desc.serialize_to_string()
)
self._dataset.batch_size = self._batch_size
self._dataset.input_names = self._get_input_names()
profile_ctx["dataset"] = self._dataset
profile_ctx["result_filename"] = result_path
return profile_ctx
def _get_input_names(self):
input_names = []
for input_spec in self._inputs_spec[:] + self._labels_spec[:]:
input_names.append(input_spec.name)
return input_names
def _launch_profile(self, ctx_path, trial_dir):
if os.environ.get("WITH_COVERAGE", "OFF") == "ON":
coverage_args = ["-m", "coverage", "run", "--branch", "-p"]
else:
coverage_args = []
profile_args = " ".join(
[
"--rank",
str(self.rank),
"--device_id",
str(self.device_id),
"--ctx_filename",
ctx_path,
"--profile_start_step",
str(self._config.profile_start_step),
"--profile_end_step",
str(self._config.profile_end_step),
]
)
cmd_args = (
"-m paddle.distributed.auto_parallel.static.tuner.profiler"
+ " "
+ profile_args
)
cmd = [sys.executable, "-u", *coverage_args, *shlex.split(cmd_args)]
parent_env = copy.copy(os.environ.copy())
# env flags need for profile
new_env = {}
new_env.update(parent_env)
# TODO if any rank hang or fail, kill all processes
self._logger.debug("Executing cmd:\n{} .".format(" ".join(cmd)))
# new_process = subprocess.Popen(cmd, env=new_env)
with (
open(
os.path.join(trial_dir, "stdout.log" + str(self.rank)), "wb"
) as out,
open(
os.path.join(trial_dir, "stderr.log" + str(self.rank)), "wb"
) as err,
):
result = subprocess.Popen(cmd, stdout=out, stderr=err, env=new_env)
result.wait()
out.flush()
err.flush()
os.fsync(out)
os.fsync(err)
def _profile_trial(self, trial):
# Making working directory
trial_dir = self._get_trial_dir(trial)
if not os.path.exists(trial_dir):
if self.rank == 0:
pathlib.Path(trial_dir).mkdir(parents=True, exist_ok=True)
else:
while not os.path.exists(trial_dir):
pass
ctx_filename = "profile_ctx." + str(self.rank)
ctx_path = os.path.join(trial_dir, ctx_filename)
result_path = os.path.join(trial_dir, "result.json")
# Prepare Profile Context
profile_ctx = self._get_profile_context(trial, result_path)
with open(ctx_path, 'wb') as f:
pickle.dump(profile_ctx, f, protocol=4)
if self._config.debug:
debug_program(trial.main_program, trial_dir, "main_program")
debug_program(trial.startup_program, trial_dir, "startup_program")
# Run
self._launch_profile(ctx_path, trial_dir)
# Load results
try:
with open(result_path, 'r') as fp:
results = json.load(fp)
return results
except FileNotFoundError:
Error_results = {"Throughput": -1, "ErrorType": 'FatalError'}
return Error_results
def _evaluate_trial(self, trial):
self._logger.info(f"Trial {trial.name} evaluation start.")
self._apply_optimization(trial)
if self._config.mode == "PROFILE":
results = self._profile_trial(trial)
elif self._config.mode == "COSTMODEL":
raise NotImplementedError(
"COSTMODEL mode for optimization tuning is not supported yet!"
)
else:
raise NotImplementedError(
f"invalid evaluation mode: {self._config.mode}"
)
self._logger.info(
f"Trial {trial.name} evaluation finish with {parse_results(results)}."
)
return results
def _update(self, i, trial, results):
self._finished_trials.append(trial)
cur_metric = get_metric(results)
if self._best_metric is None or cur_metric > self._best_metric:
self._best_metric = cur_metric
self._best_iter = i
def _get_trial_dir(self, trial):
return os.path.join(self.project_dir, trial.name)
def get_best_config(self):
"""
Return the best optimization configuration found in the tuning.
Returns:
A object of fleet.DistributedStrategy with best configuration.
"""
assert self._best_iter >= 0, "The best configuration is not found yet !"
best_trial = self._finished_trials[self._best_iter]
return self._algorithm.get_config_from_trial(best_trial)
def summary(self):
"""
Display tuning result summary.
"""
# TODO summary with the trial_name with metric_of_trial
best_trial = self._finished_trials[self._best_iter]
summary_ = f"""
Tuning Result Summary
Run total {len(self._finished_trials)} trials with {(time.time() - self._tuning_start_time) / 60} min.
The best trial is: [{best_trial.name}], whose configuration is following:
"""
summary_ += "\n" + best_trial.summary() + "\n"
self._logger.info(summary_)
with open(os.path.join(self.project_dir, "summary.txt"), "w+") as fw:
fw.writelines(line + "\n" for line in summary_.split("\n"))
# full_strategy = self.get_best_config()
# path = os.path.join(self.project_dir, "tuned_dist_strategy.yaml")
# with open(path, 'w') as outfile:
# yaml.dump(full_strategy, outfile, default_flow_style=False)
def clear(self):
"""
Clear the temporary file generated in tuning procedure.
"""
# TODO clear up zombie process created by tuning
if not self._config.debug:
for trial in self._finished_trials:
trial_dir = self._get_trial_dir(trial)
shutil.rmtree(trial_dir, ignore_errors=True)
def tune(self):
"""
Performs the search for best hyperparameter configurations
for the selected optimization pass(es).
"""
# step1: collect model info which might be used for
# pruning the search space of the algorithm
self._tuning_start_time = time.time()
self._algorithm.collect_model_info(
self._baseline_dist_context.serial_main_program,
self._baseline_dist_context.serial_startup_program,
)
# main search loop
i = 0
while i < self._config.max_num_trial:
# step2: create a new trial
trial = self._algorithm.next_trial()
if trial.status == TrialStatus.STOPPED:
break
# step3: evaluate the trial
results = self._evaluate_trial(trial)
# step4: update the algorithm with last result,
# which could be used by algorithm to pruning the
# remaining search space.
self._algorithm.update(results)
self._update(i, trial, results)
# early stop
i += 1
if (
self._config.early_stop
and self._config.early_stop <= i - self._best_iter
):
self._logger.info(
f"Early stop the Tuning since there is no better trial found within [{self._config.early_stop}] trials"
)
break
# step5: summary the best config and return
self.summary()
self.clear()