mirror of
https://github.com/SakanaAI/doc-to-lora.git
synced 2026-07-23 17:01:04 +02:00
add use_kl_loss!
This commit is contained in:
51616
parent
68e46380e7
commit
a5a64581ca
4 changed files with 128 additions and 36 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue