chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,255 @@
|
||||
# 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 sys
|
||||
import unittest
|
||||
|
||||
import paddle
|
||||
|
||||
sys.path.append("../legacy_test")
|
||||
|
||||
import paddle.nn.functional as F
|
||||
from paddle import nn, static, utils
|
||||
from paddle.base import ParamAttr
|
||||
from paddle.distributed.auto_parallel.static.dist_context import (
|
||||
DistributedContext,
|
||||
)
|
||||
from paddle.distributed.auto_parallel.static.operators.common import (
|
||||
is_data_parallel_reduce_op,
|
||||
is_data_parallel_scale_op,
|
||||
)
|
||||
from paddle.distributed.auto_parallel.static.parallelizer_v2 import Parallelizer
|
||||
from paddle.distributed.auto_parallel.static.planner_v2 import Planner
|
||||
from paddle.distributed.auto_parallel.strategy import Strategy
|
||||
from paddle.distributed.fleet import auto
|
||||
|
||||
paddle.enable_static()
|
||||
BATCH_SIZE = 4
|
||||
SEQ_LEN = 512
|
||||
HIDDEN_SIZE = 1024
|
||||
MESH_0 = auto.ProcessMesh([[0, 1, 2, 3], [4, 5, 6, 7]], dim_names=["x", "y"])
|
||||
|
||||
|
||||
class MLPLayer(nn.Layer):
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size=1024,
|
||||
intermediate_size=4 * 1024,
|
||||
dropout_ratio=0.1,
|
||||
initializer_range=0.02,
|
||||
):
|
||||
super().__init__()
|
||||
d_model = hidden_size
|
||||
dim_feedforward = intermediate_size
|
||||
weight_attr = ParamAttr(
|
||||
initializer=paddle.nn.initializer.Normal(
|
||||
mean=0.0, std=initializer_range
|
||||
)
|
||||
)
|
||||
bias_attr = ParamAttr(
|
||||
initializer=paddle.nn.initializer.Constant(value=0.0)
|
||||
)
|
||||
|
||||
self.norm = paddle.nn.LayerNorm(d_model, epsilon=1e-5)
|
||||
self.linear0 = paddle.nn.Linear(
|
||||
d_model, dim_feedforward, weight_attr, bias_attr=bias_attr
|
||||
)
|
||||
self.linear1 = paddle.nn.Linear(
|
||||
dim_feedforward, d_model, weight_attr, bias_attr=bias_attr
|
||||
)
|
||||
self.dropout = paddle.nn.Dropout(dropout_ratio, mode="upscale_in_train")
|
||||
|
||||
def forward(self, input):
|
||||
out = self.norm(input)
|
||||
out = self.linear0(out)
|
||||
out = F.gelu(out, approximate=True)
|
||||
out = self.linear1(out)
|
||||
out = self.dropout(out)
|
||||
return out
|
||||
|
||||
|
||||
class HybridParallelNet(nn.Layer):
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size=1024,
|
||||
):
|
||||
super().__init__()
|
||||
self.mlp0 = MLPLayer(hidden_size, hidden_size * 4)
|
||||
self.mlp1 = MLPLayer(hidden_size, hidden_size * 4)
|
||||
self.mlp2 = MLPLayer(hidden_size, hidden_size * 4)
|
||||
|
||||
def forward(self, input):
|
||||
# prune dp
|
||||
auto.shard_tensor(input, MESH_0, ["x", None, None])
|
||||
activation0 = self.mlp0(input)
|
||||
auto.shard_tensor(activation0, MESH_0, ["x", None, None])
|
||||
activation1 = F.gelu(activation0, approximate=True)
|
||||
|
||||
# prune sp
|
||||
auto.shard_tensor(activation1, MESH_0, [None, "y", None])
|
||||
activation2 = self.mlp1(activation1)
|
||||
auto.shard_tensor(activation2, MESH_0, [None, "y", None])
|
||||
activation3 = F.gelu(activation2, approximate=True)
|
||||
|
||||
# dp_sp
|
||||
auto.shard_tensor(activation3, MESH_0, ["x", "y", None])
|
||||
out = self.mlp2(activation3)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def get_hybrid_parallel_model(train_program, start_program):
|
||||
with (
|
||||
static.program_guard(train_program, start_program),
|
||||
utils.unique_name.guard(),
|
||||
):
|
||||
batch_size = BATCH_SIZE
|
||||
hidden_size = HIDDEN_SIZE
|
||||
sequence_len = SEQ_LEN
|
||||
|
||||
input = static.data(
|
||||
name="input",
|
||||
shape=[batch_size, sequence_len, hidden_size],
|
||||
dtype='float32',
|
||||
)
|
||||
network = HybridParallelNet(hidden_size=HIDDEN_SIZE)
|
||||
|
||||
predict = network(input)
|
||||
error_cost = paddle.sum(predict)
|
||||
|
||||
return error_cost, train_program, start_program
|
||||
|
||||
|
||||
def get_dist_prog(rank=2):
|
||||
train_program = paddle.static.Program()
|
||||
startup_program = paddle.static.Program()
|
||||
|
||||
loss, train_program, startup_program = get_hybrid_parallel_model(
|
||||
train_program, startup_program
|
||||
)
|
||||
opt = paddle.optimizer.AdamW(learning_rate=0.00001)
|
||||
strategy = Strategy()
|
||||
strategy.auto_mode = "semi"
|
||||
dist_context = DistributedContext(
|
||||
train_program, startup_program, opt, loss, strategy=strategy
|
||||
)
|
||||
planner = Planner("train", dist_context)
|
||||
planner.plan()
|
||||
|
||||
parallelizer = Parallelizer(
|
||||
"train",
|
||||
planner.completer,
|
||||
dist_context,
|
||||
)
|
||||
parallelizer.parallel(rank=rank)
|
||||
|
||||
return (
|
||||
dist_context.dist_main_programs[rank],
|
||||
dist_context.dist_startup_programs[rank],
|
||||
)
|
||||
|
||||
|
||||
class TestGradSync(unittest.TestCase):
|
||||
def test_decoder_dp_sp(self):
|
||||
dist_main_prog, dist_startup_prog = get_dist_prog()
|
||||
|
||||
ops = dist_main_prog.global_block().ops
|
||||
allreduce_count = 0
|
||||
scale_count = 0
|
||||
# Linear, Linear, LN
|
||||
dp_sync_indices = [
|
||||
0,
|
||||
2,
|
||||
4,
|
||||
6,
|
||||
8,
|
||||
9,
|
||||
18,
|
||||
19,
|
||||
20,
|
||||
21,
|
||||
22,
|
||||
23,
|
||||
] # check data parallel sync
|
||||
sp_sync_indices = [
|
||||
1,
|
||||
3,
|
||||
5,
|
||||
7,
|
||||
10,
|
||||
11,
|
||||
12,
|
||||
13,
|
||||
14,
|
||||
15,
|
||||
16,
|
||||
17,
|
||||
] # check sp parallel sync
|
||||
dp_ring_id = None
|
||||
sp_ring_id = None
|
||||
dp_scale = 0.5
|
||||
sp_scale = 0.25
|
||||
|
||||
for op in ops:
|
||||
if is_data_parallel_reduce_op(op):
|
||||
if allreduce_count in dp_sync_indices:
|
||||
if dp_ring_id is None:
|
||||
dp_ring_id = int(op.attr("ring_id"))
|
||||
else:
|
||||
assert dp_ring_id == int(op.attr("ring_id")), (
|
||||
"gradient synchronization of dp use different communication group [{}] and [{}]".format(
|
||||
dp_ring_id, int(op.attr("ring_id"))
|
||||
)
|
||||
)
|
||||
elif allreduce_count in sp_sync_indices:
|
||||
if sp_ring_id is None:
|
||||
sp_ring_id = int(op.attr("ring_id"))
|
||||
else:
|
||||
assert sp_ring_id == int(op.attr("ring_id")), (
|
||||
"gradient synchronization of sp use different communication group [{}] and [{}]".format(
|
||||
sp_ring_id, int(op.attr("ring_id"))
|
||||
)
|
||||
)
|
||||
else:
|
||||
raise AssertionError(
|
||||
f"encounter redundant gradient synchronization: [{op}]"
|
||||
)
|
||||
allreduce_count += 1
|
||||
|
||||
elif is_data_parallel_scale_op(op):
|
||||
if scale_count in dp_sync_indices:
|
||||
assert dp_scale == float(op.attr("scale")), (
|
||||
"gradient synchronization of dp use different scale [{}] and [{}]".format(
|
||||
dp_scale, int(op.attr("scale"))
|
||||
)
|
||||
)
|
||||
elif scale_count in sp_sync_indices:
|
||||
assert sp_scale == float(op.attr("scale")), (
|
||||
"gradient synchronization of sp use different scale [{}] and [{}]".format(
|
||||
sp_scale, int(op.attr("scale"))
|
||||
)
|
||||
)
|
||||
else:
|
||||
raise AssertionError(
|
||||
f"encounter redundant gradient synchronization: [{op}]"
|
||||
)
|
||||
|
||||
scale_count += 1
|
||||
|
||||
assert scale_count == 24
|
||||
assert allreduce_count == 24
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user