mirror of
https://github.com/SakanaAI/doc-to-lora.git
synced 2026-07-23 17:01:04 +02:00
update saving/loading
This commit is contained in:
51616
parent
c6f69fb438
commit
f545ec08d7
4 changed files with 45 additions and 184 deletions
|
|
@ -14,6 +14,7 @@ remove_unused_columns: false
|
|||
|
||||
optim: schedule_free_adamw
|
||||
learning_rate: 0.0001
|
||||
lr_scheduler_type: "constant_with_warmup"
|
||||
neftune_noise_alpha: 5
|
||||
weight_decay: 0.01
|
||||
warmup_ratio: 0.1
|
||||
|
|
|
|||
|
|
@ -96,7 +96,17 @@ def main():
|
|||
if ctx_args.exp_setup == ExperimentSetup.HYPER_LORA:
|
||||
logger.info("Using HyperLoRA")
|
||||
hypernet = HyperLoRA(get_hypernet_config(model)).to(model.device)
|
||||
model = ModulatedPretrainedModel(model, hypernet).to(model.device).train()
|
||||
# HACK: hardcode the embedding layer for now
|
||||
# TODO: add explicit encoder
|
||||
ctx_encoder = torch.nn.Embedding.from_pretrained(
|
||||
model.get_input_embeddings().weight.clone(),
|
||||
freeze=True,
|
||||
)
|
||||
model = (
|
||||
ModulatedPretrainedModel(model, hypernet, ctx_encoder)
|
||||
.to(model.device)
|
||||
.train()
|
||||
)
|
||||
else:
|
||||
# activate LoRA
|
||||
logger.info("Using LoRA")
|
||||
|
|
@ -193,6 +203,10 @@ def main():
|
|||
|
||||
# TODO: use SFTTrainer instead? https://huggingface.co/docs/trl/en/sft_trainer
|
||||
# TODO: use packing with SFTTrainer
|
||||
|
||||
# HACK: see transformers/trainer.py for liger-kernel patch
|
||||
# slows down training speed w/ short inputs
|
||||
# might improve/decrease training speed w/ longer inputs
|
||||
train_model(
|
||||
model,
|
||||
train_ds,
|
||||
|
|
|
|||
|
|
@ -305,7 +305,12 @@ class HyperLoRA(nn.Module):
|
|||
|
||||
|
||||
class ModulatedPretrainedModel(nn.Module):
|
||||
def __init__(self, base_model: PreTrainedModel, hypernet: HyperLoRA):
|
||||
def __init__(
|
||||
self,
|
||||
base_model: PreTrainedModel, # TODO: maybe just use the name/config?
|
||||
hypernet: HyperLoRA,
|
||||
ctx_encoder: nn.Module, # TODO: maybe just use the name/config?
|
||||
):
|
||||
super().__init__()
|
||||
# self.base_model = base_model
|
||||
self.device = base_model.device
|
||||
|
|
@ -313,13 +318,26 @@ class ModulatedPretrainedModel(nn.Module):
|
|||
# e.g., self.extractor_lm.model.layers = self.extractor_lm.model.layers[:num_extractor_layers]
|
||||
# or just use the token embeddings extractor_lm.embed_tokens(input_ids)
|
||||
|
||||
# HACK: hardcode the embedding layer for now
|
||||
# TODO: add explicit encoder
|
||||
self.encoder = base_model.get_input_embeddings()
|
||||
# self.ctx_encoder = ctx_encoder # base_model.get_input_embeddings()
|
||||
|
||||
# register base_model as a submodule
|
||||
self.register_module("base_model", base_model)
|
||||
self.register_module("hypernet", hypernet)
|
||||
self.register_module("ctx_encoder", ctx_encoder)
|
||||
|
||||
# TODO: fix model saving/loading
|
||||
def state_dict(self, *args, **kwargs):
|
||||
state_dict = super().state_dict(*args, **kwargs)
|
||||
|
||||
# remove non-trainable params
|
||||
for name, param in self.named_parameters():
|
||||
if not param.requires_grad:
|
||||
state_dict.pop(name)
|
||||
return state_dict
|
||||
|
||||
def load_state_dict(self, state_dict: dict, *args, **kwargs):
|
||||
# NOTE: might have to set `strict=False` as we don't save all the params
|
||||
return super().load_state_dict(state_dict, *args, **kwargs)
|
||||
|
||||
def get_input_embeddings(self):
|
||||
return self.base_model.get_input_embeddings()
|
||||
|
|
@ -332,10 +350,12 @@ class ModulatedPretrainedModel(nn.Module):
|
|||
input_ids = torch.tensor(examples.get("ctx_ids")).to(self.device)
|
||||
attention_mask = torch.tensor(examples.get("ctx_attn_mask")).to(self.device)
|
||||
|
||||
if isinstance(self.encoder, nn.Embedding):
|
||||
features = self.encoder(input_ids)
|
||||
if isinstance(self.ctx_encoder, nn.Embedding):
|
||||
features = self.ctx_encoder(input_ids)
|
||||
else:
|
||||
features = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
|
||||
features = self.ctx_encoder(
|
||||
input_ids=input_ids, attention_mask=attention_mask
|
||||
)
|
||||
out["ctx_features"] = features
|
||||
return out
|
||||
|
||||
|
|
|
|||
|
|
@ -80,184 +80,10 @@ def train_model(
|
|||
|
||||
print(f"Loaded from the checkpoint: {checkpoint}")
|
||||
|
||||
# TODO: save the best model based on eval loss?
|
||||
train_result = trainer.train(resume_from_checkpoint=checkpoint)
|
||||
trainer.save_model()
|
||||
trainer.log_metrics("train", train_result.metrics)
|
||||
metrics = trainer.evaluate()
|
||||
trainer.log_metrics("eval", metrics)
|
||||
trainer.save_metrics("eval", metrics)
|
||||
|
||||
|
||||
def text_extraction(input_ids, length, lm_ratio=0.0):
|
||||
|
||||
input_len = len(input_ids)
|
||||
assert input_len >= 1, f"Error: invalid input length ({input_len})"
|
||||
|
||||
# ae
|
||||
if random.random() >= lm_ratio:
|
||||
if input_len <= length: # if shorter, keep the complete text
|
||||
return input_ids, []
|
||||
else:
|
||||
last_start = input_len - length
|
||||
random_start = random.randint(0, last_start)
|
||||
return input_ids[random_start : random_start + length], []
|
||||
|
||||
# lm
|
||||
if input_len <= length:
|
||||
r = random.randint(0, input_len - 1)
|
||||
return input_ids[: r + 1], input_ids[r + 1 :]
|
||||
else:
|
||||
last_start = input_len - length
|
||||
random_start = random.randint(0, last_start)
|
||||
return (
|
||||
input_ids[random_start : random_start + length],
|
||||
input_ids[random_start + length :],
|
||||
)
|
||||
|
||||
|
||||
def pretrain_tokenize_function(examples, model, mem, lm_ratio=0.0):
|
||||
text_output = model.tokenizer(
|
||||
examples["text"], truncation=False, padding=False, return_attention_mask=False
|
||||
)
|
||||
text_output["prompt_answer_ids"] = []
|
||||
text_output["labels"] = []
|
||||
|
||||
max_len = model.training_args.model_max_length # heuristic
|
||||
|
||||
for idx in range(len(text_output["input_ids"])):
|
||||
|
||||
ae = True
|
||||
a, b = text_extraction(text_output["input_ids"][idx], max_len, lm_ratio=lm_ratio)
|
||||
length_a = len(a)
|
||||
num_segments = model.compute_num_segments(length_a)
|
||||
total_mem_length = num_segments * model.mem_size
|
||||
|
||||
if (
|
||||
len(b) > model.training_args.min_tokens_for_lm
|
||||
): # avoid too few tokens for lm, which is a waste of computing
|
||||
ae = False
|
||||
b = b[:max_len]
|
||||
|
||||
text_output["input_ids"][idx] = a
|
||||
|
||||
# decoder part: note that in v2, we add mem_tokens to the prompt_ids
|
||||
# for easy implementation; which is different from v1 implementation
|
||||
# where mem tokens are not in the prompt_ids
|
||||
if ae: # autoencoding objective
|
||||
prompt_ids = [mem[0]] * total_mem_length + [model.ae_token_id]
|
||||
answer_ids = a + [model.eos_id] # if ae, eos token
|
||||
else: # lm objective
|
||||
prompt_ids = [mem[0]] * total_mem_length
|
||||
if model.training_args.add_special_token_for_lm:
|
||||
prompt_ids += [model.lm_token_id]
|
||||
answer_ids = b # if lm, no eos token
|
||||
|
||||
text_output["prompt_answer_ids"].append(prompt_ids + answer_ids)
|
||||
if ae:
|
||||
labels = [-100] * len(prompt_ids) + answer_ids
|
||||
else:
|
||||
labels = (
|
||||
[-100] * len(prompt_ids)
|
||||
+ [-100] * model.training_args.leave_tokens_for_lm
|
||||
+ answer_ids[model.training_args.leave_tokens_for_lm :]
|
||||
) # no loss for leave_tokens_for_lm
|
||||
text_output["labels"].append(labels)
|
||||
assert len(text_output["prompt_answer_ids"][-1]) == len(labels)
|
||||
|
||||
return text_output
|
||||
|
||||
|
||||
def instruct_ft_tokenize_function(examples, model, mem):
|
||||
text_output = model.tokenizer(
|
||||
examples["context"],
|
||||
max_length=5120,
|
||||
truncation=True,
|
||||
padding=False,
|
||||
return_attention_mask=False,
|
||||
add_special_tokens=False,
|
||||
)
|
||||
prompt_output = model.tokenizer(
|
||||
examples["prompt"],
|
||||
truncation=False,
|
||||
padding=False,
|
||||
return_attention_mask=False,
|
||||
add_special_tokens=False,
|
||||
)
|
||||
label_output = model.tokenizer(
|
||||
examples["answer"],
|
||||
truncation=False,
|
||||
padding=False,
|
||||
return_attention_mask=False,
|
||||
add_special_tokens=False,
|
||||
)
|
||||
text_output["prompt_answer_ids"] = []
|
||||
text_output["labels"] = []
|
||||
|
||||
max_len = model.training_args.model_max_length # heuristic
|
||||
|
||||
for idx in range(len(text_output["input_ids"])):
|
||||
|
||||
length = len(text_output["input_ids"][idx])
|
||||
num_segments = model.compute_num_segments(length)
|
||||
total_mem_length = num_segments * model.mem_size
|
||||
|
||||
prompt_ids = (
|
||||
[mem[0]] * total_mem_length
|
||||
+ [model.ft_token_id]
|
||||
+ prompt_output["input_ids"][idx]
|
||||
)
|
||||
prompt_ids = (
|
||||
[1, 733, 16289, 28793] + prompt_ids + [733, 28748, 16289, 28793]
|
||||
) # special formats for prompt in Mistral
|
||||
answer_ids = label_output["input_ids"][idx] + [model.eos_id]
|
||||
|
||||
text_output["prompt_answer_ids"].append(prompt_ids + answer_ids)
|
||||
|
||||
labels = [-100] * len(prompt_ids) + answer_ids
|
||||
text_output["labels"].append(labels)
|
||||
|
||||
assert len(text_output["prompt_answer_ids"][-1]) == len(labels)
|
||||
|
||||
return text_output
|
||||
|
||||
|
||||
class DataCollatorForDynamicPadding:
|
||||
def __init__(self, pad_token_id, pad_to_multiple_of=None):
|
||||
self.pad_token_id = pad_token_id
|
||||
self.pad_to_multiple_of = pad_to_multiple_of
|
||||
|
||||
def __call__(self, examples):
|
||||
input_ids = [
|
||||
torch.tensor(example["input_ids"], dtype=torch.long) for example in examples
|
||||
]
|
||||
labels = [
|
||||
torch.tensor(example["labels"], dtype=torch.long) for example in examples
|
||||
]
|
||||
prompt_answer_ids = [
|
||||
torch.tensor(example["prompt_answer_ids"], dtype=torch.long)
|
||||
for example in examples
|
||||
]
|
||||
input_ids = self.dynamic_padding(input_ids, fill_value=self.pad_token_id)
|
||||
prompt_answer_ids = self.dynamic_padding(
|
||||
prompt_answer_ids, fill_value=self.pad_token_id
|
||||
)
|
||||
labels = self.dynamic_padding(labels)
|
||||
batch = {
|
||||
"input_ids": input_ids,
|
||||
"labels": labels,
|
||||
"prompt_answer_ids": prompt_answer_ids,
|
||||
}
|
||||
return batch
|
||||
|
||||
def dynamic_padding(self, sequences, fill_value=-100):
|
||||
max_length = max(len(x) for x in sequences)
|
||||
if self.pad_to_multiple_of:
|
||||
max_length = (
|
||||
(max_length - 1) // self.pad_to_multiple_of + 1
|
||||
) * self.pad_to_multiple_of
|
||||
padded_sequences = torch.full(
|
||||
(len(sequences), max_length), fill_value, dtype=torch.long
|
||||
)
|
||||
for i, seq in enumerate(sequences):
|
||||
padded_sequences[i, : len(seq)] = seq
|
||||
return padded_sequences
|
||||
trainer.save_model()
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue