chore: import upstream snapshot with attribution

tests/train/test_channel.py

@@ -0,0 +1,24 @@
+import os
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0'
+
+
+def test_channel():
+    from swift import SftArguments, sft_main
+    sft_main(
+        SftArguments(
+            model='Qwen/Qwen2.5-7B-Instruct',
+            dataset=['channel.jsonl#1000'],
+            split_dataset_ratio=0.01,
+            enable_channel_loss=True,
+            packing=True,
+            max_length=128,
+            attn_impl='flash_attn',
+            load_from_cache_file=False,
+            deepspeed='zero2',
+            eval_steps=5))
+
+
+if __name__ == '__main__':
+    test_channel()

tests/train/test_cls.py

@@ -0,0 +1,61 @@
+import os
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0'
+kwargs = {
+    'per_device_train_batch_size': 2,
+    'per_device_eval_batch_size': 2,
+    'save_steps': 50,
+    'gradient_accumulation_steps': 4,
+    'num_train_epochs': 1,
+}
+
+
+def test_llm():
+    from swift import InferArguments, SftArguments, infer_main, sft_main
+    result = sft_main(
+        SftArguments(
+            model='Qwen/Qwen2.5-1.5B-Instruct',
+            tuner_type='lora',
+            num_labels=2,
+            dataset=['DAMO_NLP/jd:cls#2000'],
+            split_dataset_ratio=0.01,
+            **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=last_model_checkpoint, load_data_args=True))
+
+
+def test_bert():
+
+    from swift import InferArguments, SftArguments, infer_main, sft_main
+    result = sft_main(
+        SftArguments(
+            model='answerdotai/ModernBERT-base',
+            # model='iic/nlp_structbert_backbone_base_std',
+            tuner_type='full',
+            num_labels=2,
+            dataset=['DAMO_NLP/jd:cls#2000'],
+            split_dataset_ratio=0.01,
+            **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(model=last_model_checkpoint, load_data_args=True))
+
+
+def test_mllm():
+    from swift import InferArguments, SftArguments, infer_main, sft_main
+    result = sft_main(
+        SftArguments(
+            model='OpenGVLab/InternVL2-1B',
+            tuner_type='lora',
+            num_labels=2,
+            dataset=['DAMO_NLP/jd:cls#500'],
+            split_dataset_ratio=0.01,
+            **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=last_model_checkpoint, load_data_args=True))
+
+
+if __name__ == '__main__':
+    # test_llm()
+    # test_bert()
+    test_mllm()

tests/train/test_embedding.py

@@ -0,0 +1,88 @@
+import os
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0'
+kwargs = {
+    'per_device_train_batch_size': 4,
+    'save_steps': 5,
+    'gradient_accumulation_steps': 4,
+    'num_train_epochs': 1,
+}
+
+
+def test_embedding():
+    from swift import SftArguments, sft_main
+    result = sft_main(
+        SftArguments(
+            model='Qwen/Qwen3-Embedding-0.6B',
+            task_type='embedding',
+            dataset=['sentence-transformers/stsb:positive'],
+            split_dataset_ratio=0.01,
+            load_from_cache_file=False,
+            loss_type='infonce',
+            attn_impl='flash_attn',
+            max_length=2048,
+            **kwargs,
+        ))
+    last_model_checkpoint = result['last_model_checkpoint']
+    print(f'last_model_checkpoint: {last_model_checkpoint}')
+
+
+def test_reranker():
+    from swift import SftArguments, sft_main
+    result = sft_main(
+        SftArguments(
+            model='Qwen/Qwen3-Reranker-4B',
+            tuner_type='lora',
+            load_from_cache_file=True,
+            task_type='generative_reranker',
+            dataset=['MTEB/scidocs-reranking#10000'],
+            split_dataset_ratio=0.05,
+            loss_type='pointwise_reranker',
+            dataloader_drop_last=True,
+            eval_strategy='steps',
+            eval_steps=10,
+            max_length=4096,
+            attn_impl='flash_attn',
+            num_train_epochs=1,
+            save_steps=200,
+            per_device_train_batch_size=2,
+            per_device_eval_batch_size=2,
+            gradient_accumulation_steps=8,
+            dataset_num_proc=2,
+        ))
+    last_model_checkpoint = result['last_model_checkpoint']
+    print(f'last_model_checkpoint: {last_model_checkpoint}')
+
+
+def test_reranker2():
+    from swift import SftArguments, sft_main
+    result = sft_main(
+        SftArguments(
+            model='Qwen/Qwen2.5-VL-3B-Instruct',
+            tuner_type='lora',
+            load_from_cache_file=True,
+            task_type='reranker',
+            dataset=['MTEB/scidocs-reranking'],
+            split_dataset_ratio=0.05,
+            loss_type='listwise_reranker',
+            dataloader_drop_last=True,
+            eval_strategy='steps',
+            eval_steps=10,
+            max_length=4096,
+            attn_impl='flash_attn',
+            padding_side='right',
+            num_train_epochs=1,
+            save_steps=200,
+            per_device_train_batch_size=2,
+            per_device_eval_batch_size=2,
+            gradient_accumulation_steps=8,
+            dataset_num_proc=1,
+        ))
+    last_model_checkpoint = result['last_model_checkpoint']
+    print(f'last_model_checkpoint: {last_model_checkpoint}')
+
+
+if __name__ == '__main__':
+    # test_embedding()
+    test_reranker()

tests/train/test_export_cached_dataset.py

@@ -0,0 +1,27 @@
+def test_export_cached_dataset():
+    from swift import ExportArguments, export_main
+    export_main(
+        ExportArguments(
+            model='Qwen/Qwen2.5-7B-Instruct',
+            dataset='swift/Chinese-Qwen3-235B-2507-Distill-data-110k-SFT',
+            to_cached_dataset=True,
+            dataset_num_proc=4,
+        ))
+    print()
+
+
+def test_sft():
+    from swift import SftArguments, sft_main
+    sft_main(
+        SftArguments(
+            model='Qwen/Qwen2.5-7B-Instruct',
+            dataset='liucong/Chinese-DeepSeek-R1-Distill-data-110k-SFT#1000',
+            dataset_num_proc=2,
+            packing=True,
+            attn_impl='flash_attn',
+        ))
+
+
+if __name__ == '__main__':
+    # test_export_cached_dataset()
+    test_sft()

tests/train/test_freeze.py

@@ -0,0 +1,89 @@
+import os
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0'
+kwargs = {
+    'per_device_train_batch_size': 2,
+    'save_steps': 5,
+    'gradient_accumulation_steps': 4,
+    'num_train_epochs': 1,
+}
+
+
+def test_full_vit():
+    os.environ['MAX_PIXELS'] = '100352'
+    os.environ['SIZE_FACTOR'] = '12'
+    os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
+    os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0,1'
+    from swift import InferArguments, SftArguments, infer_main, sft_main
+    sft_main(
+        SftArguments(
+            model='Qwen/Qwen2-VL-7B-Instruct',
+            dataset=['modelscope/coco_2014_caption:validation#20', 'AI-ModelScope/alpaca-gpt4-data-en#20'],
+            split_dataset_ratio=0.01,
+            tuner_type='full',
+            freeze_llm=True,
+            freeze_vit=False,
+            freeze_aligner=True,
+            **kwargs))
+
+
+def test_full_aligner():
+    os.environ['MAX_PIXELS'] = '100352'
+    os.environ['SIZE_FACTOR'] = '12'
+    os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
+    os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0,1'
+    from swift import InferArguments, SftArguments, infer_main, sft_main
+    sft_main(
+        SftArguments(
+            model='Qwen/Qwen2-VL-7B-Instruct',
+            dataset=['modelscope/coco_2014_caption:validation#20', 'AI-ModelScope/alpaca-gpt4-data-en#20'],
+            split_dataset_ratio=0.01,
+            tuner_type='full',
+            freeze_llm=True,
+            freeze_vit=True,
+            freeze_aligner=False,
+            **kwargs))
+
+
+def test_lora_vit():
+    os.environ['MAX_PIXELS'] = '100352'
+    os.environ['SIZE_FACTOR'] = '12'
+    os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
+    os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0,1'
+    from swift import InferArguments, SftArguments, infer_main, sft_main
+    sft_main(
+        SftArguments(
+            model='Qwen/Qwen2-VL-7B-Instruct',
+            dataset=['modelscope/coco_2014_caption:validation#20', 'AI-ModelScope/alpaca-gpt4-data-en#20'],
+            split_dataset_ratio=0.01,
+            tuner_type='lora',
+            freeze_llm=True,
+            freeze_vit=False,
+            freeze_aligner=True,
+            **kwargs))
+
+
+def test_lora_aligner():
+    os.environ['MAX_PIXELS'] = '100352'
+    os.environ['SIZE_FACTOR'] = '12'
+    os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
+    os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0,1'
+    from swift import InferArguments, SftArguments, infer_main, sft_main
+    sft_main(
+        SftArguments(
+            model='Qwen/Qwen2-VL-7B-Instruct',
+            dataset=['modelscope/coco_2014_caption:validation#20', 'AI-ModelScope/alpaca-gpt4-data-en#20'],
+            split_dataset_ratio=0.01,
+            tuner_type='lora',
+            freeze_llm=True,
+            freeze_vit=True,
+            freeze_aligner=False,
+            **kwargs))
+
+
+if __name__ == '__main__':
+    # test_full_vit()
+    test_full_aligner()
+    # test_lora_vit()
+    # test_lora_aligner()

tests/train/test_gkd.py

@@ -0,0 +1,80 @@
+import os
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0'
+kwargs = {
+    'per_device_train_batch_size': 4,
+    'save_steps': 5,
+    'gradient_accumulation_steps': 4,
+    'num_train_epochs': 1,
+}
+
+
+def test_llm():
+    from swift import InferArguments, RLHFArguments, infer_main, rlhf_main
+    result = rlhf_main(
+        RLHFArguments(
+            rlhf_type='gkd',
+            model='Qwen/Qwen2.5-0.5B',
+            teacher_model='Qwen/Qwen2.5-1.5B-Instruct',
+            dataset=['AI-ModelScope/alpaca-gpt4-data-en#2000'],
+            split_dataset_ratio=0.01,
+            load_from_cache_file=False,
+            seq_kd=True,
+            **kwargs,
+        ))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=last_model_checkpoint, load_data_args=True, merge_lora=True))
+
+
+def test_mllm():
+    from swift import InferArguments, RLHFArguments, infer_main, rlhf_main
+    result = rlhf_main(
+        RLHFArguments(
+            rlhf_type='gkd',
+            model='OpenGVLab/InternVL3-2B-Pretrained',
+            teacher_model='OpenGVLab/InternVL3-8B',
+            dataset=['AI-ModelScope/LaTeX_OCR#2000', 'AI-ModelScope/alpaca-gpt4-data-en#2000'],
+            split_dataset_ratio=0.01,
+            load_from_cache_file=False,
+            **kwargs,
+        ))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=last_model_checkpoint, load_data_args=True, merge_lora=True))
+
+
+def test_multi_turn():
+    """GKD multi-turn smoke test: verify rollout → encode → loss works with multi_turn_scheduler.
+
+    Uses the built-in ``math_tip_trick`` scheduler with max_turns=2 to keep the test
+    lightweight. The key assertion is that training completes without raising
+    NotImplementedError (the previous block) and that multi-turn response token ids
+    are correctly propagated through the GKD loss pipeline.
+    """
+    from swift import InferArguments, RLHFArguments, infer_main, rlhf_main
+    result = rlhf_main(
+        RLHFArguments(
+            rlhf_type='gkd',
+            model='Qwen/Qwen2.5-0.5B',
+            teacher_model='Qwen/Qwen2.5-1.5B-Instruct',
+            dataset=['AI-ModelScope/alpaca-gpt4-data-en#200'],
+            split_dataset_ratio=0.01,
+            load_from_cache_file=False,
+            multi_turn_scheduler='math_tip_trick',
+            max_turns=2,
+            max_completion_length=256,
+            num_generations=2,
+            per_device_train_batch_size=2,
+            gradient_accumulation_steps=1,
+            save_steps=50,
+            num_train_epochs=1,
+        ))
+    last_model_checkpoint = result['last_model_checkpoint']
+    if last_model_checkpoint is not None:
+        infer_main(InferArguments(adapters=last_model_checkpoint, load_data_args=True, merge_lora=True))
+
+
+if __name__ == '__main__':
+    # test_llm()
+    # test_mllm()
+    test_multi_turn()

tests/train/test_grounding.py

@@ -0,0 +1,9 @@
+import os
+
+from swift import SftArguments, sft_main
+
+os.environ['MAX_PIXELS'] = str(16 * 28 * 28)
+
+if __name__ == '__main__':
+    sft_main(
+        SftArguments(model='Qwen/Qwen2.5-VL-7B-Instruct', dataset='AI-ModelScope/coco#2000', split_dataset_ratio=0.01))

tests/train/test_grpo.py

@@ -0,0 +1,153 @@
+import os
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
+os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0,1'
+kwargs = {
+    'per_device_train_batch_size': 2,
+    'per_device_eval_batch_size': 2,
+    'save_steps': 50,
+    'gradient_accumulation_steps': 1,
+    'num_train_epochs': 1,
+}
+
+SYSTEM_PROMPT = ('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>')
+
+
+def test_llm():
+    from swift import InferArguments, RLHFArguments, infer_main, rlhf_main
+    result = rlhf_main(
+        RLHFArguments(
+            rlhf_type='grpo',
+            model='Qwen/Qwen2.5-1.5B-Instruct',
+            tuner_type='full',
+            dataset=['AI-MO/NuminaMath-TIR#100'],
+            split_dataset_ratio=0.1,
+            system=SYSTEM_PROMPT,
+            reward_funcs=['accuracy', 'format'],
+            max_completion_length=4096,
+            num_generations=2,
+            **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=last_model_checkpoint, load_data_args=True, merge_lora=True))
+
+
+def test_llm_zero2():
+    from swift import InferArguments, RLHFArguments, infer_main, rlhf_main
+    result = rlhf_main(
+        RLHFArguments(
+            rlhf_type='grpo',
+            model='Qwen/Qwen2.5-1.5B-Instruct',
+            tuner_type='full',
+            dataset=['AI-MO/NuminaMath-TIR#100'],
+            system=SYSTEM_PROMPT,
+            reward_funcs=['accuracy', 'format'],
+            max_completion_length=4096,
+            num_generations=2,
+            deepspeed='zero2',
+            **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=last_model_checkpoint, load_data_args=True, merge_lora=True))
+
+
+def test_llm_vllm():
+    from swift import InferArguments, RLHFArguments, infer_main, rlhf_main
+    result = rlhf_main(
+        RLHFArguments(
+            rlhf_type='grpo',
+            model='Qwen/Qwen2.5-1.5B-Instruct',
+            reward_model='AI-ModelScope/GRM_Llama3.1_8B_rewardmodel-ft',
+            tuner_type='full',
+            dataset=['AI-MO/NuminaMath-TIR#100'],
+            system=SYSTEM_PROMPT,
+            reward_funcs=['accuracy', 'format'],
+            use_vllm=True,
+            max_completion_length=4096,
+            num_generations=2,
+            **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=last_model_checkpoint, load_data_args=True, merge_lora=True))
+
+
+def test_llm_vllm_zero2():
+    from swift import InferArguments, RLHFArguments, infer_main, rlhf_main
+    result = rlhf_main(
+        RLHFArguments(
+            rlhf_type='grpo',
+            model='Qwen/Qwen2.5-1.5B-Instruct',
+            tuner_type='full',
+            dataset=['AI-MO/NuminaMath-TIR#100'],
+            system=SYSTEM_PROMPT,
+            reward_funcs=['accuracy', 'format'],
+            use_vllm=True,
+            max_completion_length=4096,
+            num_generations=2,
+            deepspeed='zero2',
+            **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=last_model_checkpoint, load_data_args=True, merge_lora=True))
+
+
+def test_mllm_pt():
+    from swift import InferArguments, RLHFArguments, infer_main, rlhf_main
+    result = rlhf_main(
+        RLHFArguments(
+            rlhf_type='grpo',
+            model='Qwen/Qwen2-VL-2B-Instruct',
+            tuner_type='full',
+            # dataset=['AI-MO/NuminaMath-TIR#100'],
+            dataset=['modelscope/coco_2014_caption:validation#100'],
+            system=SYSTEM_PROMPT,
+            reward_funcs=['format'],
+            max_completion_length=4096,
+            num_generations=2,
+            **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=last_model_checkpoint, load_data_args=True, merge_lora=True))
+
+
+def test_grpo_minimal():
+    import trl
+    from packaging import version
+    if version.parse(trl.__version__) < version.parse('0.26'):
+        print(f'Skipping test_grpo_minimal: trl>=0.26 required, found trl=={trl.__version__}')
+        return
+    from swift import InferArguments, RLHFArguments, infer_main, rlhf_main
+    result = rlhf_main(
+        RLHFArguments(
+            rlhf_type='grpo',
+            model='Qwen/Qwen2-0.5B',
+            tuner_type='lora',
+            dataset=['AI-ModelScope/alpaca-gpt4-data-zh#20'],
+            system=SYSTEM_PROMPT,
+            reward_funcs=['format'],
+            max_completion_length=128,
+            num_generations=2,
+            max_steps=2,
+            per_device_train_batch_size=2,
+            gradient_accumulation_steps=1,
+            save_steps=2,
+            split_dataset_ratio=0.01,
+            logging_steps=1,
+            use_vllm=False,
+            **{
+                k: v
+                for k, v in kwargs.items() if k not in [
+                    'per_device_train_batch_size', 'save_steps', 'gradient_accumulation_steps', 'num_train_epochs',
+                    'per_device_eval_batch_size'
+                ]
+            }))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=last_model_checkpoint, load_data_args=True))
+
+
+if __name__ == '__main__':
+    # test_llm()
+    # test_llm_zero3()
+    # test_llm_vllm()
+    # test_llm_vllm_zero2()
+    test_mllm_pt()
+    # test_grpo_minimal()

tests/train/test_kto.py

@@ -0,0 +1,41 @@
+import os
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0'
+kwargs = {
+    'per_device_train_batch_size': 2,
+    'save_steps': 5,
+    'gradient_accumulation_steps': 4,
+    'num_train_epochs': 1,
+}
+
+
+def test_llm():
+    from swift import InferArguments, RLHFArguments, infer_main, rlhf_main
+    result = rlhf_main(
+        RLHFArguments(
+            rlhf_type='kto',
+            model='Qwen/Qwen2-7B-Instruct',
+            dataset=['AI-ModelScope/ultrafeedback-binarized-preferences-cleaned-kto#100'],
+            split_dataset_ratio=0.01,
+            **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=last_model_checkpoint, load_data_args=True, merge_lora=True))
+
+
+def test_mllm():
+    from swift import InferArguments, RLHFArguments, infer_main, rlhf_main
+    result = rlhf_main(
+        RLHFArguments(
+            rlhf_type='kto',
+            model='Qwen/Qwen2-VL-7B-Instruct',
+            dataset=['AI-ModelScope/ultrafeedback-binarized-preferences-cleaned-kto#100'],
+            split_dataset_ratio=0.01,
+            **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=last_model_checkpoint, load_data_args=True, merge_lora=True))
+
+
+if __name__ == '__main__':
+    # test_llm()
+    test_mllm()

tests/train/test_liger.py

@@ -0,0 +1,46 @@
+import os
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
+os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0,1'
+kwargs = {
+    'per_device_train_batch_size': 2,
+    'save_steps': 30,
+    'gradient_accumulation_steps': 2,
+    'num_train_epochs': 1,
+}
+
+
+def test_sft():
+    from swift import InferArguments, SftArguments, infer_main, sft_main
+    result = sft_main(
+        SftArguments(
+            model='Qwen/Qwen2.5-7B-Instruct',
+            dataset=['swift/self-cognition#200'],
+            split_dataset_ratio=0.01,
+            use_liger_kernel=True,
+            **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=last_model_checkpoint, load_data_args=True))
+
+
+def test_mllm_dpo():
+    os.environ['MAX_PIXLES'] = f'{1280 * 28 * 28}'
+    from swift import InferArguments, RLHFArguments, infer_main, rlhf_main
+    result = rlhf_main(
+        RLHFArguments(
+            rlhf_type='dpo',
+            model='Qwen/Qwen2.5-VL-3B-Instruct',
+            tuner_type='full',
+            dataset=['swift/RLAIF-V-Dataset#1000'],
+            split_dataset_ratio=0.01,
+            dataset_num_proc=8,
+            deepspeed='zero3',
+            use_liger_kernel=True,
+            **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(model=last_model_checkpoint, load_data_args=True))
+
+
+if __name__ == '__main__':
+    test_sft()
+    # test_mllm_dpo()

tests/train/test_multilabel.py

@@ -0,0 +1,46 @@
+import os
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0'
+kwargs = {
+    'per_device_train_batch_size': 2,
+    'per_device_eval_batch_size': 2,
+    'save_steps': 50,
+    'gradient_accumulation_steps': 4,
+    'num_train_epochs': 1,
+}
+
+
+def test_reg_llm():
+    from swift import InferArguments, SftArguments, infer_main, sft_main
+    result = sft_main(
+        SftArguments(
+            model='Qwen/Qwen2.5-1.5B-Instruct',
+            tuner_type='lora',
+            num_labels=1,
+            dataset=['sentence-transformers/stsb:reg#200'],
+            split_dataset_ratio=0.01,
+            **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=last_model_checkpoint, load_data_args=True, metric='acc'))
+
+
+def test_reg_mllm():
+    from swift import InferArguments, SftArguments, infer_main, sft_main
+
+    # OpenGVLab/InternVL2-1B
+    result = sft_main(
+        SftArguments(
+            model='Qwen/Qwen2-VL-2B-Instruct',
+            tuner_type='lora',
+            num_labels=1,
+            dataset=['sentence-transformers/stsb:reg#200'],
+            split_dataset_ratio=0.01,
+            **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=last_model_checkpoint, load_data_args=True, metric='acc'))
+
+
+if __name__ == '__main__':
+    # test_reg_llm()
+    test_reg_mllm()

tests/train/test_opsd.py

@@ -0,0 +1,48 @@
+import os
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
+os.environ['NPROC_PER_NODE'] = '2'
+
+
+def train():
+    from swift import RLHFArguments, rlhf_main
+    result = rlhf_main(
+        RLHFArguments(
+            rlhf_type='gkd',
+            model='Qwen/Qwen3.5-4B',
+            teacher_model='Qwen/Qwen3.5-4B',
+            tuner_type='lora',
+            lora_rank=64,
+            lora_alpha=128,
+            target_modules=['all-linear'],
+            use_vllm=True,
+            vllm_mode='colocate',
+            vllm_gpu_memory_utilization=0.7,
+            vllm_max_model_len=10240,
+            sleep_level=1,
+            external_plugins=['examples/train/rlhf/opsd/opsd_plugin.py'],
+            dataset=['open-r1/OpenThoughts-114k-math'],
+            lmbda=1.0,
+            beta=0.5,
+            temperature=1.2,
+            sft_alpha=0,
+            torch_dtype='bfloat16',
+            max_steps=1000,
+            per_device_train_batch_size=4,
+            gradient_accumulation_steps=1,
+            learning_rate=2e-5,
+            save_steps=100,
+            save_total_limit=10,
+            logging_steps=1,
+            max_length=8192,
+            max_completion_length=2048,
+            save_only_model=True,
+            gradient_checkpointing=True,
+            deepspeed='zero0',
+            attn_impl='flash_attn',
+        ))
+    return result
+
+
+if __name__ == '__main__':
+    train()

tests/train/test_packing.py

@@ -0,0 +1,66 @@
+import os
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0'
+kwargs = {
+    'per_device_train_batch_size': 2,
+    'save_steps': 50,
+    'gradient_accumulation_steps': 4,
+    'num_train_epochs': 3,
+}
+
+
+def test_llm():
+    from swift import InferArguments, SftArguments, infer_main, sft_main
+    result = sft_main(
+        SftArguments(
+            model='Qwen/Qwen2-7B-Instruct',
+            dataset=['AI-ModelScope/alpaca-gpt4-data-zh#1000', 'swift/self-cognition#1000'],
+            split_dataset_ratio=0.01,
+            packing=True,
+            max_length=4096,
+            attn_impl='flash_attn',
+            logging_steps=1,
+            **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=last_model_checkpoint, load_data_args=True, merge_lora=True))
+
+
+def test_streaming():
+    from swift import InferArguments, SftArguments, infer_main, sft_main
+    result = sft_main(
+        SftArguments(
+            model='Qwen/Qwen2-7B-Instruct',
+            dataset=['AI-ModelScope/alpaca-gpt4-data-zh#10000'],
+            packing=True,
+            max_length=4096,
+            streaming=True,
+            attn_impl='flash_attn',
+            max_steps=100,
+            dataset_num_proc=1,
+            **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=last_model_checkpoint, load_data_args=True, merge_lora=True))
+
+
+def test_mllm_streaming():
+    from swift import InferArguments, SftArguments, infer_main, sft_main
+    result = sft_main(
+        SftArguments(
+            model='Qwen/Qwen2.5-VL-7B-Instruct',
+            dataset=['AI-ModelScope/LaTeX_OCR#20000'],
+            packing=True,
+            max_length=8192,
+            streaming=True,
+            attn_impl='flash_attn',
+            max_steps=100,
+            dataset_num_proc=4,
+            **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=last_model_checkpoint, load_data_args=True, merge_lora=True))
+
+
+if __name__ == '__main__':
+    # test_llm()
+    # test_streaming()
+    test_mllm_streaming()

tests/train/test_ppo.py

@@ -0,0 +1,41 @@
+import os
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0'
+kwargs = {
+    'per_device_train_batch_size': 2,
+    'save_steps': 5,
+    'gradient_accumulation_steps': 4,
+    'num_train_epochs': 1,
+}
+
+
+def test_rm():
+    from swift import InferArguments, RLHFArguments, infer_main, rlhf_main
+    result = rlhf_main(
+        RLHFArguments(
+            rlhf_type='rm',
+            model='Shanghai_AI_Laboratory/internlm2-1_8b-reward',
+            dataset=['hjh0119/shareAI-Llama3-DPO-zh-en-emoji#100'],
+            split_dataset_ratio=0.01,
+            **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=last_model_checkpoint, load_data_args=True, merge_lora=True))
+
+
+def test_ppo():
+    from swift import InferArguments, RLHFArguments, infer_main, rlhf_main
+    result = rlhf_main(
+        RLHFArguments(
+            rlhf_type='ppo',
+            model='LLM-Research/Llama-3.2-1B-Instruct',
+            reward_model='AI-ModelScope/GRM-Llama3.2-3B-rewardmodel-ft',
+            dataset=['AI-ModelScope/alpaca-gpt4-data-zh#100', 'AI-ModelScope/alpaca-gpt4-data-en#100'],
+            **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=last_model_checkpoint, load_data_args=True, merge_lora=True))
+
+
+if __name__ == '__main__':
+    # test_rm()
+    test_ppo()

tests/train/test_pt.py

@@ -0,0 +1,58 @@
+import os
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0'
+kwargs = {
+    'per_device_train_batch_size': 2,
+    'save_steps': 5,
+    'gradient_accumulation_steps': 4,
+    'num_train_epochs': 1,
+}
+
+
+def test_llm():
+    from swift import InferArguments, PretrainArguments, infer_main, pretrain_main
+    result = pretrain_main(
+        PretrainArguments(
+            model='Qwen/Qwen2-7B-Instruct', dataset=['swift/sharegpt:all#100'], split_dataset_ratio=0.01, **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=last_model_checkpoint, load_data_args=True, merge_lora=True))
+
+
+def test_mllm():
+    from swift import InferArguments, PretrainArguments, infer_main, pretrain_main
+    result = pretrain_main(
+        PretrainArguments(
+            model='Qwen/Qwen2-VL-7B-Instruct',
+            dataset=['modelscope/coco_2014_caption:validation#20', 'AI-ModelScope/alpaca-gpt4-data-en#20'],
+            split_dataset_ratio=0.01,
+            **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=last_model_checkpoint, load_data_args=True, merge_lora=True))
+
+
+def test_pretrain_minimal():
+    from swift import PretrainArguments, pretrain_main
+    result = pretrain_main(
+        PretrainArguments(
+            model='Qwen/Qwen2-0.5B',
+            dataset=['AI-ModelScope/alpaca-gpt4-data-zh#20'],
+            max_steps=2,
+            per_device_train_batch_size=1,
+            gradient_accumulation_steps=1,
+            save_steps=2,
+            split_dataset_ratio=0.01,
+            tuner_type='lora',
+            logging_steps=1,
+            **{
+                k: v
+                for k, v in kwargs.items() if k not in
+                ['per_device_train_batch_size', 'save_steps', 'gradient_accumulation_steps', 'num_train_epochs']
+            }))
+    assert os.path.isdir(result['last_model_checkpoint'])
+
+
+if __name__ == '__main__':
+    # test_llm()
+    test_mllm()
+    # test_pretrain_minimal()

tests/train/test_resume_from_checkpoint.py

@@ -0,0 +1,40 @@
+import os
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2,3'
+os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0,1,2,3'
+kwargs = {
+    'per_device_train_batch_size': 2,
+    'per_device_eval_batch_size': 2,
+    'save_steps': 5,
+    'logging_steps': 1,
+    'gradient_accumulation_steps': 4,
+    'num_train_epochs': 1,
+    'model': 'Qwen/Qwen2-0.5B',
+    'dataset': ['AI-ModelScope/alpaca-gpt4-data-zh#2000'],
+    'val_dataset': ['AI-ModelScope/alpaca-gpt4-data-zh#10'],
+    'max_steps': 10,
+    'dataset_num_proc': 4,
+    'dataloader_num_workers': 4,
+    'max_length': 2048,
+    # optional
+    # 'padding_free': True,
+    'packing': True,
+    'attn_impl': 'flash_attn',
+    # 'streaming': True,
+    'sequence_parallel_size': 2,
+}
+
+
+def test_resume_from_checkpoint():
+    from swift import InferArguments, SftArguments, infer_main, sft_main
+    result = sft_main(SftArguments(**kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    last_model_checkpoint = last_model_checkpoint.replace('checkpoint-10', 'checkpoint-5')
+    result2 = sft_main(SftArguments(**kwargs, resume_from_checkpoint=last_model_checkpoint))
+    diff = abs(result['log_history'][6]['loss'] - result2['log_history'][6]['loss'])
+    print(f'diff: {diff}')
+    assert diff < 0.01
+
+
+if __name__ == '__main__':
+    test_resume_from_checkpoint()

tests/train/test_rlhf.py

@@ -0,0 +1,86 @@
+import os
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0'
+kwargs = {
+    'per_device_train_batch_size': 2,
+    'save_steps': 5,
+    'gradient_accumulation_steps': 4,
+    'num_train_epochs': 1,
+}
+
+
+def test_llm():
+    from swift import InferArguments, RLHFArguments, infer_main, rlhf_main
+    result = rlhf_main(
+        RLHFArguments(
+            rlhf_type='dpo',
+            model='Qwen/Qwen2-7B-Instruct',
+            dataset=['hjh0119/shareAI-Llama3-DPO-zh-en-emoji#100'],
+            split_dataset_ratio=0.01,
+            **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=last_model_checkpoint, load_data_args=True, merge_lora=True))
+
+
+def test_mllm():
+    from swift import InferArguments, RLHFArguments, infer_main, rlhf_main
+    os.environ['MAX_PIXLES'] = f'{1280 * 28 * 28}'
+    result = rlhf_main(
+        RLHFArguments(
+            rlhf_type='dpo',
+            model='Qwen/Qwen2-VL-7B-Instruct',
+            dataset=['swift/RLAIF-V-Dataset#100'],
+            split_dataset_ratio=0.01,
+            dataset_num_proc=8,
+            max_pixels=512 * 512,
+            **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=last_model_checkpoint, load_data_args=True, merge_lora=True))
+
+
+def test_mllm_zero3():
+    os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
+    os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0,1'
+    os.environ['MAX_PIXLES'] = f'{1280 * 28 * 28}'
+    from swift import InferArguments, RLHFArguments, infer_main, rlhf_main
+    rlhf_main(
+        RLHFArguments(
+            rlhf_type='dpo',
+            model='Qwen/Qwen2-VL-7B-Instruct',
+            dataset=['swift/RLAIF-V-Dataset#100'],
+            split_dataset_ratio=0.01,
+            dataset_num_proc=8,
+            max_pixels=512 * 512,
+            deepspeed='zero3',
+            **kwargs))
+
+
+def test_dpo_minimal():
+    from swift import InferArguments, RLHFArguments, infer_main, rlhf_main
+    result = rlhf_main(
+        RLHFArguments(
+            rlhf_type='dpo',
+            model='Qwen/Qwen2-0.5B',
+            dataset=['AI-ModelScope/orpo-dpo-mix-40k#20'],
+            max_steps=2,
+            per_device_train_batch_size=1,
+            gradient_accumulation_steps=1,
+            save_steps=2,
+            split_dataset_ratio=0.01,
+            tuner_type='lora',
+            logging_steps=1,
+            **{
+                k: v
+                for k, v in kwargs.items() if k not in
+                ['per_device_train_batch_size', 'save_steps', 'gradient_accumulation_steps', 'num_train_epochs']
+            }))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=last_model_checkpoint, load_data_args=True))
+
+
+if __name__ == '__main__':
+    # test_llm()
+    test_mllm()
+    # test_mllm_zero3()
+    # test_dpo_minimal()

tests/train/test_sample.py

@@ -0,0 +1,14 @@
+from swift import SamplingArguments, sampling_main
+
+
+def test_sampling():
+    sampling_main(
+        SamplingArguments(
+            model='LLM-Research/Meta-Llama-3.1-8B-Instruct',
+            sampler_engine='transformers',
+            num_return_sequences=5,
+            dataset='AI-ModelScope/alpaca-gpt4-data-zh#5'))
+
+
+if __name__ == '__main__':
+    test_sampling()

tests/train/test_sft.py

@@ -0,0 +1,500 @@
+import os
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0'
+kwargs = {
+    'per_device_train_batch_size': 2,
+    'per_device_eval_batch_size': 2,
+    'save_steps': 5,
+    'gradient_accumulation_steps': 4,
+    'num_train_epochs': 1,
+}
+
+
+def test_llm_ddp():
+    os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
+    os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0,1'
+    from swift import InferArguments, SftArguments, infer_main, sft_main
+    result = sft_main(
+        SftArguments(
+            model='Qwen/Qwen2-7B-Instruct',
+            dataset=['AI-ModelScope/alpaca-gpt4-data-zh#100', 'AI-ModelScope/alpaca-gpt4-data-en#100'],
+            split_dataset_ratio=0.01,
+            # ddp_find_unused_parameters=False,
+            gradient_checkpointing_kwargs={'use_reentrant': False},
+            target_modules=['all-linear', 'all-embedding'],
+            modules_to_save=['all-embedding', 'all-norm'],
+            **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=last_model_checkpoint, load_data_args=True))
+
+
+def test_unsloth():
+    from swift import InferArguments, SftArguments, infer_main, sft_main
+    result = sft_main(
+        SftArguments(
+            model='Qwen/Qwen2-0.5B',
+            dataset=['AI-ModelScope/alpaca-gpt4-data-zh#100', 'AI-ModelScope/alpaca-gpt4-data-en#100'],
+            split_dataset_ratio=0.01,
+            max_steps=5,
+            tuner_backend='unsloth',
+            **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    result = sft_main(SftArguments(resume_from_checkpoint=last_model_checkpoint, load_data_args=True, max_steps=10))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=last_model_checkpoint, load_data_args=True))
+
+
+def test_mllm_mp():
+    os.environ['MAX_PIXELS'] = '100352'
+    os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2,3'
+    os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0,1,2,3'
+    from swift import InferArguments, SftArguments, infer_main, sft_main
+    result = sft_main(
+        SftArguments(
+            model='Qwen/Qwen2.5-VL-7B-Instruct',
+            dataset=['modelscope/coco_2014_caption:validation#20'],
+            # dataset=['modelscope/coco_2014_caption:validation#20', 'AI-ModelScope/alpaca-gpt4-data-en#20'],
+            split_dataset_ratio=0.01,
+            tuner_type='lora',
+            target_modules=['all-linear'],
+            freeze_aligner=False,
+            **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=[last_model_checkpoint], load_data_args=True, merge_lora=True))
+
+
+def test_llm_streaming():
+    from swift import InferArguments, SftArguments, infer_main, sft_main
+    result = sft_main(
+        SftArguments(
+            model='Qwen/Qwen2-7B-Instruct', dataset=['swift/chinese-c4'], streaming=True, max_steps=16, **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=[last_model_checkpoint], load_data_args=True, merge_lora=True))
+
+
+def test_mllm_streaming():
+    from swift import InferArguments, SftArguments, infer_main, sft_main
+    result = sft_main(
+        SftArguments(
+            model='Qwen/Qwen2-VL-7B-Instruct',
+            dataset=['modelscope/coco_2014_caption:validation', 'AI-ModelScope/alpaca-gpt4-data-en'],
+            streaming=True,
+            max_steps=16,
+            split_dataset_ratio=0.01,
+            **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=[last_model_checkpoint], load_data_args=True, merge_lora=True))
+
+
+def test_mllm_zero3():
+    os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
+    os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0,1'
+    from swift import SftArguments, sft_main
+    sft_main(
+        SftArguments(
+            model='Qwen/Qwen2-VL-7B-Instruct',
+            dataset=['modelscope/coco_2014_caption:validation#100', 'AI-ModelScope/alpaca-gpt4-data-en#100'],  #
+            split_dataset_ratio=0.01,
+            deepspeed='zero3',
+            **kwargs))
+
+
+def test_qwen_vl():
+    os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
+    os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0,1'
+    from swift import SftArguments, sft_main
+    sft_main(
+        SftArguments(
+            model='Qwen/Qwen-VL-Chat',
+            dataset=['AI-ModelScope/LaTeX_OCR#40', 'modelscope/coco_2014_caption:validation#40'],
+            split_dataset_ratio=0.01,
+            **kwargs))
+
+
+def test_qwen2_audio():
+    os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
+    os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0,1'
+    from swift import SftArguments, sft_main
+    sft_main(
+        SftArguments(
+            model='Qwen/Qwen2-Audio-7B-Instruct',
+            dataset=['speech_asr/speech_asr_aishell1_trainsets:validation#200'],
+            split_dataset_ratio=0.01,
+            freeze_parameters_ratio=1,
+            trainable_parameters=['audio_tower'],
+            tuner_type='full',
+            **kwargs))
+
+
+def test_llm_gptq():
+    from swift import InferArguments, SftArguments, infer_main, sft_main
+    result = sft_main(
+        SftArguments(
+            model='Qwen/Qwen2-7B-Instruct-GPTQ-Int4',
+            dataset=['AI-ModelScope/alpaca-gpt4-data-zh#100', 'AI-ModelScope/alpaca-gpt4-data-en#100'],
+            split_dataset_ratio=0.01,
+            **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=[last_model_checkpoint], load_data_args=True))
+
+
+def test_llm_awq():
+    from swift import InferArguments, SftArguments, infer_main, sft_main
+    result = sft_main(
+        SftArguments(
+            model='Qwen/Qwen2-7B-Instruct-AWQ',
+            dataset=['AI-ModelScope/alpaca-gpt4-data-zh#100', 'AI-ModelScope/alpaca-gpt4-data-en#100'],
+            split_dataset_ratio=0.01,
+            **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=[last_model_checkpoint], load_data_args=True))
+
+
+def test_mllm_streaming_zero3():
+    os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
+    os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0,1'
+    from swift import SftArguments, sft_main
+    sft_main(
+        SftArguments(
+            model='Qwen/Qwen2-VL-7B-Instruct',
+            dataset=['modelscope/coco_2014_caption:validation', 'AI-ModelScope/alpaca-gpt4-data-en'],
+            streaming=True,
+            max_steps=16,
+            deepspeed='zero3',
+            **kwargs))
+
+
+def test_mllm_streaming_mp_ddp():
+    os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2,3'
+    os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0,1,2,3'
+    from swift import SftArguments, sft_main
+    sft_main(
+        SftArguments(
+            model='Qwen/Qwen2-VL-7B-Instruct',
+            dataset=['modelscope/coco_2014_caption:validation', 'AI-ModelScope/alpaca-gpt4-data-en'],
+            streaming=True,
+            max_steps=16,
+            gradient_checkpointing_kwargs={'use_reentrant': False},
+            **kwargs))
+
+
+def test_llm_hqq():
+    from swift import InferArguments, SftArguments, infer_main, sft_main
+    result = sft_main(
+        SftArguments(
+            model='Qwen/Qwen2-7B-Instruct',
+            dataset=['AI-ModelScope/alpaca-gpt4-data-zh#100', 'AI-ModelScope/alpaca-gpt4-data-en#100'],
+            split_dataset_ratio=0.01,
+            quant_method='hqq',
+            quant_bits=4,
+            **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=[last_model_checkpoint], load_data_args=True))
+
+
+def test_llm_bnb():
+    from swift import InferArguments, SftArguments, infer_main, sft_main
+    result = sft_main(
+        SftArguments(
+            model='Qwen/Qwen2-7B-Instruct',
+            dataset=['AI-ModelScope/alpaca-gpt4-data-zh#100', 'AI-ModelScope/alpaca-gpt4-data-en#100'],
+            split_dataset_ratio=0.01,
+            quant_method='bnb',
+            quant_bits=4,
+            **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=[last_model_checkpoint], load_data_args=True))
+
+
+def test_moe():
+    from swift import InferArguments, SftArguments, infer_main, sft_main
+    result = sft_main(
+        SftArguments(
+            model='Qwen/Qwen1.5-MoE-A2.7B-Chat-GPTQ-Int4',
+            dataset=['AI-ModelScope/alpaca-gpt4-data-zh#100', 'AI-ModelScope/alpaca-gpt4-data-en#100'],
+            split_dataset_ratio=0.01,
+            **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=last_model_checkpoint, load_data_args=True))
+
+
+def test_resume_from_checkpoint():
+    from swift import InferArguments, SftArguments, infer_main, sft_main
+    result = sft_main(
+        SftArguments(
+            model='Qwen/Qwen2-0.5B',
+            dataset=['AI-ModelScope/alpaca-gpt4-data-zh#100', 'AI-ModelScope/alpaca-gpt4-data-en#100'],
+            max_steps=5,
+            streaming=True,
+            **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    result = sft_main(
+        SftArguments(
+            model='Qwen/Qwen2-0.5B',
+            resume_from_checkpoint=last_model_checkpoint,
+            dataset=['AI-ModelScope/alpaca-gpt4-data-zh#100', 'AI-ModelScope/alpaca-gpt4-data-en#100'],
+            streaming=True,
+            load_data_args=True,
+            max_steps=10,
+            **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=last_model_checkpoint, load_data_args=True))
+
+
+def test_resume_only_model():
+    from swift import SftArguments, sft_main
+    result = sft_main(
+        SftArguments(
+            model='Qwen/Qwen2-0.5B',
+            dataset=['AI-ModelScope/alpaca-gpt4-data-zh#100', 'AI-ModelScope/alpaca-gpt4-data-en#100'],
+            max_steps=5,
+            save_only_model=True,
+            deepspeed='zero3',
+            **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    result = sft_main(
+        SftArguments(
+            model='Qwen/Qwen2-0.5B',
+            resume_from_checkpoint=last_model_checkpoint,
+            dataset=['AI-ModelScope/alpaca-gpt4-data-zh#100', 'AI-ModelScope/alpaca-gpt4-data-en#100'],
+            resume_only_model=True,
+            save_only_model=True,
+            load_data_args=True,
+            max_steps=10,
+            deepspeed='zero3',
+            **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    print(f'last_model_checkpoint: {last_model_checkpoint}')
+
+
+def test_llm_transformers_4_33():
+    os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
+    os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0,1'
+    from swift import SftArguments, sft_main
+    sft_main(
+        SftArguments(
+            model='Qwen/Qwen-7B-Chat',
+            dataset=['AI-ModelScope/alpaca-gpt4-data-zh#100', 'AI-ModelScope/alpaca-gpt4-data-en#100'],
+            split_dataset_ratio=0.01,
+            **kwargs))
+
+
+def test_predict_with_generate():
+    import os
+    os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
+    os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0,1'
+    from swift import SftArguments, sft_main
+
+    # 'modelscope/coco_2014_caption:validation#100',
+    sft_main(
+        SftArguments(
+            model='Qwen/Qwen2-7B-Instruct',
+            dataset=['AI-ModelScope/alpaca-gpt4-data-en#400'],
+            predict_with_generate=True,
+            # padding_free=True,
+            max_length=512,
+            packing=True,
+            attn_impl='flash_attn',
+            split_dataset_ratio=0.01,
+            **kwargs))
+
+
+def test_predict_with_generate_zero3():
+    import os
+    os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
+    os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0,1'
+    from swift import SftArguments, sft_main
+
+    # 'modelscope/coco_2014_caption:validation#100',
+    sft_main(
+        SftArguments(
+            model='Qwen/Qwen2-VL-7B-Instruct',
+            dataset=['AI-ModelScope/LaTeX_OCR#40'],
+            split_dataset_ratio=0.01,
+            predict_with_generate=True,
+            freeze_vit=False,
+            deepspeed='zero3',
+            **kwargs))
+
+
+def test_template():
+    from swift import InferArguments, SftArguments, infer_main, sft_main
+    global kwargs
+    kwargs = kwargs.copy()
+    kwargs['num_train_epochs'] = 3
+    result = sft_main(
+        SftArguments(
+            model='Qwen/Qwen2-0.5B',
+            dataset=['swift/self-cognition#200'],
+            split_dataset_ratio=0.01,
+            model_name=['小黄'],
+            model_author=['swift'],
+            **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=[last_model_checkpoint], load_data_args=True, merge_lora=True))
+
+
+def test_emu3_gen():
+    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+    os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0'
+    os.environ['max_position_embeddings'] = '10240'
+    os.environ['image_area'] = '518400'
+    from swift import InferArguments, SftArguments, infer_main, sft_main
+    kwargs['num_train_epochs'] = 100
+    result = sft_main(
+        SftArguments(model='BAAI/Emu3-Gen', dataset=['swift/TextCaps#2'], split_dataset_ratio=0.01, **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    args = InferArguments(
+        adapters=[last_model_checkpoint],
+        infer_backend='transformers',
+        stream=False,
+        use_chat_template=False,
+        top_k=2048,
+        max_new_tokens=40960)
+    infer_main(args)
+
+
+def test_eval_strategy():
+    os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
+    os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0,1'
+    from swift import InferArguments, SftArguments, infer_main, sft_main
+    result = sft_main(
+        SftArguments(
+            model='Qwen/Qwen2-7B-Instruct',
+            eval_strategy='no',
+            dataset=['AI-ModelScope/alpaca-gpt4-data-zh#100', 'AI-ModelScope/alpaca-gpt4-data-en#100'],
+            split_dataset_ratio=0.01,
+            **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=last_model_checkpoint, load_data_args=True))
+
+
+def test_epoch():
+    os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
+    os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0,1'
+    from swift import InferArguments, SftArguments, infer_main, sft_main
+
+    train_kwargs = kwargs.copy()
+    train_kwargs['num_train_epochs'] = 3
+    # train_kwargs['save_steps'] = 2  # not use
+    result = sft_main(
+        SftArguments(
+            model='Qwen/Qwen2-7B-Instruct',
+            dataset=['AI-ModelScope/alpaca-gpt4-data-zh#50', 'AI-ModelScope/alpaca-gpt4-data-en#50'],
+            split_dataset_ratio=0.01,
+            save_strategy='epoch',
+            **train_kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=last_model_checkpoint, load_data_args=True))
+
+
+def test_agent():
+    os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
+    os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0,1'
+    from swift import InferArguments, SftArguments, infer_main, sft_main
+
+    result = sft_main(
+        SftArguments(
+            model='Qwen/Qwen2-7B-Instruct',
+            dataset=['swift/ToolBench#500'],
+            split_dataset_ratio=0.01,
+            loss_scale='react',
+            agent_template='toolbench',
+            **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=last_model_checkpoint, load_data_args=True))
+
+
+def test_grounding():
+    os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
+    os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0,1'
+    from swift import InferArguments, SftArguments, infer_main, sft_main
+
+    result = sft_main(
+        SftArguments(
+            model='Qwen/Qwen2.5-VL-7B-Instruct',
+            dataset=['AI-ModelScope/coco#200'],
+            split_dataset_ratio=0.01,
+            dataset_num_proc=4,
+            **kwargs))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=last_model_checkpoint, load_data_args=True, stream=True, max_new_tokens=2048))
+
+
+def test_lora_sft_minimal():
+    from swift import InferArguments, SftArguments, infer_main, sft_main
+    result = sft_main(
+        SftArguments(
+            model='Qwen/Qwen2-0.5B',
+            dataset=['AI-ModelScope/alpaca-gpt4-data-zh#20'],
+            max_steps=2,
+            per_device_train_batch_size=1,
+            gradient_accumulation_steps=1,
+            save_steps=2,
+            split_dataset_ratio=0.01,
+            tuner_type='lora',
+            logging_steps=1,
+            **{
+                k: v
+                for k, v in kwargs.items() if k not in [
+                    'per_device_train_batch_size', 'save_steps', 'gradient_accumulation_steps', 'num_train_epochs',
+                    'per_device_eval_batch_size'
+                ]
+            }))
+    last_model_checkpoint = result['last_model_checkpoint']
+    infer_main(InferArguments(adapters=last_model_checkpoint, load_data_args=True))
+
+
+def test_full_sft_minimal():
+    from swift import SftArguments, sft_main
+    result = sft_main(
+        SftArguments(
+            model='Qwen/Qwen2-0.5B',
+            dataset=['AI-ModelScope/alpaca-gpt4-data-zh#20'],
+            max_steps=1,
+            per_device_train_batch_size=1,
+            gradient_accumulation_steps=1,
+            save_steps=1,
+            split_dataset_ratio=0.01,
+            tuner_type='full',
+            logging_steps=1,
+            **{
+                k: v
+                for k, v in kwargs.items() if k not in [
+                    'per_device_train_batch_size', 'save_steps', 'gradient_accumulation_steps', 'num_train_epochs',
+                    'per_device_eval_batch_size'
+                ]
+            }))
+    assert os.path.isdir(result['last_model_checkpoint'])
+
+
+if __name__ == '__main__':
+    # test_llm_ddp()
+    # test_mllm_mp()
+    # test_llm_streaming()
+    # test_mllm_streaming()
+    # test_mllm_zero3()
+    # test_llm_gptq()
+    # test_llm_awq()
+    # test_mllm_streaming_zero3()
+    # test_mllm_streaming_mp_ddp()
+    # test_llm_bnb()
+    # test_llm_hqq()
+    # test_moe()
+    # test_resume_from_checkpoint()
+    test_resume_only_model()
+    # test_llm_transformers_4_33()
+    # test_predict_with_generate()
+    # test_predict_with_generate_zero3()
+    # test_template()
+    # test_qwen_vl()
+    # test_qwen2_audio()
+    # test_emu3_gen()
+    # test_unsloth()
+    # test_eval_strategy()
+    # test_epoch()
+    # test_agent()
+    # test_grounding()
+    # test_lora_sft_minimal()
+    # test_full_sft_minimal()

tests/train/test_train_eval.py

@@ -0,0 +1,34 @@
+import os
+
+kwargs = {
+    'per_device_train_batch_size': 5,
+    'save_steps': 5,
+    'gradient_accumulation_steps': 1,
+    'num_train_epochs': 1,
+}
+
+
+def test_train_eval_loop():
+    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+    os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0'
+    from swift import SftArguments, sft_main
+    sft_main(
+        SftArguments(
+            model='Qwen/Qwen2.5-0.5B-Instruct',
+            dataset=['AI-ModelScope/alpaca-gpt4-data-zh#100'],
+            target_modules=['all-linear', 'all-embedding'],
+            modules_to_save=['all-embedding', 'all-norm'],
+            eval_strategy='steps',
+            eval_steps=5,
+            per_device_eval_batch_size=5,
+            eval_use_evalscope=True,
+            eval_dataset=['gsm8k'],
+            eval_dataset_args={'gsm8k': {
+                'few_shot_num': 0
+            }},
+            eval_limit=10,
+            **kwargs))
+
+
+if __name__ == '__main__':
+    test_train_eval_loop()

tests/train/test_vit_lr.py

@@ -0,0 +1,24 @@
+import os
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0'
+
+
+def test_vit_lr():
+    # https://github.com/QwenLM/Qwen2.5-VL/tree/main/qwen-vl-finetune
+    from swift import SftArguments, sft_main
+    sft_main(
+        SftArguments(
+            model='Qwen/Qwen2.5-VL-7B-Instruct',
+            dataset=['AI-ModelScope/LaTeX_OCR#20000'],
+            split_dataset_ratio=0.01,
+            vit_lr=2e-5,
+            learning_rate=1e-5,
+            aligner_lr=1e-4,
+            freeze_llm=False,
+            freeze_vit=False,
+            freeze_aligner=False))
+
+
+if __name__ == '__main__':
+    test_vit_lr()

tests/train/test_vllm_importance_sampling_basic.py

@@ -0,0 +1,485 @@
+"""
+Basic tests for vLLM Importance Sampling implementation
+
+This test file verifies the core functionality of the vLLM IS correction,
+including the IS weight computation and metrics calculation.
+
+Reference: verl/verl/trainer/ppo/rollout_corr_helper.py
+"""
+
+import torch
+
+
+class MockAccelerator:
+    """Mock accelerator for testing metrics gathering"""
+
+    def __init__(self, device='cpu'):
+        self.device = device
+
+    def gather_for_metrics(self, tensor):
+        # In testing, just return the tensor as-is
+        return tensor
+
+
+class MockGRPOTrainer:
+    """Mock GRPO trainer for testing IS methods"""
+
+    def __init__(self, mode='token_truncate', threshold=2.0):
+        self.rollout_importance_sampling_mode = mode
+        self.rollout_importance_sampling_threshold = threshold
+        self.accelerator = MockAccelerator()
+
+    def _compute_sequence_level_ratios(self, is_ratio: torch.Tensor, completion_mask: torch.Tensor) -> torch.Tensor:
+        """
+        Helper function to compute sequence-level importance sampling ratios.
+
+        Args:
+            is_ratio: Token-level IS ratios, shape [B, T]
+            completion_mask: Boolean mask for completion tokens, shape [B, T]
+
+        Returns:
+            Sequence-level ratios as geometric mean of token-level ratios
+        """
+        log_ratio = torch.log(is_ratio.clamp(min=1e-10))
+        seq_log_ratios = (log_ratio * completion_mask).sum(-1) / completion_mask.sum(-1).clamp(min=1.0)
+        seq_ratios = torch.exp(seq_log_ratios)
+
+        return seq_ratios
+
+    def _apply_rollout_importance_sampling(self, rollout_log_ratio: torch.Tensor,
+                                           completion_mask: torch.Tensor) -> torch.Tensor:
+        """
+        Apply vLLM importance sampling correction using one of four modes.
+
+        Args:
+            rollout_log_ratio: log(π_θ / π_rollout) per token, shape [B, T]
+            completion_mask: Boolean mask for completion tokens, shape [B, T]
+
+        Returns:
+            IS weights to multiply with loss, same shape as rollout_log_ratio
+        """
+        mode = self.rollout_importance_sampling_mode
+        threshold = self.rollout_importance_sampling_threshold
+
+        # Clamp log_ratio to prevent numerical overflow from padding values (-1e10)
+        # A log_ratio of 20 corresponds to exp(20) ≈ 485 million, which is already extreme
+        SAFETY_BOUND = 20.0
+        rollout_log_ratio_safe = torch.clamp(rollout_log_ratio, min=-SAFETY_BOUND, max=SAFETY_BOUND)
+
+        # Compute importance sampling ratios: exp(log_ratio)
+        is_ratio = torch.exp(rollout_log_ratio_safe)
+
+        if mode == 'token_truncate':
+            # Token-level truncated IS: clip ratios from above at threshold
+            is_weights = torch.clamp(is_ratio, max=threshold)
+
+        elif mode == 'token_mask':
+            # Token-level masked IS: mask out tokens with ratio > threshold
+            is_weights = torch.where(is_ratio <= threshold, is_ratio, torch.zeros_like(is_ratio))
+
+        elif mode == 'sequence_truncate':
+            # Sequence-level truncated IS: compute sequence-level ratio and clip
+            seq_ratios = self._compute_sequence_level_ratios(is_ratio, completion_mask)
+            clipped_seq_ratios = torch.clamp(seq_ratios, max=threshold)
+
+            is_weights = clipped_seq_ratios.unsqueeze(-1).expand_as(is_ratio)
+
+        elif mode == 'sequence_mask':
+            # Sequence-level masked IS: mask entire sequences with ratio > threshold
+            seq_ratios = self._compute_sequence_level_ratios(is_ratio, completion_mask)
+            seq_mask = (seq_ratios <= threshold).float()
+
+            # Apply mask to original token-level ratios
+            is_weights = is_ratio * seq_mask.unsqueeze(-1)
+        else:
+            return is_ratio
+
+        return is_weights
+
+    def _compute_is_correction_metrics(
+        self,
+        vllm_log_ratio: torch.Tensor,
+        is_weights: torch.Tensor,
+        completion_mask: torch.Tensor,
+    ) -> dict:
+        """
+        Compute importance sampling correction metrics (ess, clipped_frac, is_weight_mean).
+        Only called when rollout_importance_sampling_mode is enabled.
+
+        Args:
+            vllm_log_ratio: Log ratio log(π_policy / π_rollout), shape [B, T]
+            is_weights: Importance sampling weights after correction, shape [B, T]
+            completion_mask: Boolean mask for completion tokens, shape [B, T]
+
+        Returns:
+            Dictionary with IS-specific metrics:
+                - is_weight_mean: Mean of IS weights
+                - ess: Effective Sample Size = 1 / E[(w_i / E[w_i])²]
+                - clipped_frac: Fraction of clipped/masked samples
+        """
+        metrics = {}
+        SAFETY_BOUND = 20.0
+        threshold = self.rollout_importance_sampling_threshold
+        threshold_lower = 1.0 / threshold  # Default lower threshold (reciprocal of upper)
+
+        # Helper function for masked mean
+        def masked_mean(x, mask):
+            return (x * mask).sum() / mask.sum().clamp(min=1.0)
+
+        # Compute IS ratio with safety bounds
+        log_ratio_safe = torch.clamp(vllm_log_ratio, min=-SAFETY_BOUND, max=SAFETY_BOUND)
+        is_ratio = torch.exp(log_ratio_safe)
+
+        # 1. IS weight statistics
+        mean_is_weight = masked_mean(is_weights, completion_mask)
+        metrics['is_weight_mean'] = self.accelerator.gather_for_metrics(mean_is_weight).nanmean().item()
+
+        # 2. Compute Effective Sample Size (ESS) for IS weights
+        # ESS = 1 / E[(w_i / E[w_i])²] (using clamped weights for stability)
+        # This measures how many "effective" independent samples we have after IS weighting
+        weights_for_ess = is_weights.clamp(min=threshold_lower, max=threshold)
+        mean_for_ess = masked_mean(weights_for_ess, completion_mask)
+        is_weights_normalized = weights_for_ess / (mean_for_ess + 1e-8)  # Avoid division by zero
+        ess = 1.0 / masked_mean(is_weights_normalized.square(), completion_mask).clamp(min=1e-10)
+        metrics['ess'] = self.accelerator.gather_for_metrics(ess).nanmean().item()
+
+        # 3. Fraction of clipped/masked samples
+        if self.rollout_importance_sampling_mode in ['token_truncate', 'token_mask']:
+            # Token-level
+            if self.rollout_importance_sampling_mode == 'token_truncate':
+                clipped_frac = masked_mean((is_ratio > threshold).float(), completion_mask)
+            else:  # token_mask
+                clipped_frac = masked_mean((is_weights == 0).float(), completion_mask)
+            metrics['clipped_frac'] = self.accelerator.gather_for_metrics(clipped_frac).nanmean().item()
+        else:
+            # Sequence-level (both truncate and mask)
+            seq_ratios = self._compute_sequence_level_ratios(is_ratio, completion_mask)
+            clipped_frac = (seq_ratios > threshold).float().mean()
+            metrics['clipped_frac'] = self.accelerator.gather_for_metrics(clipped_frac).nanmean().item()
+
+        return metrics
+
+
+class TestVLLMImportanceSampling:
+    """Test suite for vLLM Importance Sampling"""
+
+    def test_token_truncate_basic(self):
+        """Test token-level truncated IS"""
+        trainer = MockGRPOTrainer(mode='token_truncate', threshold=2.0)
+
+        # Create mock data: [batch=2, seq_len=4]
+        # Log ratios that will produce ratios [0.5, 1.5, 3.0, 5.0]
+        vllm_log_ratio = torch.log(torch.tensor([[0.5, 1.5, 3.0, 5.0], [0.8, 1.2, 2.5, 4.0]]))
+        completion_mask = torch.ones_like(vllm_log_ratio)
+
+        is_weights = trainer._apply_rollout_importance_sampling(vllm_log_ratio, completion_mask)
+
+        # Check truncation at threshold=2.0
+        assert is_weights.shape == vllm_log_ratio.shape
+        assert torch.allclose(is_weights[0, 0], torch.tensor(0.5), atol=1e-5)
+        assert torch.allclose(is_weights[0, 1], torch.tensor(1.5), atol=1e-5)
+        assert torch.allclose(is_weights[0, 2], torch.tensor(2.0), atol=1e-5)  # Truncated
+        assert torch.allclose(is_weights[0, 3], torch.tensor(2.0), atol=1e-5)  # Truncated
+
+    def test_token_mask_basic(self):
+        """Test token-level masked IS"""
+        trainer = MockGRPOTrainer(mode='token_mask', threshold=2.0)
+
+        vllm_log_ratio = torch.log(torch.tensor([[0.5, 1.5, 3.0, 5.0]]))
+        completion_mask = torch.ones_like(vllm_log_ratio)
+
+        is_weights = trainer._apply_rollout_importance_sampling(vllm_log_ratio, completion_mask)
+
+        # Check masking: ratio > threshold should be 0
+        assert torch.allclose(is_weights[0, 0], torch.tensor(0.5), atol=1e-5)
+        assert torch.allclose(is_weights[0, 1], torch.tensor(1.5), atol=1e-5)
+        assert torch.allclose(is_weights[0, 2], torch.tensor(0.0), atol=1e-5)  # Masked
+        assert torch.allclose(is_weights[0, 3], torch.tensor(0.0), atol=1e-5)  # Masked
+
+    def test_sequence_truncate_basic(self):
+        """Test sequence-level truncated IS"""
+        trainer = MockGRPOTrainer(mode='sequence_truncate', threshold=2.0)
+
+        # First sequence has high ratios, second has low ratios
+        vllm_log_ratio = torch.log(
+            torch.tensor([
+                [3.0, 3.0, 3.0, 3.0],  # geometric mean=3.0 > 2.0
+                [1.0, 1.0, 1.0, 1.0]
+            ]))  # geometric mean=1.0 < 2.0
+        completion_mask = torch.ones_like(vllm_log_ratio)
+
+        is_weights = trainer._apply_rollout_importance_sampling(vllm_log_ratio, completion_mask)
+
+        # First sequence should be truncated to 2.0 for all tokens
+        assert torch.allclose(is_weights[0, :], torch.tensor(2.0), atol=1e-5)
+        # Second sequence should remain 1.0
+        assert torch.allclose(is_weights[1, :], torch.tensor(1.0), atol=1e-5)
+
+    def test_sequence_mask_basic(self):
+        """Test sequence-level masked IS"""
+        trainer = MockGRPOTrainer(mode='sequence_mask', threshold=2.0)
+
+        vllm_log_ratio = torch.log(
+            torch.tensor([
+                [3.0, 3.0, 3.0, 3.0],  # geometric mean=3.0 > 2.0
+                [1.0, 1.0, 1.0, 1.0]
+            ]))  # geometric mean=1.0 < 2.0
+        completion_mask = torch.ones_like(vllm_log_ratio)
+
+        is_weights = trainer._apply_rollout_importance_sampling(vllm_log_ratio, completion_mask)
+
+        # First sequence should be completely masked (0)
+        # Note: sequence_mask multiplies is_ratio by 0, so all tokens become 0
+        assert torch.allclose(is_weights[0, :], torch.tensor(0.0), atol=1e-5)
+        # Second sequence should keep original ratios (1.0 * 1.0 = 1.0)
+        assert torch.allclose(is_weights[1, :], torch.tensor(1.0), atol=1e-5)
+
+    def test_threshold_sensitivity(self):
+        """Test different threshold values"""
+        vllm_log_ratio = torch.log(torch.tensor([[1.0, 2.0, 3.0, 4.0]]))
+        completion_mask = torch.ones_like(vllm_log_ratio)
+
+        # Test threshold=1.5
+        trainer_low = MockGRPOTrainer(mode='token_truncate', threshold=1.5)
+        is_weights_low = trainer_low._apply_rollout_importance_sampling(vllm_log_ratio, completion_mask)
+
+        # Test threshold=3.5
+        trainer_high = MockGRPOTrainer(mode='token_truncate', threshold=3.5)
+        is_weights_high = trainer_high._apply_rollout_importance_sampling(vllm_log_ratio, completion_mask)
+
+        # Lower threshold should truncate more
+        truncated_low = (is_weights_low < torch.exp(vllm_log_ratio)).sum()
+        truncated_high = (is_weights_high < torch.exp(vllm_log_ratio)).sum()
+        assert truncated_low > truncated_high
+
+    def test_completion_mask(self):
+        """Test that completion mask is respected"""
+        trainer = MockGRPOTrainer(mode='sequence_truncate', threshold=2.0)
+
+        vllm_log_ratio = torch.log(torch.tensor([[3.0, 3.0, 3.0, 3.0]]))
+        # Mask out last two tokens
+        completion_mask = torch.tensor([[1.0, 1.0, 0.0, 0.0]])
+
+        is_weights = trainer._apply_rollout_importance_sampling(vllm_log_ratio, completion_mask)
+
+        # Should only consider masked tokens for sequence ratio calculation
+        # With only first two tokens (both 3.0), geometric mean=3.0, truncated to 2.0
+        assert torch.allclose(is_weights[0, :2], torch.tensor(2.0), atol=1e-5)
+
+    def test_edge_cases(self):
+        """Test edge cases"""
+        trainer = MockGRPOTrainer(mode='token_truncate', threshold=2.0)
+
+        # Case 1: All ratios below threshold
+        vllm_log_ratio = torch.log(torch.tensor([[0.5, 1.0, 1.5]]))
+        completion_mask = torch.ones_like(vllm_log_ratio)
+        is_weights = trainer._apply_rollout_importance_sampling(vllm_log_ratio, completion_mask)
+        assert torch.allclose(is_weights, torch.exp(vllm_log_ratio), atol=1e-5)
+
+        # Case 2: All ratios above threshold
+        vllm_log_ratio = torch.log(torch.tensor([[3.0, 4.0, 5.0]]))
+        is_weights = trainer._apply_rollout_importance_sampling(vllm_log_ratio, completion_mask[:, :3])
+        assert torch.allclose(is_weights, torch.tensor(2.0), atol=1e-5)
+
+        # Case 3: Empty mask
+        vllm_log_ratio = torch.log(torch.tensor([[1.0, 2.0, 3.0]]))
+        completion_mask = torch.zeros_like(vllm_log_ratio)
+        is_weights = trainer._apply_rollout_importance_sampling(vllm_log_ratio, completion_mask)
+        # Should still compute but result may not be meaningful
+        assert is_weights.shape == vllm_log_ratio.shape
+
+    def test_safety_bound(self):
+        """Test that extreme log ratios are clamped"""
+        trainer = MockGRPOTrainer(mode='token_truncate', threshold=2.0)
+
+        # Create extreme log ratios that would overflow without clamping
+        vllm_log_ratio = torch.tensor([[100.0, -100.0, 0.0]])  # exp(100) would overflow
+        completion_mask = torch.ones_like(vllm_log_ratio)
+
+        is_weights = trainer._apply_rollout_importance_sampling(vllm_log_ratio, completion_mask)
+
+        # Should not have inf or nan
+        assert torch.isfinite(is_weights).all()
+        # Large positive log_ratio should be clamped to threshold
+        assert is_weights[0, 0] <= 2.0
+        # Large negative log_ratio should result in small positive value
+        assert is_weights[0, 1] > 0
+
+
+class TestISCorrectionMetrics:
+    """Test suite for IS correction metrics"""
+
+    def test_ess_uniform_weights(self):
+        """Test ESS with uniform weights (should be close to 1.0)"""
+        trainer = MockGRPOTrainer(mode='token_truncate', threshold=2.0)
+
+        # Uniform weights of 1.0
+        vllm_log_ratio = torch.zeros((2, 4))  # exp(0) = 1.0
+        completion_mask = torch.ones_like(vllm_log_ratio)
+        is_weights = torch.ones_like(vllm_log_ratio)
+
+        metrics = trainer._compute_is_correction_metrics(vllm_log_ratio, is_weights, completion_mask)
+
+        # ESS should be 1.0 for uniform weights
+        assert abs(metrics['ess'] - 1.0) < 0.01
+        # Mean weight should be 1.0
+        assert abs(metrics['is_weight_mean'] - 1.0) < 0.01
+        # No clipping for uniform weights
+        assert metrics['clipped_frac'] == 0.0
+
+    def test_ess_varied_weights(self):
+        """Test ESS with varied weights (should be < 1.0)"""
+        trainer = MockGRPOTrainer(mode='token_truncate', threshold=2.0)
+
+        # Varied weights
+        vllm_log_ratio = torch.log(torch.tensor([[0.5, 1.0, 1.5, 2.0]]))
+        completion_mask = torch.ones_like(vllm_log_ratio)
+        is_weights = torch.tensor([[0.5, 1.0, 1.5, 2.0]])
+
+        metrics = trainer._compute_is_correction_metrics(vllm_log_ratio, is_weights, completion_mask)
+
+        # ESS should be less than 1.0 for non-uniform weights
+        assert metrics['ess'] < 1.0
+        assert metrics['ess'] > 0.0
+
+    def test_clipped_frac_token_truncate(self):
+        """Test clipped_frac for token_truncate mode"""
+        trainer = MockGRPOTrainer(mode='token_truncate', threshold=2.0)
+
+        # 2 out of 4 tokens exceed threshold
+        vllm_log_ratio = torch.log(torch.tensor([[0.5, 1.5, 3.0, 5.0]]))
+        completion_mask = torch.ones_like(vllm_log_ratio)
+        is_weights = trainer._apply_rollout_importance_sampling(vllm_log_ratio, completion_mask)
+
+        metrics = trainer._compute_is_correction_metrics(vllm_log_ratio, is_weights, completion_mask)
+
+        # 2/4 = 0.5 tokens clipped
+        assert abs(metrics['clipped_frac'] - 0.5) < 0.01
+
+    def test_clipped_frac_token_mask(self):
+        """Test clipped_frac for token_mask mode"""
+        trainer = MockGRPOTrainer(mode='token_mask', threshold=2.0)
+
+        # 2 out of 4 tokens exceed threshold
+        vllm_log_ratio = torch.log(torch.tensor([[0.5, 1.5, 3.0, 5.0]]))
+        completion_mask = torch.ones_like(vllm_log_ratio)
+        is_weights = trainer._apply_rollout_importance_sampling(vllm_log_ratio, completion_mask)
+
+        metrics = trainer._compute_is_correction_metrics(vllm_log_ratio, is_weights, completion_mask)
+
+        # 2/4 = 0.5 tokens masked (is_weights == 0)
+        assert abs(metrics['clipped_frac'] - 0.5) < 0.01
+
+    def test_clipped_frac_sequence_level(self):
+        """Test clipped_frac for sequence-level modes"""
+        trainer = MockGRPOTrainer(mode='sequence_truncate', threshold=2.0)
+
+        # First sequence exceeds threshold, second doesn't
+        vllm_log_ratio = torch.log(torch.tensor([[3.0, 3.0, 3.0, 3.0], [1.0, 1.0, 1.0, 1.0]]))
+        completion_mask = torch.ones_like(vllm_log_ratio)
+        is_weights = trainer._apply_rollout_importance_sampling(vllm_log_ratio, completion_mask)
+
+        metrics = trainer._compute_is_correction_metrics(vllm_log_ratio, is_weights, completion_mask)
+
+        # 1/2 = 0.5 sequences clipped
+        assert abs(metrics['clipped_frac'] - 0.5) < 0.01
+
+
+class TestOffpolicyMetrics:
+    """Test suite for off-policy diagnostic metrics"""
+
+    def test_kl_divergence_same_policy(self):
+        """Test KL divergence when policies are identical"""
+        # When per_token_logps == rollout_per_token_logps, KL should be 0
+        per_token_logps = torch.tensor([[-1.0, -2.0, -1.5, -0.5]])
+        rollout_per_token_logps = per_token_logps.clone()
+        completion_mask = torch.ones_like(per_token_logps)
+
+        # Helper function for masked mean
+        def masked_mean(x, mask, axis=None):
+            if axis is None:
+                return (x * mask).sum() / mask.sum().clamp(min=1.0)
+            else:
+                return (x * mask).sum(axis) / mask.sum(axis).clamp(min=1.0)
+
+        # KL = E[log(π_rollout) - log(π_training)]
+        kl = masked_mean(rollout_per_token_logps - per_token_logps, completion_mask)
+
+        assert abs(kl.item()) < 1e-6
+
+    def test_k3_kl_estimator(self):
+        """Test K3 KL estimator"""
+        per_token_logps = torch.tensor([[-1.0, -2.0, -1.5, -0.5]])
+        rollout_per_token_logps = torch.tensor([[-1.1, -1.9, -1.6, -0.4]])
+        completion_mask = torch.ones_like(per_token_logps)
+
+        def masked_mean(x, mask, axis=None):
+            if axis is None:
+                return (x * mask).sum() / mask.sum().clamp(min=1.0)
+            else:
+                return (x * mask).sum(axis) / mask.sum(axis).clamp(min=1.0)
+
+        # K3 estimator: E[exp(log_ratio) - log_ratio - 1]
+        log_ratio = per_token_logps - rollout_per_token_logps
+        log_ratio *= completion_mask
+        k3_kl_matrix = torch.exp(log_ratio) - log_ratio - 1
+        k3_kl = masked_mean(k3_kl_matrix, completion_mask)
+
+        # K3 KL should be non-negative
+        assert k3_kl.item() >= 0
+
+    def test_chi2_divergence(self):
+        """Test χ² divergence calculation"""
+        per_token_logps = torch.tensor([[-1.0, -2.0]])
+        rollout_per_token_logps = torch.tensor([[-1.5, -1.5]])
+        completion_mask = torch.ones_like(per_token_logps)
+
+        def masked_mean(x, mask, axis=None):
+            if axis is None:
+                return (x * mask).sum() / mask.sum().clamp(min=1.0)
+            else:
+                return (x * mask).sum(axis) / mask.sum(axis).clamp(min=1.0)
+
+        SAFETY_BOUND = 20.0
+        log_ratio = per_token_logps - rollout_per_token_logps
+        log_ratio_safe = torch.clamp(log_ratio, min=-SAFETY_BOUND, max=SAFETY_BOUND)
+        rho_token = torch.exp(log_ratio_safe)
+        rho_squared_token = rho_token.square()
+        chi2_token = masked_mean(rho_squared_token, completion_mask) - 1.0
+
+        # χ² should be >= -1 (can be negative if E[ρ²] < 1)
+        assert chi2_token.item() >= -1.0
+
+
+if __name__ == '__main__':
+    # Run tests manually
+    import sys
+
+    test_classes = [
+        ('TestVLLMImportanceSampling', TestVLLMImportanceSampling),
+        ('TestISCorrectionMetrics', TestISCorrectionMetrics),
+        ('TestOffpolicyMetrics', TestOffpolicyMetrics),
+    ]
+
+    failed_tests = []
+
+    for class_name, test_class in test_classes:
+        print(f'\n=== {class_name} ===')
+        test_instance = test_class()
+
+        test_methods = [m for m in dir(test_instance) if m.startswith('test_')]
+
+        for method_name in test_methods:
+            try:
+                print(f'Running {method_name}...')
+                getattr(test_instance, method_name)()
+                print(f'✓ {method_name} passed')
+            except Exception as e:
+                print(f'✗ {method_name} failed: {e}')
+                failed_tests.append(f'{class_name}.{method_name}')
+
+    if failed_tests:
+        print(f'\nFailed tests: {failed_tests}')
+        sys.exit(1)
+    else:
+        print('\nAll tests passed!')