File size: 6,910 Bytes
4ce9939 | 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 175 176 177 178 179 180 181 182 183 184 | 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):
r"""
Sampler that tries its best to sample batches such that no batch has only
text (unimodal) data. This is necessary for training with deepspeed.
"""
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: multimodal, entry that has both text and image/video
# uni: unimodal, entry that has only text
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, {len(num_batches)} batches. "
"Not enough multimodal data in the dataset, or the batch size is too small. "
"There will be at least one batch that is text-only, which doesn't work with deepspeed. "
"Try increasing the batch size first."
)
# shuffle indices
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()]
# distribute indices into batches
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: # last batch
mega_batch.extend(mm_indices[cur_mm_index:])
assert len(mega_batch) <= self.mega_batch_size, "Last batch is too big."
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: # needed for 16-bit
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)
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
# Borrowed from peft.utils.get_peft_model_state_dict
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):
"""Collects the state dict and dump to disk."""
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) |