| import math |
| from typing import List, Dict, Optional |
|
|
| from deepspeed import zero |
| from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus |
|
|
| import torch |
| import torch.distributed as dist |
| from torch.utils.data import Sampler |
| import transformers |
| from transformers import Trainer |
| from transformers.trainer import has_length |
|
|
|
|
| class NoTextOnlyBatchSampler(Sampler): |
| def __init__( |
| self, |
| batch_size: int, |
| world_size: int, |
| is_text_only: Optional[List[bool]] = None, |
| generator=None, |
| ): |
| if is_text_only is None: |
| raise ValueError("`is_text_only` must be provided.") |
|
|
| self.batch_size = batch_size |
| self.world_size = world_size |
| self.is_text_only = is_text_only |
| self.generator = generator |
| self.mega_batch_size = batch_size * world_size |
|
|
| def __len__(self): |
| return len(self.is_text_only) |
|
|
| def __iter__(self): |
| mm_indices = [i for i, is_text_only in enumerate(self.is_text_only) if not is_text_only] |
| uni_indices = [i for i, is_text_only in enumerate(self.is_text_only) if is_text_only] |
|
|
| num_batches = math.ceil((len(mm_indices) + len(uni_indices)) / self.mega_batch_size) |
| if len(mm_indices) < num_batches: |
| raise ValueError( |
| f"{len(mm_indices)} multimodal entries, {num_batches} batches. " |
| "Not enough multimodal data in the dataset, or the batch size is too small." |
| ) |
|
|
| mm_indices = [mm_indices[i] for i in torch.randperm(len(mm_indices), generator=None).tolist()] |
| uni_indices = [uni_indices[i] for i in torch.randperm(len(uni_indices), generator=None).tolist()] |
|
|
| num_uni_indices_in_mega_batch = [len(uni_indices) // num_batches] * num_batches |
| for i in range(len(uni_indices) % num_batches): |
| num_uni_indices_in_mega_batch[i] += 1 |
|
|
| mega_batches = [] |
| cur_uni_index = 0 |
| cur_mm_index = 0 |
| for i, num_uni_indices in enumerate(num_uni_indices_in_mega_batch): |
| mega_batch = [] |
| mega_batch.extend(uni_indices[cur_uni_index:cur_uni_index + num_uni_indices]) |
| cur_uni_index += num_uni_indices |
| assert len(mega_batch) < self.mega_batch_size |
|
|
| if i < num_batches - 1: |
| increment = self.mega_batch_size - len(mega_batch) |
| mega_batch.extend( |
| mm_indices[cur_mm_index:cur_mm_index + increment] |
| ) |
| cur_mm_index += increment |
| else: |
| mega_batch.extend(mm_indices[cur_mm_index:]) |
| assert len(mega_batch) <= self.mega_batch_size |
|
|
| mega_batches.append(mega_batch) |
|
|
| mega_batch_indices = torch.randperm(len(mega_batches), generator=self.generator) |
| mega_batches = [mega_batches[i] for i in mega_batch_indices] |
| indices = [i for mega_batch in mega_batches for i in mega_batch] |
| return iter(indices) |
|
|
|
|
| class TrainerWithCustomSampler(Trainer): |
| def _get_train_sampler(self) -> Optional[torch.utils.data.Sampler]: |
| if self.train_dataset is None or not has_length(self.train_dataset): |
| return None |
|
|
| is_text_only = self.train_dataset.is_text_only |
| return NoTextOnlyBatchSampler( |
| self.args.train_batch_size, |
| world_size=self.args.world_size * self.args.gradient_accumulation_steps, |
| is_text_only=is_text_only, |
| ) |
|
|
| def _get_eval_sampler(self, eval_dataset: torch.utils.data.Dataset) -> Optional[torch.utils.data.Sampler]: |
| is_text_only = eval_dataset.is_text_only |
| return NoTextOnlyBatchSampler( |
| self.args.eval_batch_size, |
| world_size=self.args.world_size, |
| is_text_only=is_text_only, |
| ) |
|
|
|
|
| def find_all_linear_names(named_modules: Dict, target_modules: List[str]): |
| cls = torch.nn.Linear |
| lora_module_names = set() |
| for name, module in named_modules.items(): |
| if not any([module_name in name for module_name in target_modules]): |
| continue |
|
|
| if isinstance(module, cls): |
| lora_module_names.add(name) |
|
|
| for name in list(lora_module_names): |
| if 'lm_head' in name: |
| lora_module_names.remove(name) |
|
|
| return list(lora_module_names) |
|
|
|
|
| def rank0_print(*args): |
| if dist.is_initialized(): |
| if dist.get_rank() == 0: |
| print(*args) |
| else: |
| print(*args) |
|
|
|
|
| def maybe_zero_3(param): |
| if hasattr(param, "ds_id"): |
| with zero.GatheredParameters([param]): |
| param = param.data.detach().cpu().clone() |
| else: |
| param = param.detach().cpu().clone() |
| return param |
|
|
|
|
| def get_peft_state_maybe_zero_3(named_params, bias): |
| if bias == "none": |
| to_return = {k: t for k, t in named_params if "lora_" in k} |
| elif bias == "all": |
| to_return = {k: t for k, t in named_params if "lora_" in k or "bias" in k} |
| elif bias == "lora_only": |
| to_return = {} |
| maybe_lora_bias = {} |
| lora_bias_names = set() |
| for k, t in named_params: |
| if "lora_" in k: |
| to_return[k] = t |
| bias_name = k.split("lora_")[0] + "bias" |
| lora_bias_names.add(bias_name) |
| elif "bias" in k: |
| maybe_lora_bias[k] = t |
| for k, t in maybe_lora_bias: |
| if bias_name in lora_bias_names: |
| to_return[bias_name] = t |
| else: |
| raise NotImplementedError |
| to_return = {k: maybe_zero_3(v) for k, v in to_return.items()} |
| return to_return |
|
|
|
|
| def safe_save_model_for_hf_trainer(trainer: transformers.Trainer, output_dir: str): |
| if trainer.deepspeed: |
| torch.cuda.synchronize() |
| trainer.save_model(output_dir) |
| return |
|
|
| state_dict = trainer.model.state_dict() |
| if trainer.args.should_save: |
| cpu_state_dict = { |
| key: value.cpu() |
| for key, value in state_dict.items() |
| } |
| del state_dict |
| trainer._save(output_dir, state_dict=cpu_state_dict) |
|
|
