wehub-resource-sync
354 lines
11 KiB
Python
354 lines
11 KiB
Python
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import random
|
|
import unittest
|
|
|
|
import numpy as np
|
|
|
|
import paddle
|
|
import paddle.distributed as dist
|
|
from paddle.distributed import fleet
|
|
|
|
|
|
def set_random_seed(seed):
|
|
"""Set random seed for reproducibility."""
|
|
random.seed(seed)
|
|
np.random.seed(seed)
|
|
paddle.seed(seed)
|
|
fleet.meta_parallel.model_parallel_random_seed(seed)
|
|
|
|
|
|
class ColumnLinearNet(paddle.nn.Layer):
|
|
def __init__(self, input_size, output_size, global_dtype):
|
|
super().__init__()
|
|
self.parallel_linear = fleet.meta_parallel.ColumnParallelLinear(
|
|
in_features=input_size,
|
|
out_features=output_size,
|
|
weight_attr=None,
|
|
has_bias=True,
|
|
gather_output=True,
|
|
name="test_column_linear",
|
|
)
|
|
|
|
def forward(self, x):
|
|
output = self.parallel_linear(x)
|
|
return output
|
|
|
|
|
|
class RowLinearNet(paddle.nn.Layer):
|
|
def __init__(self, input_size, output_size):
|
|
super().__init__()
|
|
self.parallel_linear = fleet.meta_parallel.RowParallelLinear(
|
|
in_features=input_size,
|
|
out_features=output_size,
|
|
has_bias=True,
|
|
input_is_parallel=False,
|
|
name="test_row_linear",
|
|
)
|
|
|
|
def forward(self, x):
|
|
output = self.parallel_linear(x)
|
|
return output
|
|
|
|
|
|
class EmbeddingNet(paddle.nn.Layer):
|
|
def __init__(self, vocab_size, hidden_size):
|
|
super().__init__()
|
|
self.embedding = fleet.meta_parallel.VocabParallelEmbedding(
|
|
vocab_size, hidden_size
|
|
)
|
|
|
|
def forward(self, x):
|
|
output = self.embedding(x)
|
|
return output
|
|
|
|
|
|
class SimpleMatmul(paddle.nn.Layer):
|
|
def __init__(self, weight, output_size, global_dtype):
|
|
super().__init__()
|
|
self.weight = paddle.create_parameter(
|
|
shape=weight.shape,
|
|
dtype=global_dtype,
|
|
attr=paddle.ParamAttr(
|
|
initializer=paddle.nn.initializer.Assign(weight)
|
|
),
|
|
)
|
|
self.bias = self.create_parameter(
|
|
shape=[output_size],
|
|
dtype=global_dtype,
|
|
attr=paddle.ParamAttr(
|
|
initializer=paddle.nn.initializer.Constant(0.0)
|
|
),
|
|
)
|
|
|
|
def forward(self, x):
|
|
output = paddle.matmul(x, self.weight) + self.bias
|
|
return output
|
|
|
|
|
|
class SimpleEmbedding(paddle.nn.Layer):
|
|
def __init__(self, vocab_size, hidden_size, weight):
|
|
super().__init__()
|
|
self.embedding = paddle.nn.Embedding(
|
|
vocab_size,
|
|
hidden_size,
|
|
weight_attr=paddle.framework.ParamAttr(
|
|
name="origin_embedding",
|
|
initializer=paddle.nn.initializer.Assign(weight),
|
|
),
|
|
)
|
|
|
|
def forward(self, x):
|
|
output = self.embedding(x)
|
|
return output
|
|
|
|
|
|
class TestDistTraining(unittest.TestCase):
|
|
def setUp(self):
|
|
strategy = fleet.DistributedStrategy()
|
|
self.model_parallel_size = 2
|
|
strategy.hybrid_configs = {
|
|
"dp_degree": 1,
|
|
"mp_degree": self.model_parallel_size,
|
|
"pp_degree": 1,
|
|
}
|
|
fleet.init(is_collective=True, strategy=strategy)
|
|
|
|
def test_column_parallel_layer(self):
|
|
set_random_seed(1024)
|
|
global_dtype = "float32"
|
|
|
|
input_size_per_card = 17
|
|
input_size = input_size_per_card * self.model_parallel_size
|
|
output_size_per_card = 13
|
|
output_size = output_size_per_card * self.model_parallel_size
|
|
batch_size = 4
|
|
|
|
model_a = ColumnLinearNet(input_size, output_size, global_dtype)
|
|
|
|
# get w
|
|
check_group = dist.new_group(list(range(self.model_parallel_size)))
|
|
integral_w = []
|
|
partial_w = model_a.parallel_linear.weight.clone().detach()
|
|
paddle.distributed.all_gather(integral_w, partial_w, group=check_group)
|
|
integral_w = paddle.concat(integral_w, axis=1)
|
|
|
|
model_b = SimpleMatmul(integral_w, output_size, global_dtype)
|
|
|
|
optimizer_a = paddle.optimizer.SGD(
|
|
learning_rate=0.001, parameters=model_a.parameters()
|
|
)
|
|
optimizer_b = paddle.optimizer.SGD(
|
|
learning_rate=0.001, parameters=model_b.parameters()
|
|
)
|
|
for idx in range(5):
|
|
input = paddle.randn([batch_size, input_size], global_dtype)
|
|
input.stop_gradient = True
|
|
|
|
output_a = model_a(input)
|
|
loss_a = output_a.mean()
|
|
loss_a.backward()
|
|
|
|
output_b = model_b(input)
|
|
loss_b = output_b.mean()
|
|
loss_b.backward()
|
|
|
|
optimizer_a.step()
|
|
optimizer_b.step()
|
|
|
|
np.testing.assert_allclose(loss_a.numpy(), loss_b.numpy())
|
|
|
|
def test_row_parallel_layer(self):
|
|
global_dtype = "float32"
|
|
paddle.set_default_dtype(global_dtype)
|
|
set_random_seed(1024)
|
|
|
|
self.hcg = fleet.get_hybrid_communicate_group()
|
|
|
|
self.word_size = self.hcg.get_model_parallel_world_size()
|
|
self.rank_id = self.hcg.get_model_parallel_rank()
|
|
|
|
input_size_per_card = 11
|
|
input_size = input_size_per_card * self.model_parallel_size
|
|
output_size_per_card = 10
|
|
output_size = output_size_per_card * self.model_parallel_size
|
|
batch_size = 4
|
|
|
|
model_a = RowLinearNet(input_size, output_size)
|
|
|
|
# get w
|
|
check_group = dist.new_group(list(range(self.model_parallel_size)))
|
|
integral_w = []
|
|
partial_w = model_a.parallel_linear.weight.clone().detach()
|
|
paddle.distributed.all_gather(integral_w, partial_w, group=check_group)
|
|
integral_w = paddle.concat(integral_w, axis=0)
|
|
|
|
model_b = SimpleMatmul(integral_w, output_size, global_dtype)
|
|
|
|
optimizer_a = paddle.optimizer.SGD(
|
|
learning_rate=0.001, parameters=model_a.parameters()
|
|
)
|
|
|
|
optimizer_b = paddle.optimizer.SGD(
|
|
learning_rate=0.001, parameters=model_b.parameters()
|
|
)
|
|
|
|
for idx in range(5):
|
|
input = paddle.randn([batch_size, input_size], global_dtype)
|
|
input.stop_gradient = True
|
|
|
|
output_a = model_a(input)
|
|
loss_a = output_a.mean()
|
|
loss_a.backward()
|
|
|
|
output_b = model_b(input)
|
|
loss_b = output_b.mean()
|
|
loss_b.backward()
|
|
|
|
optimizer_a.step()
|
|
optimizer_b.step()
|
|
|
|
np.testing.assert_allclose(
|
|
loss_a.numpy(), loss_b.numpy(), rtol=5e-5
|
|
)
|
|
|
|
def test_parallel_embedding(self):
|
|
batch_size = 17
|
|
seq_length = 23
|
|
vocab_size_per_card = 2
|
|
vocab_size = vocab_size_per_card * self.model_parallel_size
|
|
hidden_size = 2
|
|
seed = 1236
|
|
|
|
set_random_seed(seed)
|
|
rank_id = dist.get_rank()
|
|
|
|
# model_a
|
|
model_a = EmbeddingNet(vocab_size, hidden_size)
|
|
|
|
# model_b
|
|
check_group = dist.new_group(list(range(self.model_parallel_size)))
|
|
integral_w = []
|
|
partial_w = model_a.embedding.weight.clone().detach()
|
|
paddle.distributed.all_gather(integral_w, partial_w, group=check_group)
|
|
result_w = []
|
|
for idx in range(len(integral_w)):
|
|
tmp = paddle.gather(
|
|
integral_w[idx],
|
|
paddle.to_tensor(list(range(vocab_size_per_card))),
|
|
)
|
|
result_w.append(tmp)
|
|
integral_w = paddle.concat(result_w, axis=0)
|
|
|
|
model_b = SimpleEmbedding(vocab_size, hidden_size, integral_w)
|
|
|
|
optimizer_a = paddle.optimizer.SGD(
|
|
learning_rate=0.001, parameters=model_a.parameters()
|
|
)
|
|
|
|
optimizer_b = paddle.optimizer.SGD(
|
|
learning_rate=0.001, parameters=model_b.parameters()
|
|
)
|
|
|
|
for _ in range(5):
|
|
np_input_data = np.random.randint(
|
|
0, vocab_size, (batch_size, seq_length)
|
|
)
|
|
input_data = paddle.to_tensor(np_input_data, dtype="int32")
|
|
|
|
output_a = model_a(input_data)
|
|
loss_a = output_a.mean()
|
|
|
|
output_b = model_b(input_data)
|
|
loss_b = output_b.mean()
|
|
|
|
loss_a.backward()
|
|
loss_b.backward()
|
|
|
|
optimizer_a.step()
|
|
optimizer_b.step()
|
|
print(loss_a.numpy(), loss_b.numpy())
|
|
|
|
np.testing.assert_allclose(loss_a.numpy(), loss_b.numpy())
|
|
|
|
def test_parallel_cross_entropy(self):
|
|
batch_size = 8
|
|
seq_length = 16
|
|
class_size_per_card = 2
|
|
vocab_size = class_size_per_card * self.model_parallel_size
|
|
seed = 100
|
|
|
|
set_random_seed(seed)
|
|
rank_id = dist.get_rank()
|
|
|
|
# model_a
|
|
model_a = fleet.meta_parallel.ParallelCrossEntropy()
|
|
|
|
model_b = paddle.nn.CrossEntropyLoss(reduction="none")
|
|
|
|
paddle.seed(rank_id * 10)
|
|
random.seed(seed)
|
|
np.random.seed(seed)
|
|
|
|
for _ in range(5):
|
|
np_label = np.random.randint(
|
|
0, vocab_size, (batch_size, seq_length)
|
|
)
|
|
label = paddle.to_tensor(np_label, dtype="int64")
|
|
|
|
data = paddle.randn(
|
|
shape=[batch_size, seq_length, class_size_per_card],
|
|
dtype='float32',
|
|
)
|
|
data.stop_gradient = False
|
|
|
|
check_group = dist.new_group(list(range(self.model_parallel_size)))
|
|
integral_data = []
|
|
partial_data = data.clone().detach()
|
|
paddle.distributed.all_gather(
|
|
integral_data, partial_data, group=check_group
|
|
)
|
|
integral_data = paddle.concat(integral_data, axis=-1)
|
|
integral_data = integral_data.detach().clone()
|
|
integral_data.stop_gradient = False
|
|
|
|
loss_a = model_a(data, label).sum() / batch_size
|
|
loss_b = model_b(integral_data, label).sum() / batch_size
|
|
print("loss_a: ", loss_a.numpy(), "loss_b: ", loss_b.numpy())
|
|
|
|
np.testing.assert_allclose(
|
|
loss_a.numpy(), loss_b.numpy(), rtol=1e-6
|
|
)
|
|
|
|
loss_a.backward()
|
|
loss_b.backward()
|
|
|
|
integral_grad = []
|
|
partial_grad = data.grad.clone().detach()
|
|
paddle.distributed.all_gather(
|
|
integral_grad, partial_grad, group=check_group
|
|
)
|
|
integral_grad = paddle.concat(integral_grad, axis=-1)
|
|
|
|
np.testing.assert_allclose(
|
|
integral_data.grad.numpy(False),
|
|
integral_grad.numpy(False),
|
|
rtol=1e-6,
|
|
)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|