wehub-resource-sync
532 lines
22 KiB
Python
532 lines
22 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 math
|
|
import os
|
|
|
|
import paddle
|
|
from paddle import nn
|
|
from paddle.distributed import ShardedWeight, fleet
|
|
from paddle.distributed.fleet.layers.mpu import (
|
|
ColumnParallelLinear,
|
|
RowParallelLinear,
|
|
VocabParallelEmbedding,
|
|
)
|
|
from paddle.distributed.fleet.meta_optimizers.dygraph_optimizer.dygraph_sharding_optimizer import (
|
|
DygraphShardingOptimizer,
|
|
DygraphShardingOptimizerV2,
|
|
)
|
|
from paddle.distributed.fleet.meta_parallel.sharding.group_sharded_optimizer_stage2 import (
|
|
GroupShardedOptimizerStage2,
|
|
)
|
|
from paddle.distributed.fleet.meta_parallel.sharding.group_sharded_stage3 import (
|
|
GroupShardedStage3,
|
|
)
|
|
from paddle.distributed.fleet.utils.sequence_parallel_utils import (
|
|
ColumnSequenceParallelLinear,
|
|
RowSequenceParallelLinear,
|
|
)
|
|
|
|
|
|
class SimpleMLP(
|
|
nn.Layer
|
|
): # embedding_weight_size=24*100=2400,it can't be divided by 256,which is using to check the padding logic
|
|
def __init__(self, hidden_size=100, has_bias=False):
|
|
super().__init__()
|
|
self.embedding = VocabParallelEmbedding(24, hidden_size)
|
|
self.linear1 = ColumnParallelLinear(
|
|
hidden_size, hidden_size, gather_output=False, has_bias=has_bias
|
|
)
|
|
self.linear2 = RowParallelLinear(
|
|
hidden_size, hidden_size, input_is_parallel=True, has_bias=has_bias
|
|
)
|
|
self.llm_head = self.embedding # test the shared weight
|
|
|
|
def forward(self, x):
|
|
x = self.embedding(x)
|
|
x = self.linear1(x)
|
|
x = self.linear2(x)
|
|
x = paddle.matmul(x, self.llm_head.weight, transpose_y=True)
|
|
return x
|
|
|
|
|
|
class TestParallelLayersLogic:
|
|
def __init__(self):
|
|
self.optimizer_var_suffix = [".moment1_0", ".moment2_0", ".w_0"]
|
|
self.test_type = os.getenv("test_type")
|
|
self.layer_type = os.getenv("layer_type")
|
|
self.tp_degree = int(os.getenv("tp", "1"))
|
|
self.dp_degree = int(os.getenv("dp", "1"))
|
|
self.sharding_degree = int(os.getenv("sharding_degree", "1"))
|
|
self.world_size = int(os.getenv("world_size"))
|
|
self.has_bias = os.getenv("has_bias", "True").lower() == "true"
|
|
self.master_weight = (
|
|
os.getenv("master_weight", "False").lower() == "true"
|
|
)
|
|
self.batch_size = 2
|
|
self.hidden_size = 32
|
|
self.vocab_size = 24
|
|
self.seq_len = 2
|
|
self.hcg = None
|
|
|
|
def run_test(self):
|
|
strategy = fleet.DistributedStrategy()
|
|
strategy.hybrid_configs = {
|
|
"dp_degree": self.dp_degree,
|
|
"mp_degree": self.tp_degree,
|
|
"sharding_degree": self.sharding_degree,
|
|
"pp_degree": 1,
|
|
}
|
|
fleet.init(is_collective=True, strategy=strategy)
|
|
self.hcg = fleet.get_hybrid_communicate_group()
|
|
if self.test_type == "layer":
|
|
self.run_layer_test()
|
|
elif self.test_type == "optimizer":
|
|
self.run_optimizer_test()
|
|
else:
|
|
raise ValueError(f"Unknown test_type: {self.test_type}")
|
|
|
|
def run_layer_test(self):
|
|
tp_group = self.hcg.get_model_parallel_group()
|
|
layer = self._get_layer()
|
|
sharded_dict = layer.sharded_state_dict()
|
|
self._verify_parallel_layer(
|
|
sharded_dict, tp_group.rank, tp_group.nranks
|
|
)
|
|
|
|
def _get_layer(self):
|
|
if self.layer_type == "ColumnParallelLinear":
|
|
return ColumnParallelLinear(
|
|
self.hidden_size, self.hidden_size * 2, has_bias=self.has_bias
|
|
)
|
|
elif self.layer_type == "RowParallelLinear":
|
|
return RowParallelLinear(
|
|
self.hidden_size * 2, self.hidden_size, has_bias=self.has_bias
|
|
)
|
|
elif self.layer_type == "VocabParallelEmbedding":
|
|
return VocabParallelEmbedding(self.vocab_size, self.hidden_size)
|
|
elif self.layer_type == "ColumnSequenceParallelLinear":
|
|
return ColumnSequenceParallelLinear(
|
|
self.hidden_size,
|
|
self.hidden_size * 2,
|
|
has_bias=self.has_bias,
|
|
gather_output=False,
|
|
)
|
|
elif self.layer_type == "RowSequenceParallelLinear":
|
|
return RowSequenceParallelLinear(
|
|
self.hidden_size * 2,
|
|
self.hidden_size,
|
|
has_bias=self.has_bias,
|
|
input_is_parallel=True,
|
|
)
|
|
raise ValueError(f"Unknown layer_type: {self.layer_type}")
|
|
|
|
def _verify_parallel_layer(self, sharded_dict, tp_rank, tp_world_size):
|
|
if self.has_bias:
|
|
assert 'bias' in sharded_dict
|
|
bias_shard = sharded_dict['bias']
|
|
assert isinstance(bias_shard, ShardedWeight)
|
|
else:
|
|
assert 'bias' not in sharded_dict
|
|
|
|
assert 'weight' in sharded_dict
|
|
weight_shard = sharded_dict['weight']
|
|
assert isinstance(weight_shard, ShardedWeight)
|
|
|
|
if self.layer_type == "ColumnParallelLinear":
|
|
in_f, out_f = self.hidden_size, self.hidden_size * 2
|
|
assert weight_shard.global_shape == (in_f, out_f)
|
|
assert weight_shard.local_shape == (in_f, out_f // tp_world_size)
|
|
assert weight_shard.global_offset == (
|
|
0,
|
|
tp_rank * (out_f // tp_world_size),
|
|
)
|
|
if self.has_bias:
|
|
assert bias_shard.global_shape == (out_f,)
|
|
assert bias_shard.local_shape == (out_f // tp_world_size,)
|
|
assert bias_shard.global_offset == (
|
|
tp_rank * (out_f // tp_world_size),
|
|
)
|
|
|
|
elif self.layer_type == "RowParallelLinear":
|
|
in_f, out_f = self.hidden_size * 2, self.hidden_size
|
|
# Weight is sharded on axis 1
|
|
assert weight_shard.global_shape == (in_f, out_f)
|
|
assert weight_shard.local_shape == (in_f // tp_world_size, out_f)
|
|
assert weight_shard.global_offset == (
|
|
tp_rank * (in_f // tp_world_size),
|
|
0,
|
|
)
|
|
|
|
if self.has_bias:
|
|
# Bias is replicated, not sharded
|
|
assert bias_shard.global_shape == (out_f,)
|
|
assert bias_shard.local_shape == bias_shard.global_shape
|
|
assert bias_shard.global_offset == (0,)
|
|
|
|
elif self.layer_type == "VocabParallelEmbedding":
|
|
assert weight_shard.global_shape == (
|
|
self.vocab_size,
|
|
self.hidden_size,
|
|
)
|
|
assert weight_shard.local_shape == (
|
|
self.vocab_size // tp_world_size,
|
|
self.hidden_size,
|
|
)
|
|
assert weight_shard.global_offset == (
|
|
tp_rank * (self.vocab_size // tp_world_size),
|
|
0,
|
|
)
|
|
|
|
elif self.layer_type == "ColumnSequenceParallelLinear":
|
|
in_f, out_f = self.hidden_size, self.hidden_size * 2
|
|
assert weight_shard.global_shape == (in_f, out_f)
|
|
assert weight_shard.local_shape == (in_f, out_f // tp_world_size)
|
|
assert weight_shard.global_offset == (
|
|
0,
|
|
tp_rank * (out_f // tp_world_size),
|
|
)
|
|
if self.has_bias:
|
|
assert bias_shard.global_shape == (out_f,)
|
|
assert bias_shard.local_shape == (out_f // tp_world_size,)
|
|
assert bias_shard.global_offset == (
|
|
tp_rank * (out_f // tp_world_size),
|
|
)
|
|
|
|
elif self.layer_type == "RowSequenceParallelLinear":
|
|
in_f, out_f = self.hidden_size * 2, self.hidden_size
|
|
assert weight_shard.global_shape == (in_f, out_f)
|
|
assert weight_shard.local_shape == (in_f // tp_world_size, out_f)
|
|
assert weight_shard.global_offset == (
|
|
tp_rank * (in_f // tp_world_size),
|
|
0,
|
|
)
|
|
if self.has_bias:
|
|
assert bias_shard.global_shape == (out_f,)
|
|
assert bias_shard.local_shape == bias_shard.global_shape
|
|
assert bias_shard.global_offset == (0,)
|
|
|
|
def run_optimizer_test(self):
|
|
model = SimpleMLP(has_bias=self.has_bias)
|
|
model = paddle.amp.decorate(
|
|
models=model, optimizers=None, level="O2", dtype="float16"
|
|
)
|
|
if self.master_weight: # test the master_weight
|
|
opt = paddle.optimizer.AdamW(
|
|
learning_rate=0.01,
|
|
parameters=model.parameters(),
|
|
multi_precision=True,
|
|
)
|
|
else:
|
|
opt = paddle.optimizer.AdamW(
|
|
learning_rate=0.01,
|
|
parameters=model.parameters(),
|
|
multi_precision=False,
|
|
)
|
|
if self.layer_type == "AdamW":
|
|
model = fleet.distributed_model(model)
|
|
model.train()
|
|
x = paddle.randint(
|
|
low=0,
|
|
high=self.vocab_size,
|
|
shape=[self.batch_size, self.seq_len, self.hidden_size],
|
|
dtype='int64',
|
|
)
|
|
y = model(x).mean()
|
|
y.backward()
|
|
opt.step()
|
|
opt.clear_grad()
|
|
|
|
model_sharded_state_dict = model.sharded_state_dict()
|
|
opt_sharded_state_dict = opt.sharded_state_dict(
|
|
model_sharded_state_dict
|
|
)
|
|
for key, value in model_sharded_state_dict.items():
|
|
for state_name in self.optimizer_var_suffix:
|
|
opt__var_name = key + state_name
|
|
if opt__var_name in opt_sharded_state_dict:
|
|
assert tuple(
|
|
opt_sharded_state_dict[opt__var_name].local_shape
|
|
) == tuple(value.local_shape)
|
|
assert tuple(
|
|
opt_sharded_state_dict[opt__var_name].global_shape
|
|
) == tuple(value.global_shape)
|
|
assert tuple(
|
|
opt_sharded_state_dict[opt__var_name].global_offset
|
|
) == tuple(value.global_offset)
|
|
elif self.layer_type == "DygraphShardingOptimizer":
|
|
opt = DygraphShardingOptimizer(opt, self.hcg)
|
|
model.train()
|
|
x = paddle.randint(
|
|
low=0,
|
|
high=self.vocab_size,
|
|
shape=[self.batch_size, self.seq_len, self.hidden_size],
|
|
dtype='int64',
|
|
)
|
|
rank = paddle.distributed.get_rank()
|
|
sharidng_x = (
|
|
x[0 : self.batch_size // 2]
|
|
if rank == 0
|
|
else x[self.batch_size // 2 :]
|
|
)
|
|
y = model(sharidng_x).mean()
|
|
y.backward()
|
|
opt.step()
|
|
opt.clear_grad()
|
|
|
|
model_sharded_state_dict = model.sharded_state_dict()
|
|
opt_sharded_state_dict = opt.sharded_state_dict(
|
|
model_sharded_state_dict
|
|
)
|
|
|
|
for key, value in model_sharded_state_dict.items():
|
|
for state_name in self.optimizer_var_suffix:
|
|
opt__var_name = key + state_name
|
|
if opt__var_name in opt_sharded_state_dict:
|
|
assert tuple(
|
|
opt_sharded_state_dict[opt__var_name].local_shape
|
|
) == tuple(value.local_shape)
|
|
assert tuple(
|
|
opt_sharded_state_dict[opt__var_name].global_shape
|
|
) == tuple(value.global_shape)
|
|
assert tuple(
|
|
opt_sharded_state_dict[opt__var_name].global_offset
|
|
) == tuple(value.global_offset)
|
|
elif self.layer_type == "DygraphShardingOptimizerV2":
|
|
opt = DygraphShardingOptimizerV2(opt, self.hcg)
|
|
model.train()
|
|
x = paddle.randint(
|
|
low=0,
|
|
high=self.vocab_size,
|
|
shape=[self.batch_size, self.seq_len, self.hidden_size],
|
|
dtype='int64',
|
|
)
|
|
rank = paddle.distributed.get_rank()
|
|
sharidng_x = (
|
|
x[0 : self.batch_size // 2]
|
|
if rank == 0
|
|
else x[self.batch_size // 2 :]
|
|
)
|
|
y = model(sharidng_x).mean()
|
|
y.backward()
|
|
opt.step()
|
|
opt.clear_grad()
|
|
|
|
model_sharded_state_dict = model.sharded_state_dict()
|
|
opt_sharded_state_dict = opt.sharded_state_dict(
|
|
model_sharded_state_dict
|
|
)
|
|
for key, value in model_sharded_state_dict.items():
|
|
for state_name in self.optimizer_var_suffix:
|
|
opt__var_name = key + state_name
|
|
if opt__var_name in opt_sharded_state_dict:
|
|
if opt_sharded_state_dict[
|
|
opt__var_name
|
|
].flattened_range.stop - opt_sharded_state_dict[
|
|
opt__var_name
|
|
].flattened_range.start != math.prod(
|
|
value.local_shape
|
|
): # check the optimizer_var which isFragment
|
|
opt_var_globle_flattened_range = []
|
|
paddle.distributed.all_gather_object(
|
|
opt_var_globle_flattened_range,
|
|
opt_sharded_state_dict[
|
|
opt__var_name
|
|
].flattened_range,
|
|
)
|
|
|
|
first_fragment = opt_var_globle_flattened_range[0]
|
|
second_fragment = opt_var_globle_flattened_range[1]
|
|
assert (
|
|
first_fragment.stop == second_fragment.start
|
|
) # the first_flattened_range_stop == the second_flattened_range_start
|
|
opt_var_globle_size_flattened = (
|
|
second_fragment.stop - first_fragment.start
|
|
)
|
|
model_var_globle_size_flattened = math.prod(
|
|
value.local_shape
|
|
)
|
|
assert (
|
|
opt_var_globle_size_flattened
|
|
== model_var_globle_size_flattened
|
|
)
|
|
|
|
assert tuple(
|
|
opt_sharded_state_dict[opt__var_name].local_shape
|
|
) == tuple(value.local_shape)
|
|
assert tuple(
|
|
opt_sharded_state_dict[opt__var_name].global_shape
|
|
) == tuple(value.global_shape)
|
|
assert tuple(
|
|
opt_sharded_state_dict[opt__var_name].global_offset
|
|
) == tuple(value.global_offset)
|
|
|
|
elif self.layer_type == "GroupShardedOptimizerStage2":
|
|
opt = GroupShardedOptimizerStage2(
|
|
opt._parameter_list, opt, self.hcg.get_sharding_parallel_group()
|
|
)
|
|
|
|
model.train()
|
|
x = paddle.randint(
|
|
low=0,
|
|
high=self.vocab_size,
|
|
shape=[self.batch_size, self.seq_len, self.hidden_size],
|
|
dtype='int64',
|
|
)
|
|
rank = paddle.distributed.get_rank()
|
|
sharidng_x = (
|
|
x[0 : self.batch_size // 2]
|
|
if rank == 0
|
|
else x[self.batch_size // 2 :]
|
|
)
|
|
y = model(sharidng_x).mean()
|
|
y.backward()
|
|
opt.step()
|
|
opt.clear_grad()
|
|
|
|
model_sharded_state_dict = model.sharded_state_dict()
|
|
opt_sharded_state_dict = opt.sharded_state_dict(
|
|
model_sharded_state_dict
|
|
)
|
|
|
|
for key, value in model_sharded_state_dict.items():
|
|
for state_name in self.optimizer_var_suffix:
|
|
opt__var_name = key + state_name
|
|
if opt__var_name in opt_sharded_state_dict:
|
|
assert tuple(
|
|
opt_sharded_state_dict[opt__var_name].local_shape
|
|
) == tuple(value.local_shape)
|
|
assert tuple(
|
|
opt_sharded_state_dict[opt__var_name].global_shape
|
|
) == tuple(value.global_shape)
|
|
assert tuple(
|
|
opt_sharded_state_dict[opt__var_name].global_offset
|
|
) == tuple(value.global_offset)
|
|
|
|
elif self.layer_type == "GroupShardedStage3":
|
|
model = fleet.distributed_model(model)
|
|
wrapped_model = GroupShardedStage3(
|
|
model, opt, segment_size=2**12
|
|
) # slice the linear1、linear2 weight
|
|
for param in opt._parameter_list:
|
|
if hasattr(param, "fw_storage"):
|
|
assert len(param.shape) != 1
|
|
wrapped_model.init_optimizer_for_slice_param()
|
|
for param in opt._parameter_list:
|
|
if hasattr(param, "fw_storage"):
|
|
assert len(param.shape) == 1
|
|
model_sharded_state_dict = model.sharded_state_dict()
|
|
for k, v in model_sharded_state_dict.items():
|
|
if (
|
|
k == "_layers.linear1.weight"
|
|
or k == "_layers.linear2.weight"
|
|
):
|
|
assert not v.local_tensor._is_initialized()
|
|
wrapped_model.init_slice_param()
|
|
for k, v in model_sharded_state_dict.items():
|
|
if (
|
|
k == "_layers.linear1.weight"
|
|
or k == "_layers.linear2.weight"
|
|
):
|
|
assert v.local_tensor._is_initialized()
|
|
wrapped_model.align_param_to_buffer_and_clear_slice_param()
|
|
for k, v in model_sharded_state_dict.items():
|
|
if (
|
|
k == "_layers.linear1.weight"
|
|
or k == "_layers.linear2.weight"
|
|
):
|
|
assert not v.local_tensor._is_initialized()
|
|
model.train()
|
|
x = paddle.randint(
|
|
low=0,
|
|
high=self.vocab_size,
|
|
shape=[self.batch_size, self.seq_len, self.hidden_size],
|
|
dtype='int64',
|
|
)
|
|
rank = paddle.distributed.get_rank()
|
|
sharidng_x = (
|
|
x[0 : self.batch_size // 2]
|
|
if rank == 0
|
|
else x[self.batch_size // 2 :]
|
|
)
|
|
y = model(sharidng_x).mean()
|
|
y.backward()
|
|
opt.step()
|
|
opt.clear_grad()
|
|
model_sharded_state_dict = model.sharded_state_dict()
|
|
for k, v in model_sharded_state_dict.items():
|
|
if (
|
|
k == "_layers.linear1.weight"
|
|
or k == "_layers.linear2.weight"
|
|
):
|
|
assert not v.local_tensor._is_initialized()
|
|
wrapped_model.get_all_parameters()
|
|
opt_sharded_state_dict = opt.sharded_state_dict(
|
|
model_sharded_state_dict
|
|
)
|
|
|
|
for k, v in model_sharded_state_dict.items():
|
|
if (
|
|
k == "_layers.linear1.weight"
|
|
or k == "_layers.linear2.weight"
|
|
):
|
|
assert v.local_tensor._is_initialized()
|
|
|
|
for key, value in model_sharded_state_dict.items():
|
|
for state_name in self.optimizer_var_suffix:
|
|
opt__var_name = key + state_name
|
|
if opt__var_name in opt_sharded_state_dict:
|
|
if hasattr(
|
|
value.local_tensor, "fw_storage"
|
|
): # check the optimizer_var which isFragment
|
|
opt_var_globle_flattened_range = []
|
|
paddle.distributed.all_gather_object(
|
|
opt_var_globle_flattened_range,
|
|
opt_sharded_state_dict[
|
|
opt__var_name
|
|
].flattened_range,
|
|
)
|
|
|
|
first_fragment = opt_var_globle_flattened_range[0]
|
|
second_fragment = opt_var_globle_flattened_range[1]
|
|
assert (
|
|
first_fragment.stop == second_fragment.start
|
|
) # the first_flattened_range_stop == the second_flattened_range_start
|
|
opt_var_globle_size_flattened = (
|
|
second_fragment.stop - first_fragment.start
|
|
)
|
|
model_var_globle_size_flattened = math.prod(
|
|
value.local_shape
|
|
)
|
|
assert (
|
|
opt_var_globle_size_flattened
|
|
== model_var_globle_size_flattened
|
|
)
|
|
|
|
assert tuple(
|
|
opt_sharded_state_dict[opt__var_name].local_shape
|
|
) == tuple(value.local_shape)
|
|
assert tuple(
|
|
opt_sharded_state_dict[opt__var_name].global_shape
|
|
) == tuple(value.global_shape)
|
|
assert tuple(
|
|
opt_sharded_state_dict[opt__var_name].global_offset
|
|
) == tuple(value.global_offset)
|
|
else:
|
|
raise ValueError(f"Unknown layer_type: {self.layer_type}")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
TestParallelLayersLogic().run_test()
|