wehub-resource-sync
554 lines
22 KiB
Python
554 lines
22 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 logging
|
|
import os
|
|
import time
|
|
|
|
from paddle.distributed.passes.pass_base import PassManager, new_pass
|
|
from paddle.framework import get_flags
|
|
from paddle.static import append_backward, program_guard
|
|
|
|
from ...utils.log_utils import get_logger
|
|
from ..random import init_auto_parallel_rng
|
|
from .partitioner import Partitioner
|
|
from .process_group import get_world_process_group
|
|
from .reshard import Resharder
|
|
from .utils import (
|
|
get_pp_stage,
|
|
is_sequential_run,
|
|
)
|
|
|
|
PIR_PASS = [
|
|
'fused_gemm_epilogue_pass',
|
|
'fused_linear_param_grad_add_pass',
|
|
'fuse_allreduce_split_to_reducescatter_pass',
|
|
'fused_dropout_add_pass',
|
|
]
|
|
|
|
PIR_PYTHON_PASS = [
|
|
'eliminate_transpose',
|
|
]
|
|
|
|
|
|
class Parallelizer:
|
|
def __init__(self, mode, completer, dist_context):
|
|
self._mode = mode
|
|
self._completer = completer
|
|
self._dist_context = dist_context
|
|
assert self._dist_context._is_initialized
|
|
self._pass_context = self._dist_context.pass_context
|
|
self._strategy = self._dist_context.strategy
|
|
self._logger = get_logger(logging.INFO)
|
|
|
|
@property
|
|
def is_train(self):
|
|
return self._mode == "train"
|
|
|
|
@property
|
|
def is_test(self):
|
|
return self._mode in ["eval", "predict"]
|
|
|
|
def parallel_all(self, parameter_list=None):
|
|
world_process_group = get_world_process_group()
|
|
all_ranks = world_process_group.ranks
|
|
for rank in all_ranks:
|
|
# self._dist_context._backup(serial=True, dist=True)
|
|
self.parallel(rank, parameter_list)
|
|
# self._dist_context._restore(serial=True, dist=True)
|
|
|
|
def parallel(self, rank, parameter_list=None):
|
|
serial_main_program = self._dist_context.serial_main_program
|
|
serial_startup_program = self._dist_context.serial_startup_program
|
|
serial_optimizer = self._dist_context.serial_optimizer
|
|
if self.is_train and serial_optimizer:
|
|
# Generate backward
|
|
serial_loss = self._dist_context.serial_loss
|
|
params_grads = self._generate_backward(
|
|
serial_main_program,
|
|
serial_startup_program,
|
|
serial_loss,
|
|
parameter_list,
|
|
)
|
|
# Apply pre optimization passes
|
|
time0 = time.time()
|
|
(
|
|
serial_main_program,
|
|
serial_startup_program,
|
|
params_grads,
|
|
) = self._apply_pre_optimization(
|
|
serial_main_program,
|
|
serial_startup_program,
|
|
serial_loss,
|
|
serial_optimizer,
|
|
params_grads,
|
|
)
|
|
self._logger.debug(
|
|
f"within parallel apply_pre_optimization time: {time.time() - time0}, mode {self._mode}"
|
|
)
|
|
# Do logical partition
|
|
time0 = time.time()
|
|
partitioner = Partitioner(self._dist_context, rank)
|
|
(
|
|
dist_main_prog,
|
|
dist_startup_prog,
|
|
dist_params_grads,
|
|
) = partitioner.partition(
|
|
serial_main_program, serial_startup_program, params_grads
|
|
)
|
|
|
|
init_auto_parallel_rng()
|
|
|
|
self._logger.debug(
|
|
f"within parallel partitioner time: {time.time() - time0}, mode {self._mode}"
|
|
)
|
|
# Generate optimizer
|
|
time0 = time.time()
|
|
self._generate_optimizer(
|
|
dist_main_prog,
|
|
dist_startup_prog,
|
|
serial_optimizer,
|
|
dist_params_grads,
|
|
)
|
|
self._logger.debug(
|
|
f"within parallel optimizer time: {time.time() - time0}, mode {self._mode}"
|
|
)
|
|
|
|
resharder = Resharder(
|
|
dist_main_prog,
|
|
dist_startup_prog,
|
|
rank,
|
|
self._dist_context,
|
|
dist_params_grads,
|
|
)
|
|
resharder.reshard()
|
|
self._logger.debug(
|
|
f"within parallel reshard time: {time.time() - time0}, mode {self._mode}"
|
|
)
|
|
# Apply post optimization passes
|
|
time0 = time.time()
|
|
self._apply_post_optimization(
|
|
dist_main_prog, dist_startup_prog, rank, dist_params_grads
|
|
)
|
|
self._logger.debug(
|
|
f"within parallel apply_post_optimization time: {time.time() - time0}, mode {self._mode}"
|
|
)
|
|
else:
|
|
# Apply pre optimization passes
|
|
time0 = time.time()
|
|
(
|
|
serial_main_program,
|
|
serial_startup_program,
|
|
params_grads,
|
|
) = self._apply_pre_optimization(
|
|
serial_main_program, serial_startup_program, None, None, []
|
|
)
|
|
self._logger.debug(
|
|
f"within parallel apply_pre_optimization time: {time.time() - time0}, mode {self._mode}"
|
|
)
|
|
# Do logical partition
|
|
time0 = time.time()
|
|
partitioner = Partitioner(self._dist_context, rank)
|
|
(
|
|
dist_main_prog,
|
|
dist_startup_prog,
|
|
dist_params_grads,
|
|
) = partitioner.partition(
|
|
serial_main_program, serial_startup_program, []
|
|
)
|
|
# Do reshard process
|
|
self._logger.debug(
|
|
f"within parallel partitioner time: {time.time() - time0}, mode {self._mode}"
|
|
)
|
|
time0 = time.time()
|
|
# Do reshard process
|
|
micro_bsz = (
|
|
1
|
|
if not self._strategy.pipeline.enable
|
|
else self._strategy.pipeline.micro_batch_size
|
|
)
|
|
resharder = Resharder(
|
|
dist_main_prog,
|
|
dist_startup_prog,
|
|
rank,
|
|
self._dist_context,
|
|
[],
|
|
micro_bsz,
|
|
)
|
|
resharder.reshard()
|
|
self._logger.debug(
|
|
f"within parallel reshard time: {time.time() - time0}, mode {self._mode}"
|
|
)
|
|
# Apply post optimization passes
|
|
time0 = time.time()
|
|
self._apply_post_optimization(
|
|
dist_main_prog, dist_startup_prog, rank, dist_params_grads
|
|
)
|
|
self._logger.debug(
|
|
f"within parallel apply_post_optimization time: {time.time() - time0}, mode {self._mode}"
|
|
)
|
|
|
|
# Clone program for test
|
|
if self.is_test:
|
|
pipeline_opt = dist_main_prog._pipeline_opt
|
|
dist_main_prog = dist_main_prog.clone(for_test=True)
|
|
dist_startup_prog = dist_startup_prog.clone(for_test=True)
|
|
dist_main_prog._pipeline_opt = pipeline_opt
|
|
|
|
# Store the distributed programs for further usages
|
|
self._dist_context.dist_main_programs[rank] = dist_main_prog
|
|
self._dist_context.dist_startup_programs[rank] = dist_startup_prog
|
|
|
|
def _generate_backward(
|
|
self, main_program, startup_program, loss, parameter_list=None
|
|
):
|
|
# NOTE(zhaoyinglia):
|
|
# Guarantee the order of params_grads is same between dynamic mode and static mode
|
|
# by making parameter_list equal to model.parameters(),
|
|
# because the order affect the result of ClipGradByGLobalNorm.
|
|
# If parameter_list is not None, the order of params_grads is same with parameter_list.
|
|
# If parameter_list is None, params_grads will be as prog.global_block().all_parameters().
|
|
with program_guard(main_program, startup_program):
|
|
params_grads = append_backward(
|
|
loss,
|
|
parameter_list=parameter_list,
|
|
distop_context=self._dist_context.dist_op_context,
|
|
)
|
|
self._completer.complete_backward_annotation(main_program)
|
|
self._dist_context.block_state.parse_backward_blocks(main_program)
|
|
return params_grads
|
|
|
|
def _generate_optimizer(
|
|
self, main_program, startup_program, optimizer, params_grads
|
|
):
|
|
# NOTE:
|
|
# 1. `apply_gradients` will add an Accumulator for a parameter only once,
|
|
# but optimizer will be called repeatedly in re-launch, so optimizer need to be copied.
|
|
# 2. lr_scheduler cannot be deepcopy, cause 'deepcopy' will lead to difference of learning_rate between executor and engine.
|
|
learning_rate = optimizer._learning_rate
|
|
new_optimizer = copy.deepcopy(optimizer)
|
|
new_optimizer._learning_rate = learning_rate
|
|
new_optimizer._sorted = False
|
|
self._dist_context._serial_optimizer = optimizer
|
|
self._dist_context._serial_optimizer._learning_rate = learning_rate
|
|
|
|
with (
|
|
program_guard(main_program, startup_program),
|
|
main_program.switch_name_generator_guard("opt_"),
|
|
):
|
|
optimizer_ops = new_optimizer.apply_gradients(params_grads)
|
|
self._completer.complete_update_annotation(main_program)
|
|
return optimizer_ops
|
|
|
|
def _apply_pre_optimization(
|
|
self, main_program, startup_program, loss, optimizer, params_grads
|
|
):
|
|
if self._strategy is None:
|
|
return
|
|
|
|
# apply amp pass on train/eval/predict
|
|
if self._strategy.amp.enable:
|
|
config = copy.deepcopy(self._strategy.amp.to_dict())
|
|
config["dist_context"] = self._dist_context
|
|
config["params_grads"] = params_grads
|
|
config["loss"] = loss
|
|
config["input_data"] = (
|
|
self._dist_context.serial_feed_vars["inputs"]
|
|
+ self._dist_context.serial_feed_vars["labels"]
|
|
)
|
|
self._logger.info(
|
|
"Applying AMP-{}-{} ...".format(
|
|
config["dtype"], config['level']
|
|
),
|
|
)
|
|
if config['level'] == "o1":
|
|
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()
|
|
elif config['level'] in ['o2', 'o3']:
|
|
config["base_opt"] = 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:
|
|
raise ValueError("AMP level should be one of o1, o2, o3")
|
|
|
|
# apply quantization pass
|
|
# The pass can be applied when mode must be 'train'
|
|
if self.is_train and self._strategy.qat.enable:
|
|
config = copy.deepcopy(self._strategy.qat.to_dict())
|
|
config["dist_context"] = self._dist_context
|
|
config["params_grads"] = params_grads
|
|
config["mode"] = self._mode
|
|
config["loss"] = loss
|
|
auto_parallel_quantization_pass = new_pass(
|
|
"auto_parallel_quantization", config
|
|
)
|
|
auto_parallel_quantization_pass.apply(
|
|
[main_program], [startup_program], self._pass_context
|
|
)
|
|
main_program = self._pass_context.get_attr("main_program")
|
|
startup_program = self._pass_context.get_attr("startup_program")
|
|
params_grads = self._pass_context.get_attr("params_grads")
|
|
loss = self._pass_context.get_attr("loss")
|
|
|
|
# apply recompute pass
|
|
# recompute is then train-only optimization
|
|
if self.is_train and self._strategy.recompute.enable:
|
|
config = copy.deepcopy(self._strategy.recompute.to_dict())
|
|
config["dist_context"] = self._dist_context
|
|
config["no_grad_set"] = None
|
|
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
|
|
)
|
|
|
|
return main_program, startup_program, params_grads
|
|
|
|
def _check_dist_attr(self, program, num_model_chunks, dist_context):
|
|
for _, block in enumerate(program.blocks):
|
|
for _, op in enumerate(block.ops):
|
|
op_dist_attr = dist_context.get_op_dist_attr_for_program(op)
|
|
if op_dist_attr is None:
|
|
raise ValueError(
|
|
f"There is not dist_attr for op[{op.type}]."
|
|
)
|
|
|
|
def _apply_post_optimization(
|
|
self, main_program, startup_program, rank, params_grads
|
|
):
|
|
if self._strategy is None:
|
|
return
|
|
|
|
# sequence parallel optimization
|
|
if self._strategy.sp_optimization.enable:
|
|
config = copy.deepcopy(self._strategy.sp_optimization.to_dict())
|
|
config["dist_context"] = self._dist_context
|
|
config["global_rank"] = rank
|
|
sp_pass = new_pass(
|
|
"auto_parallel_sequence_parallel_optimization", config
|
|
)
|
|
sp_pass.apply([main_program], [startup_program], self._pass_context)
|
|
|
|
# apply fused linear promotion pass
|
|
if (
|
|
self.is_train
|
|
and self._strategy.fused_linear_promotion.enable
|
|
and self._strategy.fused_passes.enable
|
|
):
|
|
if (
|
|
len(self._strategy.fused_passes.fused_passes_list) > 0
|
|
and "fuse_gemm_epilogue"
|
|
in self._strategy.fused_passes.fused_passes_list
|
|
):
|
|
amp_config = None
|
|
if self._strategy.amp.enable:
|
|
amp_config = copy.deepcopy(self._strategy.amp.to_dict())
|
|
config = {}
|
|
config["dist_context"] = self._dist_context
|
|
config["global_rank"] = rank
|
|
config["enable_sp"] = self._strategy.sp_optimization.enable
|
|
config["params_grads"] = params_grads
|
|
config["amp_level"] = (
|
|
amp_config['level'] if amp_config is not None else "o0"
|
|
)
|
|
fused_linear_promotion_pass = new_pass(
|
|
"auto_parallel_fused_linear_promotion", config
|
|
)
|
|
fused_linear_promotion_pass.apply(
|
|
[main_program], [startup_program], self._pass_context
|
|
)
|
|
|
|
# apply master grad pass
|
|
if self._strategy.amp.enable:
|
|
amp_config = copy.deepcopy(self._strategy.amp.to_dict())
|
|
config = {}
|
|
config["dist_context"] = self._dist_context
|
|
config["params_grads"] = params_grads
|
|
config["completer"] = self._completer
|
|
if amp_config['level'] == "o2" and amp_config["use_master_grad"]:
|
|
master_grad_pass = new_pass(
|
|
"auto_parallel_master_grad_pass", config
|
|
)
|
|
master_grad_pass.apply(
|
|
[main_program], [startup_program], self._pass_context
|
|
)
|
|
|
|
# data parallel optimization
|
|
if self._strategy.dp_optimization.enable:
|
|
config = copy.deepcopy(self._strategy.dp_optimization.to_dict())
|
|
config["dist_context"] = self._dist_context
|
|
config["global_rank"] = rank
|
|
config["use_sharding"] = self._strategy.sharding.enable
|
|
dp_pass = new_pass(
|
|
"auto_parallel_data_parallel_optimization", config
|
|
)
|
|
dp_pass.apply([main_program], [startup_program], self._pass_context)
|
|
|
|
gradient_sync_after_accumulate = (
|
|
self._strategy.dp_optimization.gradient_sync_after_accumulate
|
|
)
|
|
if gradient_sync_after_accumulate:
|
|
global_params_grads = params_grads
|
|
|
|
if self._strategy.sharding.enable:
|
|
config = copy.deepcopy(self._strategy.sharding.to_dict())
|
|
config["dist_context"] = self._dist_context
|
|
config["params_grads"] = params_grads
|
|
config["global_rank"] = rank
|
|
config["gradient_sync_after_accumulate"] = (
|
|
gradient_sync_after_accumulate
|
|
)
|
|
if self._strategy.amp.enable:
|
|
amp_config = copy.deepcopy(self._strategy.amp.to_dict())
|
|
config["amp_dtype"] = amp_config['dtype']
|
|
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")
|
|
|
|
if self._strategy.mp_optimization.allreduce_matmul_grad_overlapping:
|
|
if int(os.getenv("CUDA_DEVICE_MAX_CONNECTIONS", "0")) != 1:
|
|
self._logger.warning(
|
|
"You set mp_optimization.allreduce_matmul_grad_overlapping=True, but you did not set environment "
|
|
"variable CUDA_DEVICE_MAX_CONNECTIONS=1, which may leads to performance "
|
|
"loss. Try to export CUDA_DEVICE_MAX_CONNECTIONS=1 for better performance."
|
|
)
|
|
|
|
config = {
|
|
"dist_context": self._dist_context,
|
|
}
|
|
allreduce_matmul_grad_overlapping_pass = new_pass(
|
|
"allreduce_matmul_grad_overlapping", config
|
|
)
|
|
allreduce_matmul_grad_overlapping_pass.apply(
|
|
[main_program], [startup_program], self._pass_context
|
|
)
|
|
|
|
if self.is_train:
|
|
# GradClip is train-only optimization
|
|
config = copy.deepcopy(self._strategy.sharding.to_dict())
|
|
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 not is_sequential_run():
|
|
# deps for newexe
|
|
config = {}
|
|
config["dist_context"] = self._dist_context
|
|
APSED_pass = new_pass(
|
|
"auto_parallel_supplement_explicit_dependencies", config
|
|
)
|
|
APSED_pass.apply(
|
|
[main_program], [startup_program], self._pass_context
|
|
)
|
|
|
|
if self.is_train and self._strategy.pipeline.enable:
|
|
self._strategy.gradient_merge.enable = True
|
|
self._strategy.gradient_merge.k_steps = (
|
|
self._strategy.pipeline.accumulate_steps
|
|
)
|
|
self._strategy.gradient_merge.avg = True
|
|
|
|
# gradient_merge is then train-only optimization
|
|
grad_to_global_grad = {}
|
|
if self.is_train and self._strategy.gradient_merge.enable:
|
|
config = copy.deepcopy(self._strategy.gradient_merge.to_dict())
|
|
config["dist_context"] = self._dist_context
|
|
config["grad_to_global_grad"] = grad_to_global_grad
|
|
config["pipeline_mode"] = self._strategy.pipeline.schedule_mode
|
|
if gradient_sync_after_accumulate:
|
|
config["params_grads"] = global_params_grads
|
|
config["gradient_sync_after_accumulate"] = (
|
|
gradient_sync_after_accumulate
|
|
)
|
|
else:
|
|
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
|
|
)
|
|
|
|
self._check_dist_attr(
|
|
main_program,
|
|
self._strategy.pipeline.vpp_degree,
|
|
self._dist_context,
|
|
)
|
|
|
|
enable_ir = get_flags("FLAGS_enable_pir_in_executor")[
|
|
'FLAGS_enable_pir_in_executor'
|
|
]
|
|
ir_pass_list = []
|
|
if self.is_train and self._strategy.fused_passes.enable:
|
|
if len(self._strategy.fused_passes.fused_passes_list) > 0:
|
|
program_pass_list = []
|
|
for p in self._strategy.fused_passes.fused_passes_list:
|
|
if enable_ir and p in (PIR_PASS + PIR_PYTHON_PASS):
|
|
ir_pass_list.append(p)
|
|
else:
|
|
program_pass_list.append(new_pass(p))
|
|
pass_manager = PassManager(program_pass_list)
|
|
pass_manager.apply([main_program], [startup_program])
|
|
|
|
main_program._pass_opt = {}
|
|
main_program._pass_opt['pass_list'] = ir_pass_list
|
|
|
|
if self.is_train and self._strategy.pipeline.enable:
|
|
enable_send_recv_overlap = (
|
|
self._strategy.pipeline.enable_send_recv_overlap
|
|
)
|
|
if (
|
|
enable_send_recv_overlap
|
|
and int(os.getenv("CUDA_DEVICE_MAX_CONNECTIONS", "0")) != 1
|
|
):
|
|
self._logger.warning(
|
|
"You set pipeline.enable_send_recv_overlap=True, but you did not set environment "
|
|
"variable CUDA_DEVICE_MAX_CONNECTIONS=1, which may leads to performance "
|
|
"loss. Try to export CUDA_DEVICE_MAX_CONNECTIONS=1 for better performance."
|
|
)
|
|
|
|
main_program._pipeline_opt = {}
|
|
main_program._pipeline_opt["standalone_opt"] = {
|
|
"enable_send_recv_overlap": enable_send_recv_overlap,
|
|
"schedule_mode": self._strategy.pipeline.schedule_mode,
|
|
"num_micro_batches": self._strategy.pipeline.accumulate_steps,
|
|
"pp_degree": len(self._dist_context.process_meshes),
|
|
"pp_stage": get_pp_stage(self._dist_context, rank),
|
|
"vpp_degree": self._strategy.pipeline.vpp_degree,
|
|
"dist_context": self._dist_context,
|
|
"program_runtimes": self._strategy.pipeline.program_runtimes,
|
|
"memory_limit_times": self._strategy.pipeline.memory_limit_times,
|
|
"split_backward": self._strategy.pipeline.split_backward,
|
|
"grad_to_global_grad": grad_to_global_grad,
|
|
}
|