| |
| import os |
| from typing import List, Union |
|
|
| from megatron.core.enums import ModelType |
| from megatron.training import pretrain |
|
|
| from swift.llm.train import SwiftSft |
| from swift.utils import get_logger, is_master, plot_images |
| from ..argument import MegatronTrainArguments |
| from ..utils import patch_megatron_tokenizer |
| from .patcher import patch_megatron_data_collator, patch_training_log |
| from .utils import build_streaming_dataloader, forward_step, get_swift_datasets_provider |
|
|
| logger = get_logger() |
|
|
|
|
| class MegatronSft(SwiftSft): |
| args_class = MegatronTrainArguments |
| args: args_class |
|
|
| def __init__(self, args: Union[List[str], MegatronTrainArguments, None] = None) -> None: |
| self.train_msg = {} |
| super(SwiftSft, self).__init__(args) |
| args = self.args |
| _, self.processor = args.get_model_processor(load_model=False) |
| patch_megatron_tokenizer(self.processor) |
| args.init_model_args(self.processor.model_info.config) |
| self._prepare_template() |
| self.template.use_megatron = True |
| args.save_args(args.save) |
|
|
| def run(self): |
| args = self.args |
|
|
| train_dataset, val_dataset = self._get_dataset() |
| train_dataset, val_dataset = self._encode_dataset(train_dataset, val_dataset) |
| data_collator = self.template.data_collator |
| if args.streaming: |
| train_dataset = build_streaming_dataloader(args, train_dataset, data_collator) |
| if val_dataset is not None: |
| val_dataset = build_streaming_dataloader(args, val_dataset, data_collator) |
| datasets_provider = get_swift_datasets_provider(train_dataset, val_dataset) |
| datasets_provider.is_distributed = True |
|
|
| logging_path = os.path.join(args.save, 'logging.jsonl') |
| logger.info(f'The logging file will be saved in: {logging_path}') |
| try: |
| with patch_training_log(), patch_megatron_data_collator(data_collator): |
| pretrain( |
| datasets_provider, |
| args.megatron_model_meta.model_provider, |
| ModelType.encoder_or_decoder, |
| forward_step, |
| args_defaults=args.extra_args) |
| finally: |
| |
| if is_master(): |
| images_dir = os.path.join(args.save, 'images') |
| logger.info(f'images_dir: {images_dir}') |
| plot_images(images_dir, args.tensorboard_dir) |
|
|
|
|
| def megatron_sft_main(args: Union[List[str], MegatronTrainArguments, None] = None): |
| return MegatronSft(args).main() |
|
|
