| |
| from typing import Any |
|
|
| import datasets |
| import torch |
| import torch.distributed as dist |
| from datasets import Dataset |
| from torch.utils.data import DataLoader |
| from transformers.trainer_utils import seed_worker |
|
|
| from .base import SequenceParallel |
|
|
|
|
| class XTuner(SequenceParallel): |
|
|
| @staticmethod |
| def assert_xtuner_runtime_condition(): |
| from swift.utils import is_xtuner_available |
| assert is_xtuner_available(), \ |
| ('Please install XTuner first to pack dataset to `max_length`.' |
| '`pip install -U \'xtuner[deepspeed]\'`') |
| assert dist.is_initialized(), 'pack_to_max_length is only available with distributed training.' |
|
|
| def pack_dataset_xtuner(self, dataset: Dataset, args: Any) -> Any: |
| self.assert_xtuner_runtime_condition() |
| if dist.get_rank() == 0: |
| ds = [i[0] for i in dataset.data] |
| train_dataset = Dataset.from_list(ds) |
| from xtuner.dataset.huggingface import pack_dataset |
| train_dataset = pack_dataset( |
| train_dataset, |
| max_length=args.max_length, |
| use_varlen_attn=False, |
| shuffle_before_pack=True, |
| map_num_proc=16) |
| objects = [train_dataset] |
| train_dataset.save_to_disk('alpaca_pack') |
| else: |
| objects = [None] |
| dist.broadcast_object_list(objects, src=0) |
| train_dataset = objects[0] |
| return train_dataset |
|
|
| @property |
| def sp_group(self): |
| from xtuner.parallel.sequence import get_sequence_parallel_group |
| return get_sequence_parallel_group() |
|
|
| def init_sequence_parallel(self, size): |
| self.assert_xtuner_runtime_condition() |
| from xtuner.parallel.sequence import init_sequence_parallel |
| init_sequence_parallel(size) |
|
|
| def prepare_model(self, model, tokenizer, split_in_forward): |
| self.assert_xtuner_runtime_condition() |
| from xtuner.model.modules.dispatch import dispatch_modules |
| dispatch_modules(model) |
|
|
| def pad_and_split_inputs(self, |
| tokenizer, |
| input_ids, |
| input_embeds, |
| labels, |
| position_ids, |
| attention_mask, |
| loss_scale, |
| embed_tokens=None): |
| self.assert_xtuner_runtime_condition() |
| from xtuner.parallel.sequence import (pad_for_sequence_parallel, split_for_sequence_parallel, |
| get_sequence_parallel_group) |
| input_ids = pad_for_sequence_parallel(input_ids, padding_value=tokenizer.pad_token_id, dim=-1) |
| labels = pad_for_sequence_parallel(labels, padding_value=-100, dim=-1) |
| position_ids = pad_for_sequence_parallel(position_ids, padding_value=0, dim=-1) |
| if attention_mask is not None: |
| attention_mask = pad_for_sequence_parallel(attention_mask, padding_value=0, dim=-1) |
|
|
| sp_group = get_sequence_parallel_group() |
| input_ids = split_for_sequence_parallel(input_ids, dim=1, sp_group=sp_group) |
| labels = split_for_sequence_parallel(labels, dim=1, sp_group=sp_group) |
| position_ids = split_for_sequence_parallel(position_ids, dim=1, sp_group=sp_group) |
| if attention_mask is not None: |
| attention_mask = split_for_sequence_parallel(attention_mask, dim=-1, sp_group=sp_group) |
| if loss_scale is not None: |
| loss_scale = pad_for_sequence_parallel(loss_scale, padding_value=0., dim=-1) |
| loss_scale = split_for_sequence_parallel(loss_scale, dim=1, sp_group=sp_group) |
|
|
| return input_ids, None, labels, position_ids, attention_mask, loss_scale |
|
|
| def reduce_outputs(self, loss, labels): |
| from xtuner.parallel.sequence import (reduce_sequence_parallel_loss, get_sequence_parallel_group) |
| |
| num_tokens = (labels != -100).sum() |
| return reduce_sequence_parallel_loss(loss, num_tokens, get_sequence_parallel_group()) |
|
|
| def world_size(self): |
| self.assert_xtuner_runtime_condition() |
| from xtuner.parallel.sequence import get_sequence_parallel_world_size |
| return get_sequence_parallel_world_size() |
|
|
| def prepare_trainer(self, trainer): |
| pass |
|
|
| def get_dataloader(self, trainer, dataset, batch_size): |
| |
| |
| self.assert_xtuner_runtime_condition() |
| data_collator = trainer.data_collator |
| if isinstance(dataset, datasets.Dataset): |
| dataset = trainer._remove_unused_columns(dataset, description='training') |
| else: |
| data_collator = trainer._get_collator_with_removed_columns(data_collator, description='training') |
|
|
| dataloader_params = { |
| 'batch_size': batch_size, |
| 'collate_fn': data_collator, |
| 'num_workers': trainer.args.dataloader_num_workers, |
| 'pin_memory': trainer.args.dataloader_pin_memory, |
| 'persistent_workers': trainer.args.dataloader_persistent_workers, |
| } |
|
|
| if not isinstance(dataset, torch.utils.data.IterableDataset): |
| from xtuner.parallel import SequenceParallelSampler |
| dataloader_params['sampler'] = SequenceParallelSampler(dataset, seed=1024) |
| dataloader_params['drop_last'] = trainer.args.dataloader_drop_last |
| dataloader_params['worker_init_fn'] = seed_worker |
|
|
| return DataLoader(dataset, **dataloader_params) |
|
|
