Add files using upload-large-folder tool

tests/general/test_dataset.py

@@ -0,0 +1,90 @@
+from typing import List
+
+from swift.llm import load_dataset
+
+
+def _test_dataset(datasets: List[str], num_proc: int = 1, strict: bool = False, **kwargs):
+    dataset = load_dataset(datasets, num_proc=num_proc, strict=strict, **kwargs)
+    print(f'dataset[0]: {dataset[0]}')
+    print(f'dataset[1]: {dataset[1]}')
+
+
+def test_sft():
+    # swift/SlimOrca  swift/cosmopedia-100k
+    # _test_dataset(['lvjianjin/AdvertiseGen'])
+    # _test_dataset(['AI-ModelScope/Duet-v0.5'])
+    # _test_dataset(['swift/SlimOrca', 'swift/cosmopedia-100k'])
+    # _test_dataset(['OmniData/Zhihu-KOL-More-Than-100-Upvotes'])
+    # _test_dataset(['OmniData/Zhihu-KOL'])
+    _test_dataset([
+        'AI-ModelScope/alpaca-gpt4-data-zh#1000', 'AI-ModelScope/alpaca-gpt4-data-en#1000',
+        'AI-ModelScope/LongAlpaca-12k#1000'
+    ])
+    # _test_dataset(['swift/Infinity-Instruct:all'])
+    # _test_dataset(['swift/sharegpt:all'])
+    # _test_dataset(['AI-ModelScope/sharegpt_gpt4:all'])
+    # _test_dataset(['iic/ms_bench'])
+    # _test_dataset(['swift/tagengo-gpt4'])
+
+
+def test_mllm():
+    # _test_dataset(['AI-ModelScope/ShareGPT4V:all'])
+    # _test_dataset(['AI-ModelScope/LLaVA-Pretrain'])
+    # _test_dataset(['swift/TextCaps'])
+    # _test_dataset(['swift/RLAIF-V-Dataset:all'])
+    # _test_dataset(['swift/OK-VQA_train'])
+    # _test_dataset(['swift/OCR-VQA'])
+    # _test_dataset(['swift/A-OKVQA'])
+    # _test_dataset(['AI-ModelScope/MovieChat-1K-test'])
+    _test_dataset([
+        'AI-ModelScope/LaTeX_OCR:all', 'modelscope/coco_2014_caption:validation',
+        'speech_asr/speech_asr_aishell1_trainsets:validation'
+    ],
+                  strict=False)
+    # _test_dataset(['swift/VideoChatGPT:all'])
+    # _test_dataset(['speech_asr/speech_asr_aishell1_trainsets:validation'])
+    # _test_dataset(['AI-ModelScope/captcha-images'])
+    # _test_dataset(['swift/gpt4v-dataset:all'])
+    # _test_dataset(['modelscope/coco_2014_caption:validation'])
+    # _test_dataset(['AI-ModelScope/LLaVA-Instruct-150K'], num_proc=16)
+
+
+def test_agent():
+    _test_dataset(['swift/ToolBench'])
+    # _test_dataset(['AI-ModelScope/ms_agent_for_agentfabric:all'])
+
+
+def test_dpo():
+    _test_dataset(['AI-ModelScope/orpo-dpo-mix-40k'])
+    _test_dataset(['AI-ModelScope/hh-rlhf:all'])
+    _test_dataset(['AI-ModelScope/hh_rlhf_cn:all'])
+    _test_dataset(['hjh0119/shareAI-Llama3-DPO-zh-en-emoji:all'])
+
+
+def test_kto():
+    _test_dataset(['AI-ModelScope/ultrafeedback-binarized-preferences-cleaned-kto'])
+
+
+def test_pretrain():
+    _test_dataset(['AI-ModelScope/ruozhiba:all'])
+
+
+def test_dataset_info():
+    _test_dataset(['swift/self-cognition#500'], model_name='xiao huang', model_author='swift')
+    # _test_dataset(['codefuse-ai/CodeExercise-Python-27k'])
+
+
+def test_cls():
+    _test_dataset(['simpleai/HC3-Chinese:baike'])
+    _test_dataset(['simpleai/HC3-Chinese:baike_cls'])
+
+
+if __name__ == '__main__':
+    # test_sft()
+    # test_agent()
+    # test_dpo()
+    # test_kto()
+    test_mllm()
+    # test_pretrain()
+    # test_dataset_info()
+    # test_cls()

tests/general/test_model.py

@@ -0,0 +1,30 @@
+import os
+
+import torch
+
+from swift.utils import get_device
+
+os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
+
+
+def test_qwen2():
+    import os
+    from swift.llm import get_model_tokenizer
+    model, tokenizer = get_model_tokenizer('Qwen/Qwen2-7B-Instruct', load_model=False)
+    print(f'model: {model}, tokenizer: {tokenizer}')
+    # test hf
+    model, tokenizer = get_model_tokenizer('Qwen/Qwen2-7B-Instruct', load_model=False, use_hf=True)
+
+    model, tokenizer = get_model_tokenizer(
+        'Qwen/Qwen2-7B-Instruct', torch.float32, device_map=get_device(), attn_impl='flash_attn')
+    print(f'model: {model}, tokenizer: {tokenizer}')
+
+
+def test_modelscope_hub():
+    from swift.llm import get_model_tokenizer
+    model, tokenizer = get_model_tokenizer('Qwen/Qwen2___5-Math-1___5B-Instruct/', load_model=False)
+
+
+if __name__ == '__main__':
+    test_qwen2()
+    # test_modelscope_hub()

tests/general/test_stream.py

@@ -0,0 +1,20 @@
+from swift.llm import load_dataset
+
+
+def test_local_dataset():
+    # please use git clone
+    from swift.llm import git_clone_github
+    model_dir = git_clone_github('https://www.modelscope.cn/datasets/swift/swift-sft-mixture.git')
+    dataset = load_dataset(datasets=[f'{model_dir}:firefly'], streaming=True)[0]
+    print(next(iter(dataset)))
+
+
+def test_hub_dataset():
+    local_dataset = 'swift/swift-sft-mixture:firefly'
+    dataset = load_dataset(datasets=[local_dataset], streaming=True)[0]
+    print(next(iter(dataset)))
+
+
+if __name__ == '__main__':
+    test_local_dataset()
+    # test_hub_dataset()

tests/general/test_template.py

@@ -0,0 +1,74 @@
+from datasets import Dataset
+
+from swift.llm import EncodePreprocessor, TemplateInputs, get_model_tokenizer, get_template, load_dataset
+
+
+def test_template():
+    _, tokenizer = get_model_tokenizer('Qwen/Qwen2-7B-Instruct', load_model=False)
+    template = get_template(tokenizer.model_meta.template, tokenizer)
+    template_inputs = TemplateInputs([{
+        'role': 'system',
+        'content': 'AAA'
+    }, {
+        'role': 'user',
+        'content': 'BBB'
+    }, {
+        'role': 'assistant',
+        'content': 'CCC'
+    }, {
+        'role': 'user',
+        'content': 'DDD'
+    }])
+    inputs = template.encode(template_inputs)
+    print(f'inputs.keys(): {inputs.keys()}')
+    print(tokenizer.decode(inputs['input_ids']))
+
+
+def test_mllm():
+    _, tokenizer = get_model_tokenizer('Qwen/Qwen2-VL-7B-Instruct', load_model=False)
+    template = get_template(tokenizer.model_meta.template, tokenizer)
+    template_inputs = TemplateInputs([{
+        'role': 'system',
+        'content': 'AAA'
+    }, {
+        'role': 'user',
+        'content': '<image>BBB'
+    }, {
+        'role': 'assistant',
+        'content': 'CCC'
+    }, {
+        'role': 'user',
+        'content': 'DDD'
+    }],
+                                     images=['http://modelscope-open.oss-cn-hangzhou.aliyuncs.com/images/cat.png'])
+    inputs = template.encode(template_inputs)
+    print(f'inputs.keys(): {inputs.keys()}')
+    print(template.safe_decode(inputs['input_ids']))
+
+
+def _test_dataset_map(model_id: str, dataset_id: str):
+    _, tokenizer = get_model_tokenizer(model_id, load_model=False)
+    template = get_template(tokenizer.model_meta.template, tokenizer)
+    dataset = load_dataset([dataset_id], num_proc=2)[0]
+
+    # 1: 1500
+    # 16: 10766.36 examples/s
+    new_dataset = EncodePreprocessor(template)(dataset, num_proc=4)
+    print(f'new_dataset: {new_dataset}')
+    print(template.safe_decode(new_dataset[0]['input_ids']))
+    print(template.safe_decode(new_dataset[1]['input_ids']))
+
+
+def test_llm_dataset_map():
+    _test_dataset_map('Qwen/Qwen2-7B-Instruct', 'AI-ModelScope/alpaca-gpt4-data-zh')
+
+
+def test_mllm_dataset_map():
+    _test_dataset_map('Qwen/Qwen2-VL-7B-Instruct', 'modelscope/coco_2014_caption:validation#100')
+
+
+if __name__ == '__main__':
+    # test_template()
+    # test_mllm()
+    # test_llm_dataset_map()
+    test_mllm_dataset_map()

tests/hub/test_check_model.py

@@ -0,0 +1,24 @@
+import os
+import shutil
+import tempfile
+import unittest
+
+from modelscope import Model, check_local_model_is_latest
+
+
+class TestCheckModel(unittest.TestCase):
+
+    def setUp(self):
+        print(('Testing %s.%s' % (type(self).__name__, self._testMethodName)))
+        self.tmp_dir = tempfile.TemporaryDirectory().name
+        if not os.path.exists(self.tmp_dir):
+            os.makedirs(self.tmp_dir)
+
+    def tearDown(self):
+        import peft
+        shutil.rmtree(self.tmp_dir)
+        super().tearDown()
+
+    def test_check_model(self):
+        model = Model.from_pretrained('damo/nlp_corom_sentence-embedding_chinese-base', revision='v1.0.0')
+        self.assertFalse(check_local_model_is_latest(model.model_dir))

tests/infer/test_agent.py

@@ -0,0 +1,30 @@
+import os
+
+import torch
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+
+kwargs = {
+    'per_device_train_batch_size': 2,
+    'save_steps': 50,
+    'gradient_accumulation_steps': 4,
+    'num_train_epochs': 1,
+}
+
+
+def test_sft():
+    os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
+    from swift.llm import sft_main, TrainArguments, infer_main, InferArguments
+    sft_main(
+        TrainArguments(model='Qwen/Qwen2-7B-Instruct', dataset=['iic/ms_agent#2000'], loss_scale='react', **kwargs))
+
+
+def test_infer():
+    from swift.llm import infer_main, InferArguments
+    ckpt_dir = 'output/Qwen2-7B-Instruct/v229-20241126-133152/checkpoint-100'
+    infer_main(InferArguments(ckpt_dir=ckpt_dir))
+
+
+if __name__ == '__main__':
+    test_sft()
+    # test_infer()

tests/infer/test_infer.py

@@ -0,0 +1,73 @@
+import os
+from typing import Literal
+
+import torch
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+
+
+def _prepare(infer_backend: Literal['vllm', 'pt', 'lmdeploy']):
+    from swift.llm import InferRequest, get_template
+    if infer_backend == 'lmdeploy':
+        from swift.llm import LmdeployEngine
+        engine = LmdeployEngine('OpenGVLab/InternVL2_5-2B', torch.float32)
+    elif infer_backend == 'pt':
+        from swift.llm import PtEngine
+        engine = PtEngine('Qwen/Qwen2-7B-Instruct', max_batch_size=16)
+    elif infer_backend == 'vllm':
+        from swift.llm import VllmEngine
+        engine = VllmEngine('Qwen/Qwen2-7B-Instruct')
+    template = get_template(engine.model_meta.template, engine.tokenizer)
+    infer_requests = [
+        # InferRequest([{'role': 'user', 'content': '晚上睡不着觉怎么办'}]) for i in range(100)
+        InferRequest([{
+            'role': 'user',
+            'content': 'hello! who are you'
+        }]) for i in range(100)
+    ]
+    return engine, template, infer_requests
+
+
+def test_infer(infer_backend):
+    from swift.llm import RequestConfig
+    from swift.plugin import InferStats
+    engine, template, infer_requests = _prepare(infer_backend=infer_backend)
+    request_config = RequestConfig(temperature=0)
+    infer_stats = InferStats()
+
+    response_list = engine.infer(
+        infer_requests, template=template, request_config=request_config, metrics=[infer_stats])
+
+    for response in response_list[:2]:
+        print(response.choices[0].message.content)
+    print(infer_stats.compute())
+
+
+def test_stream(infer_backend):
+    from swift.llm import RequestConfig
+    from swift.plugin import InferStats
+    engine, template, infer_requests = _prepare(infer_backend=infer_backend)
+    infer_stats = InferStats()
+    request_config = RequestConfig(temperature=0, stream=True, logprobs=True)
+
+    gen_list = engine.infer(infer_requests, template=template, request_config=request_config, metrics=[infer_stats])
+
+    for response in gen_list[0]:
+        if response is None:
+            continue
+        print(response.choices[0].delta.content, end='', flush=True)
+    print()
+    print(infer_stats.compute())
+
+    gen_list = engine.infer(
+        infer_requests, template=template, request_config=request_config, use_tqdm=True, metrics=[infer_stats])
+
+    for response in gen_list[0]:
+        pass
+
+    print(infer_stats.compute())
+
+
+if __name__ == '__main__':
+    test_infer('pt')
+    # test_stream('pt')

tests/infer/test_logprobs.py

@@ -0,0 +1,71 @@
+import os
+from typing import Literal
+
+import torch
+
+if __name__ == '__main__':
+    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+
+
+def _prepare(infer_backend: Literal['vllm', 'pt', 'lmdeploy']):
+    from swift.llm import InferRequest, get_template
+
+    if infer_backend == 'lmdeploy':
+        from swift.llm import LmdeployEngine
+        engine = LmdeployEngine('Qwen/Qwen2-7B-Instruct', torch.float32)
+    elif infer_backend == 'pt':
+        from swift.llm import PtEngine
+        engine = PtEngine('Qwen/Qwen2-7B-Instruct')
+    elif infer_backend == 'vllm':
+        from swift.llm import VllmEngine
+        engine = VllmEngine('Qwen/Qwen2-7B-Instruct')
+    template = get_template(engine.model_meta.template, engine.tokenizer)
+    infer_requests = [
+        InferRequest([{
+            'role': 'user',
+            'content': '晚上睡不着觉怎么办'
+        }]),
+        InferRequest([{
+            'role': 'user',
+            'content': 'hello! who are you'
+        }])
+    ]
+    return engine, template, infer_requests
+
+
+def test_infer(engine, template, infer_requests):
+    from swift.llm import RequestConfig
+    from swift.plugin import InferStats
+
+    request_config = RequestConfig(temperature=0, logprobs=True, top_logprobs=2)
+    infer_stats = InferStats()
+
+    response_list = engine.infer(
+        infer_requests, template=template, request_config=request_config, metrics=[infer_stats])
+
+    for response in response_list[:2]:
+        print(response.choices[0].message.content)
+    print(infer_stats.compute())
+
+
+def test_stream(engine, template, infer_requests):
+    from swift.llm import RequestConfig
+    from swift.plugin import InferStats
+
+    infer_stats = InferStats()
+    request_config = RequestConfig(temperature=0, stream=True, logprobs=True, top_logprobs=2)
+
+    gen_list = engine.infer(infer_requests, template=template, request_config=request_config, metrics=[infer_stats])
+
+    for response in gen_list[0]:
+        if response is None:
+            continue
+        print(response.choices[0].delta.content, end='', flush=True)
+
+    print(infer_stats.compute())
+
+
+if __name__ == '__main__':
+    engine, template, infer_requests = _prepare(infer_backend='pt')
+    test_infer(engine, template, infer_requests)
+    test_stream(engine, template, infer_requests)

tests/infer/test_main.py

@@ -0,0 +1,73 @@
+import os
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+
+
+def test_cli(infer_backend):
+    from swift.llm import infer_main, InferArguments
+    args = InferArguments(model='Qwen/Qwen2-VL-7B-Instruct', infer_backend=infer_backend)
+    infer_main(args)
+
+
+def test_cli_jinja(infer_backend):
+    from swift.llm import infer_main, InferArguments
+    args = InferArguments(model='Qwen/Qwen2-VL-7B-Instruct', infer_backend=infer_backend, template_backend='jinja')
+    infer_main(args)
+
+
+def test_dataset(infer_backend):
+    from swift.llm import infer_main, InferArguments
+    args = InferArguments(
+        model='Qwen/Qwen2-7B-Instruct',
+        infer_backend=infer_backend,
+        val_dataset=['AI-ModelScope/alpaca-gpt4-data-zh#10'],
+        stream=True)
+    infer_main(args)
+
+
+def test_mllm_dataset(infer_backend):
+    from swift.llm import infer_main, InferArguments
+    args = InferArguments(
+        model='Qwen/Qwen2-VL-7B-Instruct',
+        infer_backend=infer_backend,
+        val_dataset=['modelscope/coco_2014_caption:validation#1000'],
+        stream=True)
+    infer_main(args)
+
+
+def test_dataset_ddp():
+    os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2,3'
+    from swift.llm import infer_main, InferArguments
+    args = InferArguments(
+        model='Qwen/Qwen2-7B-Instruct', max_batch_size=64, val_dataset=['AI-ModelScope/alpaca-gpt4-data-zh#1000'])
+    infer_main(args)
+
+
+def test_dataset_mp_ddp():
+    os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2,3'
+    from swift.llm import infer_main, InferArguments
+    args = InferArguments(
+        model='Qwen/Qwen2-7B-Instruct', max_batch_size=64, val_dataset=['AI-ModelScope/alpaca-gpt4-data-zh#1000'])
+    infer_main(args)
+
+
+def test_emu3_gen(infer_backend):
+    from swift.llm import infer_main, InferArguments
+    args = InferArguments(
+        model='BAAI/Emu3-Gen',
+        infer_backend=infer_backend,
+        stream=False,
+        use_chat_template=False,
+        top_k=2048,
+        max_new_tokens=40960)
+    infer_main(args)
+
+
+if __name__ == '__main__':
+    # test_cli('pt')
+    # test_cli_jinja('pt')
+    # test_dataset('pt')
+    # test_mllm_dataset('pt')
+    # test_dataset_ddp()
+    # test_dataset_mp_ddp()
+    test_emu3_gen('pt')

tests/infer/test_max_memory.py

@@ -0,0 +1,10 @@
+from swift.llm import InferArguments, infer_main
+
+
+def test_max_memory():
+    infer_main(
+        InferArguments(model='Qwen/Qwen2.5-7B-Instruct', max_memory='{0: "50GB", 1: "5GB"}', device_map='sequential'))
+
+
+if __name__ == '__main__':
+    test_max_memory()

tests/infer/test_mllm.py

@@ -0,0 +1,79 @@
+import os
+from typing import Literal
+
+import torch
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+
+
+def _prepare(infer_backend: Literal['vllm', 'pt', 'lmdeploy']):
+    from swift.llm import InferRequest, get_template
+    if infer_backend == 'lmdeploy':
+        from swift.llm import LmdeployEngine
+        engine = LmdeployEngine('Qwen/Qwen-VL-Chat', torch.float32)
+    elif infer_backend == 'pt':
+        from swift.llm import PtEngine
+        engine = PtEngine('Qwen/Qwen2-VL-7B-Instruct')
+    elif infer_backend == 'vllm':
+        from swift.llm import VllmEngine
+        engine = VllmEngine('Qwen/Qwen2-VL-7B-Instruct')
+    template = get_template(engine.model_meta.template, engine.processor)
+    infer_requests = [
+        InferRequest([{
+            'role': 'user',
+            'content': '晚上睡不着觉怎么办'
+        }]),
+        InferRequest([{
+            'role':
+            'user',
+            'content': [{
+                'type': 'image_url',
+                'image_url': 'http://modelscope-open.oss-cn-hangzhou.aliyuncs.com/images/cat.png'
+            }]
+        }])
+    ]
+    return engine, template, infer_requests
+
+
+def test_infer(engine, template, infer_requests):
+    from swift.llm import RequestConfig
+    from swift.plugin import InferStats
+    request_config = RequestConfig(temperature=0)
+    infer_stats = InferStats()
+
+    response_list = engine.infer(
+        infer_requests, template=template, request_config=request_config, metrics=[infer_stats])
+
+    for response in response_list[:2]:
+        print(response.choices[0].message.content)
+    print(infer_stats.compute())
+
+
+def test_stream(engine, template, infer_requests):
+    from swift.llm import RequestConfig
+    from swift.plugin import InferStats
+    infer_stats = InferStats()
+    request_config = RequestConfig(temperature=0, stream=True, logprobs=True)
+
+    gen_list = engine.infer(infer_requests, template=template, request_config=request_config, metrics=[infer_stats])
+
+    for response in gen_list[0]:
+        if response is None:
+            continue
+        print(response.choices[0].delta.content, end='', flush=True)
+    print()
+    print(infer_stats.compute())
+
+    gen_list = engine.infer(
+        infer_requests, template=template, request_config=request_config, use_tqdm=True, metrics=[infer_stats])
+
+    for response in gen_list[0]:
+        pass
+
+    print(infer_stats.compute())
+
+
+if __name__ == '__main__':
+    engine, template, infer_requests = _prepare(infer_backend='pt')
+    test_infer(engine, template, infer_requests)
+    test_stream(engine, template, infer_requests)

tests/llm/config/infer.json

@@ -0,0 +1,5 @@
+{
+    "ckpt_dir": "/mnt/workspace/yzhao/modelscope/swift/output/pai_test/checkpoint-6",
+    "val_dataset_sample": 2,
+    "load_dataset_config": true
+}

tests/llm/config/sft.json

@@ -0,0 +1,7 @@
+{
+    "model_type": "qwen-1_8b-chat",
+    "dataset": "jd-sentiment-zh",
+    "output_dir": "output/pai_test",
+    "train_dataset_sample": 100,
+    "eval_steps": 5
+}

tests/llm/data/alpaca.csv

@@ -0,0 +1,4 @@
+system,instruction,input,output
+00000,11111,22222,3.3
+,aaaaa,,ccccc
+,AAAAA,BBBBB,CCCCC

tests/llm/data/alpaca.jsonl

@@ -0,0 +1,3 @@
+{"instruction": "11111", "input": "22222", "output": "33333", "history": [["aaaaa", "bbbbb"]], "system": "system123"}
+{"instruction": "aaaaa", "output": "ccccc"}
+{"instruction": "AAAAA", "input": "BBBBB", "output": "CCCCC"}

tests/llm/data/alpaca2.csv

@@ -0,0 +1,4 @@
+instruction,output
+11111,33333
+aaaaa,ccccc
+AAAAA,CCCCC

tests/llm/data/chatml.jsonl

@@ -0,0 +1,3 @@
+{"messages": [{"role": "system", "content": "00000"}, {"role": "user", "content": "11111"}, {"role": "assistant", "content": "22222"}]}
+{"messages": [{"role": "user", "content": "aaaaa"}, {"role": "assistant", "content": "bbbbb"}, {"role": "user", "content": "ccccc"}, {"role": "assistant", "content": "ddddd"}]}
+{"messages": [{"role": "user", "content": "AAAAA"}, {"role": "assistant", "content": "BBBBB"}, {"role": "user", "content": "CCCCC"}, {"role": "assistant", "content": "DDDDD"}]}

tests/llm/data/conversations.jsonl

@@ -0,0 +1,3 @@
+{"conversations": [{"from": "system", "value": "00000"}, {"from": "user", "value": "11111"}, {"from": "assistant", "value": "22222"}]}
+{"conversations": [{"from": "user", "value": "aaaaa"}, {"from": "assistant", "value": "bbbbb"}, {"from": "user", "value": "ccccc"}, {"from": "assistant", "value": "ddddd"}]}
+{"conversations": [{"from": "user", "value": "AAAAA"}, {"from": "assistant", "value": "BBBBB"}, {"from": "user", "value": "CCCCC"}, {"from": "assistant", "value": "DDDDD"}]}

tests/llm/data/multi_modal_1.jsonl

@@ -0,0 +1,3 @@
+{"query": "<img>https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg</img>55555", "response": "66666"}
+{"query": "<img>https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg</img><img>https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg</img>eeeee", "response": "fffff", "history": [["hello", "123"]]}
+{"query": "EEEEE", "response": "FFFFF", "history": [["AAAAA", "BBBBB"], ["CCCCC", "DDDDD"]]}

tests/llm/data/multi_modal_2.jsonl

@@ -0,0 +1,3 @@
+{"query": "55555", "response": "66666", "images": ["https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg"]}
+{"query": "eeeee", "response": "fffff", "history": [], "images": ["https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg"]}
+{"query": "EEEEE", "response": "FFFFF", "history": [["AAAAA", "BBBBB"], ["CCCCC", "DDDDD"]], "images": ["https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg", "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg", "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg"]}

tests/llm/data/multi_modal_3.jsonl

@@ -0,0 +1,3 @@
+{"query": "55555", "response": "66666", "images": ["https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg"]}
+{"query": "eeeee", "response": "fffff", "history": [], "images": ["https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg"]}
+{"query": "EEEEE", "response": "FFFFF", "history": [["AAAAA", "BBBBB"], ["CCCCC", "DDDDD"]], "images": ["https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg"]}

tests/llm/data/sharegpt.jsonl

@@ -0,0 +1,3 @@
+{"system": "00000", "conversation": [{"human": "11111", "assistant": "22222"}]}
+{"conversation": [{"human": "aaaaa", "assistant": "bbbbb"}]}
+{"conversation": [{"human": "AAAAA", "assistant": "BBBBB"}, {"human": "CCCCC", "assistant": "DDDDD"}, {"human": "EEEEE", "assistant": "FFFFF"}]}

tests/llm/data/swift_multi.json

@@ -0,0 +1,3 @@
+[{"system": "00000", "query": "55555", "response": "66666"},
+{"query": "eeeee", "response": "fffff", "history": []},
+{"query": "EEEEE", "response": "FFFFF", "history": [["AAAAA", "BBBBB"], ["CCCCC", "DDDDD"]]}]

tests/llm/data/swift_multi.jsonl

@@ -0,0 +1,3 @@
+{"system": "00000", "query": "55555", "response": "66666"}
+{"query": "eeeee", "response": "fffff", "history": []}
+{"query": "EEEEE", "response": "FFFFF", "history": [["AAAAA", "BBBBB"], ["CCCCC", "DDDDD"]]}

tests/llm/data/swift_pre.csv

@@ -0,0 +1,4 @@
+response
+11111
+aaaaa
+AAAAA

tests/llm/data/swift_pre.jsonl

@@ -0,0 +1,3 @@
+{"response": "11111"}
+{"response": "aaaaa"}
+{"response": "AAAAA"}

tests/llm/data/swift_single.csv

@@ -0,0 +1,4 @@
+system,query,response
+00000,11111,22222
+,aaaaa,bbbbb
+,AAAAA,BBBBB

tests/llm/data/swift_single.jsonl

@@ -0,0 +1,3 @@
+{"system": "00000", "query": "11111", "response": "22222"}
+{"query": "aaaaa", "response": "bbbbb"}
+{"query": "AAAAA", "response": "BBBBB"}

tests/llm/load_model.py

@@ -0,0 +1,45 @@
+import argparse
+from dataclasses import fields
+
+import torch
+
+from swift.llm import MODEL_ARCH_MAPPING, ModelKeys, get_model_tokenizer
+
+
+def get_model_and_tokenizer(ms_model_id, model_arch=None):
+    try:
+        import transformers
+        print(f'Test model: {ms_model_id} with transformers version: {transformers.__version__}')
+        model_ins, tokenizer = get_model_tokenizer(ms_model_id)
+        model_ins: torch.nn.Module
+        if model_arch:
+            model_arch: ModelKeys = MODEL_ARCH_MAPPING[model_arch]
+            for f in fields(model_arch):
+                value = getattr(model_arch, f.name)
+                if value is not None and f.name != 'arch_name':
+                    if isinstance(value, str):
+                        value = [value]
+                    for v in value:
+                        v = v.replace('{}', '0')
+                        model_ins.get_submodule(v)
+    except Exception:
+        import traceback
+        print(traceback.format_exc())
+        raise
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        '--ms_model_id',
+        type=str,
+        required=True,
+    )
+    parser.add_argument(
+        '--model_arch',
+        type=str,
+        required=True,
+    )
+    args = parser.parse_args()
+
+    get_model_and_tokenizer(args.ms_model_id, args.model_arch)

tests/llm/load_template.py

@@ -0,0 +1,138 @@
+import argparse
+from collections.abc import Mapping
+
+import json
+import torch
+from transformers import PreTrainedTokenizerBase
+
+
+def to_list(input_ids):
+    if isinstance(input_ids, torch.Tensor):
+        input_ids = input_ids.cpu().numpy().tolist()
+    if isinstance(input_ids, list) and isinstance(input_ids[0], list):
+        input_ids = input_ids[0]
+    return input_ids
+
+
+def load_ds(ds):
+    from swift.llm import load_dataset
+    train_dataset, val_dataset = load_dataset(
+        ds,
+        split_dataset_ratio=0.0,
+        strict=False,
+        num_proc=1,
+        model_name=['小黄', 'Xiao Huang'],
+        model_author=['魔搭', 'ModelScope'])
+    return train_dataset.select(range(1))
+
+
+def load_and_tokenize(ms_model_id, template):
+    from swift.llm import EncodePreprocessor, get_model_tokenizer, get_template
+    try:
+        vl_fields = ['vl', 'video', 'minicpmv', 'llava', 'vision', 'emu', 'florence']
+        model_ins, tokenizer = get_model_tokenizer(ms_model_id, load_model='mplug' in ms_model_id.lower())
+        template_ins = get_template(template, tokenizer)
+        if template_ins.use_model:
+            model_ins, _ = get_model_tokenizer(ms_model_id, load_model=True)
+            template_ins.model = model_ins
+        template_ins.set_mode('train')
+        if 'audio' in template_ins.__class__.__name__.lower():
+            output = EncodePreprocessor(template_ins)(
+                load_ds('speech_asr/speech_asr_aishell1_trainsets:validation/test'))
+            input_ids = output[0].get('input_ids')
+        elif any([vl in template for vl in vl_fields]):
+            for row in load_ds('modelscope/coco_2014_caption:validation'):
+                output = template_ins.encode(row)
+                input_ids = output.get('input_ids')
+                # output = EncodePreprocessor(template_ins)(load_ds('swift/OK-VQA_train'))
+                if model_ins is not None and model_ins.model_meta.is_multimodal:
+                    inputs = template_ins.pre_data_collator([output], model=model_ins)
+                    _, output = template_ins.pre_forward_hook(model_ins, None, inputs)
+        else:
+            output = EncodePreprocessor(template_ins)(load_ds('modelscope/DuReader_robust-QG'))
+            input_ids = output[0].get('input_ids')
+        if isinstance(output, Mapping):
+            assert output.get('input_ids') is not None or output.get('inputs_embeds') is not None
+        else:
+            assert output[0].get('input_ids') is not None or output[0].get('inputs_embeds') is not None
+        input_ids = to_list(input_ids)
+        sent = ''
+        try:
+            if not isinstance(tokenizer, PreTrainedTokenizerBase) and hasattr(tokenizer, 'tokenizer'):
+                tokenizer = tokenizer.tokenizer
+            sent = tokenizer.decode(input_ids)
+        except Exception:
+            pass
+        return input_ids, sent
+    except Exception:
+        import traceback
+        print(traceback.format_exc())
+        raise
+
+
+def load_ds_old(ds):
+    from swift.llm import load_dataset
+    train_dataset, val_dataset = load_dataset(ds, split_dataset_ratio=0.0)
+    return train_dataset.select(range(1))
+
+
+def load_and_tokenize_old(ms_model_id, template):
+    model_type = None
+    model_info = None
+    from swift.llm import get_model_tokenizer
+    from swift.llm import get_template, MODEL_MAPPING
+    found = False
+    for model_type, model_info in MODEL_MAPPING.items():
+        if model_info['model_id_or_path'].lower() == ms_model_id.lower():
+            found = True
+            break
+
+    if not found:
+        raise ValueError(f'No model_type found: {ms_model_id}')
+
+    vl_fields = ['vl', 'video', 'minicpm-v', 'llava', 'vision', 'emu', 'florence']
+    model_ins, tokenizer = get_model_tokenizer(model_type, load_model=True)
+
+    if model_info['template'] == 'default-generation':
+        model_info['template'] = template.replace('_', '-')
+    template_ins = get_template(model_info['template'], tokenizer)
+    template_ins.model = model_ins
+    if 'audio' in model_info['template']:
+        output = template_ins.encode(load_ds_old('aishell1-zh-mini')[0])
+    elif any([vl in model_info['template'] for vl in vl_fields]):
+        output = template_ins.encode(load_ds_old('coco-en-mini')[0])
+    else:
+        output = template_ins.encode(load_ds_old('dureader-robust-zh')[0])
+    input_ids = to_list(output[0]['input_ids'])
+    sent = ''
+    try:
+        sent = tokenizer.decode(input_ids)
+    except Exception:
+        pass
+    return input_ids, sent
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        '--ms_model_id',
+        type=str,
+        required=True,
+    )
+    parser.add_argument(
+        '--template',
+        type=str,
+        required=True,
+    )
+    parser.add_argument('--new', type=str, required=False, default='1')
+    args = parser.parse_args()
+
+    is_new = args.new == '1'
+    if is_new:
+        input_ids, sent = load_and_tokenize(args.ms_model_id, args.template)
+    else:
+        input_ids, sent = load_and_tokenize_old(args.ms_model_id, args.template)
+    file = 'new_input_ids.txt' if is_new else 'old_input_ids.txt'
+    if input_ids is not None:
+        with open(file, 'w') as f:
+            json.dump({'input_ids': input_ids, 'sent': sent}, f)

tests/llm/test_custom.py

@@ -0,0 +1,74 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import unittest
+from typing import Any, Dict, Optional
+
+import torch
+
+from swift.llm import (DatasetMeta, InferRequest, Model, ModelGroup, ModelMeta, PtEngine, RequestConfig,
+                       ResponsePreprocessor, TemplateMeta, get_model_tokenizer_with_flash_attn, load_dataset,
+                       register_dataset, register_model, register_template)
+
+
+class CustomPreprocessor(ResponsePreprocessor):
+    prompt = """Task: Based on the given two sentences, provide a similarity score between 0.0 and 5.0.
+Sentence 1: {text1}
+Sentence 2: {text2}
+Similarity score: """
+
+    def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+        return super().preprocess({
+            'query': self.prompt.format(text1=row['text1'], text2=row['text2']),
+            'response': f"{row['label']:.1f}"
+        })
+
+
+register_dataset(
+    DatasetMeta(
+        ms_dataset_id='swift/stsb',
+        hf_dataset_id='SetFit/stsb',
+        preprocess_func=CustomPreprocessor(),
+    ))
+
+register_template(
+    TemplateMeta(
+        template_type='custom',
+        prefix=['<extra_id_0>System\n{{SYSTEM}}\n'],
+        prompt=['<extra_id_1>User\n{{QUERY}}\n<extra_id_1>Assistant\n'],
+        chat_sep=['\n']))
+
+register_model(
+    ModelMeta(
+        model_type='custom',
+        model_groups=[
+            ModelGroup([Model('AI-ModelScope/Nemotron-Mini-4B-Instruct', 'nvidia/Nemotron-Mini-4B-Instruct')])
+        ],
+        template='custom',
+        get_function=get_model_tokenizer_with_flash_attn,
+        ignore_patterns=['nemo']))
+
+
+class TestCustom(unittest.TestCase):
+
+    def test_custom_model(self):
+        infer_request = InferRequest(messages=[{'role': 'user', 'content': 'who are you?'}])
+        request_config = RequestConfig(max_tokens=512, temperature=0)
+        engine = PtEngine('AI-ModelScope/Nemotron-Mini-4B-Instruct', torch.float16)
+        response = engine.infer([infer_request], request_config)
+        swift_response = response[0].choices[0].message.content
+
+        engine.default_template.template_backend = 'jinja'
+        response = engine.infer([infer_request], request_config)
+        jinja_response = response[0].choices[0].message.content
+        assert swift_response == jinja_response, (f'swift_response: {swift_response}\njinja_response: {jinja_response}')
+        print(f'response: {swift_response}')
+
+    def test_custom_dataset(self):
+        dataset = load_dataset(['swift/stsb'])[0]
+        assert len(dataset) == 5749
+        assert list(dataset[0].keys()) == ['messages']
+        print(f'dataset: {dataset}')
+        print(f'dataset[0]: {dataset[0]}')
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/llm/test_dataset.py

@@ -0,0 +1,19 @@
+import unittest
+
+from swift.llm import load_dataset
+
+
+class TestDataset(unittest.TestCase):
+
+    def test_load_v_dataset(self):
+        if not __name__ == '__main__':
+            # ignore citest error in github
+            return
+
+        for ds in ['m3it#1000', 'mantis-instruct#1000', 'llava-med-zh-instruct#1000']:
+            ds = load_dataset(ds)
+            assert len(ds[0]) > 800
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/llm/test_ollama_export.py

@@ -0,0 +1,80 @@
+import os
+import shutil
+import tempfile
+import unittest
+
+import transformers
+from packaging import version
+
+from swift.llm import ExportArguments, export_main
+
+if __name__ == '__main__':
+    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+
+
+class TestTemplate(unittest.TestCase):
+
+    def setUp(self):
+        print(('Testing %s.%s' % (type(self).__name__, self._testMethodName)))
+        self.tmp_dir = tempfile.TemporaryDirectory().name
+
+    def tearDown(self):
+        if os.path.exists(self.tmp_dir):
+            shutil.rmtree(self.tmp_dir)
+        super().tearDown()
+
+    @unittest.skip('swift2.0')
+    def test_llama3(self):
+        args = ExportArguments(model_type='llama3-8b-instruct', to_ollama=True, ollama_output_dir=self.tmp_dir)
+        export_main(args)
+
+        template = ('TEMPLATE """{{ if .System }}<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n'
+                    '{{ .System }}<|eot_id|>{{ else }}<|begin_of_text|>{{ end }}{{ if .Prompt }}<|start_header_id|>user'
+                    '<|end_header_id|>\n\n{{ .Prompt }}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n'
+                    '{{ end }}{{ .Response }}<|eot_id|>"""')
+
+        stop = 'PARAMETER stop "<|eot_id|>"'
+
+        with open(os.path.join(self.tmp_dir, 'Modelfile'), 'r') as f:
+            content = f.read()
+            self.assertTrue(template in content)
+            self.assertTrue(stop in content)
+
+    @unittest.skip('swift2.0')
+    def test_glm4(self):
+        if version.parse(transformers.__version__) >= version.parse('4.45'):
+            return
+
+        args = ExportArguments(model_type='glm4-9b-chat', to_ollama=True, ollama_output_dir=self.tmp_dir)
+        export_main(args)
+
+        template = ('TEMPLATE """{{ if .System }}[gMASK] <sop><|system|>\n{{ .System }}{{ else }}'
+                    '[gMASK] <sop>{{ end }}{{ if .Prompt }}<|user|>\n{{ .Prompt }}<|assistant|>\n'
+                    '{{ end }}{{ .Response }}<|user|>"""')
+
+        stop = 'PARAMETER stop "<|user|>"'
+
+        with open(os.path.join(self.tmp_dir, 'Modelfile'), 'r') as f:
+            content = f.read()
+            self.assertTrue(template in content)
+            self.assertTrue(stop in content)
+
+    @unittest.skip('swift2.0')
+    def test_qwen2(self):
+        args = ExportArguments(model_type='qwen2-7b-instruct', to_ollama=True, ollama_output_dir=self.tmp_dir)
+        export_main(args)
+
+        template = ('TEMPLATE """{{ if .System }}<|im_start|>system\n{{ .System }}<|im_end|>\n{{ else }}{{ end }}'
+                    '{{ if .Prompt }}<|im_start|>user\n{{ .Prompt }}<|im_end|>\n<|im_start|>assistant\n'
+                    '{{ end }}{{ .Response }}<|im_end|>"""')
+
+        stop = 'PARAMETER stop "<|im_end|>"'
+
+        with open(os.path.join(self.tmp_dir, 'Modelfile'), 'r') as f:
+            content = f.read()
+            self.assertTrue(template in content)
+            self.assertTrue(stop in content)
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/llm/test_run.py

@@ -0,0 +1,458 @@
+if __name__ == '__main__':
+    import os
+    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+    os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
+
+import os
+import shutil
+import tempfile
+import unittest
+from functools import partial
+from typing import Any, Dict, List
+
+import torch
+from datasets import Dataset as HfDataset
+from modelscope import Model, MsDataset, snapshot_download
+from torch.nn.utils.rnn import pad_sequence
+from transformers import AutoTokenizer
+
+from swift import Trainer, TrainingArguments, get_logger
+from swift.llm import (InferArguments, ModelType, RLHFArguments, TrainArguments, infer_main, merge_lora, rlhf_main,
+                       sft_main)
+
+NO_EVAL_HUMAN = True
+
+logger = get_logger()
+
+kwargs = {
+    'per_device_train_batch_size': 2,
+    'per_device_eval_batch_size': 2,
+    'save_steps': 5,
+    'gradient_accumulation_steps': 4,
+    'num_train_epochs': 1,
+}
+
+
+class TestRun(unittest.TestCase):
+
+    def setUp(self):
+        print(f'Testing {type(self).__name__}.{self._testMethodName}')
+        self._tmp_dir = tempfile.TemporaryDirectory()
+        self.tmp_dir = self._tmp_dir.name
+
+    def tearDown(self):
+        shutil.rmtree(self.tmp_dir)
+
+    def test_template(self):
+        if not __name__ == '__main__':
+            # ignore citest error in github
+            return
+        torch.cuda.empty_cache()
+        output = sft_main(
+            TrainArguments(
+                model='Qwen/Qwen1.5-0.5B',
+                train_type='full',
+                dataset='DAMO_NLP/jd',
+                val_dataset='DAMO_NLP/jd#20',
+                streaming=True,
+                max_steps=12,
+                **kwargs))
+        last_model_checkpoint = output['last_model_checkpoint']
+        torch.cuda.empty_cache()
+        result = infer_main(InferArguments(model=last_model_checkpoint, load_data_args=True, val_dataset_sample=2))
+        assert len(result[0]['response']) < 20
+
+    def test_hf_hub(self):
+        if not __name__ == '__main__':
+            # ignore citest error in github
+            return
+        torch.cuda.empty_cache()
+        train_dataset_fnames = [
+            'alpaca.csv', 'chatml.jsonl', 'swift_pre.jsonl', 'swift_single.csv', 'swift_multi.jsonl',
+            'swift_multi.json#2'
+        ]
+        folder = os.path.join(os.path.dirname(__file__), 'data')
+        dataset = [
+            'llm-wizard/alpaca-gpt4-data-zh#20',
+            'shibing624/alpaca-zh#20',
+        ] + [os.path.join(folder, fname) for fname in train_dataset_fnames]
+        output = sft_main(
+            TrainArguments(
+                model='Qwen/Qwen1.5-0.5B-Chat-GPTQ-Int4', train_type='lora', dataset=dataset, use_hf=True, **kwargs))
+        last_model_checkpoint = output['last_model_checkpoint']
+        torch.cuda.empty_cache()
+        infer_main(InferArguments(adapters=last_model_checkpoint, load_data_args=True, val_dataset_sample=2))
+
+    @unittest.skip('avoid ci error')
+    def test_basic(self):
+        output_dir = 'output'
+        quant_bits_list = [0, 4]
+        train_dataset_fnames = [
+            'alpaca.csv', 'chatml.jsonl', 'swift_pre.jsonl', 'swift_single.csv', 'swift_multi.jsonl',
+            'swift_multi.json#2'
+        ]
+        folder = os.path.join(os.path.dirname(__file__), 'data')
+        dataset = [
+            'AI-ModelScope/alpaca-gpt4-data-zh#20',
+            'hurner/alpaca-gpt4-data-zh#20',
+        ] + [os.path.join(folder, fname) for fname in train_dataset_fnames]
+        if not __name__ == '__main__':
+            output_dir = self.tmp_dir
+            quant_bits_list = [4]
+            dataset = dataset[:2]
+        for quant_bits in quant_bits_list:
+            if quant_bits == 0:
+                predict_with_generate = False
+                quant_method = None
+            else:
+                predict_with_generate = True
+                quant_method = 'bnb'
+            sft_args = TrainArguments(
+                model='Qwen/Qwen2-0.5B-Instruct',
+                quant_bits=quant_bits,
+                eval_steps=5,
+                adam_beta2=0.95,
+                quant_method=quant_method,
+                predict_with_generate=predict_with_generate,
+                dataset=dataset,
+                val_dataset='DAMO_NLP/jd#20',
+                output_dir=output_dir,
+                download_mode='force_redownload',
+                include_num_input_tokens_seen=True,
+                gradient_checkpointing=True,
+                **kwargs)
+            torch.cuda.empty_cache()
+            output = sft_main(sft_args)
+            print(output)
+            best_model_checkpoint = output['best_model_checkpoint']
+            print(f'best_model_checkpoint: {best_model_checkpoint}')
+            if __name__ == '__main__':
+                infer_args = InferArguments(
+                    adapters=best_model_checkpoint,
+                    merge_lora={
+                        0: True,
+                        4: False
+                    }[quant_bits],
+                    load_data_args=NO_EVAL_HUMAN,
+                    val_dataset_sample=5)
+                torch.cuda.empty_cache()
+                result = infer_main(infer_args)
+                print(result)
+        # if __name__ == '__main__':
+        #     app_ui_main(infer_args)
+
+    def test_vl_audio(self):
+        output_dir = 'output'
+        if not __name__ == '__main__':
+            # ignore citest error in github
+            return
+        model_type_list = ['Qwen/Qwen-VL-Chat', 'Qwen/Qwen-Audio-Chat']
+        dataset_list = [
+            'modelscope/coco_2014_caption:validation#100', 'speech_asr/speech_asr_aishell1_trainsets:validation#100'
+        ]
+        for model, dataset in zip(model_type_list, dataset_list):
+            sft_args = TrainArguments(
+                model=model,
+                eval_steps=5,
+                dataset=[dataset],
+                output_dir=output_dir,
+                gradient_checkpointing=True,
+                lazy_tokenize=True,
+                disable_tqdm=True,
+                **kwargs)
+            torch.cuda.empty_cache()
+            output = sft_main(sft_args)
+            print(output)
+            best_model_checkpoint = output['best_model_checkpoint']
+            print(f'best_model_checkpoint: {best_model_checkpoint}')
+            infer_args = InferArguments(
+                adapters=best_model_checkpoint,
+                load_data_args=True,
+                stream={
+                    'Qwen/Qwen-VL-Chat': True,
+                    'Qwen/Qwen-Audio-Chat': False
+                }[model],
+                val_dataset_sample=5)
+            torch.cuda.empty_cache()
+            result = infer_main(infer_args)
+            print(result)
+
+    def test_custom_dataset(self):
+        if not __name__ == '__main__':
+            # ignore citest error in github
+            return
+        train_dataset_fnames = [
+            'alpaca.csv', 'chatml.jsonl', 'swift_pre.jsonl', 'swift_single.csv', 'swift_multi.jsonl',
+            'swift_multi.json', 'sharegpt.jsonl'
+        ]
+        val_dataset_fnames = [
+            'alpaca.jsonl',
+            'alpaca2.csv',
+            'conversations.jsonl',
+            'swift_pre.csv',
+            'swift_single.jsonl',
+            # 'swift_#:#.jsonl#3'
+        ]
+        folder = os.path.join(os.path.dirname(__file__), 'data')
+        resume_from_checkpoint = None
+        train_kwargs = kwargs.copy()
+        train_kwargs.pop('num_train_epochs')
+        for num_train_epochs in [1, 2]:
+            sft_args = TrainArguments(
+                model='Qwen/Qwen-7B-Chat',
+                dataset=['swift/self-cognition#20'] + [os.path.join(folder, fname) for fname in train_dataset_fnames],
+                val_dataset=[os.path.join(folder, fname) for fname in val_dataset_fnames],
+                resume_from_checkpoint=resume_from_checkpoint,
+                num_train_epochs=num_train_epochs,
+                model_name='小黄',
+                model_author='魔搭',
+                **train_kwargs)
+
+            torch.cuda.empty_cache()
+            result = sft_main(sft_args)
+            best_model_checkpoint = result['best_model_checkpoint']
+            resume_from_checkpoint = result['last_model_checkpoint']
+
+        for load_args in [True, False]:
+            infer_kwargs = {}
+            if load_args is False:
+                args_json = os.path.join(best_model_checkpoint, 'args.json')
+                assert os.path.exists(args_json)
+                os.remove(args_json)
+                infer_kwargs = {'model': 'Qwen/Qwen-7B-Chat'}
+            infer_args = InferArguments(
+                adapters=best_model_checkpoint,
+                load_data_args=load_args and NO_EVAL_HUMAN,
+                merge_lora=load_args,
+                val_dataset=[os.path.join(folder, fname) for fname in val_dataset_fnames],
+                **infer_kwargs)
+            torch.cuda.empty_cache()
+            infer_main(infer_args)
+
+    def test_rlhf(self):
+        if not __name__ == '__main__':
+            # ignore citest error in github
+            return
+        torch.cuda.empty_cache()
+        # llm rlhf
+        #
+        rlhf_types = ['dpo', 'orpo', 'simpo', 'kto', 'cpo', 'rm', 'ppo']
+        for rlhf_type in rlhf_types:
+            dataset = ('AI-ModelScope/hh_rlhf_cn:harmless_base_cn#100'
+                       if rlhf_type != 'kto' else 'AI-ModelScope/ultrafeedback-binarized-preferences-cleaned-kto#100')
+            train_kwargs = {}
+            if rlhf_type == 'ppo':
+                train_kwargs['reward_model'] = 'Qwen/Qwen2-1.5B-Instruct'
+            output = rlhf_main(
+                RLHFArguments(
+                    rlhf_type=rlhf_type,
+                    model='Qwen/Qwen2-1.5B-Instruct',
+                    dataset=dataset,
+                    eval_steps=5,
+                    split_dataset_ratio=0.05,
+                    **train_kwargs,
+                    **kwargs))
+            if rlhf_type == 'ppo':
+                model_checkpoint = output['last_model_checkpoint']
+            else:
+                model_checkpoint = output['best_model_checkpoint']
+
+            torch.cuda.empty_cache()
+            infer_main(InferArguments(adapters=model_checkpoint, load_data_args=True))
+
+        # mllm rlhf
+        visual_rlhf_types = ['dpo', 'orpo', 'simpo', 'cpo', 'rm']
+        test_model = [
+            'OpenGVLab/InternVL2-2B', 'Qwen/Qwen2-VL-2B-Instruct', 'llava-hf/llava-v1.6-mistral-7b-hf',
+            'AI-ModelScope/Florence-2-base-ft'
+        ]  # decoder only and encoder-decoder
+        for rlhf_type in visual_rlhf_types:
+            for model in test_model:
+                dataset_name = 'swift/RLAIF-V-Dataset#100'
+                output = rlhf_main(
+                    RLHFArguments(
+                        rlhf_type=rlhf_type,
+                        model=model,
+                        dataset=dataset_name,
+                        eval_steps=5,
+                        dataset_num_proc=16,
+                        **kwargs))
+                best_model_checkpoint = output['best_model_checkpoint']
+                torch.cuda.empty_cache()
+                infer_main(InferArguments(adapters=best_model_checkpoint, load_data_args=True, val_dataset_sample=2))
+
+    def test_loss_matching(self):
+        output_dir = 'output'
+        if not __name__ == '__main__':
+            # ignore citest error in github
+            return
+        losses = []
+        for use_swift_lora in [False, True]:
+            bool_var = use_swift_lora
+            torch.cuda.empty_cache()
+            output = sft_main([
+                '--model', 'Qwen/Qwen-7B-Chat', '--save_steps', '5', '--dataset',
+                'AI-ModelScope/leetcode-solutions-python#200', '--output_dir', output_dir, '--gradient_checkpointing',
+                'true', '--max_new_tokens', '100', '--attn_impl', 'flash_attn', '--target_modules', 'all-linear',
+                '--seed', '0', '--lora_bias', 'all', '--modules_to_save', 'lm_head', '--use_swift_lora',
+                str(use_swift_lora), '--num_train_epochs', '1', '--gradient_accumulation_steps', '16'
+            ])
+            best_model_checkpoint = output['best_model_checkpoint']
+            print(f'best_model_checkpoint: {best_model_checkpoint}')
+            load_data_args = str(bool_var or NO_EVAL_HUMAN)
+            if load_data_args:
+                val_dataset_sample = 2
+            else:
+                val_dataset_sample = -1
+            torch.cuda.empty_cache()
+            infer_main([
+                '--adapters', best_model_checkpoint, '--val_dataset_sample',
+                str(val_dataset_sample), '--max_new_tokens', '100', '--attn_impl', 'eager', '--merge_lora',
+                str(bool_var), '--load_data_args',
+                str(load_data_args)
+            ])
+            loss = output['log_history'][-1]['train_loss']
+            losses.append(loss)
+        self.assertTrue(abs(losses[0] - losses[1]) < 5e-4)
+        print(f'swift_loss: {losses[0]}')
+        print(f'peft_loss: {losses[1]}')
+        self.assertTrue(0.95 <= losses[0] <= 1)
+
+    def test_pai_compat(self):
+        if not __name__ == '__main__':
+            # ignore citest error in github
+            return
+        from swift.llm import sft_main, infer_main
+        os.environ['PAI_TRAINING_JOB_ID'] = '123456'
+        folder = os.path.join(os.path.dirname(__file__), 'config')
+        tensorboard_dir = os.path.join('output/pai_test', 'pai_tensorboard')
+        os.environ['PAI_OUTPUT_TENSORBOARD'] = tensorboard_dir
+        sft_json = os.path.join(folder, 'sft.json')
+        infer_json = os.path.join(folder, 'infer.json')
+        torch.cuda.empty_cache()
+        output = sft_main([sft_json])
+        print()
+        infer_args = {
+            'adapters': output['best_model_checkpoint'],
+            'val_dataset_sample': 2,
+            'load_data_args': True,
+        }
+        import json
+        with open(infer_json, 'w') as f:
+            json.dump(infer_args, f, ensure_ascii=False, indent=4)
+        torch.cuda.empty_cache()
+        infer_main([infer_json])
+        os.environ.pop('PAI_TRAINING_JOB_ID')
+
+
+def data_collate_fn(batch: List[Dict[str, Any]], tokenizer) -> Dict[str, torch.Tensor]:
+    # text-classification
+    assert tokenizer.pad_token_id is not None
+    input_ids = [torch.tensor(b['input_ids']) for b in batch]
+    labels = torch.tensor([b['labels'] for b in batch])
+    attention_mask = [torch.ones(len(input_ids[i]), dtype=torch.int64) for i in range(len(input_ids))]
+
+    input_ids = pad_sequence(input_ids, batch_first=True, padding_value=tokenizer.pad_token_id)
+    attention_mask = pad_sequence(attention_mask, batch_first=True, padding_value=0)
+    return {'input_ids': input_ids, 'attention_mask': attention_mask, 'labels': labels}
+
+
+class BertTrainer(Trainer):
+
+    def compute_loss(self, model, inputs, return_outputs=False):
+        outputs = model(**inputs)
+        loss = outputs.loss
+        if loss is None:
+            logits, loss = list(outputs.logits)
+        return (loss, outputs) if return_outputs else loss
+
+
+class TestTrainer(unittest.TestCase):
+
+    def setUp(self):
+        self._tmp_dir = tempfile.TemporaryDirectory()
+        self.tmp_dir = self._tmp_dir.name
+        # self.tmp_dir = 'test'
+        logger.info(f'self.tmp_dir: {self.tmp_dir}')
+
+    def tearDown(self):
+        if os.path.isdir(self.tmp_dir):
+            shutil.rmtree(self.tmp_dir)
+        # api = HubApi()
+        # api.delete_model(self.hub_model_id)
+        # logger.info(f'delete model: {self.hub_model_id}')
+
+    def test_trainer(self):
+        self.hub_model_id = 'test_trainer2'
+        logger.info(f'self.hub_model_id: {self.hub_model_id}')
+        self.tmp_dir = 'output/damo/nlp_structbert_backbone_base_std'
+        push_to_hub = True
+        if not __name__ == '__main__':
+            # ignore citest error in github
+            return
+        model_id = 'damo/nlp_structbert_backbone_base_std'
+        model_dir = snapshot_download(model_id, 'master')
+        tokenizer = AutoTokenizer.from_pretrained(model_dir)
+        dataset = MsDataset.load('clue', subset_name='tnews')
+        num_labels = max(dataset['train']['label']) + 1
+        model = Model.from_pretrained(model_dir, task='text-classification', num_labels=num_labels)
+        train_dataset, val_dataset = dataset['train'].to_hf_dataset(), dataset['validation'].to_hf_dataset()
+        train_dataset: HfDataset = train_dataset.select(range(100))
+        val_dataset: HfDataset = val_dataset.select(range(20))
+
+        #
+        def tokenize_func(examples):
+            data = tokenizer(examples['sentence'], return_attention_mask=False)
+            examples['input_ids'] = data['input_ids']
+            examples['labels'] = examples['label']
+            del examples['sentence'], examples['label']
+            return examples
+
+        train_dataset = train_dataset.map(tokenize_func)
+        val_dataset = val_dataset.map(tokenize_func)
+
+        data_collator = partial(data_collate_fn, tokenizer=tokenizer)
+        for save_only_model in [True, False]:
+            trainer_args = TrainingArguments(
+                self.tmp_dir,
+                do_train=True,
+                do_eval=True,
+                num_train_epochs=1,
+                evaluation_strategy='steps',
+                save_strategy='steps',
+                per_device_train_batch_size=4,
+                per_device_eval_batch_size=4,
+                push_to_hub=push_to_hub,
+                hub_token=None,  # use env var
+                hub_private_repo=True,
+                hub_strategy='every_save',
+                hub_model_id=self.hub_model_id,
+                overwrite_output_dir=True,
+                save_steps=10,
+                save_total_limit=2,
+                metric_for_best_model='loss',
+                greater_is_better=False,
+                report_to=['tensorboard'],
+                gradient_accumulation_steps=1,
+                logging_steps=5,
+                eval_steps=10,
+                save_safetensors=False,
+                save_only_model=save_only_model)
+        trainer_args._n_gpu = 1
+        trainer = BertTrainer(model, trainer_args, data_collator, train_dataset, val_dataset, tokenizer)
+        self.hub_model_id = trainer_args.hub_model_id
+        trainer.train()
+        if trainer_args.push_to_hub:
+            trainer.push_to_hub()
+
+
+if __name__ == '__main__':
+    # TestRun().test_template()
+    # TestRun().test_hf_hub()
+    # TestRun().test_basic()
+    # TestRun().test_custom_dataset()
+    # TestRun().test_vl_audio()
+    # TestRun().test_loss_matching()
+    #
+    # TestRun().test_rlhf()
+    unittest.main()

tests/llm/test_run3.py

@@ -0,0 +1,172 @@
+import os
+import shutil
+import tempfile
+import unittest
+
+import json
+import numpy as np
+
+from swift.llm import MODEL_MAPPING, load_dataset
+
+
+class TestRun3(unittest.TestCase):
+
+    def setUp(self):
+        print(f'Testing {type(self).__name__}.{self._testMethodName}')
+        self._tmp_dir = tempfile.TemporaryDirectory()
+        self.tmp_dir = self._tmp_dir.name
+
+    def tearDown(self):
+        shutil.rmtree(self.tmp_dir)
+
+    def load_ds(self, ds):
+        train_dataset, val_dataset = load_dataset(
+            ds,
+            split_dataset_ratio=0.0,
+            strict=False,
+            num_proc=1,
+            model_name=['小黄', 'Xiao Huang'],
+            model_author=['魔搭', 'ModelScope'])
+        return train_dataset.select(range(min(50, len(train_dataset))))
+
+    # def test_model_load(self):
+    #     if os.path.exists('./models.txt'):
+    #         with open('./models.txt', 'r') as f:
+    #             models = json.load(f)
+    #     else:
+    #         models = []
+    #     for model_name, model_meta in MODEL_MAPPING.items():
+    #         meta_requires = model_meta.requires or []
+    #         for group in model_meta.model_groups:
+    #             model = group.models[0]
+    #             if 'skip_test' in (group.tags or []) or model.ms_model_id in models:
+    #                 break
+    #             requires = meta_requires + (group.requires or [])
+    #             for req in requires:
+    #                 os.system(f'pip install "{req}"')
+    #             if not any(['transformers' in req for req in requires]):
+    #                 os.system('pip install transformers -U')
+    #             if not any(['accelerate' in req for req in requires]):
+    #                 os.system('pip install accelerate -U')
+    #             try:
+    #                 model_arch_args = ''
+    #                 if model_meta.model_arch:
+    #                     model_arch_args = f'--model_arch {model_meta.model_arch}'
+    #                 cmd = ('PYTHONPATH=. python tests/llm/load_model.py '
+    #                        f'--ms_model_id {model.ms_model_id} {model_arch_args}')
+    #                 if os.system(cmd) != 0:
+    #                     raise RuntimeError()
+    #             except Exception:
+    #                 passed = False
+    #             else:
+    #                 passed = True
+    #                 models.append(model.ms_model_id)
+    #             finally:
+    #                 if passed:
+    #                     with open('./models.txt', 'w') as f:
+    #                         json.dump(models, f)
+
+    # def test_template_load(self):
+    #     if os.path.exists('./templates.txt'):
+    #         with open('./templates.txt', 'r') as f:
+    #             templates = json.load(f)
+    #     else:
+    #         templates = []
+    #     for model_name, model_meta in MODEL_MAPPING.items():
+    #         template = model_meta.template
+    #         meta_requires = model_meta.requires or []
+    #         for group in model_meta.model_groups:
+    #             model = group.models[0]
+    #             if 'skip_test' in (group.tags or []) or template in templates:
+    #                 break
+    #             requires = meta_requires + (group.requires or [])
+    #             for req in requires:
+    #                 os.system(f'pip install "{req}"')
+    #             if not any(['transformers' in req for req in requires]):
+    #                 os.system('pip install transformers -U')
+    #             if not any(['accelerate' in req for req in requires]):
+    #                 os.system('pip install accelerate -U')
+    #             try:
+    #                 cmd = ('PYTHONPATH=. python tests/llm/load_template.py '
+    #                        f'--ms_model_id {model.ms_model_id} --template {template}')
+    #                 if os.system(cmd) != 0:
+    #                     raise RuntimeError()
+    #             except Exception:
+    #                 import traceback
+    #                 print(traceback.format_exc())
+    #                 passed = False
+    #             else:
+    #                 passed = True
+    #                 templates.append(template)
+    #             finally:
+    #                 if passed:
+    #                     with open('./templates.txt', 'w') as f:
+    #                         json.dump(templates, f)
+
+    @unittest.skip('skip')
+    def test_template_compare(self):
+        if os.path.exists('./templates.txt'):
+            with open('./templates.txt', 'r') as f:
+                templates = json.load(f)
+        else:
+            templates = []
+        skip_model_type = {
+            'grok', 'deepseek_moe', 'deepseek_v2', 'deepseek_v2_5', 'llama3_1_omni', 'llava_next_qwen_hf',
+            'llava1_6_yi', 'llava_next_qwen', 'mixtral', 'codefuse_codellama', 'wizardlm2', 'wizardlm2_awq',
+            'openbuddy_deepseek', 'sus', 'openbuddy_mixtral', 'openbuddy_llama', 'dbrx', 'nenotron', 'reflection',
+            'xverse_moe', 'qwen2_moe', 'yuan2', 'wizardlm2_moe', 'emu3_gen', 'llava1_6_mistral', 'mplug_owl3_241101',
+            'llava1_6_yi_hf'
+        }
+        for model_name, model_meta in MODEL_MAPPING.items():
+            if model_name in skip_model_type:
+                continue
+            template = model_meta.template
+            meta_requires = model_meta.requires or []
+            for group in model_meta.model_groups:
+                model = group.models[0]
+                if 'awq' in model.ms_model_id.lower() or 'gptq' in model.ms_model_id.lower():
+                    break
+                if template in templates:
+                    break
+                requires = meta_requires + (group.requires or [])
+                for req in requires:
+                    os.system(f'pip install "{req}"')
+                if not any(['transformers' in req for req in requires]):
+                    os.system('pip install transformers -U')
+                if not any(['accelerate' in req for req in requires]):
+                    os.system('pip install accelerate -U')
+                try:
+                    cmd = ('CUDA_VISIBLE_DEVICES=0 PYTHONPATH=. python tests/llm/load_template.py '
+                           f'--ms_model_id {model.ms_model_id} --template {template}')
+                    if os.system(cmd) != 0:
+                        raise RuntimeError()
+                    cmd = (
+                        'CUDA_VISIBLE_DEVICES=0 PYTHONPATH=/mnt/workspace/yzhao/tastelikefeet/swift python tests/llm/load_template.py '  # noqa
+                        f'--ms_model_id {model.ms_model_id} --template {template} --new 0')
+                    if os.system(cmd) != 0:
+                        raise RuntimeError()
+                    with open('new_input_ids.txt', 'r') as f:
+                        input_ids_new = json.load(f)
+                    with open('old_input_ids.txt', 'r') as f:
+                        input_ids_old = json.load(f)
+                    print('model_id', model.ms_model_id, 'new:', input_ids_new, 'old:', input_ids_old)
+                    self.assertTrue(np.allclose(input_ids_new['input_ids'], input_ids_old['input_ids']))
+                except Exception:
+                    import traceback
+                    print(traceback.format_exc())
+                    passed = False
+                else:
+                    passed = True
+                    templates.append(template)
+                finally:
+                    if passed:
+                        with open('./templates.txt', 'w') as f:
+                            json.dump(templates, f)
+                    if os.path.exists('new_input_ids.txt'):
+                        os.remove('new_input_ids.txt')
+                    if os.path.exists('old_input_ids.txt'):
+                        os.remove('old_input_ids.txt')
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/llm/test_template.py

@@ -0,0 +1,104 @@
+import os
+import unittest
+
+from swift.llm import PtEngine, RequestConfig, get_model_tokenizer, get_template
+from swift.utils import get_logger, seed_everything
+
+# os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+os.environ['SWIFT_DEBUG'] = '1'
+
+logger = get_logger()
+
+
+def _infer_model(pt_engine, system=None, messages=None):
+    seed_everything(42)
+    request_config = RequestConfig(max_tokens=128, temperature=0)
+    if messages is None:
+        messages = []
+        if system is not None:
+            messages += [{'role': 'system', 'content': system}]
+        messages += [{'role': 'user', 'content': '你好'}]
+        resp = pt_engine.infer([{'messages': messages}], request_config=request_config)
+        response = resp[0].choices[0].message.content
+        messages += [{'role': 'assistant', 'content': response}, {'role': 'user', 'content': '<image>这是什么'}]
+    resp = pt_engine.infer([{
+        'messages': messages,
+    }], request_config=request_config)
+    response = resp[0].choices[0].message.content
+    messages += [{'role': 'assistant', 'content': response}]
+    logger.info(f'model: {pt_engine.model_info.model_name}, messages: {messages}')
+    return response
+
+
+class TestTemplate(unittest.TestCase):
+
+    def test_template(self):
+        pt_engine = PtEngine('Qwen/Qwen2.5-3B-Instruct-GPTQ-Int4')
+        response = _infer_model(pt_engine)
+        pt_engine.default_template.template_backend = 'jinja'
+        response2 = _infer_model(pt_engine)
+        assert response == response2
+
+    def test_tool_message_join(self):
+        from copy import deepcopy
+
+        from swift.plugin import agent_templates
+
+        messages = [
+            # first round
+            {
+                'role': 'user',
+                'content': 'user1'
+            },
+            {
+                'role': 'assistant',
+                'content': 'assistant1'
+            },
+            {
+                'role': 'assistant',
+                'content': 'assistant2'
+            },
+            {
+                'role': 'tool',
+                'content': 'tool1'
+            },
+            # second round
+            {
+                'role': 'assistant',
+                'content': 'assistant3'
+            },
+            {
+                'role': 'tool',
+                'content': 'tool2'
+            },
+            {
+                'role': 'tool',
+                'content': 'tool3'
+            },
+        ]
+
+        # testing two template type.
+        tokenizer = get_model_tokenizer('Qwen/Qwen2.5-7B-Instruct', load_model=False)[1]
+        template = get_template(tokenizer.model_meta.template, tokenizer)
+        for agent_template_type in ('react_zh', 'qwen_zh'):
+            agent_template = agent_templates[agent_template_type]()
+            template.agent_template = agent_template
+            observation = agent_template.keyword.observation
+            test_messages = deepcopy(messages)
+            test_messages[2]['content'] = 'assistant2' + observation
+            test_messages[4]['content'] = (
+                agent_template.keyword.action + agent_template.keyword.action_input + 'assistant3' + observation)
+            encoded = template.encode({'messages': test_messages})
+            res = template.safe_decode(encoded['input_ids'])
+
+            ground_truth = (
+                '<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\n'
+                '<|im_start|>user\nuser1<|im_end|>\n'
+                f'<|im_start|>assistant\nassistant1assistant2{observation}tool1'
+                f'{agent_template.keyword.action}{agent_template.keyword.action_input}assistant3'
+                f'{observation}tool2\n{observation}tool3\n')
+            assert res == ground_truth
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/llm/test_utils.py

@@ -0,0 +1,28 @@
+import unittest
+
+from swift.llm import load_dataset
+from swift.utils import lower_bound
+
+
+class TestLlmUtils(unittest.TestCase):
+
+    def test_count_startswith(self):
+        arr = [-100] * 1000 + list(range(1000))
+        self.assertTrue(lower_bound(0, len(arr), lambda i: arr[i] != -100) == 1000)
+
+    def test_count_endswith(self):
+        arr = list(range(1000)) + [-100] * 1000
+        self.assertTrue(lower_bound(0, len(arr), lambda i: arr[i] == -100) == 1000)
+
+    @unittest.skip('avoid ci error')
+    def test_dataset(self):
+        dataset = load_dataset(['AI-ModelScope/alpaca-gpt4-data-zh#1000', 'AI-ModelScope/alpaca-gpt4-data-en#200'],
+                               num_proc=4,
+                               strict=False,
+                               download_mode='force_redownload')
+        print(f'dataset[0]: {dataset[0]}')
+        print(f'dataset[1]: {dataset[1]}')
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/megatron/test_align/test_llm.py

@@ -0,0 +1,94 @@
+import os
+
+import torch
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+
+
+def _test_model(model_id):
+    from swift.llm import export_main, ExportArguments
+    if model_id.endswith('mcore'):
+        export_main(
+            ExportArguments(
+                mcore_model=model_id,
+                to_hf=True,
+                exist_ok=True,
+                test_convert_precision=True,
+                torch_dtype=torch.bfloat16))
+    else:
+        export_main(
+            ExportArguments(
+                model=model_id,
+                to_mcore=True,
+                exist_ok=True,
+                test_convert_precision=True,
+                torch_dtype=torch.bfloat16,
+            ))
+
+
+def test_qwen2():
+    _test_model('Qwen/Qwen2-0.5B-Instruct')
+
+
+def test_llama2():
+    _test_model('modelscope/Llama-2-7b-chat-ms')
+
+
+def test_llama3():
+    _test_model('LLM-Research/Meta-Llama-3-8B-Instruct')
+
+
+def test_marco_o1():
+    _test_model('AIDC-AI/Marco-o1')
+
+
+def test_deepseek_r1_llama():
+    _test_model('deepseek-ai/DeepSeek-R1-Distill-Llama-8B')
+
+
+def test_deepseek_r1_qwen():
+    _test_model('deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B')
+
+
+def test_yi():
+    _test_model('01ai/Yi-1.5-6B-Chat')
+
+
+def test_megrez():
+    _test_model('InfiniAI/Megrez-3b-Instruct')
+
+
+def test_llama3_1():
+    _test_model('LLM-Research/Meta-Llama-3.1-8B-Instruct')
+
+
+def test_llama3_2():
+    _test_model('LLM-Research/Llama-3.2-1B-Instruct')
+
+
+def test_qwen3():
+    _test_model('Qwen/Qwen3-0.6B-Base')
+
+
+def test_qwen2_moe():
+    _test_model('Qwen/Qwen1.5-MoE-A2.7B-Chat')
+
+
+def test_qwen3_moe():
+    _test_model('Qwen/Qwen3-15B-A2B-Base')
+
+
+if __name__ == '__main__':
+    # test_qwen2()
+    # test_llama2()
+    # test_llama3()
+    # test_marco_o1()
+    # test_deepseek_r1_llama()
+    # test_deepseek_r1_qwen()
+    # test_yi()
+    # test_megrez()
+    # test_llama3_1()
+    # test_llama3_2()
+    # test_qwen3()
+    # test_qwen2_moe()
+    test_qwen3_moe()

tests/megatron/test_export.py

@@ -0,0 +1,64 @@
+import os
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+
+
+def _infer_model(pt_engine, system=None, messages=None):
+    from swift.utils import seed_everything, get_logger
+    from swift.llm import RequestConfig
+    logger = get_logger()
+    seed_everything(42)
+    request_config = RequestConfig(max_tokens=128, temperature=0)
+    if messages is None:
+        messages = []
+        if system is not None:
+            messages += [{'role': 'system', 'content': system}]
+        messages += [{'role': 'user', 'content': 'who are you?'}]
+        resp = pt_engine.infer([{'messages': messages}], request_config=request_config)
+        response = resp[0].choices[0].message.content
+        messages += [{'role': 'assistant', 'content': response}, {'role': 'user', 'content': '<image>这是什么'}]
+    else:
+        messages = messages.copy()
+    resp = pt_engine.infer([{
+        'messages': messages,
+    }], request_config=request_config)
+    response = resp[0].choices[0].message.content
+    messages += [{'role': 'assistant', 'content': response}]
+    logger.info(f'model: {pt_engine.model_info.model_name}, messages: {messages}')
+    return response
+
+
+model_id = 'Qwen/Qwen2-7B-Instruct'
+
+
+def hf2mcore():
+    from swift.llm import export_main, ExportArguments
+    export_main(
+        ExportArguments(
+            model=model_id, to_mcore=True, torch_dtype='bfloat16', exist_ok=True, test_convert_precision=True))
+
+
+def mcore2hf():
+    from swift.llm import export_main, ExportArguments
+    export_main(
+        ExportArguments(
+            mcore_model='Qwen2-7B-Instruct-mcore',
+            to_hf=True,
+            torch_dtype='bfloat16',
+            exist_ok=True,
+            test_convert_precision=True))
+
+
+def infer_hf_align():
+    from swift.llm import PtEngine
+    pt_engine = PtEngine(model_id)
+    response = _infer_model(pt_engine)
+    pt_engine = PtEngine('Qwen2-7B-Instruct-mcore-hf')
+    response2 = _infer_model(pt_engine)
+    assert response == response2
+
+
+if __name__ == '__main__':
+    # hf2mcore()
+    mcore2hf()
+    infer_hf_align()

tests/megatron/test_model.py

@@ -0,0 +1,65 @@
+import os
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+os.environ['MASTER_PORT'] = '29560'
+
+
+def get_mg_model_tokenizer(model_id):
+    from megatron.training.initialize import initialize_megatron
+    set_default_ddp_config()
+    hf_model, processor = get_model_tokenizer(model_id, torch_dtype=torch.float32)
+    megatron_model_meta = get_megatron_model_meta(processor.model_meta.model_type)
+    model_info = processor.model_info
+    kwargs = megatron_model_meta.convert_hf_config(model_info.config)
+    megatron_args = MegatronArguments(**kwargs, seq_length=1, use_cpu_initialization=True, no_initialization=True)
+    patch_megatron_tokenizer(processor)
+    extra_args = megatron_args.parse_to_megatron()
+    initialize_megatron(args_defaults=extra_args)
+    mg_model = megatron_model_meta.model_provider()
+    megatron_model_meta.convert_hf2mcore(hf_model, mg_model)
+    return hf_model, mg_model, processor
+
+
+def test_bf16_fp32():
+    hf_model_fp32, processor = get_model_tokenizer(model_id, torch_dtype=torch.float32)
+    hf_model_bf16, processor = get_model_tokenizer(model_id, torch_dtype=torch.bfloat16)
+    template = get_template(hf_model_fp32.model_meta.template, processor)
+    input_ids = template.encode(InferRequest(messages=[{'role': 'user', 'content': 'who are you?'}]))['input_ids']
+    input_ids = torch.tensor(input_ids)[None].to('cuda')
+    with torch.inference_mode():
+        hf_logits_fp32 = hf_model_fp32(input_ids).logits
+        hf_logits_bf16 = hf_model_bf16(input_ids).logits
+    mean_diff = (hf_logits_fp32 - hf_logits_bf16).abs().mean().item()
+    max_diff = (hf_logits_fp32 - hf_logits_bf16).abs().max().item()
+    # mean_diff: 0.13342587649822235, max_diff: 7.1983513832092285
+    print(f'mean_diff: {mean_diff}, max_diff: {max_diff}')
+
+
+def test_align(hf_model, mg_model, processor):
+    from megatron.training.utils import get_ltor_masks_and_position_ids
+    template = get_template(hf_model.model_meta.template, processor)
+    input_ids = template.encode(InferRequest(messages=[{'role': 'user', 'content': 'who are you?'}]))['input_ids']
+    input_ids = torch.tensor(input_ids)[None].to('cuda')
+    attention_mask, _, position_ids = get_ltor_masks_and_position_ids(input_ids, -100, True, True, True)
+    with torch.inference_mode():
+        hf_model.cuda()
+        mg_model.cuda()
+        hf_logits = hf_model(input_ids).logits
+        mg_logits = mg_model(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids)
+    mean_diff = (mg_logits - hf_logits).abs().mean().item()
+    max_diff = (mg_logits - hf_logits).abs().max().item()
+    print(f'mean_diff: {mean_diff}, max_diff: {max_diff}')
+
+
+model_id = 'Qwen/Qwen2-7B-Instruct'
+
+if __name__ == '__main__':
+    import torch
+    from swift.llm import InferRequest, get_model_tokenizer, get_template
+    from swift.utils import set_default_ddp_config
+    from swift.megatron.argument import MegatronArguments
+    from swift.megatron.model import get_megatron_model_meta
+    from swift.megatron.utils import patch_megatron_tokenizer
+    # test_bf16_fp32()
+    hf_model, mg_model, processor = get_mg_model_tokenizer(model_id)
+    test_align(hf_model, mg_model, processor)

tests/megatron/test_save.py

@@ -0,0 +1,61 @@
+import os
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+
+
+def get_mg_model_tokenizer():
+    model_id = 'Qwen/Qwen2.5-7B-Instruct'
+    hf_model_id = 'Qwen/Qwen2.5-7B'
+    from megatron.training.initialize import initialize_megatron
+    set_default_ddp_config()
+    hf_model, _ = get_model_tokenizer(hf_model_id, torch_dtype=torch.float32)
+    _, processor = get_model_tokenizer(model_id, load_model=False)
+    megatron_model_meta = get_megatron_model_meta(processor.model_meta.model_type)
+    model_info = processor.model_info
+    kwargs = megatron_model_meta.convert_hf_config(model_info.config)
+    megatron_args = MegatronArguments(
+        **kwargs,
+        seq_length=1,
+        use_cpu_initialization=True,
+        no_initialization=True,
+        load='Qwen2-7B-Instruct-mcore',
+        save='mcore-hf-test',
+        no_load_optim=True,
+        no_load_rng=True)
+    patch_megatron_tokenizer(processor)
+    extra_args = megatron_args.parse_to_megatron()
+    initialize_megatron(args_defaults=extra_args)
+    mg_model = megatron_model_meta.model_provider()
+    megatron_model_meta.convert_mcore2hf(hf_model, mg_model)
+    return hf_model, mg_model, processor
+
+
+def test_align(hf_model, mg_model, processor):
+    from megatron.training.utils import get_ltor_masks_and_position_ids
+    template = get_template(hf_model.model_meta.template, processor)
+    input_ids = template.encode(InferRequest(messages=[{'role': 'user', 'content': 'who are you?'}]))['input_ids']
+    input_ids = torch.tensor(input_ids)[None].to('cuda')
+    attention_mask, _, position_ids = get_ltor_masks_and_position_ids(input_ids, -100, True, True, True)
+    with torch.inference_mode():
+        hf_model.cuda()
+        mg_model.cuda()
+        hf_logits = hf_model(input_ids).logits
+        mg_logits = mg_model(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids)
+    mean_diff = (mg_logits - hf_logits).abs().mean().item()
+    max_diff = (mg_logits - hf_logits).abs().max().item()
+    print(f'mean_diff: {mean_diff}, max_diff: {max_diff}')
+
+
+def test_save():
+    hf_model, mg_model, processor = get_mg_model_tokenizer()
+    test_align(hf_model, mg_model, processor)
+
+
+if __name__ == '__main__':
+    import torch
+    from swift.llm import InferRequest, get_model_tokenizer, get_template
+    from swift.utils import set_default_ddp_config
+    from swift.megatron.argument import MegatronArguments
+    from swift.megatron.model import get_megatron_model_meta
+    from swift.megatron.utils import patch_megatron_tokenizer
+    test_save()

tests/megatron/test_train.py

@@ -0,0 +1,37 @@
+import os
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
+
+
+def test_sft():
+    from swift.megatron import megatron_sft_main, MegatronTrainArguments
+    megatron_sft_main(
+        MegatronTrainArguments(
+            load='Qwen2-7B-Instruct-mcore',
+            dataset=[
+                'AI-ModelScope/alpaca-gpt4-data-zh#500', 'swift/self-cognition#500',
+                'AI-ModelScope/alpaca-gpt4-data-en#500'
+            ],
+            tensor_model_parallel_size=2,
+            train_iters=100,
+            model_author='swift',
+            model_name='swift-robot',
+            eval_iters=5,
+            finetune=True))
+
+
+def test_pt():
+    from swift.megatron import megatron_pt_main, MegatronTrainArguments
+    megatron_pt_main(
+        MegatronTrainArguments(
+            load='Qwen2-7B-mcore',
+            dataset=['AI-ModelScope/alpaca-gpt4-data-zh#500', 'AI-ModelScope/alpaca-gpt4-data-en#500'],
+            tensor_model_parallel_size=2,
+            train_iters=200,
+            eval_iters=5,
+            finetune=True))
+
+
+if __name__ == '__main__':
+    # test_sft()
+    test_pt()

tests/models/test_flash_attn.py

@@ -0,0 +1,8 @@
+from swift.llm import get_model_tokenizer
+
+if __name__ == '__main__':
+    # model, tokenizer = get_model_tokenizer('Qwen/Qwen2-7B-Instruct', attn_impl='flash_attn')
+    # model, tokenizer = get_model_tokenizer('AIDC-AI/Ovis2-2B', attn_impl='flash_attn')
+    # model, tokenizer = get_model_tokenizer('OpenGVLab/InternVL2-2B', attn_impl='flash_attn')
+    model, tokenizer = get_model_tokenizer('Shanghai_AI_Laboratory/internlm3-8b-instruct', attn_impl='flash_attn')
+    print(model)

tests/models/test_llm.py

@@ -0,0 +1,16 @@
+import os
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
+
+
+def test_llama3():
+    from swift.llm import infer_main, InferArguments
+    infer_main(
+        InferArguments(
+            model='LLM-Research/Meta-Llama-3.1-8B-Instruct',
+            max_batch_size=2,
+            val_dataset='AI-ModelScope/alpaca-gpt4-data-en#2'))
+
+
+if __name__ == '__main__':
+    test_llama3()

tests/models/test_mllm.py

@@ -0,0 +1,16 @@
+import os
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+
+
+def test_cogvlm():
+    from swift.llm import infer_main, InferArguments, sft_main, TrainArguments
+    # infer_main(InferArguments(model='ZhipuAI/cogvlm2-video-llama3-chat'))
+    sft_main(
+        TrainArguments(
+            model='ZhipuAI/cogvlm2-video-llama3-chat',
+            dataset=['AI-ModelScope/alpaca-gpt4-data-zh#200', 'swift/VideoChatGPT:Generic#200']))
+
+
+if __name__ == '__main__':
+    test_cogvlm()

tests/sample/test_client.py

@@ -0,0 +1,35 @@
+import os
+
+
+def test_client():
+    from swift.llm import sampling_main, SamplingArguments
+    import json
+    base_url = 'https://dashscope.aliyuncs.com/compatible-mode/v1'
+    api_key = os.environ.get('OPENAI_API_KEY')
+    engine_kwargs = json.dumps({
+        'base_url': base_url,
+        'api_key': api_key,
+    })
+    dataset = 'tastelikefeet/competition_math#5'
+    system = """A conversation between User and Assistant. The user asks a question, and the Assistant solves it.
+        The assistant first thinks about the reasoning process in the mind and then provides the user
+        with the answer. The reasoning process and answer are enclosed
+        within <think> </think> and <answer> </answer> tags, respectively,
+        i.e., <think> reasoning process here </think> <answer> answer here </answer>."""
+    args = SamplingArguments(
+        sampler_type='distill',
+        sampler_engine='client',
+        model='deepseek-r1',
+        dataset=dataset,
+        num_return_sequences=1,
+        stream=True,
+        system=system,
+        temperature=0.6,
+        top_p=0.95,
+        engine_kwargs=engine_kwargs,
+    )
+    sampling_main(args)
+
+
+if __name__ == '__main__':
+    test_client()

tests/test_align/test_cls.py

@@ -0,0 +1,60 @@
+import os
+from pprint import pprint
+
+import torch
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '1'
+kwargs = {
+    'per_device_train_batch_size': 4,
+    'per_device_eval_batch_size': 4,
+    'gradient_accumulation_steps': 4,
+    'num_train_epochs': 1,
+    'save_steps': 100,
+    'max_length': 512,
+    'task_type': 'seq_cls',
+    'num_labels': 2,
+}
+
+
+def calc_acc(infer_result):
+    n_correct = 0
+    for res in infer_result:
+        if res['response'] == res['labels']:
+            n_correct += 1
+    return f'acc: {n_correct/len(infer_result)}, n_correct: {n_correct}, len(res): {len(infer_result)}'
+
+
+def test_llm():
+    from swift.llm import sft_main, TrainArguments, infer_main, InferArguments, Template
+    res = []
+    for model in ['Qwen/Qwen2.5-0.5B-Instruct', 'Qwen/Qwen2.5-0.5B', 'AI-ModelScope/bert-base-chinese']:
+        dataset = ['DAMO_NLP/jd:cls#2000']
+        result = sft_main(TrainArguments(model=model, dataset=dataset, split_dataset_ratio=0.1, **kwargs))
+        last_model_checkpoint = result['last_model_checkpoint']
+        infer_result = infer_main(
+            InferArguments(ckpt_dir=last_model_checkpoint, load_data_args=True, truncation_strategy='right'))
+        res.append(calc_acc(infer_result))
+        infer_result2 = infer_main(
+            InferArguments(
+                ckpt_dir=last_model_checkpoint, load_data_args=True, max_batch_size=16, truncation_strategy='right'))
+        res.append(calc_acc(infer_result2))
+
+    model = 'Qwen/Qwen2.5-0.5B-Instruct'
+    dataset = ['DAMO_NLP/jd#2000']
+    train_kwargs = kwargs.copy()
+    train_kwargs.pop('task_type')
+    train_kwargs.pop('num_labels')
+    result = sft_main(TrainArguments(model=model, dataset=dataset, split_dataset_ratio=0.1, **train_kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_result = infer_main(
+        InferArguments(ckpt_dir=last_model_checkpoint, load_data_args=True, truncation_strategy='right'))
+    res.append(calc_acc(infer_result))
+    infer_result2 = infer_main(
+        InferArguments(
+            ckpt_dir=last_model_checkpoint, load_data_args=True, max_batch_size=16, truncation_strategy='right'))
+    res.append(calc_acc(infer_result2))
+    pprint(res)
+
+
+if __name__ == '__main__':
+    test_llm()

tests/test_align/test_lmdeploy_vlm.py

@@ -0,0 +1,80 @@
+import os
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+
+
+def _infer_image(model, system=None, images=None):
+    engine = LmdeployEngine(model)
+    if images is None:
+        images = ['http://modelscope-open.oss-cn-hangzhou.aliyuncs.com/images/cat.png']
+    messages = []
+    if system is not None:
+        messages += [{'role': 'system', 'content': system}]
+    messages.append({'role': 'user', 'content': 'describe the image.'})
+    resp_list = engine.infer([InferRequest(messages=messages, images=images)],
+                             RequestConfig(temperature=0, max_tokens=64, repetition_penalty=1.))
+    return resp_list[0].choices[0].message.content
+
+
+def _infer_image_pipeline(model, images=None, prefix='<IMAGE_TOKEN>\n'):
+    from lmdeploy import pipeline, GenerationConfig
+    from lmdeploy.vl import load_image
+    from swift.llm import safe_snapshot_download
+    gen_config = GenerationConfig(temperature=0., repetition_penalty=1., max_new_tokens=64)
+    pipe = pipeline(safe_snapshot_download(model))
+
+    image = load_image('http://modelscope-open.oss-cn-hangzhou.aliyuncs.com/images/cat.png')
+    response = pipe((f'{prefix}describe the image.', image), gen_config=gen_config)
+    return response.text
+
+
+def test_internvl2_5():
+    model = 'OpenGVLab/InternVL2_5-4B'
+    response = _infer_image(model)
+    response2 = _infer_image_pipeline(model)
+    assert response == response2
+
+
+def test_internvl2():
+    model = 'OpenGVLab/InternVL2-2B'
+    response = _infer_image(model)
+    response2 = _infer_image_pipeline(model)  # Missing '\n' after '<|im_end|>'
+    assert response == response2
+
+
+def test_deepseek_vl():
+    model = 'deepseek-ai/deepseek-vl-1.3b-chat'
+    response = _infer_image(model)
+    response2 = _infer_image_pipeline(model, prefix='<IMAGE_TOKEN>')
+    assert response == response2
+
+
+def test_qwen_vl():
+    model = 'Qwen/Qwen-VL-Chat'
+    response = _infer_image_pipeline(model)  # Missing: 'Picture 1: '
+    response2 = _infer_image(model)
+    assert response == response2
+
+
+def test_qwen2_vl():
+    model = 'Qwen/Qwen2-VL-2B-Instruct'
+    response = _infer_image_pipeline(model, prefix='<IMAGE_TOKEN>')
+    response2 = _infer_image(model)
+    assert response == response2
+
+
+def test_qwen2_5_vl():
+    model = 'Qwen/Qwen2.5-VL-3B-Instruct'
+    response = _infer_image(model)
+    response2 = _infer_image_pipeline(model, prefix='<IMAGE_TOKEN>')
+    assert response == response2
+
+
+if __name__ == '__main__':
+    from swift.llm import LmdeployEngine, InferRequest, RequestConfig
+    # test_internvl2()
+    # test_internvl2_5()
+    # test_deepseek_vl()
+    # test_qwen_vl()
+    # test_qwen2_vl()
+    test_qwen2_5_vl()