# Copyright (c) ModelScope Contributors. All rights reserved. import os from datasets import Dataset as HfDataset from typing import List, Optional, Union from swift.arguments import SftArguments from swift.dataset import (AddLengthPreprocessor, DatasetLoader, EncodePreprocessor, IterablePackingDataset, LazyLLMDataset, PackingDataset) from swift.infer_engine import prepare_generation_config from swift.ray_utils import RayHelper from swift.sequence_parallel import sequence_parallel from swift.trainers import TrainerFactory from swift.utils import append_to_jsonl, get_logger, get_model_parameter_info, is_master, plot_images, stat_array from ..base import SwiftPipeline from ..utils import get_cached_dataset from .tuner import TunerMixin logger = get_logger() @RayHelper.worker(group=['default']) class SwiftSft(SwiftPipeline, TunerMixin): args_class = SftArguments args: args_class def __init__(self, args: Optional[Union[List[str], SftArguments]] = None) -> None: super().__init__(args) self.train_msg = {} self._prepare_model_tokenizer() self._prepare_template() self._prepare_flash_ckpt() @RayHelper.function(group='default') def _prepare_flash_ckpt(self): if self.args.use_flash_ckpt: try: import dlrover.trainer.torch.flash_checkpoint.hf_trainer except ImportError: raise ValueError('Please install dlrover to use flash ckpt `pip install dlrover[k8s,torch]') def _prepare_generation_config(self): args = self.args self.model.origin_generation_config = self.model.generation_config self.model.generation_config = prepare_generation_config(self.model.generation_config, args.get_request_config(), self.tokenizer) logger.info(f'model.generation_config: {self.model.generation_config}') @RayHelper.function(group='default') def _prepare_model_tokenizer(self, **kwargs): args = self.args self.model, self.processor = args.get_model_processor(**kwargs) if args.sequence_parallel_size > 1: sequence_parallel.prepare( args.sequence_parallel_size, model=self.model, tokenizer=self.processor, padding_free=args.padding_free) if self.model is None: return if hasattr(self.model, 'hf_device_map'): logger.info(f'model.hf_device_map: {self.model.hf_device_map}') logger.info(f'model_info: {self.model.model_info}') self._prepare_generation_config() @RayHelper.function(group='default') def _prepare_template(self) -> None: args = self.args template = args.get_template(self.processor) template.set_mode('train') if template.use_model: template.model = self.model support_padding_free = template.support_padding_free if support_padding_free is None: support_padding_free = not args.model_meta.is_multimodal if (args.padding_free or args.packing) and not support_padding_free: raise ValueError(f'Template `{args.template}` does not support padding free or packing.') self.template = template def _get_dataset(self): # The random shuffling of the training set occurs in the dataloader of the trainer. args = self.args train_dataset, val_dataset = args.load_dataset() if args.truncation_strategy == 'split': logger.info(f'train_dataset: {train_dataset}') logger.info(f'val_dataset: {val_dataset}') return train_dataset, val_dataset def _save_val_dataset(self, val_dataset): args = self.args output_dir = getattr(args, 'output_dir', None) or getattr(args, 'save') if is_master() and isinstance(val_dataset, HfDataset) and not args.val_dataset: os.makedirs(output_dir, exist_ok=True) val_dataset_path = os.path.join(output_dir, 'val_dataset.jsonl') append_to_jsonl(val_dataset_path, val_dataset.to_list()) logger.info(f'The split dataset from the training set will be saved at: `{val_dataset_path}`.') @RayHelper.function(group='default') def _prepare_dataset(self): args = self.args # Defer encoding to the training phase pre_process = not (hasattr(args, 'rlhf_type') and args.rlhf_type in ['grpo', 'gkd']) if args.cached_dataset or args.cached_val_dataset: assert not args.streaming, 'Cached dataset does not support streaming.' train_datasets, val_datasets = get_cached_dataset(self.args) else: train_datasets, val_datasets = [], [] if args.dataset or args.val_dataset: train_dataset, val_dataset = self._get_dataset() train_dataset, val_dataset = self._encode_dataset(train_dataset, val_dataset, pre_process=pre_process) if train_dataset is not None: train_datasets.append(train_dataset) if val_dataset is not None: val_datasets.append(val_dataset) train_dataset = DatasetLoader.concat_datasets(train_datasets) val_dataset = DatasetLoader.concat_datasets(val_datasets) if args.truncation_strategy != 'split': logger.info(f'train_dataset: {train_dataset}') logger.info(f'val_dataset: {val_dataset}') datasets = [train_dataset, val_dataset] if not pre_process: return datasets datasets = self._post_process_datasets(datasets) self._show_dataset(*datasets) return datasets def _post_process_datasets(self, datasets: List) -> List: args = self.args predict_with_generate = getattr(args, 'predict_with_generate', False) template = self.template for i, dataset in enumerate(datasets): if dataset is None: continue if i == 1 and predict_with_generate: # val_dataset continue if not args.streaming and args.truncation_strategy != 'split': dataset = LazyLLMDataset(dataset, template.encode, strict=args.strict, random_state=args.data_seed) if args.packing: packing_dataset_cls = IterablePackingDataset if args.streaming else PackingDataset dataset = packing_dataset_cls( template, dataset, num_proc=args.dataset_num_proc, packing_length=args.packing_length, packing_num_proc=args.packing_num_proc, packing_strategy=args.packing_strategy, strict=args.strict, load_from_cache_file=args.load_from_cache_file) elif args.streaming: preprocessor = EncodePreprocessor(template=template) dataset = preprocessor( dataset, num_proc=args.dataset_num_proc, load_from_cache_file=args.load_from_cache_file, strict=args.strict) datasets[i] = dataset return datasets @RayHelper.function(group='default') def run(self): args = self.args train_dataset, val_dataset = self._prepare_dataset() if args.task_type == 'seq_cls': args.problem_type = args.problem_type or getattr(self.model.config, 'problem_type', None) logger.info(f'args.problem_type: {args.problem_type}') args.save_args() # Some tuners require train_dataset and data_collator for preparation: LoRA-GA self.model = self.prepare_model(self.args, self.model, template=self.template, train_dataset=train_dataset) logger.info(f'model: {self.model}') model_parameter_info = get_model_parameter_info(self.model) self.train_msg['model_parameter_info'] = model_parameter_info logger.info(f'model_parameter_info: {model_parameter_info}') trainer_cls = TrainerFactory.get_trainer_cls(args) trainer = trainer_cls( model=self.model, args=self.args.training_args, template=self.template, train_dataset=train_dataset, eval_dataset=val_dataset, **self._get_trainer_kwargs(), ) return self.train(trainer) def _get_trainer_kwargs(self): return {} def _handle_trainer_state(self, trainer, is_write_rank: bool): state = trainer.state if hasattr(state, 'last_model_checkpoint'): if self.args.create_checkpoint_symlink: last_checkpoint = os.path.join(self.args.output_dir, 'last') best_checkpoint = os.path.join(self.args.output_dir, 'best') if is_write_rank: os.symlink(state.last_model_checkpoint, last_checkpoint) os.symlink(state.best_model_checkpoint, best_checkpoint) state.last_model_checkpoint = last_checkpoint state.best_model_checkpoint = best_checkpoint else: state.last_model_checkpoint = None logger.info_if(f'last_model_checkpoint: {state.last_model_checkpoint}', cond=is_write_rank) logger.info_if(f'best_model_checkpoint: {state.best_model_checkpoint}', cond=is_write_rank) def _save_trainer_state(self, trainer): training_args = trainer.args state = trainer.state self._handle_trainer_state(trainer, is_master()) if is_master(): # Visualization if 'tensorboard' in training_args.report_to: images_dir = os.path.join(training_args.output_dir, 'images') logger.info(f'images_dir: {images_dir}') plot_images(images_dir, training_args.logging_dir, ['train/loss'], 0.9) if training_args.push_to_hub: trainer.push_to_hub() self.train_msg.update({ 'last_model_checkpoint': state.last_model_checkpoint, 'best_model_checkpoint': state.best_model_checkpoint, 'best_metric': state.best_metric, 'global_step': state.global_step, 'log_history': state.log_history, 'memory': getattr(state, 'max_memory', None), }) if is_master(): jsonl_path = os.path.join(training_args.output_dir, 'logging.jsonl') append_to_jsonl(jsonl_path, self.train_msg, strict=False) return self.train_msg def _get_resume_checkpoint(self, trainer): args = trainer.args if args.resume_from_checkpoint: return args.resume_from_checkpoint resume_checkpoint = None # If flash checkpoint is enabled, try to resume from the last complete checkpoint. # If the previous training finished, resume_checkpoint stays None. if args.use_flash_ckpt: # resume_checkpoint = /checkpoint- resume_checkpoint = trainer.get_resume_checkpoint() # Elastic runs require a universal checkpoint; fall back when missing or incomplete. callbacks = set(getattr(args, 'callbacks', [])) elastic_enabled = 'deepspeed_elastic' in callbacks if elastic_enabled and (resume_checkpoint is None or not os.path.exists(os.path.join(resume_checkpoint, 'latest_universal'))): # get_resume_checkpoint_until_find_ucp returns /checkpoint- with latest_universal, # or None; when None, no universal checkpoint exists and training starts from scratch. resume_checkpoint = trainer.get_resume_checkpoint_until_find_ucp() return resume_checkpoint def train(self, trainer): logging_path = os.path.join(trainer.args.output_dir, 'logging.jsonl') logger.info(f'The logging file will be saved in: {logging_path}') resume_checkpoint = self._get_resume_checkpoint(trainer) try: trainer.train(resume_checkpoint) finally: res = self._save_trainer_state(trainer) if self.args.use_flash_ckpt and hasattr(trainer, 'flash_checkpointer'): trainer.wait_latest_checkpoint(trainer.FLASH_CKPT_WAIT_TIMEOUT, trainer.state.global_step) return res @staticmethod def _stat_dataset(dataset: Union[HfDataset, PackingDataset, LazyLLMDataset]): if isinstance(dataset, LazyLLMDataset): dataset = dataset.dataset if isinstance(dataset, HfDataset): lengths = dataset['lengths'] lengths = [max(length) if isinstance(length, list) else length for length in lengths] else: lengths = dataset.packed_length _, stat_str = stat_array(lengths) logger.info(f'Dataset Token Length: {stat_str}') return stat_str def _show_dataset(self, train_dataset, val_dataset): args = self.args predict_with_generate = getattr(args, 'predict_with_generate', False) if is_master(): inputs = train_dataset[0] if hasattr(train_dataset, '__len__') else next(iter(train_dataset)) if isinstance(inputs, list): inputs = inputs[0] self.template.print_inputs(inputs) elif hasattr(train_dataset, '__len__'): # Avoid the random mismatch issue in LazyLLMDataset. inputs = train_dataset[0] if val_dataset is not None and hasattr(val_dataset, '__len__') and len(val_dataset) == 0: val_dataset = None if not args.lazy_tokenize and not args.streaming: self.train_msg['train_dataset'] = self._stat_dataset(train_dataset) if val_dataset is not None and not predict_with_generate: self.train_msg['val_dataset'] = self._stat_dataset(val_dataset) def _encode_dataset(self, train_dataset, val_dataset, pre_process=True): template = self.template args = self.args self._save_val_dataset(val_dataset) predict_with_generate = getattr(args, 'predict_with_generate', False) datasets = [train_dataset, val_dataset] if not pre_process: return datasets origin_template_model = template.model template.model = None # Avoid serializing the model. if args.truncation_strategy == 'split': if args.task_type != 'causal_lm' or template.mode != 'train' or args.use_chat_template: raise ValueError('`--truncation_strategy split` is currently only supported for pre-training.') assert not args.lazy_tokenize, '`--truncation_strategy split` does not support lazy_tokenize' for i, dataset in enumerate(datasets): if dataset is None: continue if i == 1 and predict_with_generate: # val_dataset continue if not args.lazy_tokenize and not args.streaming: # Compatible with cached_dataset, only additionally write length here. preprocessor_cls = EncodePreprocessor if args.truncation_strategy == 'split' else AddLengthPreprocessor preprocessor = preprocessor_cls(template=template) batch_size = 100 if args.model_meta.is_multimodal else 1000 dataset = preprocessor( dataset, num_proc=args.dataset_num_proc, load_from_cache_file=args.load_from_cache_file, strict=args.strict, batch_size=batch_size) if len(dataset) == 0: dataset = None datasets[i] = dataset template.model = origin_template_model return datasets def sft_main(args: Optional[Union[List[str], SftArguments]] = None): return SwiftSft(args).main()