chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,232 @@
|
||||
# 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.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], dim_names=["x"])
|
||||
|
||||
|
||||
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 = False
|
||||
|
||||
self.norm0 = paddle.nn.LayerNorm(d_model, epsilon=1e-5)
|
||||
self.norm0.bias.stop_gradient = True
|
||||
self.norm1 = paddle.nn.LayerNorm(d_model, epsilon=1e-5)
|
||||
self.norm1.bias.stop_gradient = True
|
||||
self.linear0 = paddle.nn.Linear(
|
||||
d_model, dim_feedforward, weight_attr, bias_attr=bias_attr
|
||||
)
|
||||
auto.shard_tensor(self.linear0.weight, MESH_0, [None, "x"])
|
||||
self.linear1 = paddle.nn.Linear(
|
||||
dim_feedforward, d_model, weight_attr, bias_attr=bias_attr
|
||||
)
|
||||
auto.shard_tensor(self.linear1.weight, MESH_0, ["x", None])
|
||||
self.dropout = paddle.nn.Dropout(dropout_ratio, mode="upscale_in_train")
|
||||
|
||||
def forward(self, input):
|
||||
# sp region
|
||||
auto.shard_tensor(input, MESH_0, ["x", None, None])
|
||||
out = self.norm0(input)
|
||||
auto.shard_tensor(input, MESH_0, ["x", None, None])
|
||||
out = F.gelu(out, approximate=True)
|
||||
|
||||
# tp region
|
||||
auto.shard_tensor(out, MESH_0, [None, None, None])
|
||||
out = self.linear0(out)
|
||||
out = F.gelu(out, approximate=True)
|
||||
out = self.linear1(out)
|
||||
auto.shard_tensor(out, MESH_0, [None, None, None])
|
||||
|
||||
# sp region
|
||||
out = self.dropout(out)
|
||||
auto.shard_tensor(out, MESH_0, ["x", None, None])
|
||||
out = F.gelu(out, approximate=True)
|
||||
out = self.norm1(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)
|
||||
|
||||
def forward(self, input):
|
||||
out = self.mlp0(input)
|
||||
out = self.mlp1(out)
|
||||
|
||||
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"
|
||||
strategy.sp_optimization.enable = True
|
||||
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()
|
||||
|
||||
with open("test_static_sequence_parallel.txt", "w+") as f:
|
||||
f.write(str(dist_main_prog))
|
||||
|
||||
ops = dist_main_prog.global_block().ops
|
||||
sp_ring_id = None
|
||||
allgather_count = 0
|
||||
reducescatter_count = 0
|
||||
allreduce_count = 0
|
||||
|
||||
for op in ops:
|
||||
# check sequence parallel allgather
|
||||
if op.type == "all_gather":
|
||||
assert int(op.attr("nranks")) == 4, (
|
||||
"sequence parallel allgather error with nranks [{}]".format(
|
||||
op.attr("nranks")
|
||||
)
|
||||
)
|
||||
if sp_ring_id is None:
|
||||
sp_ring_id = int(op.attr("ring_id"))
|
||||
else:
|
||||
assert sp_ring_id == int(op.attr("ring_id")), (
|
||||
"sequence parallel allgather error with ring_id [{}]".format(
|
||||
op.attr("ring_id")
|
||||
)
|
||||
)
|
||||
allgather_count += 1
|
||||
|
||||
# check sequence parallel reducescatter
|
||||
elif op.type == "reduce_scatter":
|
||||
assert int(op.attr("nranks")) == 4, (
|
||||
"sequence parallel reducescatter error with nranks [{}]".format(
|
||||
op.attr("nranks")
|
||||
)
|
||||
)
|
||||
assert sp_ring_id == int(op.attr("ring_id")), (
|
||||
"sequence parallel reducescatter error with ring_id [{}]".format(
|
||||
op.attr("ring_id")
|
||||
)
|
||||
)
|
||||
reducescatter_count += 1
|
||||
|
||||
# check sequence parallel grad sync
|
||||
elif op.type == "c_allreduce_sum":
|
||||
assert "layer_norm" in op.output_arg_names[0], (
|
||||
f"sequence parallel reducescatter error grad sync var [{op.output_arg_names[0]}]"
|
||||
)
|
||||
assert sp_ring_id == int(op.attr("ring_id")), (
|
||||
"sequence parallel reducescatter error with ring_id [{}]".format(
|
||||
op.attr("ring_id")
|
||||
)
|
||||
)
|
||||
allreduce_count += 1
|
||||
|
||||
assert allgather_count == 4, (
|
||||
f"sequence parallel should have 4 allgather, but got [{allgather_count}]"
|
||||
)
|
||||
assert reducescatter_count == 4, (
|
||||
f"sequence parallel should have 4 allgather, but got [{reducescatter_count}]"
|
||||
)
|
||||
assert allreduce_count == 4, (
|
||||
f"sequence parallel should have 4 allgather, but got [{allreduce_count}]"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user