From f545ec08d778bb00965f59b20a442a4c01e3c929 Mon Sep 17 00:00:00 2001 From: 51616 Date: Sun, 22 Dec 2024 16:07:11 +0000 Subject: [PATCH] update saving/loading --- configs/default.yaml | 1 + hyperlora/intx_sft.py | 16 +++- hyperlora/modeling_utils.py | 34 +++++-- hyperlora/training_utils.py | 178 +----------------------------------- 4 files changed, 45 insertions(+), 184 deletions(-) diff --git a/configs/default.yaml b/configs/default.yaml index 606acb3..d8173cd 100644 --- a/configs/default.yaml +++ b/configs/default.yaml @@ -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 diff --git a/hyperlora/intx_sft.py b/hyperlora/intx_sft.py index dba475c..5bafd32 100644 --- a/hyperlora/intx_sft.py +++ b/hyperlora/intx_sft.py @@ -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, diff --git a/hyperlora/modeling_utils.py b/hyperlora/modeling_utils.py index ef771a7..872ab97 100644 --- a/hyperlora/modeling_utils.py +++ b/hyperlora/modeling_utils.py @@ -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 diff --git a/hyperlora/training_utils.py b/hyperlora/training_utils.py index 230278f..d3c39ad 100644 --- a/hyperlora/training_utils.py +++ b/hyperlora/training_utils.py @@ -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()