wehub-resource-sync
540 lines
21 KiB
Python
540 lines
21 KiB
Python
# Copyright (c) 2021 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 os
|
|
import pathlib
|
|
import pickle
|
|
import shlex
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
|
|
import paddle
|
|
from paddle.distributed.passes import PassContext, new_pass
|
|
from paddle.distributed.utils.log_utils import get_logger
|
|
from paddle.framework import core
|
|
from paddle.static import append_backward, program_guard
|
|
|
|
from .cluster import Cluster
|
|
from .completion import Completer
|
|
from .dist_context import DistributedContext, set_default_distributed_context
|
|
from .dist_op import DistributedOperator
|
|
from .dist_tensor import DistributedTensor
|
|
from .mapper import mapping
|
|
from .partitioner import Partitioner
|
|
from .planner import Planner
|
|
from .process_group import (
|
|
ProcessGroup,
|
|
_g_process_group_map,
|
|
get_all_process_groups,
|
|
get_process_group,
|
|
get_world_process_group,
|
|
)
|
|
from .reshard import Resharder
|
|
from .utils import SerialProgramInfo, make_data_unshard
|
|
|
|
_logger = get_logger(logging.INFO)
|
|
|
|
|
|
class AutoParallelizer:
|
|
"""
|
|
AutoParallelizer is the main controller class to do the auto parallel process.
|
|
And the auto parallel process will be triggered in the wrapped parallelize function.
|
|
To facilitate the auto parallelization, it will contain information about program, cluster and the
|
|
related context. In this basic version, the program information will be retrieved from
|
|
Fleet object, and the cluster information can be retrieved in the new created Cluster object,
|
|
and the context information can be retrieved in the new created DistributedContext.
|
|
"""
|
|
|
|
def __init__(self, fleet):
|
|
self._fleet = fleet
|
|
self._optimizer = self._fleet.user_defined_optimizer
|
|
self._dist_strategy = self._fleet._user_defined_strategy
|
|
self._dist_context = DistributedContext()
|
|
self._cluster = None
|
|
self._cluster_topo_path = os.getenv("PADDLE_CLUSTER_TOPO_PATH", None)
|
|
if self._cluster_topo_path is not None:
|
|
self._cluster = Cluster()
|
|
self._cluster.build_from_file(self._cluster_topo_path)
|
|
# Prepare information for auto mapping
|
|
self._rank_mapping_path = os.getenv("PADDLE_RANK_MAPPING_PATH", None)
|
|
enable_auto_mapping_env = os.getenv("PADDLE_ENABLE_AUTO_MAPPING", None)
|
|
if enable_auto_mapping_env is None:
|
|
self._enable_auto_mapping = False
|
|
else:
|
|
self._enable_auto_mapping = True
|
|
self._pass_context = PassContext()
|
|
|
|
self._need_rank_mapping = os.getenv("PADDLE_NEED_RANK_MAPPING")
|
|
self._need_rank_mapping = (
|
|
True
|
|
if self._need_rank_mapping
|
|
and self._need_rank_mapping.lower() == 'true'
|
|
else False
|
|
)
|
|
# self._pass_context = None
|
|
|
|
def _remove_distributed_attrs(self, main_program):
|
|
suffix = core.kAutoParallelSuffix()
|
|
# distributed attributes for variable have been removed
|
|
# in previous process.
|
|
for block in main_program.blocks:
|
|
for op in block.ops:
|
|
for attr_name in op.attr_names:
|
|
if suffix in attr_name:
|
|
op._remove_attr(attr_name)
|
|
|
|
def _apply_pre_optimization_passes(
|
|
self, main_program, startup_program, loss, params_grads, no_grad_set
|
|
):
|
|
# apply amp pass
|
|
if self._dist_strategy.amp:
|
|
config = copy.deepcopy(self._dist_strategy.amp_configs)
|
|
config["dist_context"] = self._dist_context
|
|
config["params_grads"] = params_grads
|
|
config["loss"] = loss
|
|
if config["use_pure_fp16"]:
|
|
config["base_opt"] = self._optimizer
|
|
auto_parallel_fp16_pass = new_pass("auto_parallel_fp16", config)
|
|
auto_parallel_fp16_pass.apply(
|
|
[main_program], [startup_program], self._pass_context
|
|
)
|
|
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], self._pass_context
|
|
)
|
|
loss = auto_parallel_amp_pass.get_loss()
|
|
|
|
# apply recompute pass
|
|
if self._dist_strategy.recompute:
|
|
config = copy.deepcopy(self._dist_strategy.recompute_configs)
|
|
config["dist_context"] = self._dist_context
|
|
config["no_grad_set"] = copy.deepcopy(no_grad_set)
|
|
config["loss"] = loss
|
|
auto_parallel_recompute_pass = new_pass(
|
|
"auto_parallel_recompute", config
|
|
)
|
|
auto_parallel_recompute_pass.apply(
|
|
[main_program], [startup_program], self._pass_context
|
|
)
|
|
|
|
def _generate_backward(
|
|
self,
|
|
main_program,
|
|
startup_program,
|
|
loss,
|
|
parameter_list,
|
|
no_grad_set,
|
|
callbacks,
|
|
):
|
|
with program_guard(main_program, startup_program):
|
|
params_grads = append_backward(
|
|
loss,
|
|
parameter_list,
|
|
no_grad_set,
|
|
callbacks,
|
|
distop_context=self._dist_context.dist_op_context,
|
|
)
|
|
self._completer = Completer(self._dist_context)
|
|
self._completer.complete_backward_annotation(main_program)
|
|
self._dist_context.block_state.parse_backward_blocks(main_program)
|
|
return params_grads
|
|
|
|
def _apply_optimize(self, main_program, startup_program, params_grads):
|
|
optimizer = copy.deepcopy(self._optimizer)
|
|
with program_guard(main_program, startup_program):
|
|
optimize_ops = optimizer.apply_gradients(params_grads)
|
|
|
|
self._dist_context._serial_optimizer = optimizer
|
|
# update completion
|
|
self._completer = Completer(self._dist_context)
|
|
self._completer.complete_update_annotation(main_program)
|
|
|
|
return optimize_ops
|
|
|
|
def _apply_post_optimization_passes(
|
|
self, main_program, startup_program, rank, params_grads
|
|
):
|
|
if self._dist_strategy.sharding:
|
|
config = copy.deepcopy(self._dist_strategy.sharding_configs)
|
|
config["dist_context"] = self._dist_context
|
|
config["params_grads"] = params_grads
|
|
config["global_rank"] = rank
|
|
auto_parallel_sharding_pass = new_pass(
|
|
"auto_parallel_sharding", config
|
|
)
|
|
auto_parallel_sharding_pass.apply(
|
|
[main_program], [startup_program], self._pass_context
|
|
)
|
|
params_grads = self._pass_context.get_attr("params_grads")
|
|
|
|
config = copy.deepcopy(self._dist_strategy.sharding_configs)
|
|
config["dist_context"] = self._dist_context
|
|
config["params_grads"] = params_grads
|
|
config["rank_id"] = rank
|
|
auto_parallel_clip_pass = new_pass("auto_parallel_grad_clip", config)
|
|
auto_parallel_clip_pass.apply(
|
|
[main_program], [startup_program], self._pass_context
|
|
)
|
|
|
|
if self._dist_strategy.gradient_merge:
|
|
config = copy.deepcopy(self._dist_strategy.gradient_merge_configs)
|
|
config["dist_context"] = self._dist_context
|
|
config["params_grads"] = params_grads
|
|
auto_parallel_gradient_merge_pass = new_pass(
|
|
"auto_parallel_gradient_merge_pass", config
|
|
)
|
|
auto_parallel_gradient_merge_pass.apply(
|
|
[main_program], [startup_program], self._pass_context
|
|
)
|
|
|
|
def _get_dist_program(self, rank, dist_context=None, relaunch_phase=False):
|
|
completed_main_program = None
|
|
serial_main_program = self._main_program.clone()
|
|
serial_startup_program = self._startup_program.clone()
|
|
serial_loss = serial_main_program.global_block().var(self._loss.name)
|
|
|
|
# generating serial
|
|
if dist_context is None:
|
|
# Annotation completion
|
|
self._dist_context = DistributedContext()
|
|
_logger.info("Start annotation dist attr.")
|
|
self._completer = Completer(self._dist_context)
|
|
completed_main_program = (
|
|
self._completer.complete_forward_annotation(serial_main_program)
|
|
)
|
|
else:
|
|
completed_main_program = serial_main_program
|
|
self._dist_context = copy.deepcopy(dist_context)
|
|
|
|
# parse forward sub block
|
|
self._dist_context.block_state.parse_forward_blocks(serial_main_program)
|
|
|
|
# serial backward pass
|
|
params_grads = self._generate_backward(
|
|
completed_main_program,
|
|
serial_startup_program,
|
|
serial_loss,
|
|
self._parameter_list,
|
|
self._no_grad_set,
|
|
self._callbacks,
|
|
)
|
|
|
|
# serial forward pass
|
|
self._apply_pre_optimization_passes(
|
|
completed_main_program,
|
|
serial_startup_program,
|
|
serial_loss,
|
|
params_grads,
|
|
self._no_grad_set,
|
|
)
|
|
# Logical partition
|
|
partitioner = Partitioner(self._dist_context, rank)
|
|
(
|
|
dist_main_prog,
|
|
dist_startup_prog,
|
|
dist_params_grads,
|
|
) = partitioner.partition(
|
|
completed_main_program, serial_startup_program, params_grads
|
|
)
|
|
|
|
# TODO refactor the placement of optimizer
|
|
# generate optimize program
|
|
dist_optimize_ops = self._apply_optimize(
|
|
dist_main_prog, dist_startup_prog, dist_params_grads
|
|
)
|
|
|
|
make_data_unshard(dist_main_prog, dist_startup_prog, self._dist_context)
|
|
|
|
resharder = Resharder(
|
|
dist_main_prog,
|
|
dist_startup_prog,
|
|
rank,
|
|
self._dist_context,
|
|
dist_params_grads,
|
|
)
|
|
resharder.reshard()
|
|
|
|
self._apply_post_optimization_passes(
|
|
dist_main_prog, dist_startup_prog, rank, dist_params_grads
|
|
)
|
|
g_process_group_map = None
|
|
if not relaunch_phase:
|
|
g_process_group_map = copy.deepcopy(_g_process_group_map)
|
|
_g_process_group_map.clear()
|
|
_g_process_group_map[0] = ProcessGroup(0, [])
|
|
for process_mesh in self._dist_context._process_meshes:
|
|
_g_process_group_map[0].add_ranks(process_mesh.process_ids)
|
|
return (
|
|
dist_optimize_ops,
|
|
dist_params_grads,
|
|
dist_startup_prog,
|
|
dist_main_prog,
|
|
g_process_group_map,
|
|
)
|
|
|
|
def parallelize(
|
|
self,
|
|
loss,
|
|
startup_program,
|
|
parameter_list=None,
|
|
no_grad_set=None,
|
|
callbacks=None,
|
|
):
|
|
assert startup_program is not None
|
|
self._loss = loss
|
|
self._startup_program = startup_program
|
|
self._main_program = loss.block.program
|
|
self._parameter_list = parameter_list
|
|
self._no_grad_set = no_grad_set
|
|
self._callbacks = callbacks
|
|
|
|
if self._enable_auto_mapping and self._need_rank_mapping:
|
|
# Do the mapping pass before parallelization
|
|
assert self._cluster is not None, (
|
|
"The cluster must not be none when using auto mapping."
|
|
)
|
|
dist_programs = {}
|
|
world_process_group = get_world_process_group()
|
|
dist_context = None
|
|
# auto search
|
|
if self._dist_strategy.auto_search:
|
|
logging.info("Start searching dist attr.")
|
|
serial_program_info = SerialProgramInfo(
|
|
self._main_program,
|
|
self._startup_program,
|
|
self._loss,
|
|
self._optimizer,
|
|
self._cluster,
|
|
)
|
|
planner = Planner(
|
|
serial_program_info,
|
|
self,
|
|
algorithm_config={"name": "mcmc", "max_search_times": 5},
|
|
)
|
|
dist_context, _ = planner.search()
|
|
logging.info("End searching dist attr.")
|
|
|
|
# serialize the dist context by planner
|
|
if dist_context is not None:
|
|
logging.info("Start serialize searched dist attr")
|
|
cwd = pathlib.Path().cwd()
|
|
searched_dist_context_path = os.path.join(
|
|
cwd, f"searched_dist_context_{time.time()}.pkl"
|
|
)
|
|
saved_dist_context = {}
|
|
ops_dist_attr = {}
|
|
tensors_dist_attr = {}
|
|
for key, dist_op in dist_context._dist_ops_for_program.items():
|
|
ops_dist_attr[key] = dist_op.dist_attr
|
|
for (
|
|
key,
|
|
dist_tensor,
|
|
) in dist_context._dist_tensors_for_program.items():
|
|
tensors_dist_attr[key] = dist_tensor.dist_attr
|
|
saved_dist_context["ops_dist_attr"] = ops_dist_attr
|
|
saved_dist_context["tensors_dist_attr"] = tensors_dist_attr
|
|
saved_dist_context["process_meshes"] = (
|
|
dist_context._process_meshes
|
|
)
|
|
with open(
|
|
searched_dist_context_path, "wb"
|
|
) as dist_context_file:
|
|
pickle.dump(saved_dist_context, dist_context_file)
|
|
os.environ['PADDLE_SEARCHED_DIST_CONTEXT_PATH'] = (
|
|
searched_dist_context_path
|
|
)
|
|
logging.info(
|
|
f"End serialize searched dist attr to {searched_dist_context_path}"
|
|
)
|
|
|
|
for rank in world_process_group.ranks:
|
|
(
|
|
dist_optimize_ops,
|
|
dist_params_grads,
|
|
dist_startup_prog,
|
|
dist_main_prog,
|
|
g_process_group_map,
|
|
) = self._get_dist_program(rank, dist_context)
|
|
dist_programs[rank] = [dist_main_prog, g_process_group_map]
|
|
|
|
# Do the mapping between the distributed program graph and the cluster graph
|
|
rank_mapping_dict = mapping(dist_programs, self._cluster)
|
|
rank_mapping = list(rank_mapping_dict.values())
|
|
|
|
# Relaunch the training by using the rank mapping file
|
|
with open(self._rank_mapping_path, "w") as rank_mapping_file:
|
|
json.dump(rank_mapping, rank_mapping_file)
|
|
|
|
enable_elastic = os.getenv("PADDLE_ENABLE_ELASTIC")
|
|
enable_elastic = (
|
|
True
|
|
if enable_elastic and enable_elastic.lower() == 'true'
|
|
else False
|
|
)
|
|
if enable_elastic:
|
|
print("Auto mapping finished, now do elastic re-launch")
|
|
sys.exit(
|
|
paddle.distributed.fleet.elastic.manager.ELASTIC_AUTO_PARALLEL_EXIT_CODE
|
|
)
|
|
|
|
original_cmd_args = os.getenv("PADDLE_ORIGINAL_CMD_ARGS")
|
|
rank_mapping_args = " ".join(
|
|
["--rank_mapping_path", self._rank_mapping_path]
|
|
)
|
|
if os.environ.get("WITH_COVERAGE", "OFF") == "ON":
|
|
coverage_args = ["-m", "coverage", "run", "--branch", "-p"]
|
|
else:
|
|
coverage_args = []
|
|
new_cmd_args = (
|
|
"-m paddle.distributed.fleet.launch"
|
|
+ " "
|
|
+ rank_mapping_args
|
|
+ " "
|
|
+ original_cmd_args
|
|
)
|
|
new_cmd = [
|
|
sys.executable,
|
|
"-u",
|
|
*coverage_args,
|
|
*shlex.split(new_cmd_args),
|
|
]
|
|
new_process = subprocess.Popen(new_cmd)
|
|
new_process.wait()
|
|
assert new_process.returncode == 0, (
|
|
"Launch failed with rank mapping"
|
|
)
|
|
print("Successfully do the second launch for auto mapping!")
|
|
sys.exit(0)
|
|
else:
|
|
# Parallelization after the mapping pass
|
|
rank = paddle.distributed.get_rank()
|
|
dist_context = None
|
|
searched_dist_context_path = os.getenv(
|
|
"PADDLE_SEARCHED_DIST_CONTEXT_PATH", None
|
|
)
|
|
if searched_dist_context_path is not None:
|
|
with open(
|
|
searched_dist_context_path, "rb"
|
|
) as dist_context_file:
|
|
from paddle.framework.restricted_unpickler import (
|
|
safe_load_pickle,
|
|
)
|
|
|
|
saved_dist_context = safe_load_pickle(dist_context_file)
|
|
dist_context = DistributedContext()
|
|
for op in self._main_program.global_block().ops:
|
|
dist_attr = saved_dist_context["ops_dist_attr"][
|
|
op.desc.id()
|
|
]
|
|
dist_op = DistributedOperator(op, dist_attr)
|
|
dist_context.add_dist_op_for_program(dist_op)
|
|
|
|
vars = self._main_program.global_block().vars
|
|
for var in vars.values():
|
|
dist_attr = saved_dist_context["tensors_dist_attr"][
|
|
var.desc.id()
|
|
]
|
|
dist_tensor = DistributedTensor(var, dist_attr)
|
|
dist_context.add_dist_tensor_for_program(dist_tensor)
|
|
|
|
dist_context._process_meshes = saved_dist_context[
|
|
"process_meshes"
|
|
]
|
|
|
|
else:
|
|
if self._dist_strategy.auto_search:
|
|
serial_program_info = SerialProgramInfo(
|
|
self._main_program,
|
|
self._startup_program,
|
|
self._loss,
|
|
self._optimizer,
|
|
cluster=self._cluster,
|
|
)
|
|
planner = Planner(
|
|
serial_program_info,
|
|
self,
|
|
algorithm_config={
|
|
"name": "mcmc",
|
|
"max_search_times": 5,
|
|
},
|
|
)
|
|
dist_context, _ = planner.search()
|
|
|
|
# rebuild g_process_group
|
|
if dist_context is not None:
|
|
pg0 = get_process_group(0)
|
|
for process_mesh in dist_context._process_meshes:
|
|
pg0.add_ranks(process_mesh.process_ids)
|
|
(
|
|
dist_optimize_ops,
|
|
dist_params_grads,
|
|
dist_startup_prog,
|
|
dist_main_prog,
|
|
_,
|
|
) = self._get_dist_program(rank, dist_context, relaunch_phase=True)
|
|
|
|
# NOTE: This is a trick to fix hang in pipeline mode when dist context is searched by planner
|
|
if self._dist_strategy.auto_search:
|
|
is_pipeline = False
|
|
for op in dist_main_prog.global_block().ops:
|
|
if op.type == "send_v2" or op.type == "recv_v2":
|
|
is_pipeline = True
|
|
break
|
|
if is_pipeline:
|
|
with paddle.static.program_guard(dist_main_prog):
|
|
paddle.distributed.barrier()
|
|
|
|
# Traverse different rank programs and traverse each op of them,
|
|
# instantiate communication by process_mapping.
|
|
all_process_groups = get_all_process_groups()
|
|
for process_group in all_process_groups:
|
|
process_group.instantiate()
|
|
|
|
# Copy distributed info to the default context
|
|
set_default_distributed_context(self._dist_context)
|
|
|
|
# The last step: remove all distributed attributes to be compatible
|
|
# with inference.
|
|
self._remove_distributed_attrs(dist_main_prog)
|
|
|
|
return (
|
|
dist_optimize_ops,
|
|
dist_params_grads,
|
|
dist_startup_prog,
|
|
dist_main_prog,
|
|
)
|
|
|
|
def __deepcopy__(self, memo):
|
|
cls = self.__class__
|
|
result = cls.__new__(cls)
|
|
memo[id(self)] = result
|
|
for k, v in self.__dict__.items():
|
|
if (
|
|
k == "_main_program"
|
|
or k == "_startup_program"
|
|
or k == "_dist_context"
|
|
or k == "_fleet"
|
|
or k == "_loss"
|
|
):
|
|
setattr(result, k, v)
|
|
else:
|
|
setattr(result, k, copy.deepcopy(v, memo))
|
|
return result
|