From 4dbeac1340267f7fd3d9211470e6092fb229988f Mon Sep 17 00:00:00 2001 From: 51616 Date: Thu, 9 Jan 2025 20:19:14 +0000 Subject: [PATCH] readd remove cols after adding samples + no grad for base chat output --- hyperlora/data_utils.py | 14 ++++++++++++-- hyperlora/modeling_utils.py | 8 +++++--- 2 files changed, 17 insertions(+), 5 deletions(-) diff --git a/hyperlora/data_utils.py b/hyperlora/data_utils.py index 12ddf9a..5159d6e 100644 --- a/hyperlora/data_utils.py +++ b/hyperlora/data_utils.py @@ -126,9 +126,19 @@ def get_tokenized_dataset( ds = ds.map(get_preprocessing_fn(ds_name)) ds = ds.filter(filter_long_samples, batched=True) if add_negative_prompt: - ds = ds.map(add_negative_prompt_fn, batched=True) + 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) if add_repeat_prompt and "context_numbers" not in ds_name: - ds = ds.map(add_repeat_prompt_fn, batched=True) + 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) # for sft + chat_model, we need to convert the dataset to chat format # add "messages" field diff --git a/hyperlora/modeling_utils.py b/hyperlora/modeling_utils.py index 93f15e4..e9b9f36 100644 --- a/hyperlora/modeling_utils.py +++ b/hyperlora/modeling_utils.py @@ -688,9 +688,11 @@ class ModulatedPretrainedModel(nn.Module): 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 torch.no_grad(): + chat_outputs = self.base_model( + input_ids=chat_ids, attention_mask=chat_attn_mask + ) + with apply_generated_loras( self.base_model, generated_loras,