File size: 6,039 Bytes
01cf98b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 | 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)
|