wehub-resource-sync
157 lines
4.9 KiB
Python
157 lines
4.9 KiB
Python
# 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 glob
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
import paddle
|
|
from paddle.distributed import fleet
|
|
|
|
sys.path.append(str(Path(__file__).parent.parent.parent))
|
|
from tests.parallel_launch import TestMultipleGpus
|
|
from tests.testing_utils import require_paddle_at_least_2_gpu
|
|
|
|
tp_size = paddle.distributed.get_world_size()
|
|
tp_rank = 0
|
|
if tp_size > 1:
|
|
strategy = fleet.DistributedStrategy()
|
|
strategy.hybrid_configs = {
|
|
"dp_degree": 1,
|
|
"mp_degree": tp_size,
|
|
"pp_degree": 1,
|
|
"sharding_degree": 1,
|
|
}
|
|
fleet.init(is_collective=True, strategy=strategy)
|
|
hcg = fleet.get_hybrid_communicate_group()
|
|
tp_rank = hcg.get_model_parallel_rank()
|
|
mp_group = hcg.get_model_parallel_group()
|
|
|
|
|
|
def prepare_config(config):
|
|
config.hidden_size = 512
|
|
config.num_layers = 2
|
|
config.num_hidden_layers = 2
|
|
config.num_attention_heads = 16
|
|
config.num_key_value_heads = 16
|
|
config.intermediate_size = config.hidden_size * 3
|
|
config.tensor_parallel_degree = tp_size
|
|
config.tensor_parallel_rank = tp_rank
|
|
return config
|
|
|
|
|
|
def common_test_load(model_class, tempdir):
|
|
paddle.distributed.barrier()
|
|
if model_class is not None:
|
|
model_class.from_pretrained(tempdir)
|
|
paddle.distributed.barrier()
|
|
if paddle.distributed.get_rank() == 0:
|
|
files = glob.glob(tempdir + "/*")
|
|
for f in files:
|
|
os.remove(f)
|
|
paddle.distributed.barrier()
|
|
|
|
|
|
def common_test_merge(model, model_class=None):
|
|
rank = paddle.distributed.get_rank()
|
|
is_main_process = rank == 0
|
|
object_list = []
|
|
with tempfile.TemporaryDirectory() as tempdir:
|
|
paddle.distributed.all_gather_object(object_list, tempdir, group=mp_group)
|
|
tempdir = object_list[0]
|
|
# test merge one
|
|
model.save_pretrained(save_dir=tempdir, merge_tensor_parallel=True, is_main_process=is_main_process)
|
|
common_test_load(model_class, tempdir)
|
|
# test merge shard
|
|
model.save_pretrained(
|
|
tempdir,
|
|
merge_tensor_parallel=True,
|
|
variant=f"tp{rank:0>2d}",
|
|
max_shard_size="5MB",
|
|
is_main_process=is_main_process,
|
|
)
|
|
common_test_load(model_class, tempdir)
|
|
# test save tp
|
|
model.save_pretrained(tempdir, max_shard_size="5MB", is_main_process=is_main_process)
|
|
common_test_load(model_class, tempdir)
|
|
# test save shard safe
|
|
model.save_pretrained(tempdir, max_shard_size="5MB", safe_serialization=True, is_main_process=is_main_process)
|
|
common_test_load(model_class, tempdir)
|
|
# test save safe tensor
|
|
model.save_pretrained(tempdir, safe_serialization=True, is_main_process=is_main_process)
|
|
common_test_load(model_class, tempdir)
|
|
paddle.distributed.barrier()
|
|
|
|
|
|
def _test_llama():
|
|
from paddlenlp.transformers import LlamaConfig, LlamaForCausalLM
|
|
|
|
config = LlamaConfig()
|
|
config = prepare_config(config)
|
|
model = LlamaForCausalLM.from_config(config)
|
|
common_test_merge(model, LlamaForCausalLM)
|
|
|
|
|
|
def _test_chatglm():
|
|
from paddlenlp.transformers import ChatGLMConfig, ChatGLMForCausalLM
|
|
|
|
config = ChatGLMConfig()
|
|
config = prepare_config(config)
|
|
model = ChatGLMForCausalLM.from_config(config)
|
|
common_test_merge(model, ChatGLMForCausalLM)
|
|
|
|
|
|
def _test_bloom():
|
|
from paddlenlp.transformers import BloomConfig, BloomForCausalLM
|
|
|
|
config = BloomConfig()
|
|
config = prepare_config(config)
|
|
model = BloomForCausalLM.from_config(config)
|
|
common_test_merge(model, BloomForCausalLM)
|
|
|
|
|
|
def _test_qwen2():
|
|
from paddlenlp.transformers import Qwen2Config, Qwen2ForCausalLM
|
|
|
|
config = Qwen2Config()
|
|
config = prepare_config(config)
|
|
model = Qwen2ForCausalLM.from_config(config)
|
|
common_test_merge(model, Qwen2ForCausalLM)
|
|
|
|
|
|
def _test_gemma():
|
|
from paddlenlp.transformers import GemmaConfig, GemmaForCausalLM
|
|
|
|
config = GemmaConfig()
|
|
config = prepare_config(config)
|
|
model = GemmaForCausalLM.from_config(config)
|
|
common_test_merge(model, GemmaForCausalLM)
|
|
|
|
|
|
@require_paddle_at_least_2_gpu
|
|
class TestTensorParallel(TestMultipleGpus):
|
|
def test_model_load_merge(self):
|
|
self.run_2gpu(__file__)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
_test_llama()
|
|
_test_chatglm()
|
|
_test_bloom()
|
|
_test_gemma()
|
|
_test_qwen2()
|