chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,153 @@
|
||||
# Copyright (c) 2020 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.
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from paddle.distributed import fleet
|
||||
|
||||
from . import ( # noqa: F401
|
||||
hybrid_parallel_util,
|
||||
log_util,
|
||||
mix_precision_utils,
|
||||
sequence_parallel_utils,
|
||||
tensor_parallel_utils,
|
||||
)
|
||||
from .fs import HDFSClient, LocalFS
|
||||
from .ps_util import DistributedInfer
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
|
||||
from paddle.nn import Layer
|
||||
|
||||
__all__ = ["LocalFS", "recompute", "DistributedInfer", "HDFSClient"]
|
||||
|
||||
|
||||
def recompute(
|
||||
function: Layer | Callable[..., Any], *args: Any, **kwargs: Any
|
||||
) -> Any:
|
||||
"""
|
||||
recompute intermediate activations to save the memory.
|
||||
|
||||
Parameters:
|
||||
function(paddle.nn.Layer): layer of sequence of layers that describes part of forward pass of the model
|
||||
whose intermediate activations will be released to save memory in forward stage and will be recomputed
|
||||
in backward stage for gradient calculation.
|
||||
*args(Tensor): inputs to the function.
|
||||
**kwargs(Dict): Kwargs should only contain two kinds of key-value params, the one is part of function's key-value params,
|
||||
and the other contains ``preserve_rng_state`` and ``use_reentrant``. the key-value pair of ``preserve_rng_state``,
|
||||
which is used to indicate whether to save the forward rng. If it is True, then the last forward rng value
|
||||
will be restored when the forward recalculation of backpropagation is performed, its default value is True.
|
||||
the key-value pair of ``use_reentrant`` is used to indicate which implementation of recompute you will be used.
|
||||
``use_reentrant=True`` means to use the PyLayer implementation of recompute, ``use_reentrant=False`` means to
|
||||
use the Hook implementation of recompute, its default value is True.
|
||||
Returns:
|
||||
Output of function on args.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +REQUIRES(env:DISTRIBUTED, env:GPU)
|
||||
>>> import paddle
|
||||
>>> from paddle.distributed.fleet.utils import recompute
|
||||
>>> import random
|
||||
>>> paddle.seed(2023)
|
||||
>>> def get_fc_block(block_idx, input_size, is_last=False):
|
||||
... block_name = "block_" + str(block_idx)
|
||||
... block = paddle.nn.Sequential(
|
||||
... (block_name + "_fc_0", paddle.nn.Linear(input_size, input_size, bias_attr=False)),
|
||||
... (block_name + "_dropout", paddle.nn.Dropout(p=0.5)),
|
||||
... (block_name + "_relu_1", paddle.nn.ReLU()),
|
||||
... (block_name + "_fc_1", paddle.nn.Linear(input_size, input_size, bias_attr=False)),
|
||||
... (block_name + "_relu_2", paddle.nn.ReLU()),
|
||||
... )
|
||||
... if is_last:
|
||||
... block.add_sublayer(
|
||||
... block_name + "_fc_2",
|
||||
... paddle.nn.Linear(input_size, 1, bias_attr=False),
|
||||
... )
|
||||
... else:
|
||||
... block.add_sublayer(
|
||||
... block_name + "_fc_2",
|
||||
... paddle.nn.Linear(input_size, input_size, bias_attr=False),
|
||||
... )
|
||||
... return block
|
||||
|
||||
>>> class Naive_fc_net(paddle.nn.Layer):
|
||||
... def __init__(
|
||||
... self,
|
||||
... input_size=10,
|
||||
... recompute_blocks=[1, 3],
|
||||
... recompute_kwargs={},
|
||||
... ):
|
||||
... super().__init__()
|
||||
... self.recompute_blocks = recompute_blocks
|
||||
... self.recompute_kwargs = recompute_kwargs
|
||||
... self.runfunc0 = get_fc_block(0, input_size, is_last=False)
|
||||
... self.runfunc1 = get_fc_block(1, input_size, is_last=False)
|
||||
... self.runfunc2 = get_fc_block(2, input_size, is_last=False)
|
||||
... self.runfunc3 = get_fc_block(3, input_size, is_last=False)
|
||||
... self.runfunc4 = get_fc_block(4, input_size, is_last=True)
|
||||
... self.total_func = [self.runfunc0, self.runfunc1, self.runfunc2, self.runfunc3, self.runfunc4]
|
||||
...
|
||||
... def forward(self, inputs):
|
||||
... nums = len(self.total_func)
|
||||
... for i in range(nums):
|
||||
... if i in self.recompute_blocks:
|
||||
... inputs = recompute(self.total_func[i], inputs, **{"preserve_rng_state": True})
|
||||
... else:
|
||||
... inputs = self.total_func[i](inputs)
|
||||
... return inputs
|
||||
|
||||
>>> def run_model(cuda_state, recompute_block=[], recompute_kwargs={}):
|
||||
... gen = paddle.seed(10)
|
||||
... gen.manual_seed(10)
|
||||
... random.seed(10)
|
||||
... if cuda_state:
|
||||
... paddle.set_cuda_rng_state(cuda_state)
|
||||
... batch_size, input_size = 1, 10
|
||||
... model = Naive_fc_net(
|
||||
... input_size,
|
||||
... recompute_blocks=recompute_block,
|
||||
... recompute_kwargs=recompute_kwargs,
|
||||
... )
|
||||
... optimizer = paddle.optimizer.SGD(learning_rate=0.01, parameters=model.parameters())
|
||||
... loss_ = []
|
||||
... param_ = []
|
||||
... grad_ = []
|
||||
... for _ in range(5):
|
||||
... x = paddle.rand(shape=[batch_size, input_size], dtype="float32")
|
||||
... y_pred = model(x)
|
||||
... loss = y_pred.mean()
|
||||
... loss_.append(loss.item())
|
||||
... loss.backward()
|
||||
... optimizer.step()
|
||||
... param_.append(model.parameters()[9])
|
||||
... grad_.append(model.parameters()[3]._grad_ivar())
|
||||
... optimizer.clear_grad()
|
||||
... return loss_, param_, grad_
|
||||
|
||||
>>> cuda_state = paddle.get_cuda_rng_state()
|
||||
>>> # without recompute
|
||||
>>> loss_ref, param_ref, grad_ref = run_model(cuda_state, recompute_block=[])
|
||||
|
||||
>>> loss, param, grad = run_model(cuda_state, recompute_block=[1, 2])
|
||||
>>> print("normal_loss: {}, recompute_loss: {}".format(loss_ref, loss))
|
||||
>>> # The result of the recompute_loss should be the same as the normal_loss.
|
||||
normal_loss: [0.0018744759727269411, 0.0, 0.035971127450466156, 0.0, 0.0], recompute_loss: [0.0018744759727269411, 0.0, 0.035971127450466156, 0.0, 0.0]
|
||||
|
||||
"""
|
||||
|
||||
return fleet.recompute.recompute(function, *args, **kwargs)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,190 @@
|
||||
# Copyright (c) 2020 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.
|
||||
"""Http Server."""
|
||||
|
||||
import http.server as SimpleHTTPServer
|
||||
import logging
|
||||
import threading
|
||||
from http.server import HTTPServer
|
||||
|
||||
__all__ = []
|
||||
|
||||
|
||||
def get_logger(name, level, fmt):
|
||||
logger = logging.getLogger(name)
|
||||
logger.setLevel(level)
|
||||
handler = logging.FileHandler('http.log', mode='w')
|
||||
formatter = logging.Formatter(fmt=fmt)
|
||||
handler.setFormatter(formatter)
|
||||
logger.addHandler(handler)
|
||||
logger.propagate = False
|
||||
return logger
|
||||
|
||||
|
||||
_http_server_logger = get_logger(
|
||||
__name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
|
||||
)
|
||||
|
||||
|
||||
class KVHandler(SimpleHTTPServer.SimpleHTTPRequestHandler):
|
||||
"""
|
||||
kv handler class for kv http server,
|
||||
it defines the way to get/set kv in server.
|
||||
"""
|
||||
|
||||
def do_GET(self):
|
||||
"""
|
||||
get method for kv handler, get value according to key.
|
||||
"""
|
||||
log_str = "GET " + self.address_string() + self.path
|
||||
paths = self.path.split('/')
|
||||
if len(paths) < 3:
|
||||
print('len of request path must be 3: ' + self.path)
|
||||
self.send_status_code(400)
|
||||
return
|
||||
_, scope, key = paths
|
||||
with self.server.kv_lock:
|
||||
value = self.server.kv.get(scope, {}).get(key)
|
||||
if value is None:
|
||||
log_str += ' , key not found: ' + key
|
||||
self.send_status_code(404)
|
||||
else:
|
||||
log_str += ' , key found: ' + key
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Length", str(len(value)))
|
||||
self.end_headers()
|
||||
self.wfile.write(value)
|
||||
_http_server_logger.info(log_str)
|
||||
|
||||
def do_PUT(self):
|
||||
"""
|
||||
put method for kv handler, set value according to key.
|
||||
"""
|
||||
log_str = "PUT " + self.address_string() + self.path
|
||||
paths = self.path.split('/')
|
||||
if len(paths) < 3:
|
||||
print('len of request path must be 3: ' + self.path)
|
||||
self.send_status_code(400)
|
||||
return
|
||||
_, scope, key = paths
|
||||
content_length = int(self.headers['Content-Length'])
|
||||
try:
|
||||
value = self.rfile.read(content_length)
|
||||
except:
|
||||
print("receive error invalid request")
|
||||
self.send_status_code(404)
|
||||
return
|
||||
with self.server.kv_lock:
|
||||
if self.server.kv.get(scope) is None:
|
||||
self.server.kv[scope] = {}
|
||||
self.server.kv[scope][key] = value
|
||||
self.send_status_code(200)
|
||||
_http_server_logger.info(log_str)
|
||||
|
||||
def do_DELETE(self):
|
||||
"""
|
||||
delete method for kv handler, set value according to key.
|
||||
"""
|
||||
log_str = "DELETE " + self.address_string() + self.path
|
||||
paths = self.path.split('/')
|
||||
if len(paths) < 3:
|
||||
print('len of request path must be 3: ' + self.path)
|
||||
self.send_status_code(400)
|
||||
return
|
||||
_, scope, key = paths
|
||||
with self.server.delete_kv_lock:
|
||||
if self.server.delete_kv.get(scope) is None:
|
||||
self.server.delete_kv[scope] = set()
|
||||
self.server.delete_kv[scope].add(key)
|
||||
self.send_status_code(200)
|
||||
_http_server_logger.info(log_str)
|
||||
|
||||
def log_message(self, format, *args):
|
||||
"""
|
||||
ignore all logging messages in kv handler.
|
||||
"""
|
||||
pass
|
||||
|
||||
def send_status_code(self, code):
|
||||
"""
|
||||
send status code back to client.
|
||||
"""
|
||||
self.send_response(code)
|
||||
self.send_header("Content-Length", 0)
|
||||
self.end_headers()
|
||||
|
||||
|
||||
class KVHTTPServer(HTTPServer):
|
||||
"""
|
||||
it is a http server storing kv pairs.
|
||||
"""
|
||||
|
||||
def __init__(self, port, handler):
|
||||
"""Init."""
|
||||
super().__init__(('', port), handler)
|
||||
self.delete_kv_lock = threading.Lock()
|
||||
self.delete_kv = {}
|
||||
self.kv_lock = threading.Lock()
|
||||
self.kv = {}
|
||||
|
||||
def get_deleted_size(self, key):
|
||||
"""
|
||||
get deleted size in key.
|
||||
"""
|
||||
ret = 0
|
||||
with self.delete_kv_lock:
|
||||
ret = len(self.delete_kv.get(key, set()))
|
||||
return ret
|
||||
|
||||
|
||||
class KVServer:
|
||||
"""
|
||||
it is a server storing kv pairs, has a http server inside.
|
||||
"""
|
||||
|
||||
def __init__(self, port, size={}):
|
||||
"""Init."""
|
||||
self.http_server = KVHTTPServer(port, KVHandler)
|
||||
self.listen_thread = None
|
||||
self.size = size
|
||||
|
||||
def start(self):
|
||||
"""
|
||||
start server until user calls stop to let it quit.
|
||||
"""
|
||||
self.listen_thread = threading.Thread(
|
||||
target=lambda: self.http_server.serve_forever()
|
||||
)
|
||||
self.listen_thread.start()
|
||||
|
||||
def stop(self):
|
||||
"""
|
||||
stop server and clear its resources.
|
||||
"""
|
||||
self.http_server.shutdown()
|
||||
self.listen_thread.join()
|
||||
self.http_server.server_close()
|
||||
|
||||
def should_stop(self):
|
||||
"""
|
||||
return whether the server should stop.
|
||||
|
||||
Returns:
|
||||
ret(bool): whether the server should stop
|
||||
"""
|
||||
for key in self.size:
|
||||
s = self.http_server.get_deleted_size(key)
|
||||
if s != self.size.get(key, 0):
|
||||
return False
|
||||
return True
|
||||
@@ -0,0 +1,842 @@
|
||||
# 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.
|
||||
|
||||
from collections import defaultdict
|
||||
|
||||
import numpy as np
|
||||
|
||||
# (TODO: GhostScreaming) It will be removed later.
|
||||
from paddle.base import core
|
||||
from paddle.distributed import fleet
|
||||
from paddle.framework import Block, Program, in_dynamic_mode
|
||||
|
||||
|
||||
class HybridParallelInferenceHelper:
|
||||
"""
|
||||
A helper class to split program for inference with hybrid parallelism.
|
||||
|
||||
Args:
|
||||
startup_program (Program): the startup program.
|
||||
main_program (Program): the main program.
|
||||
num_mp (int): number of model parallel degree. Default ``1``.
|
||||
num_pp (int): number of pipeline parallel degree. Default ``1``.
|
||||
micro_batch_size (int): number of micro batch size. Default ``1``.
|
||||
beam_size (int): number of beam search size. Default ``1``.
|
||||
init_comm (bool): whether if initialize communication group. Default ``True``.
|
||||
role_maker (RoleMakerBase or subclass): user custom define RoleMakerBase.
|
||||
If ``role_maker==None``, then use PaddleCloudRoleMaker. Default ``None``.
|
||||
|
||||
Returns:
|
||||
None.
|
||||
|
||||
Write Paradigm:
|
||||
|
||||
.. code-block:: text
|
||||
:name: text-example1
|
||||
|
||||
>>> # doctest: +REQUIRES(env:DISTRIBUTED, env:GPU)
|
||||
>>> import paddle
|
||||
>>> # while op pattern
|
||||
>>> with paddle.base.device_guard(f'{device}:all'):
|
||||
... # init global cond
|
||||
... max_len = paddle.full(shape=[1], dtype="int64", fill_value=10)
|
||||
... step_idx = paddle.full(shape=[1], dtype="int64", fill_value=0)
|
||||
... cond_int = paddle.full(shape=[1], dtype="int64", fill_value=0, name="cond_int")
|
||||
... cond = layers.cast(step_idx < max_len, dtype="bool")
|
||||
... while_op = layers.While(cond, is_test=True)
|
||||
|
||||
... # init global lod_tensor_array for generation task
|
||||
... arr = paddle.tensor.array_write(data, step_idx)
|
||||
|
||||
>>> with while_op.block():
|
||||
... with paddle.base.device_guard(f'{device}:all'):
|
||||
... # read data from global lod_tensor_array
|
||||
... element_in_arr = paddle.tensor.array_read(array=arr, i=step_idx)
|
||||
... # write placeholder data to global lod_tensor_array,
|
||||
... # it need for send_v2 of lod_tensor_array
|
||||
... paddle.increment(x=step_idx, value=1.0)
|
||||
... paddle.tensor.array_write(element_in_arr, i=step_idx, array=arr)
|
||||
... with paddle.base.device_guard(f'{device}:0'):
|
||||
... pass # some code
|
||||
... with paddle.base.device_guard(f'{device}:1'):
|
||||
... pass # some code
|
||||
... with paddle.base.device_guard(f'{device}:{num_pp - 1}'):
|
||||
... # generate some data in while block and write to global lod_tensor_array
|
||||
... # that they are read in next while step.
|
||||
... # we will using send_v2 to send global lod_tensor_array to other pipeline and sync
|
||||
... paddle.tensor.array_write(other_var, i=step_idx, array=arr)
|
||||
... # update cond and assign to cond_int, we will sync cond_int
|
||||
... layers.assign(layers.cast(cond, dtype="int32"), cond_int)
|
||||
... with paddle.base.device_guard(f'{model._device}:all'):
|
||||
... # the code below must at end of while block and exists in device:all
|
||||
... layers.assign(layers.cast(cond_int, dtype='bool'), cond)
|
||||
|
||||
>>> with paddle.base.device_guard(f'{model._device}:all'):
|
||||
... # use a empty lod_tensor_array to clear lod_tensor_array
|
||||
... layers.assign(layers.create_array(data.dtype), arr)
|
||||
|
||||
Examples:
|
||||
|
||||
.. code-block:: pycon
|
||||
:name: code-example1
|
||||
|
||||
>>> # doctest: +REQUIRES(env:DISTRIBUTED, env:GPU)
|
||||
>>> import os
|
||||
>>> import numpy as np
|
||||
>>> import paddle
|
||||
>>> import paddle.distributed.fleet as fleet
|
||||
>>> from paddle.distributed.fleet.utils import hybrid_parallel_inference
|
||||
>>> paddle.enable_static()
|
||||
>>> nranks = int(os.getenv("PADDLE_TRAINERS_NUM", 1))
|
||||
>>> rank = int(os.getenv("PADDLE_TRAINER_ID", 0))
|
||||
>>> dev_id = int(os.getenv("FLAGS_selected_gpus", 0))
|
||||
>>> main_program = paddle.static.Program()
|
||||
>>> startup_program = paddle.static.Program()
|
||||
>>> if nranks > 1:
|
||||
... dist_strategy = fleet.DistributedStrategy()
|
||||
... dist_strategy.without_graph_optimization = True
|
||||
... fleet.init(is_collective=True, strategy=dist_strategy)
|
||||
>>> device = "gpu"
|
||||
>>> with paddle.static.program_guard(main_program, startup_program):
|
||||
... with paddle.base.device_guard(f'{device}:0'):
|
||||
... X = paddle.static.data(name='X', shape=[None, 2], dtype='float32')
|
||||
... with paddle.base.device_guard(f'{device}:all'):
|
||||
... max_len = paddle.full(shape=[1], dtype="int64", fill_value=5, name="n")
|
||||
... step_idx = paddle.full(shape=[1], dtype="int64", fill_value=0, name="i")
|
||||
... data = paddle.tensor.array_write(X, step_idx)
|
||||
... cond_int = paddle.full(shape=[1], dtype="int64", fill_value=0, name="cond_int")
|
||||
... cond = paddle.less_than(x=step_idx, y=max_len)
|
||||
... while_op = paddle.static.nn.control_flow.While(cond, is_test=True)
|
||||
... with while_op.block():
|
||||
... with paddle.base.device_guard(f'{device}:all'):
|
||||
... input = paddle.tensor.array_read(array=data, i=step_idx)
|
||||
... paddle.increment(x=step_idx, value=1.0)
|
||||
... paddle.tensor.array_write(input, i=step_idx, array=data)
|
||||
... with paddle.base.device_guard(f'{device}:0'):
|
||||
... param_attr = paddle.ParamAttr(initializer=paddle.nn.initializer.Constant(1.0))
|
||||
... weight1 = paddle.static.create_parameter(shape=[2, 5], dtype='float32', attr=param_attr, is_bias=False)
|
||||
... hidden1 = paddle.matmul(input, weight1)
|
||||
... with paddle.base.device_guard(f'{device}:1'):
|
||||
... param_attr = paddle.ParamAttr(initializer=paddle.nn.initializer.Constant(2.0))
|
||||
... weight2 = paddle.static.create_parameter(shape=[5, 2], dtype='float32', attr=param_attr, is_bias=False)
|
||||
... hidden2 = paddle.matmul(hidden1, weight2)
|
||||
... paddle.tensor.array_write(hidden2, i=step_idx, array=data)
|
||||
... # update cond and assign to cond_int, we will sync cond_int
|
||||
... paddle.assign(paddle.less_than(x=step_idx, y=max_len), cond)
|
||||
... paddle.assign(paddle.cast(cond, dtype="int32"), cond_int)
|
||||
... with paddle.base.device_guard(f'{device}:all'):
|
||||
... # the code below must at end of while block and exists in device:all
|
||||
... paddle.assign(paddle.cast(cond_int, dtype='bool'), cond)
|
||||
... with paddle.base.device_guard(f'{device}:all'):
|
||||
... out = paddle.tensor.create_array(data.dtype)
|
||||
... paddle.assign(data, out)
|
||||
... with paddle.base.device_guard(f'{device}:all'):
|
||||
... # use a empty lod_tensor_array to clear lod_tensor_array
|
||||
... paddle.assign(paddle.tensor.create_array(data.dtype), data)
|
||||
>>> helper = hybrid_parallel_inference.HybridParallelInferenceHelper(
|
||||
... startup_program,
|
||||
... main_program,
|
||||
... micro_batch_size=2,
|
||||
... num_pp=2,
|
||||
... init_comm=nranks > 1,
|
||||
... )
|
||||
>>> helper.gen_infer_program(['array_write_0.out'], ['cond_int.tmp_0'])
|
||||
>>> exe = paddle.static.Executor(paddle.CUDAPlace(dev_id))
|
||||
>>> exe.run(startup_program)
|
||||
>>> np.random.seed(2333)
|
||||
>>> for step in range(5):
|
||||
... init_data = np.random.uniform(low=0.0, high=1.0, size=[2, 2]).astype('float32')
|
||||
... [res] = exe.run(main_program, feed={"X": init_data}, fetch_list=[out])
|
||||
... print('-------- step', step, ' --------')
|
||||
... print(res)
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
startup_program,
|
||||
main_program,
|
||||
num_mp=1,
|
||||
num_pp=1,
|
||||
micro_batch_size=1,
|
||||
beam_size=1,
|
||||
init_comm=True,
|
||||
role_maker=None,
|
||||
):
|
||||
assert isinstance(startup_program, Program)
|
||||
assert isinstance(main_program, Program)
|
||||
|
||||
self._device = None
|
||||
if core.is_compiled_with_cuda():
|
||||
self._device = "gpu"
|
||||
assert self._device, "Only gpu are supported."
|
||||
|
||||
assert not in_dynamic_mode(), "Only static graph mode is supported."
|
||||
|
||||
op_maker = core.op_proto_and_checker_maker
|
||||
self._op_role = op_maker.OpRole
|
||||
self._op_role_key = op_maker.kOpRoleAttrName()
|
||||
self._op_device_key = op_maker.kOpDeviceAttrName()
|
||||
|
||||
self._param_device_map = {}
|
||||
|
||||
self._pipeline_pair = []
|
||||
self._pipeline_pair_in_while = []
|
||||
self._pp_ring_map = {}
|
||||
self.ring_id = 20 # Just a magic number
|
||||
|
||||
self.micro_batch_size = micro_batch_size
|
||||
self.beam_size = beam_size
|
||||
self.init_comm = init_comm
|
||||
|
||||
self._output_var_to_op = None
|
||||
self._input_var_to_op = None
|
||||
self._main_program = main_program
|
||||
self._startup_program = startup_program
|
||||
|
||||
if role_maker is None:
|
||||
self.role_maker = fleet.base.role_maker.PaddleCloudRoleMaker(
|
||||
is_collective=True
|
||||
)
|
||||
else:
|
||||
if isinstance(role_maker, fleet.base.role_maker.RoleMakerBase):
|
||||
assert role_maker._is_collective
|
||||
self.role_maker = role_maker
|
||||
|
||||
# communication_group info
|
||||
self.mp_ring_id = 0
|
||||
self.global_ring_id = 1
|
||||
|
||||
self.endpoints = self.role_maker._get_trainer_endpoints()
|
||||
self.current_endpoint = self.endpoints[self.role_maker._worker_index()]
|
||||
self.rank = self.role_maker._worker_index()
|
||||
self.nranks = self.role_maker._worker_num()
|
||||
assert num_mp * num_pp == self.nranks
|
||||
self.num_pp = num_pp
|
||||
self.num_mp = num_mp
|
||||
|
||||
# global ring info
|
||||
self.global_endpoints = self.endpoints
|
||||
self.global_rank = self.rank
|
||||
self.global_nranks = self.nranks
|
||||
|
||||
arr = np.arange(0, self.num_pp * self.num_mp).reshape(
|
||||
[self.num_pp, self.num_mp]
|
||||
)
|
||||
ipp, imp = np.where(arr == self.rank)
|
||||
ipp = ipp[0]
|
||||
imp = imp[0]
|
||||
self.mp_group = arr[ipp, :]
|
||||
self.pp_group = arr[:, imp]
|
||||
|
||||
self._stage = ipp
|
||||
|
||||
def _init_communication_group(self):
|
||||
dev_ids = []
|
||||
for pair in self._pipeline_pair:
|
||||
prev_id, cur_id = pair
|
||||
if prev_id not in dev_ids:
|
||||
dev_ids.append(prev_id)
|
||||
if cur_id not in dev_ids:
|
||||
dev_ids.append(cur_id)
|
||||
num_pp = len(dev_ids)
|
||||
num_pp = max(1, num_pp)
|
||||
assert num_pp == self.num_pp, (
|
||||
f'num_pp: {num_pp}, self.num_pp: {self.num_pp}'
|
||||
)
|
||||
|
||||
collective_helper = fleet.meta_optimizers.common.CollectiveHelper(
|
||||
self.role_maker, wait_port=False
|
||||
)
|
||||
|
||||
# Create global rings
|
||||
collective_helper._init_communicator(
|
||||
self._startup_program,
|
||||
self.current_endpoint,
|
||||
self.global_endpoints,
|
||||
self.global_rank,
|
||||
self.global_ring_id,
|
||||
True,
|
||||
self.global_ring_id,
|
||||
True,
|
||||
)
|
||||
|
||||
# Create mp rings
|
||||
if self.num_mp > 1:
|
||||
mp_endpoints = [self.endpoints[mp_idx] for mp_idx in self.mp_group]
|
||||
mp_rank = next(
|
||||
idx
|
||||
for idx, mp_idx in enumerate(self.mp_group)
|
||||
if mp_idx == self.rank
|
||||
)
|
||||
collective_helper._init_communicator(
|
||||
self._startup_program,
|
||||
self.current_endpoint,
|
||||
mp_endpoints,
|
||||
mp_rank,
|
||||
self.mp_ring_id,
|
||||
True,
|
||||
self.global_ring_id,
|
||||
True,
|
||||
)
|
||||
|
||||
# Create pipeline rings
|
||||
if self.num_pp > 1:
|
||||
for pair in self._pipeline_pair:
|
||||
pair_key = pair[0] * 1000 + pair[1]
|
||||
ring_id = self._pp_ring_map[pair_key]
|
||||
|
||||
first_node = self.pp_group[pair[0]]
|
||||
second_node = self.pp_group[pair[1]]
|
||||
if self.rank != first_node and self.rank != second_node:
|
||||
collective_helper._init_communicator(
|
||||
self._startup_program,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
False,
|
||||
self.global_ring_id,
|
||||
True,
|
||||
)
|
||||
continue
|
||||
|
||||
pipeline_endpoints = [
|
||||
self.endpoints[first_node],
|
||||
self.endpoints[second_node],
|
||||
]
|
||||
pipeline_rank = 0 if self.rank == first_node else 1
|
||||
collective_helper._init_communicator(
|
||||
self._startup_program,
|
||||
self.current_endpoint,
|
||||
pipeline_endpoints,
|
||||
pipeline_rank,
|
||||
ring_id,
|
||||
False,
|
||||
self.global_ring_id,
|
||||
True,
|
||||
)
|
||||
|
||||
def _get_input_output_info(self, block):
|
||||
'''
|
||||
Get info of op input and output.
|
||||
'''
|
||||
# A map from output var to op which generate it.
|
||||
output_var_to_op = defaultdict(list)
|
||||
# A map from var to op which takes it as input.
|
||||
input_var_to_op = defaultdict(list)
|
||||
|
||||
for index, op in enumerate(block.ops):
|
||||
for var_name in op.input_arg_names:
|
||||
input_var_to_op[var_name].append([op, index])
|
||||
for var_name in op.output_arg_names:
|
||||
output_var_to_op[var_name].append([op, index])
|
||||
|
||||
return output_var_to_op, input_var_to_op
|
||||
|
||||
def _update_param_device_map(self):
|
||||
"""
|
||||
Get the device info for parameters.
|
||||
"""
|
||||
params = [param.name for param in self._main_program.all_parameters()]
|
||||
for each_block in self._main_program.blocks:
|
||||
for op in each_block.ops:
|
||||
for var_name in op.input_arg_names:
|
||||
if (
|
||||
var_name not in params
|
||||
or var_name in self._param_device_map
|
||||
):
|
||||
continue
|
||||
device = op.attr(self._op_device_key)
|
||||
|
||||
self._param_device_map[var_name] = device
|
||||
|
||||
def _split_program(self, program, stage, block_idx):
|
||||
"""
|
||||
Split a program and get the one with the given pipeline stage.
|
||||
|
||||
Args:
|
||||
stage (int): pipeline stage
|
||||
block_idx (int): block index
|
||||
|
||||
Returns:
|
||||
used_var_names (set): used var names in block_idx block
|
||||
"""
|
||||
|
||||
used_var_names = set()
|
||||
block = program.block(block_idx)
|
||||
op_idx = 0
|
||||
for op in list(block.ops):
|
||||
op_stage = op.attr(self._op_device_key).split(':')[1]
|
||||
# Copy ops whose op_device set to "gpu:all" to all sections.
|
||||
if op_stage == "all" or int(op_stage) == stage:
|
||||
op_idx += 1
|
||||
if op.type == "while":
|
||||
sub_block_id = int(op.attr('sub_block').id)
|
||||
sub_used_var_names = self._split_program(
|
||||
program, stage, sub_block_id
|
||||
)
|
||||
|
||||
used_var_names.update(sub_used_var_names)
|
||||
|
||||
input_idxs = []
|
||||
input_arg_names = op.input("X")
|
||||
for i, name in enumerate(input_arg_names):
|
||||
if name not in sub_used_var_names:
|
||||
input_idxs.append(i)
|
||||
if len(input_idxs) > 0:
|
||||
for i in reversed(input_idxs):
|
||||
input_arg_names.pop(i)
|
||||
op.desc.set_input("X", input_arg_names)
|
||||
|
||||
output_idxs = []
|
||||
output_arg_names = op.output("Out")
|
||||
for i, name in enumerate(output_arg_names):
|
||||
if name not in sub_used_var_names:
|
||||
output_idxs.append(i)
|
||||
if len(output_idxs) > 0:
|
||||
for i in reversed(output_idxs):
|
||||
output_arg_names.pop(i)
|
||||
op.desc.set_output("Out", output_arg_names)
|
||||
|
||||
for var_name in op.input_arg_names + op.output_arg_names:
|
||||
used_var_names.add(var_name)
|
||||
else:
|
||||
block._remove_op(op_idx)
|
||||
|
||||
for var_name in list(block.vars.keys()):
|
||||
if var_name not in used_var_names:
|
||||
block._remove_var(var_name)
|
||||
|
||||
return used_var_names
|
||||
|
||||
# def _find_post_op(self, index, var_name):
|
||||
# """
|
||||
# Find the post op that has variable named var_name as input.
|
||||
# """
|
||||
# # bugfix for uniform hybrid parallelism
|
||||
# if '.cast_fp32' in var_name:
|
||||
# var_name = var_name.replace('.cast_fp32', '')
|
||||
# if '.cast_fp16' in var_name:
|
||||
# var_name = var_name.replace('.cast_fp16', '')
|
||||
|
||||
# post_ops = self._input_var_to_op[var_name]
|
||||
# if post_ops == None: return None
|
||||
# result_op = None
|
||||
# for post_op, post_idx in reversed(post_ops):
|
||||
# if post_idx > index:
|
||||
# result_op = post_op
|
||||
# break
|
||||
# return result_op
|
||||
|
||||
def _find_prev_op(self, index, var_name):
|
||||
"""
|
||||
Find the previous op of op with index that outputs
|
||||
variable named var_name.
|
||||
"""
|
||||
prev_ops = self._output_var_to_op[var_name]
|
||||
if prev_ops is None:
|
||||
return None
|
||||
result_op = None
|
||||
for prev_op, prev_idx in reversed(prev_ops):
|
||||
if prev_idx < index:
|
||||
result_op = prev_op
|
||||
break
|
||||
return result_op
|
||||
|
||||
def _add_op_device_attr(self, block):
|
||||
"""
|
||||
Add op_device attribute for ops in block that have
|
||||
not that attribute set.
|
||||
|
||||
Args:
|
||||
block (Block): the block to process.
|
||||
"""
|
||||
assert isinstance(block, Block)
|
||||
|
||||
# Ops should be copied to all pipeline stages.
|
||||
device_all_ops = [
|
||||
"create_py_reader",
|
||||
"read",
|
||||
"create_double_buffer_reader",
|
||||
"while",
|
||||
]
|
||||
|
||||
for op in block.ops:
|
||||
if op.type in device_all_ops:
|
||||
# We use "gpu:all" to represent an op should be put on all
|
||||
# pipeline stages, such as read ops. Note that: "gpu:all"
|
||||
# is only used by pipeline as an indicator.
|
||||
op._set_attr(self._op_device_key, self._device + ":all")
|
||||
if op.type == "while":
|
||||
sub_block_id = op.attr('sub_block').id
|
||||
sub_block = block.program.block(sub_block_id)
|
||||
self._add_op_device_attr(sub_block)
|
||||
|
||||
def _check_validation(self, block):
|
||||
"""
|
||||
Check whether ops in a block have both the op_device and the
|
||||
op_role attributes set.
|
||||
"""
|
||||
assert isinstance(block, Block)
|
||||
|
||||
pre_stage_id = None
|
||||
for op in block.ops:
|
||||
assert op.has_attr(self._op_role_key), (
|
||||
f"{op.type} has no {self._op_role_key} set ."
|
||||
)
|
||||
op_role = op.attr(self._op_role_key)
|
||||
assert op_role == int(self._op_role.Forward), (
|
||||
"Only forward is supported for inference."
|
||||
)
|
||||
if not op._has_kernel(op.type):
|
||||
assert op.type in [
|
||||
"while",
|
||||
"conditional_block",
|
||||
], "The only supported op without kernel is while."
|
||||
sub_block_id = op.attr('sub_block').id
|
||||
sub_block = block.program.block(sub_block_id)
|
||||
self._check_validation(sub_block)
|
||||
assert op.has_attr(self._op_device_key), (
|
||||
f"{op.type} has no {self._op_device_key} set."
|
||||
)
|
||||
|
||||
device = op.attr(self._op_device_key)
|
||||
assert device, f"{op.type} has no {self._op_device_key} set."
|
||||
if device.split(':')[1] == "all":
|
||||
continue
|
||||
|
||||
dev_type = device.split(':')[0]
|
||||
assert dev_type == self._device
|
||||
stage_id = int(device.split(':')[1])
|
||||
pre_stage_id = stage_id
|
||||
|
||||
def _insert_sendrecv_ops_for_boundaries(self, block, is_while_block):
|
||||
"""
|
||||
Insert a pair of send and recv ops for every two
|
||||
consecutive ops on different devices.
|
||||
"""
|
||||
# A map from var to device where op takes it as input,
|
||||
# avoiding multiple send and recv ops.
|
||||
input_var_to_device = {}
|
||||
|
||||
extra_index_info = {
|
||||
'index': 0,
|
||||
}
|
||||
|
||||
for index, op in enumerate(list(block.ops)):
|
||||
cur_device = op.attr(self._op_device_key)
|
||||
if cur_device.split(':')[-1] == "all":
|
||||
continue
|
||||
for var_name in op.input_arg_names:
|
||||
if not block.has_var(var_name) and block._find_var_recursive(
|
||||
var_name
|
||||
):
|
||||
continue
|
||||
var = block.var(var_name)
|
||||
# skip data var
|
||||
if var.is_data:
|
||||
continue
|
||||
prev_device = None
|
||||
generate_ops = self._output_var_to_op.get(var_name)
|
||||
if generate_ops is None:
|
||||
if var_name not in self._param_device_map:
|
||||
continue
|
||||
prev_device = self._param_device_map[var_name]
|
||||
|
||||
prev_op = self._find_prev_op(index, var_name)
|
||||
|
||||
if not prev_device:
|
||||
prev_device = (
|
||||
prev_op.attr(self._op_device_key) if prev_op else None
|
||||
)
|
||||
|
||||
if prev_device is None or prev_device.split(":")[-1] == "all":
|
||||
continue
|
||||
|
||||
if prev_device == cur_device:
|
||||
continue
|
||||
|
||||
if var_name not in input_var_to_device:
|
||||
input_var_to_device[var_name] = []
|
||||
if (cur_device, prev_device) in input_var_to_device[var_name]:
|
||||
continue
|
||||
|
||||
assert self._device == cur_device.split(':')[0], (
|
||||
"More than one device type found."
|
||||
)
|
||||
device_type = cur_device.split(':')[0] + ':'
|
||||
|
||||
def _insert_send_recv(cur_id, prev_id):
|
||||
assert cur_id > prev_id
|
||||
cur_dev = device_type + str(cur_id)
|
||||
prev_dev = device_type + str(prev_id)
|
||||
if (cur_dev, prev_dev) in input_var_to_device[var_name]:
|
||||
return
|
||||
|
||||
if cur_id - prev_id > 1:
|
||||
_insert_send_recv(cur_id - 1, prev_id)
|
||||
_insert_send_recv(cur_id, cur_id - 1)
|
||||
input_var_to_device[var_name].append(
|
||||
(cur_dev, prev_dev)
|
||||
)
|
||||
return
|
||||
|
||||
assert cur_id - prev_id == 1
|
||||
input_var_to_device[var_name].append((cur_dev, prev_dev))
|
||||
|
||||
op_role = op.attr(self._op_role_key)
|
||||
var = block.vars[var_name]
|
||||
pair = (prev_id, cur_id)
|
||||
if (
|
||||
is_while_block
|
||||
and pair not in self._pipeline_pair_in_while
|
||||
):
|
||||
self._pipeline_pair_in_while.append(pair)
|
||||
|
||||
# 1000 is just a magic number
|
||||
pair_key = prev_id * 1000 + cur_id
|
||||
if pair not in self._pipeline_pair:
|
||||
self._pipeline_pair.append(pair)
|
||||
self._pp_ring_map[pair_key] = self.ring_id
|
||||
ring_id = self.ring_id
|
||||
self.ring_id += 1
|
||||
else:
|
||||
ring_id = self._pp_ring_map[pair_key]
|
||||
|
||||
block._insert_op_without_sync(
|
||||
index=index + extra_index_info['index'],
|
||||
type='send_v2',
|
||||
inputs={'X': var},
|
||||
attrs={
|
||||
self._op_device_key: prev_dev,
|
||||
self._op_role_key: op_role,
|
||||
'use_calc_stream': True,
|
||||
'peer': 1,
|
||||
'ring_id': ring_id,
|
||||
},
|
||||
)
|
||||
extra_index_info['index'] += 1
|
||||
var_shape = list(var.shape)
|
||||
if var_shape[0] < 0:
|
||||
if is_while_block:
|
||||
var_shape[0] = (
|
||||
self.micro_batch_size * self.beam_size
|
||||
)
|
||||
else:
|
||||
var_shape[0] = self.micro_batch_size
|
||||
|
||||
block._insert_op_without_sync(
|
||||
index=index + extra_index_info['index'],
|
||||
type='recv_v2',
|
||||
outputs={'Out': [var]},
|
||||
attrs={
|
||||
'out_shape': var_shape,
|
||||
'dtype': var.dtype,
|
||||
self._op_device_key: cur_dev,
|
||||
self._op_role_key: op_role,
|
||||
'use_calc_stream': True,
|
||||
'peer': 0,
|
||||
'ring_id': ring_id,
|
||||
},
|
||||
)
|
||||
extra_index_info['index'] += 1
|
||||
|
||||
_insert_send_recv(
|
||||
int(cur_device.split(':')[1]),
|
||||
int(prev_device.split(':')[1]),
|
||||
)
|
||||
block._sync_with_cpp()
|
||||
|
||||
def _insert_sendrecv_ops_in_while_block(
|
||||
self,
|
||||
block,
|
||||
sync_in_while_lastpp2firstpp_var_names,
|
||||
sync_in_while_var_names,
|
||||
stage,
|
||||
):
|
||||
dev_ids = []
|
||||
for pair in self._pipeline_pair_in_while:
|
||||
prev_id, cur_id = pair
|
||||
if prev_id not in dev_ids:
|
||||
dev_ids.append(prev_id)
|
||||
if cur_id not in dev_ids:
|
||||
dev_ids.append(cur_id)
|
||||
|
||||
if len(dev_ids) == 0:
|
||||
return
|
||||
|
||||
first_id = min(dev_ids)
|
||||
last_id = max(dev_ids)
|
||||
|
||||
assert len(block.ops) > 2, (
|
||||
"It must have more than 2 ops in while sub block, "
|
||||
"layers.assign(layers.cast(cond_int, dtype='bool'), cond) must at end of while block, "
|
||||
"because nccl cannot send bool dtype var"
|
||||
)
|
||||
index = len(block.ops) - 2
|
||||
|
||||
for prev_id in dev_ids:
|
||||
if prev_id == cur_id:
|
||||
continue
|
||||
assert cur_id > prev_id
|
||||
|
||||
pair = (prev_id, cur_id)
|
||||
# 1000 is just a magic number
|
||||
pair_key = prev_id * 1000 + cur_id
|
||||
if pair not in self._pipeline_pair:
|
||||
self._pipeline_pair.append(pair)
|
||||
self._pp_ring_map[pair_key] = self.ring_id
|
||||
ring_id = self.ring_id
|
||||
self.ring_id += 1
|
||||
else:
|
||||
ring_id = self._pp_ring_map[pair_key]
|
||||
|
||||
if cur_id == last_id and prev_id == first_id:
|
||||
var_names = (
|
||||
sync_in_while_lastpp2firstpp_var_names
|
||||
+ sync_in_while_var_names
|
||||
)
|
||||
else:
|
||||
var_names = sync_in_while_var_names
|
||||
|
||||
for var_name in var_names:
|
||||
var = block._var_recursive(var_name)
|
||||
if stage == cur_id:
|
||||
block._insert_op_without_sync(
|
||||
index=index,
|
||||
type='send_v2',
|
||||
inputs={'X': var},
|
||||
attrs={
|
||||
self._op_device_key: self._device
|
||||
+ ':'
|
||||
+ str(cur_id),
|
||||
self._op_role_key: int(self._op_role.Forward),
|
||||
'use_calc_stream': True,
|
||||
'peer': 0,
|
||||
'ring_id': ring_id,
|
||||
},
|
||||
)
|
||||
else:
|
||||
var_shape = list(var.shape)
|
||||
print(var_name)
|
||||
if len(var.shape) > 0:
|
||||
var_shape[0] = (
|
||||
self.micro_batch_size
|
||||
if var_shape[0] < 0
|
||||
else var_shape[0]
|
||||
)
|
||||
block._insert_op_without_sync(
|
||||
index=index,
|
||||
type='recv_v2',
|
||||
outputs={'Out': [var]},
|
||||
attrs={
|
||||
'out_shape': var_shape,
|
||||
'dtype': var.dtype,
|
||||
self._op_device_key: self._device
|
||||
+ ':'
|
||||
+ str(prev_id),
|
||||
self._op_role_key: int(self._op_role.Forward),
|
||||
'use_calc_stream': True,
|
||||
'peer': 1,
|
||||
'ring_id': ring_id,
|
||||
},
|
||||
)
|
||||
index += 1
|
||||
block._sync_with_cpp()
|
||||
|
||||
def _get_while_block(self):
|
||||
"""
|
||||
Get the while sub-block.
|
||||
"""
|
||||
main_block = self._main_program.global_block()
|
||||
num_while = 0
|
||||
sub_block_id = None
|
||||
for op in main_block.ops:
|
||||
assert num_while < 2, "More than one while op found."
|
||||
if op.type == 'while':
|
||||
sub_block_id = op.attr('sub_block').id
|
||||
num_while += 1
|
||||
if sub_block_id:
|
||||
return op, self._main_program.block(sub_block_id)
|
||||
return None, None
|
||||
|
||||
def gen_infer_program(
|
||||
self,
|
||||
sync_in_while_lastpp2firstpp_var_names=None,
|
||||
sync_in_while_var_names=None,
|
||||
debug=False,
|
||||
):
|
||||
"""
|
||||
Generate inference program.
|
||||
Params:
|
||||
sync_in_while_lastpp2firstpp_var_names (list(str)): the vars in the last pipeline
|
||||
that need to send var to first pipeline and exclude bool dtype var
|
||||
sync_in_while_var_names (list(str)): the vars sync among all pipeline in while block
|
||||
e.g cond. Note that cond cannot be bool dtype.
|
||||
debug (bool): the flag indicate debug
|
||||
"""
|
||||
main_block = self._main_program.global_block()
|
||||
startup_block = self._startup_program.global_block()
|
||||
|
||||
if debug:
|
||||
with open('main_program.txt', 'w') as f:
|
||||
f.write(str(self._main_program))
|
||||
with open('startup_program.txt', 'w') as f:
|
||||
f.write(str(self._startup_program))
|
||||
|
||||
# step1: add op_device attribute for all ops
|
||||
self._add_op_device_attr(startup_block)
|
||||
self._check_validation(startup_block)
|
||||
self._add_op_device_attr(main_block)
|
||||
self._check_validation(main_block)
|
||||
|
||||
# step2: add send/recv ops
|
||||
self._update_param_device_map()
|
||||
# step2.1: add send/recv for main_block
|
||||
out_var_to_op, in_var_to_op = self._get_input_output_info(main_block)
|
||||
self._output_var_to_op = out_var_to_op
|
||||
self._input_var_to_op = in_var_to_op
|
||||
self._insert_sendrecv_ops_for_boundaries(main_block, False)
|
||||
|
||||
# step2.2: add send/recv for while_block
|
||||
while_op, while_block = self._get_while_block()
|
||||
if while_block:
|
||||
out_var_to_op, in_var_to_op = self._get_input_output_info(
|
||||
while_block
|
||||
)
|
||||
self._output_var_to_op = out_var_to_op
|
||||
self._input_var_to_op = in_var_to_op
|
||||
|
||||
self._insert_sendrecv_ops_for_boundaries(while_block, True)
|
||||
|
||||
self._insert_sendrecv_ops_in_while_block(
|
||||
while_block,
|
||||
sync_in_while_lastpp2firstpp_var_names,
|
||||
sync_in_while_var_names,
|
||||
self._stage,
|
||||
)
|
||||
|
||||
# step3: split programs
|
||||
self._split_program(self._startup_program, self._stage, 0)
|
||||
self._split_program(self._main_program, self._stage, 0)
|
||||
|
||||
if debug:
|
||||
with open(f'main_program.txt.{self.rank}', 'w') as f:
|
||||
f.write(str(self._main_program))
|
||||
with open(f'startup_program.txt.{self.rank}', 'w') as f:
|
||||
f.write(str(self._startup_program))
|
||||
|
||||
if self.init_comm:
|
||||
self._init_communication_group()
|
||||
@@ -0,0 +1,340 @@
|
||||
# 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 paddle
|
||||
from paddle import framework
|
||||
|
||||
# (TODO: GhostScreaming) It will be removed later.
|
||||
from paddle.base import core
|
||||
from paddle.distributed.parallel import (
|
||||
_split_tensors,
|
||||
build_groups,
|
||||
in_dynamic_mode,
|
||||
sync_params_buffers,
|
||||
)
|
||||
|
||||
from .log_util import logger
|
||||
|
||||
__all__ = []
|
||||
|
||||
|
||||
def obtain_optimizer_parameters_list(optimizer):
|
||||
if getattr(optimizer, '_param_groups', None) and isinstance(
|
||||
optimizer._param_groups[0], dict
|
||||
):
|
||||
parameters_list = []
|
||||
for group in optimizer._param_groups:
|
||||
for param in group['params']:
|
||||
parameters_list.append(param)
|
||||
else:
|
||||
parameters_list = list(optimizer._parameter_list)
|
||||
|
||||
return parameters_list
|
||||
|
||||
|
||||
def _apply_collective_grads(parameters, comm_group, bucket_size, scale=None):
|
||||
grad_var_set = set()
|
||||
grad_vars = []
|
||||
sparse_grad_vars = []
|
||||
|
||||
for param in parameters:
|
||||
if param.trainable and (param._grad_ivar() is not None):
|
||||
g_var = param._grad_ivar()
|
||||
assert not g_var._is_sparse(), (
|
||||
"Now, it doesn't support sparse parameters"
|
||||
)
|
||||
grad_vars.append(g_var)
|
||||
assert g_var not in grad_var_set
|
||||
grad_var_set.add(g_var)
|
||||
|
||||
coalesced_grads_and_vars = build_groups(grad_vars, bucket_size)
|
||||
|
||||
nranks = (
|
||||
paddle.distributed.get_world_size()
|
||||
if comm_group is None
|
||||
else comm_group.nranks
|
||||
)
|
||||
|
||||
scale = nranks if scale is None else 1.0 / scale
|
||||
scale = None if scale == 1.0 else scale
|
||||
|
||||
for coalesced_grad, _, _ in coalesced_grads_and_vars:
|
||||
# need to div nranks
|
||||
if scale is not None:
|
||||
div_factor = paddle.to_tensor(scale, dtype=coalesced_grad.dtype)
|
||||
paddle.base.framework._dygraph_tracer().trace_op(
|
||||
type="elementwise_div",
|
||||
inputs={'X': coalesced_grad, 'Y': div_factor},
|
||||
outputs={'Out': coalesced_grad},
|
||||
attrs={'axis': -1},
|
||||
)
|
||||
paddle.distributed.all_reduce(coalesced_grad, group=comm_group)
|
||||
|
||||
_split_tensors(coalesced_grads_and_vars)
|
||||
|
||||
|
||||
def _apply_collective_grads_eager(
|
||||
parameters, comm_group, bucket_size, scale=None
|
||||
):
|
||||
grad_var_set = set()
|
||||
grad_vars = []
|
||||
|
||||
for param in parameters:
|
||||
g_var = None
|
||||
if param.trainable and (param._grad_ivar() is not None):
|
||||
g_var = param._grad_ivar()
|
||||
if param.trainable and hasattr(param, "main_grad"):
|
||||
assert param._grad_ivar() is None, "param.grad is not None"
|
||||
g_var = param.main_grad
|
||||
if g_var is not None:
|
||||
assert not g_var.is_sparse(), (
|
||||
"Now, it doesn't support sparse parameters"
|
||||
)
|
||||
grad_vars.append(g_var)
|
||||
assert g_var not in grad_var_set
|
||||
grad_var_set.add(g_var)
|
||||
|
||||
if len(grad_vars) == 0:
|
||||
return
|
||||
|
||||
coalesced_grads_and_vars = build_groups(grad_vars, bucket_size)
|
||||
|
||||
nranks = (
|
||||
paddle.distributed.get_world_size()
|
||||
if comm_group is None
|
||||
else comm_group.nranks
|
||||
)
|
||||
|
||||
scale = 1.0 / nranks if scale is None else scale
|
||||
scale = None if scale == 1.0 else scale
|
||||
|
||||
for coalesced_grad, _, _ in coalesced_grads_and_vars:
|
||||
# need to div nranks
|
||||
if scale is not None:
|
||||
coalesced_grad.scale_(scale)
|
||||
paddle.distributed.all_reduce(coalesced_grad, group=comm_group)
|
||||
|
||||
_split_tensors(coalesced_grads_and_vars)
|
||||
|
||||
|
||||
def _broadcast_data_help(data, shape, dtype, hcg):
|
||||
model_parallel_group = hcg.get_model_parallel_group()
|
||||
src_rank = hcg.get_model_parallel_group_src_rank()
|
||||
mp_rank = hcg.get_model_parallel_rank()
|
||||
|
||||
shape_gpu = paddle.to_tensor(shape, dtype="int32")
|
||||
paddle.distributed.broadcast(
|
||||
shape_gpu, src=src_rank, group=model_parallel_group, sync_op=True
|
||||
)
|
||||
|
||||
if mp_rank != 0:
|
||||
input_data = paddle.zeros(shape_gpu, dtype=dtype)
|
||||
else:
|
||||
input_data = data
|
||||
|
||||
paddle.distributed.broadcast(
|
||||
input_data, src=src_rank, group=model_parallel_group, sync_op=True
|
||||
)
|
||||
|
||||
if mp_rank != 0:
|
||||
if in_dynamic_mode():
|
||||
data._clear_data()
|
||||
input_data._share_buffer_to(data)
|
||||
else:
|
||||
data.value().get_tensor()._clear()
|
||||
data.value().get_tensor()._share_data_with(
|
||||
input_data.value().get_tensor()
|
||||
)
|
||||
|
||||
|
||||
def _broadcast_object_list_help(object_list, hcg):
|
||||
model_parallel_group = hcg.get_model_parallel_group()
|
||||
src_rank = hcg.get_model_parallel_group_src_rank()
|
||||
mp_rank = hcg.get_model_parallel_rank()
|
||||
|
||||
paddle.distributed.broadcast_object_list(
|
||||
object_list, src=src_rank, group=model_parallel_group
|
||||
)
|
||||
|
||||
|
||||
def _process_element(hcg, dev, place, element):
|
||||
if isinstance(element, core.eager.Tensor):
|
||||
with framework.no_grad():
|
||||
if (
|
||||
in_dynamic_mode()
|
||||
and not eval(f"element.place.is_{dev}_place")()
|
||||
):
|
||||
element_gpu = element._copy_to(place, True)
|
||||
element._clear_data()
|
||||
element_gpu._share_buffer_to(element)
|
||||
_broadcast_data_help(element, element.shape, element.dtype, hcg)
|
||||
elif isinstance(element, (dict, list, tuple)):
|
||||
return _broadcast_nested_data(hcg, dev, place, element)
|
||||
else:
|
||||
_broadcast_object_list_help([element], hcg)
|
||||
|
||||
|
||||
def _broadcast_nested_data(hcg, dev, place, data):
|
||||
if isinstance(data, dict):
|
||||
return {
|
||||
key: _process_element(hcg, dev, place, value)
|
||||
for key, value in data.items()
|
||||
}
|
||||
elif isinstance(data, list):
|
||||
return [_process_element(hcg, dev, place, item) for item in data]
|
||||
elif isinstance(data, tuple):
|
||||
return tuple(_process_element(hcg, dev, place, item) for item in data)
|
||||
else:
|
||||
raise TypeError(f"Unsupported data type: {type(data)}")
|
||||
|
||||
|
||||
def broadcast_input_data(hcg, *inputs, **kwargs):
|
||||
cur_device = paddle.get_device()
|
||||
dev = cur_device.split(":")[0]
|
||||
assert (
|
||||
dev
|
||||
in [
|
||||
"xpu",
|
||||
"gpu",
|
||||
]
|
||||
or dev in paddle.device.get_all_custom_device_type()
|
||||
), f"Only support xpu, gpu and custom_device now, but this is {dev}"
|
||||
dev_idx = int(cur_device.split(':')[1])
|
||||
if dev == "gpu":
|
||||
place = paddle.CUDAPlace(dev_idx)
|
||||
elif dev in paddle.device.get_all_custom_device_type():
|
||||
place = paddle.CustomPlace(dev, dev_idx)
|
||||
dev = 'custom'
|
||||
else:
|
||||
place = eval(f"paddle.{dev.upper()}Place")(dev_idx)
|
||||
|
||||
if len(inputs) > 0:
|
||||
inputs = _broadcast_nested_data(hcg, dev, place, inputs)
|
||||
if len(kwargs) > 0:
|
||||
kwargs = _broadcast_nested_data(hcg, dev, place, kwargs)
|
||||
return inputs, kwargs
|
||||
|
||||
|
||||
def broadcast_mp_parameters(model, hcg, fuse_params=True):
|
||||
model_parallel_group = hcg.get_model_parallel_group()
|
||||
src_rank = hcg.get_model_parallel_group_src_rank()
|
||||
sync_params_buffers(
|
||||
model,
|
||||
model_parallel_group,
|
||||
src_rank,
|
||||
is_model_parallel=True,
|
||||
fuse_params=fuse_params,
|
||||
)
|
||||
|
||||
|
||||
def broadcast_dp_parameters(model, hcg, fuse_params=True):
|
||||
data_parallel_group = hcg.get_data_parallel_group()
|
||||
src_rank = hcg.get_data_parallel_group_src_rank()
|
||||
sync_params_buffers(
|
||||
model,
|
||||
data_parallel_group,
|
||||
src_rank,
|
||||
is_model_parallel=False,
|
||||
fuse_params=fuse_params,
|
||||
)
|
||||
|
||||
|
||||
def fused_allreduce_gradients_with_group(
|
||||
parameter_list, group, bucket_size=128 * 1024 * 1024, scale=None
|
||||
):
|
||||
apply_func = (
|
||||
_apply_collective_grads_eager
|
||||
if in_dynamic_mode()
|
||||
else _apply_collective_grads
|
||||
)
|
||||
with framework.no_grad():
|
||||
apply_func(parameter_list, group, bucket_size, scale)
|
||||
|
||||
|
||||
def fused_allreduce_gradients(parameter_list, hcg):
|
||||
group = None
|
||||
scale = None
|
||||
if hcg is not None:
|
||||
dp_enabled = hcg.get_data_parallel_world_size() > 1
|
||||
sep_enabled = hcg.get_sep_parallel_world_size() > 1
|
||||
assert dp_enabled or sep_enabled, (
|
||||
f"dp_enabled {dp_enabled}; sep_enabled {sep_enabled}"
|
||||
)
|
||||
group = None
|
||||
# sep all reduce is not scaled
|
||||
scale = 1.0
|
||||
if dp_enabled:
|
||||
group = hcg.get_data_parallel_group()
|
||||
scale = scale / group.nranks
|
||||
if sep_enabled:
|
||||
sep_group = hcg.get_sep_parallel_group()
|
||||
dp_sep_group = hcg.get_dp_sep_parallel_group()
|
||||
group = sep_group if group is None else dp_sep_group
|
||||
|
||||
logger.debug("dp or sep start fuse allreduce gradients")
|
||||
from paddle.distributed import in_auto_parallel_align_mode
|
||||
|
||||
if in_auto_parallel_align_mode():
|
||||
scale = 1.0
|
||||
fused_allreduce_gradients_with_group(parameter_list, group, scale=scale)
|
||||
|
||||
|
||||
def broadcast_sharding_parameters(model, hcg, fuse_params=True):
|
||||
# TODO TO save memory, use un-fused broadcast to avoid potential OOM
|
||||
logger.debug("sharding start init parameters sync")
|
||||
sharding_parallel_group = hcg.get_sharding_parallel_group()
|
||||
src_rank = hcg.get_sharding_parallel_group_src_rank()
|
||||
sync_params_buffers(
|
||||
model,
|
||||
sharding_parallel_group,
|
||||
src_rank,
|
||||
is_model_parallel=False,
|
||||
fuse_params=fuse_params,
|
||||
)
|
||||
|
||||
|
||||
def broadcast_sep_parameters(model, hcg, fuse_params=True):
|
||||
# TODO TO save memory, use un-fused broadcast to avoid potential OOM
|
||||
logger.debug("sep start init parameters sync")
|
||||
sep_group = hcg.get_sep_parallel_group()
|
||||
src_rank = hcg.get_sep_parallel_group_src_rank()
|
||||
sync_params_buffers(
|
||||
model,
|
||||
sep_group,
|
||||
src_rank,
|
||||
is_model_parallel=False,
|
||||
fuse_params=fuse_params,
|
||||
)
|
||||
|
||||
|
||||
def broadcast_moe_sharding_parameters(model, hcg, fuse_params=True):
|
||||
# TODO TO save memory, use un-fused broadcast to avoid potential OOM
|
||||
logger.debug("moe sharding start init parameters sync")
|
||||
moe_sharding_group = hcg.get_moe_sharding_parallel_group()
|
||||
src_rank = hcg.get_moe_sharding_parallel_group_src_rank()
|
||||
sync_params_buffers(
|
||||
model,
|
||||
moe_sharding_group,
|
||||
src_rank,
|
||||
is_model_parallel=False,
|
||||
fuse_params=fuse_params,
|
||||
is_moe_sharding_parallel=True,
|
||||
)
|
||||
|
||||
|
||||
def unwrap_optimizer(optimizer, optimizer_instances=()):
|
||||
_inner_opt = optimizer
|
||||
while isinstance(_inner_opt, optimizer_instances):
|
||||
_inner_opt = _inner_opt._inner_opt
|
||||
return _inner_opt
|
||||
@@ -0,0 +1,174 @@
|
||||
# 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 logging
|
||||
import os
|
||||
import subprocess
|
||||
from logging.handlers import RotatingFileHandler
|
||||
|
||||
import paddle
|
||||
from paddle.distributed.utils.log_utils import get_logger
|
||||
|
||||
logger = get_logger("INFO", __name__)
|
||||
|
||||
|
||||
def set_log_level(level):
|
||||
"""
|
||||
Set log level
|
||||
|
||||
Args:
|
||||
level (str|int): a specified level
|
||||
|
||||
Example 1:
|
||||
import paddle
|
||||
import paddle.distributed.fleet as fleet
|
||||
fleet.init()
|
||||
fleet.setLogLevel("DEBUG")
|
||||
|
||||
Example 2:
|
||||
import paddle
|
||||
import paddle.distributed.fleet as fleet
|
||||
fleet.init()
|
||||
fleet.setLogLevel(1)
|
||||
|
||||
"""
|
||||
assert isinstance(level, (str, int)), "level's type must be str or int"
|
||||
if isinstance(level, int):
|
||||
logger.setLevel(level)
|
||||
else:
|
||||
logger.setLevel(level.upper())
|
||||
|
||||
|
||||
def get_log_level_code():
|
||||
"""
|
||||
Return current log level code
|
||||
"""
|
||||
return logger.getEffectiveLevel()
|
||||
|
||||
|
||||
def get_log_level_name():
|
||||
"""
|
||||
Return current log level name
|
||||
"""
|
||||
return logging.getLevelName(get_log_level_code())
|
||||
|
||||
|
||||
def layer_to_str(base, *args, **kwargs):
|
||||
name = base + "("
|
||||
if args:
|
||||
name += ", ".join(str(arg) for arg in args)
|
||||
if kwargs:
|
||||
name += ", "
|
||||
if kwargs:
|
||||
name += ", ".join(f"{key}={value}" for key, value in kwargs.items())
|
||||
name += ")"
|
||||
return name
|
||||
|
||||
|
||||
class DistributedLogger(logging.Logger):
|
||||
def __init__(self, name, level=logging.NOTSET):
|
||||
super().__init__(name, level)
|
||||
|
||||
def info(self, msg, *args, **kwargs):
|
||||
paddle.device.synchronize()
|
||||
super().info(f"Distributed Debug: {msg}", *args, **kwargs)
|
||||
|
||||
|
||||
def get_rotate_file_logger(log_level, name='root'):
|
||||
distributed_logger = DistributedLogger(name + '_rotate', level=log_level)
|
||||
distributed_logger.propagate = False
|
||||
|
||||
device_id = int(os.getenv("FLAGS_selected_gpus", "0"))
|
||||
log_dir = os.path.join(os.getcwd(), "hybrid_parallel")
|
||||
os.makedirs(log_dir, exist_ok=True)
|
||||
|
||||
path = os.path.join(log_dir, f"worker_{device_id}.log")
|
||||
handler = RotatingFileHandler(
|
||||
path,
|
||||
maxBytes=2 * 1024 * 1024 * 1024,
|
||||
backupCount=3, # 2GB
|
||||
)
|
||||
|
||||
log_format = logging.Formatter(
|
||||
'[%(asctime)-15s] [%(levelname)8s] %(filename)s:%(lineno)s - %(message)s'
|
||||
)
|
||||
handler.setFormatter(log_format)
|
||||
distributed_logger.addHandler(handler)
|
||||
return distributed_logger
|
||||
|
||||
|
||||
g_sync_rotate_logger = None
|
||||
|
||||
|
||||
def get_sync_logger():
|
||||
global logger
|
||||
paddle.device.synchronize()
|
||||
return logger
|
||||
|
||||
|
||||
def sync_rotate_logger():
|
||||
global g_sync_rotate_logger
|
||||
if g_sync_rotate_logger is None:
|
||||
g_sync_rotate_logger = get_rotate_file_logger("INFO", __name__)
|
||||
return g_sync_rotate_logger
|
||||
|
||||
|
||||
def check_memory_usage(msg=""):
|
||||
GB = 1024.0 * 1024.0 * 1024.0
|
||||
mem_dict = {}
|
||||
mem_dict['max_memory_allocated_size'] = (
|
||||
paddle.device.cuda.max_memory_allocated() / GB
|
||||
)
|
||||
mem_dict['max_memory_reserved_size'] = (
|
||||
paddle.device.cuda.max_memory_reserved() / GB
|
||||
)
|
||||
mem_dict['memory_allocated_size'] = (
|
||||
paddle.device.cuda.memory_allocated() / GB
|
||||
)
|
||||
mem_dict['memory_reserved_size'] = paddle.device.cuda.memory_reserved() / GB
|
||||
mem_msg = f"checking gpu memory usage {msg}:"
|
||||
for key in mem_dict:
|
||||
mem_msg += f"\n{key}: {mem_dict[key]}GB"
|
||||
logger.info(mem_msg)
|
||||
|
||||
if hasattr(paddle.device.cuda, 'max_pinned_memory_allocated'):
|
||||
mem_dict = {}
|
||||
mem_dict['max_memory_allocated_size'] = (
|
||||
paddle.device.cuda.max_pinned_memory_allocated() / GB
|
||||
)
|
||||
mem_dict['max_memory_reserved_size'] = (
|
||||
paddle.device.cuda.max_pinned_memory_reserved() / GB
|
||||
)
|
||||
mem_dict['memory_allocated_size'] = (
|
||||
paddle.device.cuda.pinned_memory_allocated() / GB
|
||||
)
|
||||
mem_dict['memory_reserved_size'] = (
|
||||
paddle.device.cuda.pinned_memory_reserved() / GB
|
||||
)
|
||||
mem_msg = f"checking pinned memory usage {msg}:"
|
||||
for key in mem_dict:
|
||||
mem_msg += f"\n{key}: {mem_dict[key]}GB"
|
||||
logger.info(mem_msg)
|
||||
|
||||
# Execute the command and get the output
|
||||
result = subprocess.run(["free", "-h"], capture_output=True, text=True)
|
||||
lines = result.stdout.strip().split('\n')
|
||||
|
||||
# Extract data
|
||||
mem_data = lines[1].split()
|
||||
swap_data = lines[2].split()
|
||||
|
||||
# Format and print
|
||||
formatted_output = f"checking CPU memory usage: {msg} Memory - Total: {mem_data[1]}, Used: {mem_data[2]}, Free: {mem_data[3]} Available:{mem_data[-1]}"
|
||||
logger.info(formatted_output)
|
||||
@@ -0,0 +1,260 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
|
||||
from collections import defaultdict
|
||||
from types import MethodType
|
||||
|
||||
import numpy as np
|
||||
|
||||
import paddle
|
||||
from paddle import _legacy_C_ops, nn
|
||||
from paddle.base import framework
|
||||
from paddle.base.dygraph import (
|
||||
base as imperative_base,
|
||||
)
|
||||
from paddle.distributed import fleet
|
||||
from paddle.distributed.fleet.utils.hybrid_parallel_util import (
|
||||
obtain_optimizer_parameters_list,
|
||||
)
|
||||
from paddle.framework import core
|
||||
from paddle.utils import deprecated
|
||||
|
||||
|
||||
class MixPrecisionLayer(nn.Layer):
|
||||
def __init__(self, layers, dtype="float16"):
|
||||
super().__init__(layers.full_name() + "_mix_precision")
|
||||
|
||||
self._layers = layers
|
||||
self._dtype = dtype
|
||||
|
||||
assert self._dtype in ["float16", "bfloat16"]
|
||||
|
||||
for param in self._layers.parameters():
|
||||
if not hasattr(param, "main_grad"):
|
||||
param.main_grad = None
|
||||
param._register_grad_hook(self._update_main_grad_hook(param))
|
||||
|
||||
def _update_main_grad_hook(self, param):
|
||||
"""Create the update_main_grad hook for back-prop."""
|
||||
|
||||
# Hook used for back-prop and grad-merge.
|
||||
@paddle.autograd.no_grad()
|
||||
def param_hook(tmp_grad):
|
||||
assert param.grad is None, (
|
||||
f"In main_grad node, param.grad should be None, but find param[{param.name}] has grad."
|
||||
)
|
||||
if tmp_grad is not None and tmp_grad._is_initialized():
|
||||
# Some previous pylayer may return None, should check grad validation.
|
||||
if param.main_grad is None:
|
||||
param.main_grad = core.eager.Tensor(
|
||||
value=tmp_grad.cast(paddle.float32).value(),
|
||||
place=tmp_grad.place,
|
||||
name="main_grad@" + param.name,
|
||||
)
|
||||
else:
|
||||
param.main_grad.add_(tmp_grad)
|
||||
|
||||
tmp_grad._clear_data()
|
||||
|
||||
return param_hook
|
||||
|
||||
def forward(self, *inputs, **kwargs):
|
||||
outputs = self._layers(*inputs, **kwargs)
|
||||
|
||||
return outputs
|
||||
|
||||
def state_dict(
|
||||
self,
|
||||
destination=None,
|
||||
include_sublayers=True,
|
||||
structured_name_prefix="",
|
||||
):
|
||||
return self._layers.state_dict(
|
||||
destination=destination,
|
||||
include_sublayers=include_sublayers,
|
||||
structured_name_prefix=structured_name_prefix,
|
||||
)
|
||||
|
||||
@framework.deprecate_stat_dict
|
||||
def set_state_dict(self, state_dict, use_structured_name=True):
|
||||
self._layers.set_state_dict(
|
||||
state_dict, use_structured_name=use_structured_name
|
||||
)
|
||||
|
||||
|
||||
class MixPrecisionOptimizer:
|
||||
def __init__(self, optimizer):
|
||||
self._inner_opt = optimizer
|
||||
self._parameter_list = obtain_optimizer_parameters_list(optimizer)
|
||||
|
||||
@imperative_base.no_grad
|
||||
@framework.dygraph_only
|
||||
def step(self):
|
||||
need_shard = any(
|
||||
hasattr(p, '_need_shard') for p in self._parameter_list
|
||||
)
|
||||
if need_shard:
|
||||
fleet.meta_parallel.sharding.group_sharded_fully_shard.FullyShardOptimizer(
|
||||
self
|
||||
)
|
||||
self.step()
|
||||
return
|
||||
|
||||
if not isinstance(self._parameter_list[0], dict):
|
||||
params_grads = []
|
||||
for param in self._parameter_list:
|
||||
if param.stop_gradient:
|
||||
continue
|
||||
grad_var = param.main_grad
|
||||
if grad_var is None:
|
||||
continue
|
||||
if paddle.in_dynamic_mode():
|
||||
if (
|
||||
hasattr(grad_var, "is_selected_rows")
|
||||
and grad_var.is_selected_rows()
|
||||
and self._inner_opt.regularization is not None
|
||||
):
|
||||
raise RuntimeError(
|
||||
"AdamW don't support weight_decay with sparse parameters, please set it to None."
|
||||
)
|
||||
else:
|
||||
if (
|
||||
hasattr(grad_var, "_is_sparse")
|
||||
and grad_var._is_sparse()
|
||||
and self._inner_opt.regularization is not None
|
||||
):
|
||||
raise RuntimeError(
|
||||
"AdamW don't support weight_decay with sparse parameters, please set it to None."
|
||||
)
|
||||
params_grads.append((param, grad_var))
|
||||
|
||||
optimize_ops = self._inner_opt._apply_optimize(
|
||||
loss=None, startup_program=None, params_grads=params_grads
|
||||
)
|
||||
else:
|
||||
# optimize parameters in groups
|
||||
for param_group in self._inner_opt._param_groups:
|
||||
params_grads = defaultdict(lambda: [])
|
||||
for param in param_group['params']:
|
||||
if param.stop_gradient:
|
||||
continue
|
||||
grad_var = param.main_grad
|
||||
if grad_var is None:
|
||||
continue
|
||||
if paddle.in_dynamic_mode():
|
||||
if (
|
||||
hasattr(grad_var, "is_selected_rows")
|
||||
and grad_var.is_selected_rows()
|
||||
and self._inner_opt.regularization is not None
|
||||
):
|
||||
raise RuntimeError(
|
||||
"AdamW don't support weight_decay with sparse parameters, please set it to None."
|
||||
)
|
||||
else:
|
||||
if (
|
||||
hasattr(grad_var, "_is_sparse")
|
||||
and grad_var._is_sparse()
|
||||
and self._inner_opt.regularization is not None
|
||||
):
|
||||
raise RuntimeError(
|
||||
"AdamW don't support weight_decay with sparse parameters, please set it to None."
|
||||
)
|
||||
params_grads['params'].append((param, grad_var))
|
||||
params_grads.update(
|
||||
{k: v for k, v in param_group.items() if k != 'params'}
|
||||
)
|
||||
self._apply_optimize(
|
||||
loss=None, startup_program=None, params_grads=params_grads
|
||||
)
|
||||
|
||||
@framework.dygraph_only
|
||||
def clear_grad(self, set_to_zero=True):
|
||||
param_list = []
|
||||
if self._parameter_list is None or not isinstance(
|
||||
self._parameter_list[0], dict
|
||||
):
|
||||
for p in self._parameter_list:
|
||||
if not p.stop_gradient:
|
||||
param_list.append(p)
|
||||
else:
|
||||
for param_group in self._param_groups:
|
||||
for p in param_group['params']:
|
||||
if not p.stop_gradient:
|
||||
param_list.append(p)
|
||||
|
||||
for p in param_list:
|
||||
if hasattr(p, "main_grad") and p.main_grad is not None:
|
||||
if set_to_zero:
|
||||
p.main_grad.zero_()
|
||||
else:
|
||||
p.main_grad._clear()
|
||||
p.main_grad = None
|
||||
elif not hasattr(p, "main_grad"):
|
||||
p.clear_gradient(set_to_zero)
|
||||
|
||||
def __getattr__(self, item):
|
||||
return getattr(self._inner_opt, item)
|
||||
|
||||
|
||||
def unscale_method(self, optimizer):
|
||||
if not self._enable:
|
||||
return
|
||||
param_grads = []
|
||||
if getattr(optimizer, '_param_groups', None) and isinstance(
|
||||
optimizer._param_groups[0], dict
|
||||
):
|
||||
for group in optimizer._param_groups:
|
||||
for param in group['params']:
|
||||
if param.main_grad is not None:
|
||||
assert param.main_grad.dtype == paddle.float32
|
||||
param_grads.append(param.main_grad)
|
||||
else:
|
||||
for param in optimizer._parameter_list:
|
||||
if param.main_grad is not None:
|
||||
assert param.main_grad.dtype == paddle.float32
|
||||
param_grads.append(param.main_grad)
|
||||
|
||||
temp_found_inf = paddle.to_tensor(np.array([0]).astype(np.bool_))
|
||||
if len(param_grads):
|
||||
_legacy_C_ops.check_finite_and_unscale(
|
||||
param_grads,
|
||||
self._scale,
|
||||
param_grads,
|
||||
temp_found_inf,
|
||||
)
|
||||
|
||||
self._found_inf = 1 if temp_found_inf else 0
|
||||
|
||||
hcg = fleet.get_hybrid_communicate_group()
|
||||
if hcg is not None and hcg.nranks > hcg.get_data_parallel_world_size():
|
||||
is_found_inf = paddle.to_tensor([self._found_inf], dtype="int32")
|
||||
paddle.distributed.all_reduce(
|
||||
is_found_inf, op=paddle.distributed.ReduceOp.MAX, group=None
|
||||
)
|
||||
self._found_inf = int(is_found_inf)
|
||||
|
||||
|
||||
@deprecated(
|
||||
since="2.5.0",
|
||||
update_to="paddle.distributed_scaler",
|
||||
level=1,
|
||||
)
|
||||
class MixPrecisionScaler:
|
||||
def __init__(self, scaler):
|
||||
self._inner_scaler = scaler
|
||||
self._inner_scaler._unscale = MethodType(unscale_method, scaler)
|
||||
|
||||
def __getattr__(self, item):
|
||||
return getattr(self._inner_scaler, item)
|
||||
@@ -0,0 +1,604 @@
|
||||
# Copyright (c) 2023 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 argparse
|
||||
import math
|
||||
import re
|
||||
import shutil
|
||||
from collections import OrderedDict
|
||||
|
||||
import paddle
|
||||
|
||||
|
||||
class ParallelConfig:
|
||||
def __init__(self, mp: int, pp: int, vpp: int = 1, sharding: int = 1):
|
||||
self.mp = mp
|
||||
self.pp = pp
|
||||
self.vpp = vpp
|
||||
self.sharding = sharding
|
||||
|
||||
def pipe_parallel_group(self, i: int, j: int):
|
||||
ans = []
|
||||
for k in range(self.pp):
|
||||
ans.append((i, j, k))
|
||||
return ans
|
||||
|
||||
|
||||
class LayerReNamingHelper:
|
||||
def __init__(self, template: str):
|
||||
self._template = template
|
||||
self._i = -1
|
||||
self._last_old_layer_name = None
|
||||
|
||||
def get_new_layer_name(self, old_layer_name: str):
|
||||
old_layer_name = old_layer_name.split(".")[0]
|
||||
if (
|
||||
self._last_old_layer_name is None
|
||||
or old_layer_name != self._last_old_layer_name
|
||||
):
|
||||
self._i = self._i + 1
|
||||
self._last_old_layer_name = old_layer_name
|
||||
return self._template.format(self._i)
|
||||
|
||||
|
||||
class LayerReNamingManager:
|
||||
def __init__(self):
|
||||
self._renaming_helpers = OrderedDict()
|
||||
self._renaming_helpers["linear"] = LayerReNamingHelper("linear_{}")
|
||||
self._renaming_helpers["layer_norm"] = LayerReNamingHelper(
|
||||
"layer_norm_{}"
|
||||
)
|
||||
self._renaming_helpers["embedding"] = LayerReNamingHelper(
|
||||
"embedding_{}"
|
||||
)
|
||||
|
||||
def get_new_layer_name(self, old_name: str):
|
||||
layer_name = ""
|
||||
for k, v in self._renaming_helpers.items():
|
||||
if old_name.startswith(k):
|
||||
layer_name = v.get_new_layer_name(old_name)
|
||||
break
|
||||
return layer_name
|
||||
|
||||
def get_new_param_name(self, old_name: str):
|
||||
names = old_name.split(".")
|
||||
layer_name = self.get_new_layer_name(names[0])
|
||||
assert layer_name, f"can not rename layer {names[0]}"
|
||||
names[0] = layer_name
|
||||
return ".".join(names)
|
||||
|
||||
|
||||
class PipeLineModelAdaptor:
|
||||
def __init__(
|
||||
self,
|
||||
src_parallel_config: ParallelConfig,
|
||||
dst_parallel_config: ParallelConfig,
|
||||
transformer_layer_num: int,
|
||||
segment_method: str = "layer",
|
||||
):
|
||||
self._src_parallel_config = src_parallel_config
|
||||
self._dst_parallel_config = dst_parallel_config
|
||||
self._transformer_layer_num = transformer_layer_num
|
||||
self._segment_method = segment_method
|
||||
|
||||
def apply(self, src_model_path: str, dst_model_path: str):
|
||||
for i in range(self._src_parallel_config.mp):
|
||||
for j in range(self._src_parallel_config.sharding):
|
||||
# TODO(liuzhenhai): use multiple process
|
||||
layers = []
|
||||
|
||||
# 1、extract layers in the same pp group
|
||||
group = self._src_parallel_config.pipe_parallel_group(i, j)
|
||||
src_dirs = [
|
||||
"{}/mp_{:0>2d}_sharding_{:0>2d}_pp_{:0>2d}".format(
|
||||
src_model_path, *e
|
||||
)
|
||||
for e in group
|
||||
]
|
||||
# first rank extract shared layer
|
||||
with_shared = True
|
||||
for dir in src_dirs:
|
||||
print(f"extract layer params in dir {dir}")
|
||||
layers.extend(self.extract_layers(dir, with_shared))
|
||||
with_shared = False
|
||||
# 2、sort and unique layers
|
||||
layers = self.sort_layers(layers)
|
||||
|
||||
# 3、resplit layers among pp group according new pp config
|
||||
layer_segments = self.segment_layers(
|
||||
layers, self._dst_parallel_config, self._segment_method
|
||||
)
|
||||
dst_group = self._dst_parallel_config.pipe_parallel_group(i, j)
|
||||
dst_dirs = [
|
||||
"{}/mp_{:0>2d}_sharding_{:0>2d}_pp_{:0>2d}".format(
|
||||
dst_model_path, *e
|
||||
)
|
||||
for e in dst_group
|
||||
]
|
||||
|
||||
# 4、merge layers belonging to the same node
|
||||
for layer_segment, dir_ in zip(layer_segments, dst_dirs):
|
||||
print(f"merge {len(layer_segment)} layers to {dir_}")
|
||||
self.merge_layers(layer_segment, dir_)
|
||||
|
||||
# 5、copy meta_state.pdopt
|
||||
for src_dir, dst_dir in zip(src_dirs, dst_dirs):
|
||||
shutil.copyfile(
|
||||
f"{src_dir}/meta_state.pdopt",
|
||||
f"{dst_dir}/meta_state.pdopt",
|
||||
)
|
||||
|
||||
def peek_model(self, model_dir: str):
|
||||
for i in range(self._src_parallel_config.mp):
|
||||
for j in range(self._src_parallel_config.sharding):
|
||||
group = self._src_parallel_config.pipe_parallel_group(i, j)
|
||||
dirs = [
|
||||
"{}/mp_{:0>2d}_sharding_{:0>2d}_pp_{:0>2d}".format(
|
||||
model_dir, *e
|
||||
)
|
||||
for e in group
|
||||
]
|
||||
for dir in dirs:
|
||||
print(f"peek partial model in {dir}:")
|
||||
self.peek_partial_model(dir)
|
||||
|
||||
def peek_partial_model(self, sub_dir: str):
|
||||
state_dict = paddle.load(f"{sub_dir}/model.pdparams")
|
||||
for k, v in state_dict.items():
|
||||
print(f"\t{k} -> {v.name}")
|
||||
|
||||
def extract_layers(self, dir: str, with_shared: bool):
|
||||
opt = paddle.load(dir + "/model_state.pdopt")
|
||||
params = paddle.load(dir + "/model.pdparams")
|
||||
shared_layer_parsed = False
|
||||
# tname -> (layer, param_name)
|
||||
tname_to_layer_and_pname = {}
|
||||
for k, v in params.items():
|
||||
layer = self._extract_layer_name(k)
|
||||
assert layer
|
||||
# special treatment for embedding layer, skip duplicated shared layer
|
||||
# shared layer may exist or not, if it exist it share weight with _layers.0
|
||||
# _layers.shared_layers.embed.word_embeddings.weight -> embedding_0.w_0
|
||||
# _layers.shared_layers.embed.position_embeddings.weight -> embedding_1.w_0
|
||||
# _layers.0.word_embeddings.weight -> embedding_0.w_0
|
||||
# _layers.0.position_embeddings.weight -> embedding_1.w_0
|
||||
shared_layer_parsed = shared_layer_parsed or (
|
||||
"_layers.shared_layers" in layer
|
||||
)
|
||||
if (
|
||||
"_layers.shared_layers" not in layer
|
||||
and ("word_embeddings" in k or "position_embeddings" in k)
|
||||
and shared_layer_parsed
|
||||
):
|
||||
continue
|
||||
tname_to_layer_and_pname[v.name] = (layer, k)
|
||||
|
||||
# get opt-> param mapping
|
||||
tensor_names = list(tname_to_layer_and_pname.keys())
|
||||
opt_names = [
|
||||
e for e in opt.keys() if e not in ["master_weights", "LR_Scheduler"]
|
||||
]
|
||||
opt_to_t = self._opt_name_to_tname(tensor_names, opt_names)
|
||||
# gather tensors belonging to one layer together
|
||||
layers = OrderedDict()
|
||||
for k, v in params.items():
|
||||
layer, p = tname_to_layer_and_pname[v.name]
|
||||
if layer not in layers:
|
||||
layers[layer] = {}
|
||||
layers[layer]["opt"] = OrderedDict()
|
||||
layers[layer]["params"] = OrderedDict()
|
||||
layers[layer]["master_weights"] = OrderedDict()
|
||||
layers[layer]["params"][p] = v
|
||||
|
||||
for k, v in opt.items():
|
||||
if k in ["master_weights", "LR_Scheduler"]:
|
||||
continue
|
||||
layer, _ = tname_to_layer_and_pname[opt_to_t[v.name]]
|
||||
layers[layer]["opt"][k] = v
|
||||
|
||||
if "master_weights" in opt:
|
||||
for k, v in opt["master_weights"].items():
|
||||
layer, _ = tname_to_layer_and_pname[k]
|
||||
layers[layer]["master_weights"][k] = v
|
||||
|
||||
if "LR_Scheduler" in opt:
|
||||
for layer in layers:
|
||||
layers[layer]["LR_Scheduler"] = opt["LR_Scheduler"]
|
||||
|
||||
ans = []
|
||||
|
||||
for layer_name, layer in layers.items():
|
||||
# special treatment for embedding layer
|
||||
if (not with_shared) and "shared_layers" in layer_name:
|
||||
continue
|
||||
file_name = f"./tmp_layer_files/{layer_name}.tmp"
|
||||
paddle.save(layer, file_name)
|
||||
ans.append((layer_name, file_name))
|
||||
print(f"save layer {layer_name} to {file_name}")
|
||||
return ans
|
||||
|
||||
def sort_layers(self, layers: list):
|
||||
def priority(elem):
|
||||
layer_name = elem[0]
|
||||
if "shared_layers" in layer_name:
|
||||
return -0.5
|
||||
match = re.search(
|
||||
r"^_layers((\.\d+)+|(\.shared_layers\.[^\.]+))", layer_name
|
||||
)
|
||||
assert match, f"{layer_name} not a valid layer name"
|
||||
return float(match.group(1).lstrip("."))
|
||||
|
||||
# strictly sort layers
|
||||
print("before sort {}".format("|".join([e[0] for e in layers])))
|
||||
layers.sort(key=priority)
|
||||
# unique
|
||||
unique_layers = []
|
||||
for e in layers:
|
||||
if unique_layers and e[0] == unique_layers[-1][0]:
|
||||
continue
|
||||
unique_layers.append(e)
|
||||
print("after sort {} ".format("|".join([e[0] for e in unique_layers])))
|
||||
return unique_layers
|
||||
|
||||
def segment_layers(
|
||||
self,
|
||||
layers: list,
|
||||
config: ParallelConfig,
|
||||
segment_method: str = "layer",
|
||||
):
|
||||
layer_num = len(layers)
|
||||
stage_num = config.pp * config.vpp
|
||||
|
||||
# segment by weights
|
||||
def segment_by_layer():
|
||||
# assume model is of the structure below
|
||||
# embedding -> n*(transformer layer) -> [optional output layer]
|
||||
# segment index
|
||||
weights = [0 for _ in range(layer_num)]
|
||||
non_zero_layers = range(1, layer_num - 1)
|
||||
# input layer is embedding
|
||||
if self._transformer_layer_num:
|
||||
assert self._transformer_layer_num < layer_num
|
||||
non_zero_layers = range(1, 1 + self._transformer_layer_num)
|
||||
for i in non_zero_layers:
|
||||
weights[i] = 1
|
||||
|
||||
part_size = sum(weights) // stage_num
|
||||
result = [0 for _ in range(stage_num + 1)]
|
||||
memory_counter = 0
|
||||
result_idx = 1
|
||||
for idx, weight in enumerate(weights):
|
||||
memory_counter += weight
|
||||
if memory_counter == part_size:
|
||||
result[result_idx] = idx + 1
|
||||
result_idx += 1
|
||||
memory_counter = 0
|
||||
result[stage_num] = layer_num
|
||||
return result
|
||||
|
||||
def segment_uniform():
|
||||
result = [0 for _ in range(stage_num + 1)]
|
||||
part_size = math.floor(layer_num / stage_num)
|
||||
extra_layers = layer_num % stage_num
|
||||
for i in range(1, stage_num):
|
||||
offset = 1 if i > (stage_num - extra_layers) else 0
|
||||
result[i] = int(
|
||||
min(result[i - 1] + part_size + offset, layer_num)
|
||||
)
|
||||
result[stage_num] = layer_num
|
||||
return result
|
||||
|
||||
result = (
|
||||
segment_uniform()
|
||||
if (segment_method == "uniform")
|
||||
else segment_by_layer()
|
||||
)
|
||||
index_segments = [[] for _ in range(config.pp)]
|
||||
for i in range(stage_num):
|
||||
index_segments[i % config.pp].append((result[i], result[i + 1]))
|
||||
|
||||
# name layers
|
||||
segments = [[] for i in range(config.pp)]
|
||||
for i in range(config.pp):
|
||||
for start, end in index_segments[i]:
|
||||
for j in range(start, end):
|
||||
if config.vpp > 1:
|
||||
segments[i].append(
|
||||
(
|
||||
[f"_layers.{start}.{j - start}"],
|
||||
layers[j][1],
|
||||
)
|
||||
)
|
||||
else:
|
||||
segments[i].append(([f"_layers.{j}"], layers[j][1]))
|
||||
|
||||
shared_layer_exist = any(
|
||||
"_layers.shared_layers" in e[0] for e in layers
|
||||
)
|
||||
if shared_layer_exist:
|
||||
# special treatment for shared layer
|
||||
if config.vpp > 1:
|
||||
segments[0] = [
|
||||
([layers[0][0], segments[0][0][0][0]], layers[0][1]),
|
||||
*segments[0][1:],
|
||||
]
|
||||
else:
|
||||
segments[0] = [([layers[0][0]], layers[0][1]), *segments[0][1:]]
|
||||
|
||||
for i in range(1, config.pp):
|
||||
segments[i] = [([layers[0][0]], layers[0][1])] + segments[i]
|
||||
|
||||
for pp_rank, segs in enumerate(segments):
|
||||
print(f"segment result for pp_rank {pp_rank}:")
|
||||
print(50 * "=")
|
||||
for seg in segs:
|
||||
print(f"{seg[0]} => {seg[1]}")
|
||||
return segments
|
||||
|
||||
def merge_layers(self, layers_segment: list, save_dir: str):
|
||||
params = OrderedDict()
|
||||
opt = OrderedDict()
|
||||
master_weights = OrderedDict()
|
||||
renaming_manager = LayerReNamingManager()
|
||||
|
||||
def merge(src, dst, map_k=None):
|
||||
for k, v in src.items():
|
||||
k = map_k(k) if map_k is not None else k
|
||||
dst[k] = v
|
||||
|
||||
lr_scheduler = None
|
||||
for layer_names, file_path in layers_segment:
|
||||
print(f"load {file_path}")
|
||||
layer = paddle.load(file_path)
|
||||
|
||||
def get_param_name_mapper(layer_name):
|
||||
# replace layer name
|
||||
def map_param_name(param_name):
|
||||
layer_pre = self._extract_layer_name(param_name)
|
||||
return layer_name + param_name[len(layer_pre) :]
|
||||
|
||||
return map_param_name
|
||||
|
||||
(
|
||||
layer_params,
|
||||
layer_opt,
|
||||
layer_master_weight,
|
||||
) = self._map_tensor_names(
|
||||
layer["params"],
|
||||
layer["opt"],
|
||||
layer["master_weights"],
|
||||
renaming_manager,
|
||||
)
|
||||
for layer_name in layer_names:
|
||||
merge(layer_params, params, get_param_name_mapper(layer_name))
|
||||
merge(layer_opt, opt)
|
||||
merge(layer_master_weight, master_weights)
|
||||
lr_scheduler = layer["LR_Scheduler"]
|
||||
|
||||
opt = self._pack_opt_state_dict(opt, master_weights, lr_scheduler)
|
||||
paddle.save(params, save_dir + "/model.pdparams")
|
||||
paddle.save(opt, save_dir + "/model_state.pdopt")
|
||||
|
||||
def _pack_opt_state_dict(self, opt, master_weights, lr_scheduler):
|
||||
opt["master_weights"] = master_weights
|
||||
opt["LR_Scheduler"] = lr_scheduler
|
||||
return opt
|
||||
|
||||
def _extract_layer_name(self, param_name: str):
|
||||
match = re.search(
|
||||
r"^_layers((\.\d+)+|(\.shared_layers\.[^\.]+))", param_name
|
||||
)
|
||||
layer_name = ""
|
||||
return "" if (not match) else match.group()
|
||||
|
||||
# map opt names to tensor name
|
||||
def _opt_name_to_tname(self, tensor_names, opt_names):
|
||||
tensor_names = set(tensor_names)
|
||||
all_names = []
|
||||
all_names.extend(list(tensor_names))
|
||||
all_names.extend(opt_names)
|
||||
all_names.sort()
|
||||
pre_t_name = ""
|
||||
opt_to_t = {}
|
||||
for n in all_names:
|
||||
if n in tensor_names:
|
||||
# we get a param
|
||||
pre_t_name = n
|
||||
else:
|
||||
assert pre_t_name
|
||||
opt_to_t[n] = pre_t_name
|
||||
return opt_to_t
|
||||
|
||||
def _map_tensor_names(self, params, opt, master_weights, renaming_manager):
|
||||
opt_renamed = OrderedDict()
|
||||
master_weights_renamed = OrderedDict()
|
||||
# old name to new name
|
||||
t_name_mapping = {}
|
||||
# map tensor names
|
||||
for k, v in params.items():
|
||||
t_name_mapping[v.name] = renaming_manager.get_new_param_name(v.name)
|
||||
v.name = t_name_mapping[v.name]
|
||||
# map opt names
|
||||
opt_to_tname = self._opt_name_to_tname(
|
||||
t_name_mapping.keys(), opt.keys()
|
||||
)
|
||||
for k, v in opt.items():
|
||||
old_t_name = opt_to_tname[k]
|
||||
t_name = t_name_mapping[old_t_name]
|
||||
opt_name = t_name + k[len(old_t_name) :]
|
||||
v.name = opt_name
|
||||
opt_renamed[opt_name] = v
|
||||
|
||||
# map master names
|
||||
for k, v in master_weights.items():
|
||||
t_name = t_name_mapping[k]
|
||||
v.name = t_name + v.name[len(k) :]
|
||||
master_weights_renamed[t_name] = v
|
||||
return (params, opt_renamed, master_weights_renamed)
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(
|
||||
prog='model converter', description='converter a model'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--src_path',
|
||||
type=str,
|
||||
default="./output/epoch_0_step_30",
|
||||
help='path of the model to convert',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--dst_path',
|
||||
type=str,
|
||||
default="./test_adapt",
|
||||
help='path to saved the converted model',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--src_mp',
|
||||
type=int,
|
||||
default=2,
|
||||
help='mp degree of the origin training task that dumped this model',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--src_pp',
|
||||
type=int,
|
||||
default=2,
|
||||
help='pp degree of the origin training task that dumped this model',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--src_vp',
|
||||
type=int,
|
||||
default=2,
|
||||
help='vp degree of the origin training task that dumped this model',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--dst_mp',
|
||||
type=int,
|
||||
default=None,
|
||||
help='mp degree of the origin training task that dumped this model',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--dst_pp',
|
||||
type=int,
|
||||
default=None,
|
||||
help='pp degree of the expected training task that would recover this model',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--dst_vp',
|
||||
type=int,
|
||||
default=2,
|
||||
help='vp degree of the expected training task that would recover this model',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--sharding',
|
||||
type=int,
|
||||
default=1,
|
||||
help=" sharding degree of both the origin training task that dumped this model and the expected training task that would recover this model",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--method',
|
||||
type=str,
|
||||
default="adapt_model",
|
||||
help='vp degree of the expected training task that would recover this model',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--segment_method',
|
||||
type=str,
|
||||
default="layer",
|
||||
help='method to segment layers to pp or vp stages',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--transformer_layer_num',
|
||||
type=int,
|
||||
default=0,
|
||||
help='transformer_layer_num of the model',
|
||||
)
|
||||
# assume model is of the structure below
|
||||
# embedding -> n*[transformer layer] -> optional output layer
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.dst_mp is None:
|
||||
args.dst_mp = args.src_mp
|
||||
if args.dst_pp is None:
|
||||
args.dst_pp = args.src_pp
|
||||
|
||||
assert args.src_mp == args.dst_mp, (
|
||||
f"src mp {args.src_mp} dst mp {args.dst_mp}"
|
||||
)
|
||||
|
||||
assert args.method in [
|
||||
'peek_model',
|
||||
'adapt_model',
|
||||
], "method should be in ['peek_model', 'adapt_model']"
|
||||
assert args.segment_method in [
|
||||
"uniform",
|
||||
"layer",
|
||||
], "segment_method should be 'uniform' or 'layer"
|
||||
|
||||
print(
|
||||
f"adapt model dumped by task with pp degree:{args.src_pp}, vp degree:{args.src_vp}, mp degree:{args.src_mp} to task with pp degree:{args.dst_pp}, vp degree:{args.dst_vp}, mp degree:{args.dst_mp}"
|
||||
)
|
||||
|
||||
return args
|
||||
|
||||
|
||||
def adaptor_from_args(args):
|
||||
src_parallel_config = ParallelConfig(
|
||||
args.src_mp, args.src_pp, args.src_vp, args.sharding
|
||||
)
|
||||
|
||||
dst_parallel_config = ParallelConfig(
|
||||
args.dst_mp, args.dst_pp, args.dst_vp, args.sharding
|
||||
)
|
||||
|
||||
adaptor = PipeLineModelAdaptor(
|
||||
src_parallel_config,
|
||||
dst_parallel_config,
|
||||
args.transformer_layer_num,
|
||||
args.segment_method,
|
||||
)
|
||||
return adaptor
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
adaptor = adaptor_from_args(args)
|
||||
if args.method == "peek_model":
|
||||
adaptor.peek_model(args.dst_path)
|
||||
elif args.method == "adapt_model":
|
||||
adaptor.apply(args.src_path, args.dst_path)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
"""
|
||||
Usage:
|
||||
python pp_parallel_adaptor.py --src_mp xxx --src_path xxx --method \
|
||||
adapt_model/peek_model --dst_path xxx --sharding xxx --segment_method xxx --transformer_layer_num xxx
|
||||
|
||||
for the meaning of a specific arg, please use:
|
||||
python pp_parallel_adaptor.py -h
|
||||
"""
|
||||
main()
|
||||
@@ -0,0 +1,354 @@
|
||||
# Copyright (c) 2020 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.
|
||||
"""Parameter Server utils"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import warnings
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import paddle
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from paddle import Tensor
|
||||
from paddle.distributed.fleet.base.role_maker import RoleMakerBase
|
||||
from paddle.static import Executor, Program
|
||||
|
||||
__all__ = []
|
||||
|
||||
|
||||
class DistributedInfer:
|
||||
"""
|
||||
Utility class for distributed infer of PaddlePaddle.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
main_program: Program | None = None,
|
||||
startup_program: Program | None = None,
|
||||
) -> None:
|
||||
if main_program:
|
||||
self.origin_main_program = main_program.clone()
|
||||
else:
|
||||
self.origin_main_program = (
|
||||
paddle.static.default_main_program().clone()
|
||||
)
|
||||
|
||||
if startup_program:
|
||||
self.origin_startup_program = startup_program
|
||||
else:
|
||||
self.origin_startup_program = (
|
||||
paddle.static.default_startup_program()
|
||||
)
|
||||
self.sparse_table_maps = None
|
||||
|
||||
def init_distributed_infer_env(
|
||||
self,
|
||||
exe: Executor,
|
||||
loss: Tensor,
|
||||
role_maker: RoleMakerBase | None = None,
|
||||
dirname: str | None = None,
|
||||
) -> None:
|
||||
from paddle.distributed import fleet
|
||||
|
||||
if fleet.fleet._runtime_handle is None:
|
||||
fleet.init(role_maker=role_maker)
|
||||
|
||||
fake_optimizer = paddle.optimizer.SGD()
|
||||
strategy = fleet.DistributedStrategy()
|
||||
strategy.a_sync = True
|
||||
optimizer = fleet.distributed_optimizer(
|
||||
fake_optimizer, strategy=strategy
|
||||
)
|
||||
optimizer.minimize(
|
||||
loss, startup_program=self.origin_startup_program
|
||||
)
|
||||
|
||||
if fleet.is_server():
|
||||
fleet.init_server(dirname=dirname)
|
||||
fleet.run_server()
|
||||
else:
|
||||
exe.run(paddle.static.default_startup_program())
|
||||
fleet.init_worker()
|
||||
self._init_dense_params(exe, dirname)
|
||||
global_startup_program = paddle.static.default_startup_program()
|
||||
global_startup_program = self.origin_startup_program
|
||||
global_main_program = paddle.static.default_main_program()
|
||||
global_main_program = self.origin_main_program
|
||||
|
||||
def _get_sparse_table_map(self):
|
||||
from paddle.distributed import fleet
|
||||
|
||||
if self.sparse_table_maps is None:
|
||||
self.sparse_table_maps = {}
|
||||
send_ctx = fleet.fleet._runtime_handle._send_ctx
|
||||
for gradname, ctx in send_ctx.items():
|
||||
if ctx.is_sparse:
|
||||
param = gradname.strip("@GRAD")
|
||||
self.sparse_table_maps[param] = ctx.table_id()
|
||||
else:
|
||||
continue
|
||||
return self.sparse_table_maps
|
||||
|
||||
def _init_dense_params(self, exe=None, dirname=None):
|
||||
sparse_table_maps = self._get_sparse_table_map()
|
||||
|
||||
if dirname is not None and exe is not None:
|
||||
all_persist_vars = [
|
||||
v
|
||||
for v in self.origin_main_program.list_vars()
|
||||
if paddle.static.io.is_persistable(v)
|
||||
]
|
||||
dense_persist_vars = [
|
||||
(v.name, v)
|
||||
for v in all_persist_vars
|
||||
if v.name not in sparse_table_maps
|
||||
]
|
||||
need_load_vars = [
|
||||
v[1]
|
||||
for v in dense_persist_vars
|
||||
if os.path.isfile(os.path.join(dirname, v[0]))
|
||||
]
|
||||
paddle.static.load_vars(
|
||||
exe,
|
||||
dirname,
|
||||
main_program=self.origin_main_program,
|
||||
vars=need_load_vars,
|
||||
)
|
||||
|
||||
def get_dist_infer_program(self) -> Program:
|
||||
varname2tables = self._get_sparse_table_map()
|
||||
convert_program = self._convert_program(
|
||||
self.origin_main_program, varname2tables
|
||||
)
|
||||
return convert_program
|
||||
|
||||
def _convert_program(self, main_program, varname2tables):
|
||||
def distributed_ops_pass(program):
|
||||
SPARSE_OP_TYPE_DICT = {"lookup_table": "W", "lookup_table_v2": "W"}
|
||||
|
||||
def _get_pull_sparse_ops(_program):
|
||||
pull_sparse_ops = {}
|
||||
for op in _program.global_block().ops:
|
||||
if (
|
||||
op.type in SPARSE_OP_TYPE_DICT.keys()
|
||||
and op.attr('remote_prefetch') is True
|
||||
):
|
||||
param_name = op.input(SPARSE_OP_TYPE_DICT[op.type])[0]
|
||||
ops = pull_sparse_ops.get(param_name, [])
|
||||
ops.append(op)
|
||||
pull_sparse_ops[param_name] = ops
|
||||
return pull_sparse_ops
|
||||
|
||||
def _pull_sparse_fuse(_program, pull_sparse_ops):
|
||||
def dag_check_up_and_reorder(program, inputs, outputs):
|
||||
global_block = program.global_block()
|
||||
min_output_index = len(global_block.ops)
|
||||
max_input_index = -1
|
||||
input_indexes = [0] * len(global_block.ops)
|
||||
output_indexes = [0] * len(global_block.ops)
|
||||
for idx, op in enumerate(global_block.ops):
|
||||
for i in range(0, len(op.output_names)):
|
||||
if input_indexes[idx] == 1:
|
||||
break
|
||||
outs = op.output(op.output_names[i])
|
||||
for in_id, in_var in enumerate(inputs):
|
||||
if in_var.name in outs:
|
||||
input_indexes[idx] = 1
|
||||
max_input_index = max(max_input_index, idx)
|
||||
break
|
||||
|
||||
for i in range(0, len(op.input_names)):
|
||||
if output_indexes[idx] == 1:
|
||||
break
|
||||
ins = op.input(op.input_names[i])
|
||||
for out_id, out_var in enumerate(outputs):
|
||||
if out_var.name in ins:
|
||||
output_indexes[idx] = 1
|
||||
min_output_index = min(
|
||||
min_output_index, idx
|
||||
)
|
||||
|
||||
for i in range(len(global_block.ops)):
|
||||
if input_indexes[i] == 1 and output_indexes[i] == 1:
|
||||
warnings.warn(
|
||||
"unable to re-arrange dags order to combine distributed embedding ops because a op both needs embedding table's output as input and produces ids as the same embedding table's input"
|
||||
)
|
||||
return
|
||||
|
||||
if min_output_index < max_input_index:
|
||||
move_ops = []
|
||||
for i in range(
|
||||
min_output_index + 1, len(input_indexes)
|
||||
):
|
||||
if input_indexes[i] == 1:
|
||||
move_ops.append((global_block.ops[i], i))
|
||||
for i, op in enumerate(move_ops):
|
||||
queue = []
|
||||
visited = set()
|
||||
queue.append(op[1])
|
||||
visited.add(op[0])
|
||||
start = 0
|
||||
while start < len(queue):
|
||||
pos = queue[start]
|
||||
op = global_block.ops[pos]
|
||||
op_inputs = []
|
||||
for k in range(0, len(op.input_names)):
|
||||
ins = op.input(op.input_names[k])
|
||||
op_inputs.append(ins)
|
||||
for j in range(
|
||||
pos - 1, min_output_index - 1, -1
|
||||
):
|
||||
op1 = global_block.ops[j]
|
||||
if op1 in visited:
|
||||
continue
|
||||
found = False
|
||||
for k in range(0, len(op1.output_names)):
|
||||
outs = op1.output(op1.output_names[k])
|
||||
for t in range(len(op_inputs)):
|
||||
for y in op_inputs[t]:
|
||||
if y in outs:
|
||||
found = True
|
||||
break
|
||||
if found:
|
||||
break
|
||||
if found:
|
||||
break
|
||||
if found:
|
||||
if output_indexes[j]:
|
||||
warnings.warn(
|
||||
"unable to re-arrange dags order to combine distributed embedding ops"
|
||||
)
|
||||
return
|
||||
queue.append(j)
|
||||
visited.add(global_block.ops[j])
|
||||
start = start + 1
|
||||
|
||||
queue.sort()
|
||||
for index in queue:
|
||||
desc = global_block.desc._insert_op(
|
||||
min_output_index
|
||||
)
|
||||
desc.copy_from(global_block.ops[index].desc)
|
||||
global_block.desc._remove_op(
|
||||
index + 1, index + 2
|
||||
)
|
||||
global_block.ops[index].desc = desc
|
||||
insert_op = global_block.ops.pop(index)
|
||||
input_state = input_indexes.pop(index)
|
||||
output_state = output_indexes.pop(index)
|
||||
global_block.ops.insert(
|
||||
min_output_index, insert_op
|
||||
)
|
||||
input_indexes.insert(
|
||||
min_output_index, input_state
|
||||
)
|
||||
output_indexes.insert(
|
||||
min_output_index, output_state
|
||||
)
|
||||
min_output_index = min_output_index + 1
|
||||
|
||||
assert global_block.desc.op_size() == len(
|
||||
global_block.ops
|
||||
)
|
||||
for i in range(len(global_block.ops)):
|
||||
assert (
|
||||
global_block.desc.op(i)
|
||||
== global_block.ops[i].desc
|
||||
)
|
||||
|
||||
for param, ops in pull_sparse_ops.items():
|
||||
all_ops = program.global_block().ops
|
||||
|
||||
inputs = [
|
||||
program.global_block().vars[op.input("Ids")[0]]
|
||||
for op in ops
|
||||
]
|
||||
|
||||
w = program.global_block().vars[ops[0].input("W")[0]]
|
||||
|
||||
if w.name not in varname2tables.keys():
|
||||
raise ValueError(
|
||||
f"can not find variable {w.name}, please check your configuration"
|
||||
)
|
||||
|
||||
table_id = varname2tables[w.name]
|
||||
|
||||
padding_idx = ops[0].attr("padding_idx")
|
||||
is_distributed = ops[0].attr("is_distributed")
|
||||
op_type = ops[0].type
|
||||
|
||||
outputs = [
|
||||
program.global_block().vars[op.output("Out")[0]]
|
||||
for op in ops
|
||||
]
|
||||
|
||||
dag_check_up_and_reorder(program, inputs, outputs)
|
||||
op_idxs = [all_ops.index(op) for op in ops]
|
||||
|
||||
for idx in op_idxs[::-1]:
|
||||
program.global_block()._remove_op(idx)
|
||||
|
||||
inputs_idxs = [-1] * len(inputs)
|
||||
outputs_idxs = [len(program.global_block().ops) + 1] * len(
|
||||
outputs
|
||||
)
|
||||
|
||||
for idx, op in enumerate(program.global_block().ops):
|
||||
for i in range(0, len(op.output_names)):
|
||||
outs = op.output(op.output_names[i])
|
||||
for in_id, in_var in enumerate(inputs):
|
||||
if in_var.name in outs:
|
||||
inputs_idxs[in_id] = max(
|
||||
idx, inputs_idxs[in_id]
|
||||
)
|
||||
for i in range(0, len(op.input_names)):
|
||||
ins = op.input(op.input_names[i])
|
||||
for out_id, out_var in enumerate(outputs):
|
||||
if out_var.name in ins:
|
||||
outputs_idxs[out_id] = min(
|
||||
idx, outputs_idxs[out_id]
|
||||
)
|
||||
|
||||
if min(outputs_idxs) - max(inputs_idxs) >= 1:
|
||||
distributed_idx = max(inputs_idxs) + 1
|
||||
|
||||
program.global_block()._insert_op(
|
||||
index=distributed_idx,
|
||||
type="distributed_lookup_table",
|
||||
inputs={"Ids": inputs, 'W': w},
|
||||
outputs={"Outputs": outputs},
|
||||
attrs={
|
||||
"is_distributed": is_distributed,
|
||||
"padding_idx": padding_idx,
|
||||
"table_id": table_id,
|
||||
"is_test": True,
|
||||
"lookup_table_version": op_type,
|
||||
},
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
"something wrong with Fleet, submit a issue is recommended"
|
||||
)
|
||||
|
||||
pull_sparse_ops = _get_pull_sparse_ops(program)
|
||||
warnings.warn(
|
||||
"lookup_table will be forced to test mode when use DistributedInfer"
|
||||
)
|
||||
_pull_sparse_fuse(program, pull_sparse_ops)
|
||||
return program
|
||||
|
||||
covert_program = distributed_ops_pass(main_program)
|
||||
return covert_program
|
||||
@@ -0,0 +1,711 @@
|
||||
# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
|
||||
# Copyright 2018 The OpenAI Team Authors and HuggingFace Inc. team.
|
||||
# Copyright (c) 2018, NVIDIA CORPORATION. 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 os
|
||||
|
||||
import paddle
|
||||
from paddle import distributed as dist
|
||||
from paddle.autograd import PyLayer
|
||||
from paddle.base import core
|
||||
from paddle.distributed import fleet
|
||||
from paddle.distributed.fleet.meta_parallel import get_rng_state_tracker
|
||||
from paddle.distributed.fleet.utils.hybrid_parallel_util import (
|
||||
fused_allreduce_gradients_with_group,
|
||||
)
|
||||
from paddle.distributed.flex_checkpoint.dcp.sharded_weight import (
|
||||
build_sharded_state_dict,
|
||||
)
|
||||
from paddle.nn import (
|
||||
Layer,
|
||||
functional as F,
|
||||
)
|
||||
|
||||
from .log_util import logger
|
||||
|
||||
####################################################
|
||||
# #
|
||||
# Distributed Communication Operator #
|
||||
# #
|
||||
####################################################
|
||||
|
||||
|
||||
def scatter(input):
|
||||
hcg = fleet.get_hybrid_communicate_group()
|
||||
group = hcg.get_model_parallel_group()
|
||||
parallelism = group.nranks
|
||||
rank = group.rank
|
||||
seq_len = input.shape[0]
|
||||
assert seq_len % parallelism == 0, (
|
||||
f"Input sequence length {seq_len} can't be divided exactly by sequence parallelism {parallelism}"
|
||||
)
|
||||
interval = seq_len // parallelism
|
||||
input = paddle.slice(
|
||||
input, axes=[0], starts=[interval * rank], ends=[interval * (rank + 1)]
|
||||
)
|
||||
return input
|
||||
|
||||
|
||||
def all_gather(input):
|
||||
hcg = fleet.get_hybrid_communicate_group()
|
||||
group = hcg.get_model_parallel_group()
|
||||
parallelism = group.nranks
|
||||
output_shape = input.shape
|
||||
output_shape[0] = output_shape[0] * parallelism
|
||||
output = paddle.empty(shape=output_shape, dtype=input.dtype)
|
||||
group.process_group.all_gather(input, output).wait()
|
||||
return output
|
||||
|
||||
|
||||
def reduce_scatter(input):
|
||||
hcg = fleet.get_hybrid_communicate_group()
|
||||
group = hcg.get_model_parallel_group()
|
||||
parallelism = group.nranks
|
||||
output_shape = input.shape
|
||||
assert input.shape[0] % parallelism == 0, (
|
||||
f"Input sequence length {input.shape[0]} can't be divided exactly by sequence parallelism {parallelism}"
|
||||
)
|
||||
output_shape[0] = output_shape[0] // parallelism
|
||||
output = paddle.empty(shape=output_shape, dtype=input.dtype)
|
||||
dist.stream.reduce_scatter(
|
||||
output, input, op=dist.ReduceOp.SUM, group=group, sync_op=True
|
||||
)
|
||||
return output
|
||||
|
||||
|
||||
class ScatterOp(PyLayer):
|
||||
# input shape: [s, b, h], n is mp parallelism
|
||||
# after forward shape: [s/n, b, h]
|
||||
@staticmethod
|
||||
def forward(ctx, input):
|
||||
return scatter(input)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad):
|
||||
return all_gather(grad)
|
||||
|
||||
|
||||
class GatherOp(PyLayer):
|
||||
# input shape: [s/n, b, h], n is mp parallelism
|
||||
# after forward shape: [s, b, h]
|
||||
@staticmethod
|
||||
def forward(ctx, input):
|
||||
return all_gather(input)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad):
|
||||
return scatter(grad)
|
||||
|
||||
|
||||
# All gather along the first dim during forward pass
|
||||
# All reduce and scatter along the first dim during backward pass
|
||||
class AllGatherOp(PyLayer):
|
||||
# input shape: [s/n, b, h], n is mp parallelism
|
||||
# after forward shape: [s, b, h]
|
||||
@staticmethod
|
||||
def forward(ctx, input):
|
||||
return all_gather(input)
|
||||
|
||||
# grad shape: [s, b, h], n is mp parallelism
|
||||
# after forward shape: [s/n, b, h]
|
||||
@staticmethod
|
||||
def backward(ctx, grad):
|
||||
return reduce_scatter(grad)
|
||||
|
||||
|
||||
# All reduce and scatter along the first dim during forward pass
|
||||
# All gather along the first dim during backward pass
|
||||
class ReduceScatterOp(PyLayer):
|
||||
# input shape: [s, b, h], n is mp parallelism
|
||||
# after forward shape: [s/n, b, h]
|
||||
@staticmethod
|
||||
def forward(ctx, input):
|
||||
return reduce_scatter(input)
|
||||
|
||||
# grad shape: [s/n, b, h], n is mp parallelism
|
||||
# after forward shape: [s, b, h]
|
||||
@staticmethod
|
||||
def backward(ctx, grad):
|
||||
return all_gather(grad)
|
||||
|
||||
|
||||
###################################################
|
||||
# #
|
||||
# Modified Parallel Linear Operator #
|
||||
# #
|
||||
###################################################
|
||||
|
||||
|
||||
def mark_as_sequence_parallel_parameter(parameter):
|
||||
parameter.sequence_parallel = True
|
||||
|
||||
|
||||
def is_sequence_parallel_parameter(parameter):
|
||||
return getattr(parameter, "sequence_parallel", False)
|
||||
|
||||
|
||||
def create_fused_allreduce_gradient_hook(parameter_list, accumulation_steps):
|
||||
hcg = fleet.get_hybrid_communicate_group()
|
||||
group = hcg.get_model_parallel_group()
|
||||
|
||||
step = [0]
|
||||
accumulation_steps *= len(parameter_list)
|
||||
|
||||
def __impl__(grad):
|
||||
step[0] += 1
|
||||
if step[0] == accumulation_steps:
|
||||
step[0] = 0
|
||||
fused_allreduce_gradients_with_group(
|
||||
parameter_list, group=group, scale=1.0
|
||||
)
|
||||
return grad
|
||||
|
||||
return __impl__
|
||||
|
||||
|
||||
def create_non_fused_allreduce_gradient_hook(param, accumulation_steps):
|
||||
hcg = fleet.get_hybrid_communicate_group()
|
||||
pg = hcg.get_model_parallel_group().process_group
|
||||
step = [0]
|
||||
|
||||
@paddle.autograd.no_grad()
|
||||
def __impl__():
|
||||
step[0] += 1
|
||||
if (step[0] % accumulation_steps) == 0:
|
||||
if hasattr(param, "main_grad"):
|
||||
pg.allreduce(param.main_grad).wait()
|
||||
else:
|
||||
pg.allreduce(param.grad).wait()
|
||||
|
||||
return __impl__
|
||||
|
||||
|
||||
def register_sequence_parallel_allreduce_hooks(
|
||||
model, accumulation_steps, fuse_sequence_parallel_allreduce
|
||||
):
|
||||
if accumulation_steps <= 0 or not paddle.distributed.is_initialized():
|
||||
return
|
||||
|
||||
mp_group = fleet.get_hybrid_communicate_group().get_model_parallel_group()
|
||||
if mp_group.nranks <= 1:
|
||||
return
|
||||
|
||||
params = []
|
||||
for p in model.parameters():
|
||||
if is_sequence_parallel_parameter(p) and not p.stop_gradient:
|
||||
params.append(p)
|
||||
|
||||
if fuse_sequence_parallel_allreduce:
|
||||
hook = create_fused_allreduce_gradient_hook(params, accumulation_steps)
|
||||
for p in params:
|
||||
p._register_backward_hook(hook)
|
||||
else:
|
||||
for p in params:
|
||||
hook = create_non_fused_allreduce_gradient_hook(
|
||||
p, accumulation_steps
|
||||
)
|
||||
p._register_backward_hook(hook)
|
||||
|
||||
|
||||
def is_fused_matmul_bias_supported():
|
||||
if (
|
||||
paddle.is_compiled_with_cuda()
|
||||
and not paddle.is_compiled_with_rocm()
|
||||
or paddle.is_compiled_with_xpu()
|
||||
):
|
||||
return hasattr(core.eager.ops.legacy, "fused_gemm_epilogue")
|
||||
else:
|
||||
return False
|
||||
|
||||
|
||||
def is_fused_linear_param_grad_add_supported():
|
||||
if (
|
||||
paddle.is_compiled_with_cuda() and not paddle.is_compiled_with_rocm()
|
||||
) or paddle.is_compiled_with_xpu():
|
||||
return hasattr(paddle._C_ops, 'fused_linear_param_grad_add')
|
||||
else:
|
||||
return False
|
||||
|
||||
|
||||
_raise_cuda_env_unset_warning_for_sp = True
|
||||
|
||||
|
||||
def _check_environment_for_overlap():
|
||||
if int(os.getenv("CUDA_DEVICE_MAX_CONNECTIONS", "0")) != 1:
|
||||
global _raise_cuda_env_unset_warning_for_sp
|
||||
if _raise_cuda_env_unset_warning_for_sp:
|
||||
logger.warning(
|
||||
"You set mp_async_allreduce=True or recompute_allgather=True, but you forget to 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."
|
||||
)
|
||||
_raise_cuda_env_unset_warning_for_sp = False
|
||||
|
||||
# Using small operation to preempt GPU SMs for all_gather or reduce_scatter to achieve overlap.
|
||||
tmp = paddle.ones([512])
|
||||
|
||||
|
||||
class SPInnerOverlapLinear(paddle.autograd.PyLayer):
|
||||
@staticmethod
|
||||
def forward(
|
||||
ctx,
|
||||
x,
|
||||
weight,
|
||||
bias,
|
||||
fuse_matmul_bias,
|
||||
recompute_allgather,
|
||||
mp_fused_linear_param_grad_add,
|
||||
model_parallel_group,
|
||||
):
|
||||
ctx.recompute_allgather = recompute_allgather
|
||||
ctx.mp_fused_linear_param_grad_add = mp_fused_linear_param_grad_add
|
||||
ctx.model_parallel_group = model_parallel_group
|
||||
|
||||
world_size = model_parallel_group.nranks
|
||||
input_parallel = all_gather(x)
|
||||
|
||||
if not recompute_allgather:
|
||||
ctx.save_for_backward(x, weight, bias, input_parallel)
|
||||
else:
|
||||
ctx.save_for_backward(x, weight, bias)
|
||||
|
||||
if not fuse_matmul_bias:
|
||||
output = paddle._C_ops.linear(input_parallel, weight, bias)
|
||||
else:
|
||||
output = paddle._legacy_C_ops.fused_gemm_epilogue(
|
||||
input_parallel, weight, bias
|
||||
)
|
||||
return output
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, dy):
|
||||
group = ctx.model_parallel_group
|
||||
parallelism = group.nranks
|
||||
|
||||
if not ctx.recompute_allgather:
|
||||
x, weight, bias, input_parallel = ctx.saved_tensor()
|
||||
else:
|
||||
x, weight, bias = ctx.saved_tensor()
|
||||
|
||||
# all-gather x
|
||||
input_parallel_shape = x.shape
|
||||
input_parallel_shape[0] = input_parallel_shape[0] * parallelism
|
||||
input_parallel = paddle.empty(
|
||||
shape=input_parallel_shape, dtype=x.dtype
|
||||
)
|
||||
allgather_task = dist.all_gather(
|
||||
input_parallel, x, group=group, sync_op=False
|
||||
)
|
||||
|
||||
# compute dx
|
||||
_check_environment_for_overlap()
|
||||
if dy.dtype == weight.dtype:
|
||||
dinput_parallel = paddle.matmul(dy, weight, transpose_y=True)
|
||||
else:
|
||||
dinput_parallel = paddle.matmul(
|
||||
dy, paddle.cast(weight, dtype=dy.dtype), transpose_y=True
|
||||
)
|
||||
|
||||
assert dinput_parallel.shape[0] % parallelism == 0, (
|
||||
f"Input sequence length {dinput_parallel.shape[0]} can't be divided exactly by sequence parallelism {parallelism}"
|
||||
)
|
||||
|
||||
if ctx.recompute_allgather:
|
||||
# wait the finish of all-gather of x
|
||||
allgather_task.wait()
|
||||
|
||||
# reduce-scatter dx
|
||||
dx_shape = dinput_parallel.shape
|
||||
dx_shape[0] = dx_shape[0] // parallelism
|
||||
dx = paddle.empty(shape=dx_shape, dtype=dinput_parallel.dtype)
|
||||
task = dist.stream.reduce_scatter(
|
||||
dx,
|
||||
dinput_parallel,
|
||||
op=dist.ReduceOp.SUM,
|
||||
group=group,
|
||||
sync_op=False,
|
||||
)
|
||||
|
||||
# compute dw and dbias
|
||||
_check_environment_for_overlap()
|
||||
if ctx.mp_fused_linear_param_grad_add:
|
||||
if not is_fused_linear_param_grad_add_supported():
|
||||
raise NotImplementedError(
|
||||
"You set mp_fused_linear_param_grad_add=True, "
|
||||
"however, the paddle you are using not support this operation. "
|
||||
"Please unset fused_linear_param_grad_add or use paddle compiled "
|
||||
"with cuda 11.6 or higher."
|
||||
)
|
||||
if bias is None:
|
||||
if hasattr(weight, "main_grad"):
|
||||
(
|
||||
weight.main_grad,
|
||||
_,
|
||||
) = paddle._C_ops.fused_linear_param_grad_add(
|
||||
input_parallel, dy, weight.main_grad, None, True, False
|
||||
)
|
||||
task.wait()
|
||||
return dx, None
|
||||
else:
|
||||
if weight.grad is not None:
|
||||
(
|
||||
weight.grad,
|
||||
_,
|
||||
) = paddle._C_ops.fused_linear_param_grad_add(
|
||||
input_parallel, dy, weight.grad, None, False, False
|
||||
)
|
||||
task.wait()
|
||||
return dx, None
|
||||
else:
|
||||
(
|
||||
dw,
|
||||
_,
|
||||
) = paddle._C_ops.fused_linear_param_grad_add(
|
||||
input_parallel, dy, None, None, False, False
|
||||
)
|
||||
task.wait()
|
||||
return dx, dw
|
||||
|
||||
if hasattr(weight, "main_grad") and hasattr(bias, "main_grad"):
|
||||
(
|
||||
weight.main_grad,
|
||||
bias.main_grad,
|
||||
) = paddle._C_ops.fused_linear_param_grad_add(
|
||||
input_parallel,
|
||||
dy,
|
||||
weight.main_grad,
|
||||
bias.main_grad,
|
||||
True,
|
||||
True,
|
||||
)
|
||||
task.wait()
|
||||
return dx, None, None
|
||||
else:
|
||||
if weight.grad is not None:
|
||||
assert bias.grad is not None
|
||||
(
|
||||
weight.grad,
|
||||
bias.grad,
|
||||
) = paddle._C_ops.fused_linear_param_grad_add(
|
||||
input_parallel, dy, weight.grad, bias.grad, False, True
|
||||
)
|
||||
task.wait()
|
||||
return dx, None, None
|
||||
else:
|
||||
# When main_grad is not enabled and gradient_accumulation is used, the grad is not initialized for the first acc step.
|
||||
(
|
||||
dw,
|
||||
dbias,
|
||||
) = paddle._C_ops.fused_linear_param_grad_add(
|
||||
input_parallel, dy, None, None, False, True
|
||||
)
|
||||
task.wait()
|
||||
return dx, dw, dbias
|
||||
else:
|
||||
dy = dy.reshape([-1, dy.shape[-1]])
|
||||
dw = paddle.matmul(
|
||||
input_parallel.reshape([-1, input_parallel.shape[-1]]),
|
||||
dy,
|
||||
transpose_x=True,
|
||||
)
|
||||
if bias is None:
|
||||
task.wait()
|
||||
return dx, dw
|
||||
else:
|
||||
dbias = paddle.sum(dy, axis=0)
|
||||
task.wait()
|
||||
return dx, dw, dbias
|
||||
|
||||
|
||||
class ColumnSequenceParallelLinear(Layer):
|
||||
def __init__(
|
||||
self,
|
||||
in_features,
|
||||
out_features,
|
||||
weight_attr=None,
|
||||
has_bias=None,
|
||||
gather_output=True,
|
||||
fuse_matmul_bias=False,
|
||||
mp_group=None,
|
||||
name=None,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
hcg = fleet.get_hybrid_communicate_group()
|
||||
self.model_parallel_group = (
|
||||
hcg.get_model_parallel_group() if mp_group is None else mp_group
|
||||
)
|
||||
self.world_size = (
|
||||
hcg.get_model_parallel_group().nranks
|
||||
if mp_group is None
|
||||
else mp_group.nranks
|
||||
)
|
||||
assert self.world_size > 1, (
|
||||
"tensor parallel degree must be greater than 1 in sequence parallel"
|
||||
)
|
||||
|
||||
self._name = name
|
||||
self.is_mp = self.world_size > 1
|
||||
assert gather_output is False, (
|
||||
"If sequence_parallel is True, gather_output is False"
|
||||
)
|
||||
|
||||
self.gather_output = gather_output
|
||||
assert out_features % self.world_size == 0, (
|
||||
f"Number of column of the weight for linear ({out_features}) must be"
|
||||
f" divisible by model parallel size ({self.world_size})"
|
||||
)
|
||||
self.output_size_per_partition = out_features // self.world_size
|
||||
|
||||
self._weight_attr = weight_attr
|
||||
self._dtype = self._helper.get_default_dtype()
|
||||
|
||||
if self.is_mp and paddle.in_dynamic_mode():
|
||||
with get_rng_state_tracker().rng_state():
|
||||
self.weight = self.create_parameter(
|
||||
shape=[in_features, self.output_size_per_partition],
|
||||
attr=self._weight_attr,
|
||||
dtype=self._dtype,
|
||||
is_bias=False,
|
||||
)
|
||||
else:
|
||||
self.weight = self.create_parameter(
|
||||
shape=[in_features, self.output_size_per_partition],
|
||||
attr=self._weight_attr,
|
||||
dtype=self._dtype,
|
||||
is_bias=False,
|
||||
)
|
||||
|
||||
self.weight.is_distributed = True if self.is_mp else False
|
||||
self.fuse_matmul_bias = fuse_matmul_bias
|
||||
|
||||
if has_bias:
|
||||
# initialize bias to zero like Megatron
|
||||
self.bias = self.create_parameter(
|
||||
shape=[self.output_size_per_partition],
|
||||
attr=paddle.nn.initializer.Constant(value=0.0),
|
||||
dtype=self._dtype,
|
||||
is_bias=True,
|
||||
)
|
||||
self.bias.is_distributed = True if self.is_mp else False
|
||||
else:
|
||||
self.bias = None
|
||||
|
||||
if self.weight.is_distributed:
|
||||
self.weight.split_axis = 1
|
||||
|
||||
if has_bias and self.bias.is_distributed:
|
||||
self.bias.split_axis = 0
|
||||
|
||||
self.linear = F.linear
|
||||
|
||||
if fuse_matmul_bias:
|
||||
if not is_fused_matmul_bias_supported():
|
||||
raise NotImplementedError(
|
||||
"You set fuse_matmul_bias=True in ColumnSequenceParallelLinear, "
|
||||
"however, the paddle you are using not support this operation. "
|
||||
"Please set fuse_matmul_bias=False or use paddle compiled "
|
||||
"with cuda 11.6 or higher, or use xpu version."
|
||||
)
|
||||
from paddle.incubate.nn.functional import fused_linear
|
||||
|
||||
self.linear = fused_linear
|
||||
|
||||
mp_configs = fleet.fleet._user_defined_strategy.hybrid_configs[
|
||||
"mp_configs"
|
||||
]
|
||||
self.mp_async_allreduce = mp_configs.mp_async_allreduce
|
||||
self.sp_async_reduce_scatter = mp_configs.sp_async_reduce_scatter
|
||||
self.recompute_allgather = mp_configs.recompute_allgather
|
||||
|
||||
self.mp_fused_linear_param_grad_add = (
|
||||
self.mp_async_allreduce
|
||||
and mp_configs.mp_fused_linear_param_grad_add
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
# sequence parallel is same as tensor parallel, if sequence parallel is true, input shape is [s, b, h], else input shape is [b, s, h]
|
||||
if self.sp_async_reduce_scatter:
|
||||
output = SPInnerOverlapLinear.apply(
|
||||
x,
|
||||
self.weight,
|
||||
self.bias,
|
||||
self.fuse_matmul_bias,
|
||||
self.recompute_allgather,
|
||||
self.mp_fused_linear_param_grad_add,
|
||||
self.model_parallel_group,
|
||||
)
|
||||
else:
|
||||
if self.is_mp:
|
||||
input_parallel = AllGatherOp.apply(x)
|
||||
else:
|
||||
input_parallel = x
|
||||
output = self.linear(
|
||||
input_parallel, self.weight, self.bias, name=self._name
|
||||
)
|
||||
return output
|
||||
|
||||
def sharded_state_dict(
|
||||
self,
|
||||
structured_name_prefix: str = "",
|
||||
):
|
||||
state_dict = self.state_dict(structured_name_prefix="")
|
||||
return build_sharded_state_dict(
|
||||
state_dict, {"weight": 1, "bias": 0}, structured_name_prefix
|
||||
)
|
||||
|
||||
|
||||
class MPScale(PyLayer):
|
||||
@staticmethod
|
||||
def forward(ctx, x, mp_degree):
|
||||
out = paddle.scale(x, 1.0 / mp_degree)
|
||||
return out
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, dout):
|
||||
return dout
|
||||
|
||||
|
||||
class RowSequenceParallelLinear(Layer):
|
||||
def __init__(
|
||||
self,
|
||||
in_features,
|
||||
out_features,
|
||||
weight_attr=None,
|
||||
has_bias=True,
|
||||
input_is_parallel=False,
|
||||
fuse_matmul_bias=False,
|
||||
mp_group=None,
|
||||
name=None,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.in_features = in_features
|
||||
self.out_features = out_features
|
||||
assert input_is_parallel is True, (
|
||||
"If sequence_parallel is True, input_is_parallel should be true."
|
||||
)
|
||||
|
||||
self.input_is_parallel = input_is_parallel
|
||||
self._weight_attr = weight_attr
|
||||
self._dtype = self._helper.get_default_dtype()
|
||||
self._name = name
|
||||
|
||||
hcg = fleet.get_hybrid_communicate_group()
|
||||
self.model_parallel_group = (
|
||||
hcg.get_model_parallel_group() if mp_group is None else mp_group
|
||||
)
|
||||
self.world_size = (
|
||||
hcg.get_model_parallel_group().nranks
|
||||
if mp_group is None
|
||||
else mp_group.nranks
|
||||
)
|
||||
self.rank = (
|
||||
hcg.get_model_parallel_group().rank
|
||||
if mp_group is None
|
||||
else mp_group.rank
|
||||
)
|
||||
|
||||
self.is_mp = self.world_size > 1
|
||||
assert in_features % self.world_size == 0, (
|
||||
f"Number of row of the weight for linear ({in_features}) must be"
|
||||
f" divisible by model parallel size ({self.world_size})"
|
||||
)
|
||||
|
||||
self.input_size_per_partition = in_features // self.world_size
|
||||
|
||||
if self.is_mp and paddle.in_dynamic_mode():
|
||||
with get_rng_state_tracker().rng_state():
|
||||
self.weight = self.create_parameter(
|
||||
shape=[self.input_size_per_partition, self.out_features],
|
||||
attr=self._weight_attr,
|
||||
dtype=self._dtype,
|
||||
is_bias=False,
|
||||
)
|
||||
else:
|
||||
self.weight = self.create_parameter(
|
||||
shape=[self.input_size_per_partition, self.out_features],
|
||||
attr=self._weight_attr,
|
||||
dtype=self._dtype,
|
||||
is_bias=False,
|
||||
)
|
||||
|
||||
self.weight.is_distributed = True if self.is_mp else False
|
||||
|
||||
# if sequence parallel is true,
|
||||
# register hook to all_reduce gradient of weight and bias
|
||||
if has_bias:
|
||||
self.bias = self.create_parameter(
|
||||
shape=[self.out_features],
|
||||
attr=paddle.nn.initializer.Constant(value=0.0),
|
||||
dtype=self._dtype,
|
||||
is_bias=True,
|
||||
)
|
||||
if self.is_mp:
|
||||
mark_as_sequence_parallel_parameter(self.bias)
|
||||
else:
|
||||
self.bias = None
|
||||
|
||||
if self.weight.is_distributed:
|
||||
self.weight.split_axis = 0
|
||||
|
||||
self.linear = F.linear
|
||||
|
||||
self.mp_scale = None
|
||||
if fuse_matmul_bias:
|
||||
if not is_fused_matmul_bias_supported():
|
||||
raise NotImplementedError(
|
||||
"You set fuse_matmul_bias=True in RowParallelLinear, "
|
||||
"however, the paddle you are using not support this operation. "
|
||||
"Please set fuse_matmul_bias=False or use paddle compiled "
|
||||
"with cuda 11.6 or higher."
|
||||
)
|
||||
from paddle.incubate.nn.functional import fused_linear
|
||||
|
||||
self.linear = fused_linear
|
||||
if self.is_mp and has_bias:
|
||||
self.mp_scale = MPScale.apply
|
||||
|
||||
def forward(self, x):
|
||||
input_parallel = x
|
||||
if self.is_mp:
|
||||
if self.mp_scale is not None:
|
||||
bias = self.mp_scale(self.bias, self.world_size)
|
||||
else:
|
||||
bias = None
|
||||
output_parallel = self.linear(
|
||||
input_parallel, self.weight, bias, name=self._name
|
||||
)
|
||||
output_ = ReduceScatterOp.apply(output_parallel)
|
||||
# if self.bias is not none, sequence parallel will use
|
||||
# register_hook to all_reduce self.bias
|
||||
if bias is None and self.bias is not None:
|
||||
output = output_ + self.bias
|
||||
else:
|
||||
output = output_
|
||||
else:
|
||||
output = self.linear(
|
||||
input_parallel, self.weight, self.bias, name=self._name
|
||||
)
|
||||
return output
|
||||
|
||||
def sharded_state_dict(
|
||||
self,
|
||||
structured_name_prefix: str = "",
|
||||
):
|
||||
state_dict = self.state_dict(structured_name_prefix="")
|
||||
return build_sharded_state_dict(
|
||||
state_dict, {"weight": 0}, structured_name_prefix
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,351 @@
|
||||
# Copyright (c) 2023 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 logging
|
||||
|
||||
import paddle
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
formatter = logging.Formatter(
|
||||
fmt='%(asctime)s %(levelname)-8s %(message)s', datefmt='%Y-%m-%d %H:%M:%S'
|
||||
)
|
||||
ch = logging.StreamHandler()
|
||||
ch.setFormatter(formatter)
|
||||
logger.addHandler(ch)
|
||||
|
||||
from paddle.base import core
|
||||
from paddle.distributed.fleet.meta_optimizers.common import OP_ROLE_KEY
|
||||
from paddle.static import Parameter
|
||||
|
||||
_supported_optimizer_type = [
|
||||
"adam",
|
||||
"adamax",
|
||||
"adamw",
|
||||
"decayed_adagrad",
|
||||
"momentum",
|
||||
"dgc_momentum",
|
||||
"lars_momentum",
|
||||
"merged_momentum",
|
||||
"lamb",
|
||||
"sgd",
|
||||
]
|
||||
|
||||
|
||||
def tensor_parallel_sync_filter_fn(
|
||||
param, pos_emb=True, layer_norm=True, bias=True
|
||||
):
|
||||
"""
|
||||
Layer filter function for tensor parallelism transformer.
|
||||
|
||||
In tensor parallelism of transformer like model, there is 4 kind of param
|
||||
that are supposed to be the same in all tensor parallel peers:
|
||||
* position embedding
|
||||
* scale of layer norm
|
||||
* bias of layer norm
|
||||
* bias of row parallel linear
|
||||
|
||||
set corresponding input args to select specific layers.
|
||||
NOTE adopting the param name pattern for different transformer blocks.
|
||||
"""
|
||||
p_name = param.name
|
||||
if pos_emb and p_name.startswith("pos_embedding"):
|
||||
return True
|
||||
|
||||
elif layer_norm and p_name.endswith("_layer_norm_bias"):
|
||||
return True
|
||||
|
||||
elif layer_norm and p_name.endswith("_layer_norm_scale"):
|
||||
return True
|
||||
|
||||
elif bias and ".b_" in p_name and (param.is_distributed is False):
|
||||
return True
|
||||
|
||||
else:
|
||||
return False
|
||||
|
||||
|
||||
def resolute_tensor_parallel_ring_id(program):
|
||||
ops = program.global_block().ops
|
||||
ring_id = None
|
||||
|
||||
for op in ops:
|
||||
if op.type == "c_identity":
|
||||
if ring_id is None:
|
||||
ring_id = int(op.attr("ring_id"))
|
||||
else:
|
||||
assert ring_id == int(op.attr("ring_id")), (
|
||||
"Found two different ring_id for Tensor Parallel: ring_id={} and ring_id={}.".format(
|
||||
ring_id, int(op.attr("ring_id"))
|
||||
)
|
||||
)
|
||||
assert ring_id is not None, "Could NOT found ring_id for Tensor Parallel."
|
||||
|
||||
return ring_id
|
||||
|
||||
|
||||
def copy_parameters(block_, params):
|
||||
for param in params:
|
||||
new_p = Parameter(
|
||||
block=block_,
|
||||
shape=param.shape,
|
||||
dtype=param.dtype,
|
||||
type=param.type,
|
||||
lod_level=(
|
||||
param.lod_level
|
||||
if param.type == core.VarDesc.VarType.DENSE_TENSOR
|
||||
else None
|
||||
),
|
||||
stop_gradient=param.stop_gradient,
|
||||
trainable=param.trainable,
|
||||
optimize_attr=param.optimize_attr,
|
||||
regularizer=param.regularizer,
|
||||
error_clip=param.error_clip,
|
||||
name=param.name,
|
||||
)
|
||||
assert param.is_distributed is False, (
|
||||
f"Try to sync Distributed Parameter: {param}"
|
||||
)
|
||||
new_p.is_distributed = False
|
||||
|
||||
block_.vars[new_p.name] = new_p
|
||||
|
||||
|
||||
def insert_sync_op(
|
||||
block, idx, tp_degree, sync_mode, sync_ring_id, src_rank, varname, op_role
|
||||
):
|
||||
if sync_mode == "broadcast":
|
||||
block._insert_op_without_sync(
|
||||
idx,
|
||||
type='broadcast',
|
||||
inputs={'x': varname},
|
||||
outputs={'out': varname},
|
||||
attrs={
|
||||
'ring_id': sync_ring_id,
|
||||
'root': src_rank,
|
||||
OP_ROLE_KEY: op_role,
|
||||
},
|
||||
)
|
||||
|
||||
elif sync_mode == "average":
|
||||
block._insert_op_without_sync(
|
||||
idx,
|
||||
type='scale',
|
||||
inputs={'X': varname},
|
||||
outputs={'Out': varname},
|
||||
attrs={'scale': 1.0 / tp_degree, OP_ROLE_KEY: op_role},
|
||||
)
|
||||
block._insert_op_without_sync(
|
||||
idx,
|
||||
type='all_reduce',
|
||||
inputs={'x': varname},
|
||||
outputs={'out': varname},
|
||||
attrs={
|
||||
'ring_id': sync_ring_id,
|
||||
'reduce_type': paddle.distributed.ReduceOp.SUM,
|
||||
OP_ROLE_KEY: op_role,
|
||||
},
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f'Sync mode of [{sync_mode}] is NOT supported.'
|
||||
)
|
||||
|
||||
|
||||
def insert_synchronization(
|
||||
block,
|
||||
params_to_sync,
|
||||
tp_degree,
|
||||
sync_ring_id,
|
||||
sync_param,
|
||||
sync_grad,
|
||||
sync_moment,
|
||||
sync_mode,
|
||||
src_rank,
|
||||
):
|
||||
unsync_param_names = [p.name for p in params_to_sync]
|
||||
|
||||
for idx, op in reversed(list(enumerate(block.ops))):
|
||||
if op.type in _supported_optimizer_type:
|
||||
assert "Param" in op.input_names
|
||||
assert len(op.input("Param")) == 1
|
||||
param_name = op.input("Param")[0]
|
||||
op_role = op.attr(OP_ROLE_KEY)
|
||||
|
||||
if param_name in unsync_param_names:
|
||||
unsync_param_names.remove(param_name)
|
||||
|
||||
# Param sync after opt
|
||||
if sync_param:
|
||||
assert (
|
||||
"ParamOut" in op.output_names
|
||||
and op.output("ParamOut")[0] == param_name
|
||||
)
|
||||
insert_sync_op(
|
||||
block,
|
||||
idx + 1,
|
||||
tp_degree,
|
||||
sync_mode,
|
||||
sync_ring_id,
|
||||
src_rank,
|
||||
param_name,
|
||||
op_role,
|
||||
)
|
||||
|
||||
if (
|
||||
"MasterParamOut" in op.output_names
|
||||
and len(op.output("MasterParamOut")) == 1
|
||||
):
|
||||
sync_var = op.output("MasterParamOut")[0]
|
||||
insert_sync_op(
|
||||
block,
|
||||
idx + 1,
|
||||
tp_degree,
|
||||
sync_mode,
|
||||
sync_ring_id,
|
||||
src_rank,
|
||||
sync_var,
|
||||
op_role,
|
||||
)
|
||||
|
||||
# Moment sync after opt
|
||||
if sync_moment:
|
||||
if (
|
||||
"Moment1Out" in op.output_names
|
||||
and len(op.output("Moment1Out")) == 1
|
||||
):
|
||||
sync_var = op.output("Moment1Out")[0]
|
||||
insert_sync_op(
|
||||
block,
|
||||
idx + 1,
|
||||
tp_degree,
|
||||
sync_mode,
|
||||
sync_ring_id,
|
||||
src_rank,
|
||||
sync_var,
|
||||
op_role,
|
||||
)
|
||||
|
||||
if (
|
||||
"Moment2Out" in op.output_names
|
||||
and len(op.output("Moment2Out")) == 1
|
||||
):
|
||||
sync_var = op.output("Moment2Out")[0]
|
||||
insert_sync_op(
|
||||
block,
|
||||
idx + 1,
|
||||
tp_degree,
|
||||
sync_mode,
|
||||
sync_ring_id,
|
||||
src_rank,
|
||||
sync_var,
|
||||
op_role,
|
||||
)
|
||||
|
||||
# Grad sync before opt
|
||||
if sync_grad:
|
||||
assert (
|
||||
"Grad" in op.input_names and len(op.input("Grad")) == 1
|
||||
)
|
||||
sync_var = op.input("Grad")[0]
|
||||
insert_sync_op(
|
||||
block,
|
||||
idx,
|
||||
tp_degree,
|
||||
sync_mode,
|
||||
sync_ring_id,
|
||||
src_rank,
|
||||
sync_var,
|
||||
op_role,
|
||||
)
|
||||
|
||||
assert len(unsync_param_names) == 0, (
|
||||
f"The following param is unsync by some error: {unsync_param_names}"
|
||||
)
|
||||
|
||||
|
||||
def add_extra_synchronization(
|
||||
program,
|
||||
params_filter_fn=tensor_parallel_sync_filter_fn,
|
||||
tp_degree=8,
|
||||
sync_mode="broadcast",
|
||||
sync_param=True,
|
||||
sync_grad=False,
|
||||
sync_moment=False,
|
||||
src_rank=0,
|
||||
sync_ring_id=None,
|
||||
):
|
||||
"""
|
||||
Inplace add extra synchronization for input program.
|
||||
|
||||
program(Paddle.Program): distributed train program.
|
||||
|
||||
params_filter_fn(callable): function to filter out parameter for synchronization.
|
||||
|
||||
sync_mode(string): select from
|
||||
"broadcast": parameter is sync by broadcasted from 'src_rank' to all other ranks.
|
||||
"average": parameter is sync by average among all ranks
|
||||
|
||||
src_rank(int): the src used in broadcast sync_mode.
|
||||
|
||||
sync_param(bool): extra synchronize parameters.
|
||||
|
||||
sync_grad(bool): extra synchronize gradients.
|
||||
|
||||
sync_grad(bool): extra synchronize optimizer momentum.
|
||||
|
||||
sync_ring_id(int): communicator id use for synchronization, if it is None, use the ring_id of tensor parallel.
|
||||
"""
|
||||
|
||||
logger.info("Constructing Extra Parameter Synchronization.")
|
||||
logger.info(
|
||||
f"Tensor Parallel Degree: {tp_degree}, Synchronization mode: {sync_mode}"
|
||||
)
|
||||
|
||||
# adopt for pipeline opt
|
||||
if program._pipeline_opt is not None:
|
||||
assert program._pipeline_opt['section_program'] is not None, (
|
||||
"Pipeline is enable but section_program is None"
|
||||
)
|
||||
program = program._pipeline_opt['section_program']
|
||||
|
||||
# step1: collect the param that need to be sync
|
||||
params_to_sync = []
|
||||
# TODO support multiple blocks with different parameter.
|
||||
all_params = program.global_block().all_parameters()
|
||||
for param in all_params:
|
||||
if params_filter_fn(param):
|
||||
params_to_sync.append(param)
|
||||
logger.info(
|
||||
"The following param are going to be synchronization everytime the optimizer update phase of the program is run: "
|
||||
)
|
||||
logger.info([p.name for p in params_to_sync])
|
||||
|
||||
# step2: resolute synchronization communicator group (ring_id)
|
||||
if sync_ring_id is None:
|
||||
sync_ring_id = resolute_tensor_parallel_ring_id(program)
|
||||
|
||||
# step3: insert synchronization
|
||||
# TODO support gradient merge with different update block
|
||||
block = program.global_block()
|
||||
insert_synchronization(
|
||||
block,
|
||||
params_to_sync,
|
||||
tp_degree,
|
||||
sync_ring_id,
|
||||
sync_param,
|
||||
sync_grad,
|
||||
sync_moment,
|
||||
sync_mode,
|
||||
src_rank,
|
||||
)
|
||||
@@ -0,0 +1,134 @@
|
||||
# Copyright (c) 2023 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 os
|
||||
import time
|
||||
|
||||
import paddle
|
||||
from paddle.base import core
|
||||
|
||||
_GLOBAL_TIMERS = None
|
||||
|
||||
|
||||
def is_timer_initialized():
|
||||
return _GLOBAL_TIMERS is not None
|
||||
|
||||
|
||||
def _ensure_var_is_not_initialized(var, name):
|
||||
"""Make sure the input variable is not None."""
|
||||
assert var is None, f"{name} has been already initialized."
|
||||
|
||||
|
||||
def _ensure_var_is_initialized(var, name):
|
||||
"""Make sure the input variable is not None."""
|
||||
assert var is not None, f"{name} is not initialized."
|
||||
|
||||
|
||||
def get_timers():
|
||||
_ensure_var_is_initialized(_GLOBAL_TIMERS, "timers")
|
||||
return _GLOBAL_TIMERS
|
||||
|
||||
|
||||
def set_timers():
|
||||
"""Initialize timers."""
|
||||
global _GLOBAL_TIMERS
|
||||
_ensure_var_is_not_initialized(_GLOBAL_TIMERS, "timers")
|
||||
_GLOBAL_TIMERS = Timers()
|
||||
|
||||
|
||||
class _Timer:
|
||||
"""Timer."""
|
||||
|
||||
def __init__(self, name):
|
||||
self.name = name
|
||||
self.elapsed_ = 0.0
|
||||
self.started_ = False
|
||||
self.start_time = time.time()
|
||||
|
||||
def start(self):
|
||||
"""Start the timer."""
|
||||
assert not self.started_, "timer has already started"
|
||||
paddle.device.cuda.synchronize()
|
||||
self.start_time = time.time()
|
||||
self.started_ = True
|
||||
|
||||
def stop(self):
|
||||
"""Stop the timers."""
|
||||
assert self.started_, "timer is not started."
|
||||
paddle.device.cuda.synchronize()
|
||||
self.elapsed_ += time.time() - self.start_time
|
||||
self.started_ = False
|
||||
|
||||
def reset(self):
|
||||
"""Reset timer."""
|
||||
self.elapsed_ = 0.0
|
||||
self.started_ = False
|
||||
|
||||
def elapsed(self, reset=True):
|
||||
"""Calculate the elapsed time."""
|
||||
started_ = self.started_
|
||||
# If the timing in progress, end it first.
|
||||
if self.started_:
|
||||
self.stop()
|
||||
# Get the elapsed time.
|
||||
elapsed_ = self.elapsed_
|
||||
# Reset the elapsed time
|
||||
if reset:
|
||||
self.reset()
|
||||
# If timing was in progress, set it back.
|
||||
if started_:
|
||||
self.start()
|
||||
return elapsed_
|
||||
|
||||
|
||||
class _GPUEventTimer:
|
||||
"""GPUEventTimer."""
|
||||
|
||||
def __init__(self, name):
|
||||
self.name = name
|
||||
dev_id = int(os.getenv("FLAGS_selected_gpus", "0"))
|
||||
self.timer = core.GPUEventTimer(core.CUDAPlace(dev_id))
|
||||
|
||||
def __getattr__(self, name):
|
||||
return getattr(self.timer, name)
|
||||
|
||||
|
||||
class Timers:
|
||||
"""Group of timers."""
|
||||
|
||||
def __init__(self):
|
||||
self.timers = {}
|
||||
|
||||
def __call__(self, name, use_event=False):
|
||||
clazz = _GPUEventTimer if use_event else _Timer
|
||||
timer = self.timers.get(name)
|
||||
if timer is None:
|
||||
timer = clazz(name)
|
||||
self.timers[name] = timer
|
||||
else:
|
||||
assert type(timer) == clazz, (
|
||||
f"Invalid timer type: {clazz} vs {type(timer)}"
|
||||
)
|
||||
return timer
|
||||
|
||||
def log(self, names, normalizer=1.0, reset=True):
|
||||
"""Log a group of timers."""
|
||||
assert normalizer > 0.0
|
||||
string = "time (ms)"
|
||||
for name in names:
|
||||
elapsed_time = (
|
||||
self.timers[name].elapsed(reset=reset) * 1000.0 / normalizer
|
||||
)
|
||||
string += f" | {name}: {elapsed_time:.2f}"
|
||||
print(string, flush=True)
|
||||
Reference in New Issue
Block a user