From a5a64581cac3a403ae31e5774ec69d961c111903 Mon Sep 17 00:00:00 2001 From: 51616 Date: Thu, 9 Jan 2025 19:56:08 +0000 Subject: [PATCH] add use_kl_loss! --- hyperlora/configs.py | 4 +++ hyperlora/data_utils.py | 54 ++++++++++++++++------------ hyperlora/intx_sft.py | 35 +++++++++++++++++- hyperlora/modeling_utils.py | 71 ++++++++++++++++++++++++++++++------- 4 files changed, 128 insertions(+), 36 deletions(-) diff --git a/hyperlora/configs.py b/hyperlora/configs.py index bddf53c..05d18c0 100644 --- a/hyperlora/configs.py +++ b/hyperlora/configs.py @@ -244,6 +244,10 @@ class CtxTrainingArguments: default=True, metadata={"help": "Whether to add negative prompt to the dataset."}, ) + use_kl_loss: bool = field( + default=False, + metadata={"help": "Whether to use KL loss."}, + ) @dataclass diff --git a/hyperlora/data_utils.py b/hyperlora/data_utils.py index efecfe6..12ddf9a 100644 --- a/hyperlora/data_utils.py +++ b/hyperlora/data_utils.py @@ -18,6 +18,8 @@ def validate_columns(tokenized_ds): cols = ["input_ids", "attention_mask", "labels"] if "ctx_ids" in tokenized_ds.column_names: cols += ["ctx_ids", "ctx_attn_mask"] + if "chat_ids" in tokenized_ds.column_names: + cols += ["chat_ids", "chat_attn_mask", "chat_labels"] ref_cols = set(cols) assert ( set(tokenized_ds.column_names) == ref_cols @@ -110,6 +112,7 @@ def get_tokenized_dataset( add_ctx_to_chat: bool, add_repeat_prompt: bool, add_negative_prompt: bool, + use_kl_loss: bool, ) -> dict[str, Any]: need_ctx_ids = not add_ctx_to_chat @@ -123,19 +126,10 @@ def get_tokenized_dataset( ds = ds.map(get_preprocessing_fn(ds_name)) ds = ds.filter(filter_long_samples, batched=True) if add_negative_prompt: - cols_to_remove = [ - col - for col in ds.column_names - if col not in ["context", "prompt", "response"] - ] - ds = ds.map(add_negative_prompt_fn, batched=True, remove_columns=cols_to_remove) + ds = ds.map(add_negative_prompt_fn, batched=True) if add_repeat_prompt and "context_numbers" not in ds_name: - cols_to_remove = [ - col - for col in ds.column_names - if col not in ["context", "prompt", "response"] - ] - ds = ds.map(add_repeat_prompt_fn, batched=True, remove_columns=cols_to_remove) + ds = ds.map(add_repeat_prompt_fn, batched=True) + # for sft + chat_model, we need to convert the dataset to chat format # add "messages" field ds = ds.map( @@ -146,6 +140,7 @@ def get_tokenized_dataset( ds = ds.map(get_sft_prompt_formatting_fn(TRAINING_TASK.COMPLETION, tokenizer)) # tokenize the chat + mask the assistant inputs + # add "input_ids", "attention_mask", "labels" tokenized_ds = ds.map( tokenize_chat_messages, fn_kwargs={ @@ -153,7 +148,6 @@ def get_tokenized_dataset( "mask_assistant_inputs": True, "tokenizer_kwargs": tokenizer_kwargs, }, - remove_columns=["messages", "prompt", "response", "chat"], ) other_cols = [ @@ -162,7 +156,25 @@ def get_tokenized_dataset( if col not in ["input_ids", "attention_mask", "labels"] ] - # computes ctx_features offline when using hyperlora + # for use_kl_loss, we need "chat_ids" and "chat_attn_mask" + if use_kl_loss: + tokenized_ds = tokenized_ds.map( + convert_ctx_prompt_response_to_messages, + fn_kwargs={"add_ctx_to_chat": True}, + ) + tokenized_ds = tokenized_ds.map( + get_sft_prompt_formatting_fn(TRAINING_TASK.COMPLETION, tokenizer) + ) + tokenized_ds = tokenized_ds.map( + tokenize_chat_messages, + fn_kwargs={ + "tokenizer": tokenizer, + "mask_assistant_inputs": True, + "tokenizer_kwargs": tokenizer_kwargs, + "for_kl_loss": True, + }, + ) + if need_ctx_ids: # TODO: can we batch this? # TODO: can we cache this? @@ -172,14 +184,6 @@ def get_tokenized_dataset( fn_kwargs={"tokenizer": tokenizer}, ) - # # FIXME: arrow cant consturct this for very long ctx_len (>1k tokens) - # # might have to really do this online - # tokenized_ds = tokenized_ds.map( - # # TODO: truncate the ctx_ids to the max_ctx_len - # get_ctx_features_fn, - # remove_columns=["ctx_ids"], - # ) - tokenized_ds = tokenized_ds.remove_columns(other_cols) tokenized_ds.set_format(type="pt") @@ -413,6 +417,7 @@ def tokenize_chat_messages( example: dict[str, Any], tokenizer: PreTrainedTokenizerBase, mask_assistant_inputs: bool = True, + for_kl_loss: bool = False, tokenizer_kwargs: Optional[dict[str, Any]] = None, ) -> dict[str, list[int]]: """ @@ -422,6 +427,7 @@ def tokenize_chat_messages( example: Dictionary containing 'chat' and 'messages' keys tokenizer: Tokenizer to use mask_assistant_inputs: Whether to mask non-assistant tokens + for_kl_loss: change the column names to "chat_ids" and "chat_attn_mask" tokenizer_kwargs: Additional arguments to pass to tokenizer Returns: @@ -447,6 +453,10 @@ def tokenize_chat_messages( del conversation_ids["offset_mapping"] else: conversation_ids["labels"] = conversation_ids["input_ids"] + if for_kl_loss: + conversation_ids["chat_ids"] = conversation_ids.pop("input_ids") + conversation_ids["chat_attn_mask"] = conversation_ids.pop("attention_mask") + conversation_ids["chat_labels"] = conversation_ids.pop("labels") return conversation_ids diff --git a/hyperlora/intx_sft.py b/hyperlora/intx_sft.py index c8031ae..52a840f 100644 --- a/hyperlora/intx_sft.py +++ b/hyperlora/intx_sft.py @@ -254,7 +254,12 @@ def main(): # ).to(model.device) # ctx_encoder = EarlyExit(get_base_model(model), ctx_encoder_args.layer_idx) - model = ModulatedPretrainedModel(model, hypernet_config, ctx_encoder_args) + model = ModulatedPretrainedModel( + model, + hypernet_config, + ctx_encoder_args, + ctx_args.use_kl_loss, + ) else: # activate LoRA logger.info("Using LoRA") @@ -277,6 +282,7 @@ def main(): add_ctx_to_chat=add_ctx_to_chat, add_repeat_prompt=ctx_args.add_repeat_prompt, add_negative_prompt=ctx_args.add_negative_prompt, + use_kl_loss=ctx_args.use_kl_loss, ) tokenized_ds = {} for split, ds_names in zip( @@ -354,6 +360,29 @@ def main(): batch_first=True, padding_value=0, ) + + chat_ids = None + if "chat_ids" in inp_list[0]: + chat_ids = [x.pop("chat_ids") for x in inp_list] + chat_ids = torch.nn.utils.rnn.pad_sequence( + chat_ids, + batch_first=True, + padding_value=0, + ) + chat_attn_mask = [x.pop("chat_attn_mask") for x in inp_list] + chat_attn_mask = torch.nn.utils.rnn.pad_sequence( + chat_attn_mask, + batch_first=True, + padding_value=0, + ) + chat_labels = [x.pop("chat_labels") for x in inp_list] + chat_labels = torch.nn.utils.rnn.pad_sequence( + chat_labels, + batch_first=True, + padding_value=-100, + ) + chat_labels = torch.where(chat_attn_mask == 0, -100, chat_labels) + labels = [x.pop("labels") for x in inp_list] padded_seq = tokenizer.pad(inp_list, **padding_kwargs) @@ -364,6 +393,10 @@ def main(): if ctx_ids is not None: out["ctx_ids"] = ctx_ids out["ctx_attn_mask"] = ctx_attn_mask + if chat_ids is not None: + out["chat_ids"] = chat_ids + out["chat_attn_mask"] = chat_attn_mask + out["chat_labels"] = chat_labels return out def generation_collator(inp_list, tokenizer): diff --git a/hyperlora/modeling_utils.py b/hyperlora/modeling_utils.py index 89d8b51..93f15e4 100644 --- a/hyperlora/modeling_utils.py +++ b/hyperlora/modeling_utils.py @@ -441,7 +441,7 @@ class HyperLoRA(nn.Module): self.d_in = {k: self.config.light_weight_latent_size for k in self.d_in} self.d_out = {k: self.config.light_weight_latent_size for k in self.d_out} - logger.info(f"Using light-weight LoRA with d_lora = {d_lora}.") + logger.info(f"Using light-weight LoRA with d_lora = {d_lora // 2}") else: d_lora = max(self.d_in[m] + self.d_out[m] for m in self.target_modules) @@ -525,11 +525,14 @@ class ModulatedPretrainedModel(nn.Module): base_model: PreTrainedModel, hypernet_config: HypernetConfig, ctx_encoder_args: CtxEncoderArguments, + # use_kl_loss is only used for training + use_kl_loss: bool = False, ): super().__init__() self.device = base_model.device self.hypernet_config = hypernet_config self.ctx_encoder_args = ctx_encoder_args + self.use_kl_loss = use_kl_loss self.register_module("base_model", base_model) self._init_model() @@ -641,8 +644,11 @@ class ModulatedPretrainedModel(nn.Module): def forward( self, # ctx_features: Optional[Float[Tensor, "bs ctx_length feature_dim"]] = None, - ctx_ids: Optional[Integer[Tensor, "bs ctx_length"]] = None, - ctx_attn_mask: Optional[Integer[Tensor, "bs ctx_length"]] = None, + ctx_ids: Optional[Integer[Tensor, "bs ctx_len"]] = None, + ctx_attn_mask: Optional[Integer[Tensor, "bs ctx_len"]] = None, + chat_ids: Optional[Integer[Tensor, "bs chat_len"]] = None, + chat_attn_mask: Optional[Integer[Tensor, "bs chat_len"]] = None, + chat_labels: Optional[Integer[Tensor, "bs chat_len"]] = None, **model_inputs_kwargs: dict[str, Any], ) -> Union[tuple, ModelOutput]: """Forward pass of the modulated model.""" @@ -667,16 +673,55 @@ class ModulatedPretrainedModel(nn.Module): ) generated_loras = self.hypernet.generate_loras(ctx_features, ctx_attn_mask) - # apply lora hook to the base model - # self.apply_generated_loras(generated_loras) - with apply_generated_loras( - self.base_model, - generated_loras, - self.hypernet.layer_indices, - self.training, - ): - model_outputs = self.base_model(**model_inputs_kwargs) - # model_outputs.generated_loras = generated_loras + # compute kl loss + # - compute logits from the base model from tokenized chat [bs, chat_len, vocab_size] + # - compute logits from model w/ generated loras [bs, answer_len, vocab_size] + # - select only the position of the label tokens (but we're actually not using the labels themselves) + # - compute kl loss between the two logits + if self.use_kl_loss and "labels" in model_inputs_kwargs: + bs = chat_ids.shape[0] + labels = model_inputs_kwargs.pop("labels") # [bs, prompt_res_len] + first_label_pos = torch.argmax((labels != -100).float(), dim=-1) + # just a place holder, we dont use labels as groundtruth anyway + labels[torch.arange(bs), first_label_pos - 1] = 1 + # same for chat_labels + first_label_pos = torch.argmax((chat_labels != -100).float(), dim=-1) + chat_labels[torch.arange(bs), first_label_pos - 1] = 1 + + chat_outputs = self.base_model( + input_ids=chat_ids, attention_mask=chat_attn_mask + ) + with apply_generated_loras( + self.base_model, + generated_loras, + self.hypernet.layer_indices, + self.training, + ): + model_outputs = self.base_model(**model_inputs_kwargs) + pred_logits = model_outputs.logits # [bs, prompt_res_len, vocab_size] + + # [bs, ctx_prompt_res_len, vocab_size] + base_chat_logits = chat_outputs.logits + base_chat_logits = base_chat_logits[torch.where(chat_labels != -100)] + pred_logits = pred_logits[torch.where(labels != -100)] + kl_loss = F.kl_div( + F.log_softmax(pred_logits, dim=-1), + F.softmax(base_chat_logits, dim=-1), + reduction="none", + ) + # only compute kl loss on tokens where labels != -100 + kl_loss = kl_loss.sum(dim=-1).mean() + model_outputs = ModelOutput(loss=kl_loss, **model_outputs) + else: + # input_ids in model_inputs_kwargs contains only + # prompt + response (for hypernet training) + with apply_generated_loras( + self.base_model, + generated_loras, + self.hypernet.layer_indices, + self.training, + ): + model_outputs = self.base_model(**model_inputs_kwargs) return model_outputs