[bugfix] fix mulitmodal cached_dataset (#5671)

Author: Jintao-Huang

examples/export/cached_dataset/mcore.sh

@@ -1,3 +1,4 @@
+# Note: cached_dataset does not support CP temporarily.
 swift export \
     --model Qwen/Qwen3-30B-A3B-Base \
     --dataset 'swift/Chinese-Qwen3-235B-2507-Distill-data-110k-SFT' \

examples/export/cached_dataset/vlm.sh

@@ -0,0 +1,70 @@
+OMP_NUM_THREADS=14 \
+MAX_PIXELS=1003520 \
+VIDEO_MAX_PIXELS=50176 \
+FPS_MAX_FRAMES=12 \
+swift export \
+    --model Qwen/Qwen2.5-Omni-7B \
+    --dataset 'AI-ModelScope/alpaca-gpt4-data-zh#10000' \
+              'AI-ModelScope/LaTeX_OCR:human_handwrite#5000' \
+              'speech_asr/speech_asr_aishell1_trainsets:validation#5000' \
+    --max_length 4096 \
+    --split_dataset_ratio 0.01 \
+    --dataset_num_proc 16 \
+    --to_cached_dataset true \
+    --lazy_tokenize false \
+    --output_dir ./qwen2_5_omni_cached_dataset
+
+# 4 * 70GiB
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+MAX_PIXELS=1003520 \
+VIDEO_MAX_PIXELS=50176 \
+FPS_MAX_FRAMES=12 \
+NPROC_PER_NODE=4 \
+ENABLE_AUDIO_OUTPUT=0 \
+CUDA_VISIBLE_DEVICES=0,1,2,3 \
+swift sft \
+    --model Qwen/Qwen2.5-Omni-7B \
+    --train_type full \
+    --cached_dataset './qwen2_5_omni_cached_dataset' \
+    --num_train_epochs 1 \
+    --split_dataset_ratio 0.01 \
+    --torch_dtype bfloat16 \
+    --per_device_train_batch_size 1 \
+    --per_device_eval_batch_size 1 \
+    --learning_rate 1e-5 \
+    --gradient_accumulation_steps 1 \
+    --packing true \
+    --freeze_llm false \
+    --freeze_vit true \
+    --freeze_aligner true \
+    --eval_steps 200 \
+    --save_steps 200 \
+    --logging_steps 5 \
+    --max_length 4096 \
+    --warmup_ratio 0.05 \
+    --dataloader_num_workers 8 \
+    --dataset_num_proc 8 \
+    --save_total_limit 2 \
+    --save_only_model true \
+    --output_dir output/Qwen2.5-Omni-7B \
+    --deepspeed zero2 \
+    --use_liger_kernel true \
+    --attn_impl flash_attn
+
+# Use the validation set
+CUDA_VISIBLE_DEVICES=0 \
+MAX_PIXELS=1003520 \
+VIDEO_MAX_PIXELS=50176 \
+FPS_MAX_FRAMES=12 \
+ENABLE_AUDIO_OUTPUT=0 \
+swift infer \
+    --model output/Qwen2.5-Omni-7B/vx-xxx/checkpoint-xxx \
+    --dataset 'AI-ModelScope/alpaca-gpt4-data-zh#10000' \
+              'AI-ModelScope/LaTeX_OCR:human_handwrite#5000' \
+              'speech_asr/speech_asr_aishell1_trainsets:validation#5000' \
+    --max_length 4096 \
+    --split_dataset_ratio 0.01 \
+    --attn_impl flash_attn \
+    --stream true \
+    --temperature 0 \
+    --max_new_tokens 512

examples/train/multimodal/omni/sft.sh

@@ -2,13 +2,13 @@
 # A demo for four modalities that can be run directly
 nproc_per_node=4
 
+# If using zero3, please set `ENABLE_AUDIO_OUTPUT=0`.
 CUDA_VISIBLE_DEVICES=0,1,2,3 \
 ENABLE_AUDIO_OUTPUT=1 \
 NPROC_PER_NODE=$nproc_per_node \
 VIDEO_MAX_PIXELS=50176 \
 FPS_MAX_FRAMES=12 \
 MAX_PIXELS=1003520 \
-ENABLE_AUDIO_OUTPUT=0 \
 swift sft \
     --model Qwen/Qwen2.5-Omni-7B \
     --dataset 'AI-ModelScope/alpaca-gpt4-data-zh#2000' \

swift/llm/argument/base_args/data_args.py

@@ -27,10 +27,8 @@ class DataArguments:
         custom_dataset_info (Optional[str]): Path to custom dataset_info.json file. Default is None.
     """
     # dataset_id or dataset_dir or dataset_path
-    dataset: List[str] = field(
-        default_factory=list, metadata={'help': f'dataset choices: {list(DATASET_MAPPING.keys())}'})
-    val_dataset: List[str] = field(
-        default_factory=list, metadata={'help': f'dataset choices: {list(DATASET_MAPPING.keys())}'})
+    dataset: List[str] = field(default_factory=list)
+    val_dataset: List[str] = field(default_factory=list)
     split_dataset_ratio: float = 0.
 
     data_seed: int = 42

swift/llm/argument/export_args.py

@@ -119,7 +119,8 @@ def __post_init__(self):
         if self.quant_method in {'gptq', 'awq'} and len(self.dataset) == 0:
             raise ValueError(f'self.dataset: {self.dataset}, Please input the quant dataset.')
         if self.to_cached_dataset:
+            self.lazy_tokenize = False
             if self.packing:
                 raise ValueError('Packing will be handled during training; here we only perform tokenization '
                                  'in advance, so you do not need to set up packing separately.')
-            assert not self.streaming and not self.lazy_tokenize, 'not supported'
+            assert not self.streaming, 'not supported'

swift/llm/export/cached_dataset.py

@@ -2,7 +2,9 @@
 import os
 from typing import List, Optional, Union
 
-from swift.llm import ExportArguments
+import torch
+
+from swift.llm import TEMPLATE_MAPPING, ExportArguments
 from swift.llm.train import SwiftSft
 from swift.utils import get_logger
 
@@ -16,10 +18,14 @@ class ExportCachedDataset(SwiftSft):
     def __init__(self, args: Optional[Union[List[str], ExportArguments]] = None) -> None:
         super(SwiftSft, self).__init__(args)
         self.train_msg = {}  # dummy
-        self.processor = None
+        template_cls = TEMPLATE_MAPPING[args.template].template_cls
+        if template_cls and template_cls.use_model:
+            kwargs = {'return_dummy_model': True}
+        else:
+            kwargs = {'load_model': False}
+        with torch.device('meta'):
+            self._prepare_model_tokenizer(**kwargs)
         self._prepare_template()
-        self._prepare_model_tokenizer(load_model=self.template.use_model)
-        self.template.init_processor(self.processor)
 
     def main(self):
         train_dataset, val_dataset = self._get_dataset()

swift/llm/train/sft.py

@@ -42,12 +42,12 @@ def _prepare_generation_config(self):
                                                                  args.get_request_config(), self.tokenizer)
         logger.info(f'model.generation_config: {self.model.generation_config}')
 
-    def _prepare_model_tokenizer(self, load_model=True):
+    def _prepare_model_tokenizer(self, **kwargs):
         args = self.args
         if args.sequence_parallel_size > 1:
             from swift.trainers.sequence_parallel import sequence_parallel
             sequence_parallel.init_sequence_parallel(args.sequence_parallel_size)
-        self.model, self.processor = args.get_model_processor(load_model=load_model)
+        self.model, self.processor = args.get_model_processor(**kwargs)
         if self.model is None:
             return
         if hasattr(self.model, 'hf_device_map'):