wehub-resource-sync
521 lines
13 KiB
Python
521 lines
13 KiB
Python
# Copyright (c) 2025 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 numpy as np
|
|
|
|
import paddle
|
|
import paddle.distributed as dist
|
|
from paddle import nn
|
|
|
|
dist.init_parallel_env()
|
|
|
|
|
|
class SimpleConvNet(nn.Layer):
|
|
def __init__(
|
|
self,
|
|
in_channel,
|
|
out_channel,
|
|
kernel_size,
|
|
padding,
|
|
bias_attr,
|
|
stride,
|
|
data_format="NCHW",
|
|
):
|
|
super().__init__()
|
|
self.conv1 = nn.Conv2D(
|
|
in_channel,
|
|
out_channel,
|
|
kernel_size=kernel_size,
|
|
padding=padding,
|
|
data_format=data_format,
|
|
bias_attr=bias_attr,
|
|
stride=stride,
|
|
)
|
|
self.relu = nn.ReLU()
|
|
|
|
def forward(self, x):
|
|
x = self.conv1(x)
|
|
return self.relu(x)
|
|
|
|
|
|
class TestTPConv:
|
|
def __init__(self):
|
|
self.rank = dist.get_rank()
|
|
self.world_size = dist.get_world_size()
|
|
self._tp_mesh = dist.ProcessMesh(
|
|
list(range(self.world_size)), dim_names=["tp"]
|
|
)
|
|
|
|
def set_seed(self, seed):
|
|
paddle.seed(seed)
|
|
np.random.seed(seed)
|
|
random.seed(seed)
|
|
|
|
def _test_intermediate(
|
|
self,
|
|
N,
|
|
C,
|
|
H,
|
|
W,
|
|
kernel_size,
|
|
padding,
|
|
bias_attr,
|
|
mesh,
|
|
test_name="conv_test",
|
|
dtype_str="float32",
|
|
data_format="NCHW",
|
|
stride=1,
|
|
):
|
|
self.set_seed(2025)
|
|
|
|
dist.auto_parallel.set_mesh(mesh)
|
|
|
|
conv_layer = SimpleConvNet(
|
|
C,
|
|
C,
|
|
kernel_size=kernel_size,
|
|
padding=padding,
|
|
data_format=data_format,
|
|
bias_attr=bias_attr,
|
|
stride=stride,
|
|
)
|
|
original_weight = conv_layer.conv1.weight
|
|
conv_layer.conv1.weight = original_weight
|
|
|
|
if data_format == "NCHW":
|
|
input_tensor = paddle.randn([N, C, H, W])
|
|
shard_axis_input = 3
|
|
else:
|
|
input_tensor = paddle.randn([N, H, W, C])
|
|
shard_axis_input = 2
|
|
|
|
input_placements = [
|
|
dist.Replicate() for _ in range(len(mesh.dim_names))
|
|
]
|
|
mp_dim_index = mesh.dim_names.index("mp")
|
|
input_placements[mp_dim_index] = dist.Shard(shard_axis_input)
|
|
|
|
sharded_input = dist.shard_tensor(input_tensor, mesh, input_placements)
|
|
output_ref = conv_layer(input_tensor)
|
|
loss_ref = output_ref.mean()
|
|
loss_ref.backward()
|
|
weight_grad_ref = conv_layer.conv1.weight.grad.clone()
|
|
|
|
if (
|
|
conv_layer.conv1.bias is not None
|
|
and conv_layer.conv1.bias.grad is not None
|
|
):
|
|
bias_grad_ref = conv_layer.conv1.bias.grad.clone()
|
|
|
|
conv_layer.clear_gradients()
|
|
conv_layer.conv1.weight = original_weight
|
|
|
|
opt = paddle.optimizer.AdamW(
|
|
learning_rate=0.001, parameters=conv_layer.parameters()
|
|
)
|
|
mp_config = {"parallelize_plan": {"conv1": dist.ConvParallel()}}
|
|
|
|
parallel_config = {
|
|
"mp_config": mp_config,
|
|
}
|
|
|
|
dist_model, dist_opt = dist.parallelize(
|
|
conv_layer, opt, config=parallel_config
|
|
)
|
|
|
|
output_intermediate = dist_model(input_tensor)
|
|
loss_intermediate = paddle.mean(output_intermediate)
|
|
loss_intermediate.backward()
|
|
weight_grad_intermediate = dist_model.conv1.weight.grad.clone()
|
|
|
|
if (
|
|
dist_model.conv1.bias is not None
|
|
and dist_model.conv1.bias.grad is not None
|
|
):
|
|
bias_grad_intermediate = dist_model.conv1.bias.grad.clone()
|
|
|
|
def compare_tensors(name, tensor1, tensor2):
|
|
np.testing.assert_allclose(
|
|
tensor1.numpy(), tensor2.numpy(), rtol=1e-8, atol=1e-8
|
|
)
|
|
|
|
def compare_grads(name, grad1, grad2):
|
|
np.testing.assert_allclose(
|
|
grad1.numpy(), grad2.numpy(), rtol=1e-6, atol=1e-6
|
|
)
|
|
|
|
if data_format == "NCHW":
|
|
w_size = output_ref.shape[-1]
|
|
else:
|
|
w_size = output_ref.shape[-2]
|
|
|
|
if dist.get_rank() == 0:
|
|
start_index = 0
|
|
end_index = w_size // 2
|
|
else:
|
|
start_index = w_size // 2
|
|
end_index = w_size
|
|
|
|
if data_format == "NCHW":
|
|
compare_tensors(
|
|
"output",
|
|
output_ref[:, :, :, start_index:end_index],
|
|
output_intermediate._local_value(),
|
|
)
|
|
else:
|
|
compare_tensors(
|
|
"output",
|
|
output_ref[:, :, start_index:end_index, :],
|
|
output_intermediate._local_value(),
|
|
)
|
|
|
|
compare_grads("w", weight_grad_ref, weight_grad_intermediate)
|
|
|
|
if (
|
|
conv_layer.conv1.bias is not None
|
|
and conv_layer.conv1.bias.grad is not None
|
|
):
|
|
compare_grads("b", bias_grad_ref, bias_grad_intermediate)
|
|
|
|
def _test_conv_case(
|
|
self,
|
|
N,
|
|
C,
|
|
H,
|
|
W,
|
|
kernel_size,
|
|
padding,
|
|
bias_attr,
|
|
mesh,
|
|
test_name="conv_test",
|
|
dtype_str="float32",
|
|
data_format="NCHW",
|
|
stride=1,
|
|
):
|
|
self.set_seed(2025)
|
|
|
|
conv_layer = nn.Conv2D(
|
|
C,
|
|
C,
|
|
kernel_size=kernel_size,
|
|
padding=padding,
|
|
bias_attr=bias_attr,
|
|
data_format=data_format,
|
|
stride=stride,
|
|
)
|
|
original_weight = conv_layer.weight
|
|
conv_layer.weight = original_weight
|
|
|
|
if data_format == "NCHW":
|
|
input_tensor = paddle.randn([N, C, H, W])
|
|
shard_axis_input = 3
|
|
else:
|
|
input_tensor = paddle.randn([N, H, W, C])
|
|
shard_axis_input = 2
|
|
|
|
output_ref = conv_layer(input_tensor)
|
|
loss_ref = output_ref.mean()
|
|
loss_ref.backward()
|
|
weight_grad_ref = conv_layer.weight.grad.clone()
|
|
|
|
if conv_layer.bias is not None and conv_layer.bias.grad is not None:
|
|
bias_grad_ref = conv_layer.bias.grad.clone()
|
|
|
|
conv_layer.clear_gradients()
|
|
conv_layer.weight = original_weight
|
|
|
|
rank = dist.get_rank()
|
|
input_placements = [
|
|
dist.Replicate() for _ in range(len(mesh.dim_names))
|
|
]
|
|
mp_dim_index = mesh.dim_names.index("mp")
|
|
input_placements[mp_dim_index] = dist.Shard(shard_axis_input)
|
|
|
|
sharded_input = dist.shard_tensor(input_tensor, mesh, input_placements)
|
|
|
|
output_sharded = conv_layer(sharded_input)
|
|
loss_sharded = paddle.mean(output_sharded)
|
|
loss_sharded.backward()
|
|
|
|
weight_grad_shard = conv_layer.weight.grad.clone()
|
|
if conv_layer.bias is not None and conv_layer.bias.grad is not None:
|
|
bias_grad_shard = conv_layer.bias.grad.clone()
|
|
|
|
def compare_grads(name, grad1, grad2):
|
|
np.testing.assert_allclose(
|
|
grad1.numpy(), grad2.numpy(), rtol=1e-6, atol=1e-7
|
|
)
|
|
|
|
def compare_tensors(name, tensor1, tensor2):
|
|
np.testing.assert_allclose(
|
|
tensor1.numpy(), tensor2.numpy(), rtol=1e-8, atol=1e-8
|
|
)
|
|
|
|
if data_format == "NCHW":
|
|
w_size = output_ref.shape[-1]
|
|
else:
|
|
w_size = output_ref.shape[-2]
|
|
|
|
if dist.get_rank() == 0:
|
|
start_index = 0
|
|
end_index = w_size // 2
|
|
else:
|
|
start_index = w_size // 2
|
|
end_index = w_size
|
|
|
|
if data_format == "NCHW":
|
|
compare_tensors(
|
|
"output",
|
|
output_ref[:, :, :, start_index:end_index],
|
|
output_sharded._local_value(),
|
|
)
|
|
else:
|
|
compare_tensors(
|
|
"output",
|
|
output_ref[:, :, start_index:end_index, :],
|
|
output_sharded._local_value(),
|
|
)
|
|
|
|
compare_grads("w", weight_grad_ref, weight_grad_shard)
|
|
|
|
if conv_layer.bias is not None and conv_layer.bias.grad is not None:
|
|
compare_grads("b", bias_grad_ref, bias_grad_shard)
|
|
|
|
def run_test_cases(self):
|
|
mesh1 = dist.ProcessMesh([0, 1], dim_names=['mp'])
|
|
|
|
# ========= Case 1: padding > 0, stride = 1 =========
|
|
# Typical convolution with halo exchange required.
|
|
self._test_conv_case(
|
|
N=1,
|
|
C=10,
|
|
H=32,
|
|
W=32,
|
|
kernel_size=3,
|
|
padding=1,
|
|
bias_attr=True,
|
|
mesh=mesh1,
|
|
)
|
|
self._test_conv_case(
|
|
N=2,
|
|
C=8,
|
|
H=16,
|
|
W=32,
|
|
kernel_size=5,
|
|
padding=2,
|
|
bias_attr=False,
|
|
mesh=mesh1,
|
|
)
|
|
self._test_conv_case(
|
|
N=4,
|
|
C=6,
|
|
H=28,
|
|
W=28,
|
|
kernel_size=3,
|
|
padding=1,
|
|
bias_attr=True,
|
|
mesh=mesh1,
|
|
)
|
|
|
|
# NHWC format with padding > 0
|
|
self._test_conv_case(
|
|
N=2,
|
|
C=8,
|
|
H=16,
|
|
W=32,
|
|
kernel_size=3,
|
|
padding=1,
|
|
bias_attr=True,
|
|
mesh=mesh1,
|
|
data_format="NHWC",
|
|
)
|
|
self._test_conv_case(
|
|
N=4,
|
|
C=6,
|
|
H=28,
|
|
W=28,
|
|
kernel_size=5,
|
|
padding=2,
|
|
bias_attr=False,
|
|
mesh=mesh1,
|
|
data_format="NHWC",
|
|
)
|
|
|
|
# ========= Case 2: padding = 0, stride == kernel_size =========
|
|
# No halo exchange needed, input width must be divisible by stride.
|
|
self._test_conv_case(
|
|
N=1,
|
|
C=10,
|
|
H=32,
|
|
W=32,
|
|
kernel_size=1,
|
|
padding=0,
|
|
bias_attr=True,
|
|
mesh=mesh1,
|
|
stride=1,
|
|
)
|
|
self._test_conv_case(
|
|
N=4,
|
|
C=6,
|
|
H=32,
|
|
W=32,
|
|
kernel_size=2,
|
|
padding=0,
|
|
bias_attr=False,
|
|
mesh=mesh1,
|
|
stride=2,
|
|
)
|
|
self._test_conv_case(
|
|
N=2,
|
|
C=8,
|
|
H=16,
|
|
W=32,
|
|
kernel_size=4,
|
|
padding=0,
|
|
bias_attr=True,
|
|
mesh=mesh1,
|
|
stride=4,
|
|
)
|
|
|
|
# NHWC format with padding = 0
|
|
self._test_conv_case(
|
|
N=1,
|
|
C=10,
|
|
H=32,
|
|
W=32,
|
|
kernel_size=2,
|
|
padding=0,
|
|
bias_attr=True,
|
|
mesh=mesh1,
|
|
stride=2,
|
|
data_format="NHWC",
|
|
)
|
|
self._test_conv_case(
|
|
N=4,
|
|
C=6,
|
|
H=32,
|
|
W=32,
|
|
kernel_size=4,
|
|
padding=0,
|
|
bias_attr=False,
|
|
mesh=mesh1,
|
|
stride=4,
|
|
data_format="NHWC",
|
|
)
|
|
|
|
# ========= Case 3: 2D ProcessMesh (dp + tp) =========
|
|
mesh2 = dist.ProcessMesh([[0, 1]], dim_names=['dp', 'mp'])
|
|
|
|
# padding > 0
|
|
self._test_conv_case(
|
|
N=2,
|
|
C=8,
|
|
H=32,
|
|
W=32,
|
|
kernel_size=3,
|
|
padding=1,
|
|
bias_attr=True,
|
|
mesh=mesh2,
|
|
)
|
|
self._test_conv_case(
|
|
N=4,
|
|
C=6,
|
|
H=28,
|
|
W=28,
|
|
kernel_size=5,
|
|
padding=2,
|
|
bias_attr=False,
|
|
mesh=mesh2,
|
|
)
|
|
|
|
# padding = 0
|
|
self._test_conv_case(
|
|
N=2,
|
|
C=8,
|
|
H=16,
|
|
W=32,
|
|
kernel_size=1,
|
|
padding=0,
|
|
bias_attr=True,
|
|
mesh=mesh2,
|
|
stride=1,
|
|
)
|
|
|
|
# NHWC format, both padding > 0 and = 0
|
|
self._test_conv_case(
|
|
N=4,
|
|
C=6,
|
|
H=28,
|
|
W=28,
|
|
kernel_size=3,
|
|
padding=1,
|
|
bias_attr=True,
|
|
mesh=mesh2,
|
|
data_format="NHWC",
|
|
)
|
|
self._test_conv_case(
|
|
N=1,
|
|
C=10,
|
|
H=32,
|
|
W=32,
|
|
kernel_size=4,
|
|
padding=0,
|
|
bias_attr=True,
|
|
mesh=mesh2,
|
|
stride=4,
|
|
data_format="NHWC",
|
|
)
|
|
|
|
self._test_intermediate(
|
|
N=1,
|
|
C=10,
|
|
H=32,
|
|
W=32,
|
|
kernel_size=3,
|
|
padding=1,
|
|
bias_attr=True,
|
|
mesh=mesh1,
|
|
)
|
|
self._test_intermediate(
|
|
N=1,
|
|
C=10,
|
|
H=32,
|
|
W=32,
|
|
kernel_size=2,
|
|
padding=0,
|
|
bias_attr=True,
|
|
mesh=mesh1,
|
|
stride=2,
|
|
data_format="NHWC",
|
|
)
|
|
self._test_intermediate(
|
|
N=4,
|
|
C=6,
|
|
H=28,
|
|
W=28,
|
|
kernel_size=3,
|
|
padding=1,
|
|
bias_attr=True,
|
|
mesh=mesh2,
|
|
data_format="NHWC",
|
|
)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
tester = TestTPConv()
|
|
tester.run_test_cases()
|