mirror of
https://github.com/SakanaAI/doc-to-lora.git
synced 2026-07-23 17:01:04 +02:00
per module head
This commit is contained in:
51616
parent
302ed4dbdd
commit
59370b175f
2 changed files with 100 additions and 67 deletions
|
|
@ -445,17 +445,37 @@ class HyperLoRA(nn.Module):
|
|||
else:
|
||||
d_lora = max(self.d_in[m] + self.d_out[m] for m in self.target_modules)
|
||||
|
||||
# each module processes d -> r d_out
|
||||
self.head = Mix(
|
||||
"bs n_layers n_modules d_latent -> bs n_layers n_modules r d_lora",
|
||||
weight_shape="d_latent r d_lora",
|
||||
# bias_shape=None, # no bias
|
||||
bias_shape="r d_lora",
|
||||
n_modules=len(self.target_modules),
|
||||
d_latent=self.config.latent_size,
|
||||
r=self.config.lora_config.r,
|
||||
d_lora=d_lora,
|
||||
)
|
||||
n_modules = len(self.target_modules)
|
||||
# have to do this otherwise doesnt work with adamw_torch_fused
|
||||
# has something to do with the bias shape (n_modules r d_lora)
|
||||
# when n_modules == 1, adamw_torch_fused complains about device/layout
|
||||
# but when n_modules > 1, it works fine
|
||||
if n_modules == 1:
|
||||
self.head = Mix(
|
||||
"bs n_layers n_modules d_latent -> bs n_layers n_modules r d_lora",
|
||||
weight_shape="d_latent r d_lora",
|
||||
# bias_shape=None, # no bias
|
||||
bias_shape="r d_lora",
|
||||
d_latent=self.config.latent_size,
|
||||
r=self.config.lora_config.r,
|
||||
d_lora=d_lora,
|
||||
)
|
||||
else:
|
||||
# each module processes d -> r d_out independently
|
||||
self.head = Mix(
|
||||
"bs n_layers n_modules d_latent -> bs n_layers n_modules r d_lora",
|
||||
weight_shape="n_modules d_latent r d_lora",
|
||||
# bias_shape=None, # no bias
|
||||
bias_shape="n_modules r d_lora",
|
||||
n_modules=len(self.target_modules),
|
||||
d_latent=self.config.latent_size,
|
||||
r=self.config.lora_config.r,
|
||||
d_lora=d_lora,
|
||||
)
|
||||
# print(self.head)
|
||||
# print(self.head.weight.shape)
|
||||
# print(self.head.bias.shape)
|
||||
# breakpoint()
|
||||
|
||||
def _to_lora_dict(
|
||||
self, flat_loras: Float[Tensor, "bs n_layers n_modules r max_io_dim"]
|
||||
|
|
@ -583,18 +603,34 @@ class ModulatedPretrainedModel(nn.Module):
|
|||
logger.debug(f"peft_weights: {peft_weights}")
|
||||
self.hypernet.head.weight.data[:] = 0
|
||||
self.hypernet.head.bias.data[:] = 0
|
||||
d_in = self.hypernet.d_in
|
||||
d_out = self.hypernet.d_out
|
||||
m = max(self.hypernet.target_modules, key=lambda m: d_in[m] + d_out[m])
|
||||
# d_in = self.hypernet.d_in
|
||||
# d_out = self.hypernet.d_out
|
||||
# m = max(self.hypernet.target_modules, key=lambda m: d_in[m] + d_out[m])
|
||||
|
||||
A = peft_weights[m]["lora_A"].weight.clone()[:, : d_in[m]] # [r, d_in]
|
||||
B = peft_weights[m]["lora_B"].weight.clone()[: d_out[m]] # [d_out, r]
|
||||
biases = [A, B.T]
|
||||
# A = peft_weights[m]["lora_A"].weight.clone()[:, : d_in[m]] # [r, d_in]
|
||||
# B = peft_weights[m]["lora_B"].weight.clone()[: d_out[m]] # [d_out, r]
|
||||
# biases = [A, B.T]
|
||||
|
||||
# bias-hyperinit
|
||||
# init weights to zeros and bias to the base weights
|
||||
bias_cat = torch.cat(biases, dim=1)
|
||||
self.hypernet.head.bias.data[..., :, : bias_cat.shape[1]] = bias_cat
|
||||
# # bias-hyperinit
|
||||
# # init weights to zeros and bias to the base weights
|
||||
# bias_cat = torch.cat(biases, dim=1)
|
||||
# self.hypernet.head.bias.data[..., :, : bias_cat.shape[1]] = bias_cat
|
||||
for i, m in enumerate(self.hypernet.target_modules):
|
||||
A = peft_weights[m]["lora_A"].weight.clone() # [r, in_d]
|
||||
B = peft_weights[m]["lora_B"].weight.clone() # [out_d, r]
|
||||
if self.hypernet.config.use_light_weight_lora:
|
||||
A = A[:, : self.hypernet.config.light_weight_latent_size]
|
||||
B = B[: self.hypernet.config.light_weight_latent_size]
|
||||
biases = [A, B.T]
|
||||
|
||||
# bias-hyperinit
|
||||
# init weights to zeros and bias to the base weights
|
||||
bias_cat = torch.cat(biases, dim=1)
|
||||
self.hypernet.head.bias.data[..., i, :, : bias_cat.shape[1]] = bias_cat
|
||||
# if len(self.hypernet.target_modules) > 1:
|
||||
# self.hypernet.head.bias.data[..., :, : bias_cat.shape[1]] = bias_cat
|
||||
# else:
|
||||
# self.hypernet.head.bias.data[..., :, : bias_cat.shape[1]] = bias_cat
|
||||
# self.hypernet.head.bias.requires_grad = False
|
||||
|
||||
def state_dict(self, *args, **kwargs):
|
||||
|
|
|
|||
|
|
@ -147,60 +147,57 @@ def train_model(
|
|||
trainer.save_model()
|
||||
clear_gpu()
|
||||
|
||||
############## Evaluation
|
||||
# TODO: eval does not work when using with deepspeed
|
||||
# make a separate eval script
|
||||
# ############## Evaluation
|
||||
# # TODO: eval does not work when using with deepspeed
|
||||
# # make a separate eval script
|
||||
|
||||
# max_input_len=2**13 # for input truncation
|
||||
# # max_input_len=2**13 # for input truncation
|
||||
|
||||
gen_kwargs = dict(do_sample=False, max_new_tokens=max_new_tokens)
|
||||
# pad_token_id=tokenizer.pad_token_id,
|
||||
# eos_token_id=?
|
||||
# gen_kwargs = dict(do_sample=False, max_new_tokens=max_new_tokens)
|
||||
# # pad_token_id=tokenizer.pad_token_id,
|
||||
# # eos_token_id=?
|
||||
|
||||
eval_trainer_args = {}
|
||||
# eval_trainer_args = {}
|
||||
|
||||
# Copy only necessary attributes from training_args to eval_trainer_args
|
||||
seq2seq_training_args_fields = {f.name for f in fields(Seq2SeqTrainingArguments)}
|
||||
for attr, value in training_args.to_dict().items():
|
||||
if attr in seq2seq_training_args_fields:
|
||||
eval_trainer_args[attr] = value
|
||||
# # Copy only necessary attributes from training_args to eval_trainer_args
|
||||
# seq2seq_training_args_fields = {f.name for f in fields(Seq2SeqTrainingArguments)}
|
||||
# for attr, value in training_args.to_dict().items():
|
||||
# if attr in seq2seq_training_args_fields:
|
||||
# eval_trainer_args[attr] = value
|
||||
|
||||
eval_trainer_args["eval_strategy"] = "no"
|
||||
eval_trainer_args["save_strategy"] = "no"
|
||||
eval_trainer_args["overwrite_output_dir"] = True
|
||||
eval_trainer_args["per_device_eval_batch_size"] = gen_per_device_eval_batch_size
|
||||
# eval_trainer_args["eval_strategy"] = "no"
|
||||
# eval_trainer_args["save_strategy"] = "no"
|
||||
# eval_trainer_args["overwrite_output_dir"] = True
|
||||
# eval_trainer_args["per_device_eval_batch_size"] = gen_per_device_eval_batch_size
|
||||
|
||||
# NOTE: could also set kv_cache implementation here
|
||||
eval_trainer_args = Seq2SeqTrainingArguments(
|
||||
**eval_trainer_args,
|
||||
predict_with_generate=True,
|
||||
generation_config=GenerationConfig(**gen_kwargs),
|
||||
)
|
||||
# # NOTE: could also set kv_cache implementation here
|
||||
# eval_trainer_args = Seq2SeqTrainingArguments(
|
||||
# **eval_trainer_args,
|
||||
# predict_with_generate=True,
|
||||
# generation_config=GenerationConfig(**gen_kwargs),
|
||||
# )
|
||||
|
||||
# Seq2SeqTrainer is actually just the same as Trainer
|
||||
# (although it uses a different data collator, i.e., explicit prompt/answer separation)
|
||||
# it just allows `predict_with_generate`
|
||||
# allowing us to compute metrics on the generated outputs
|
||||
# no clue why they call this seq2seq...
|
||||
# # Seq2SeqTrainer is actually just the same as Trainer
|
||||
# # (although it uses a different data collator, i.e., explicit prompt/answer separation)
|
||||
# # it just allows `predict_with_generate`
|
||||
# # allowing us to compute metrics on the generated outputs
|
||||
# # no clue why they call this seq2seq...
|
||||
|
||||
logger.info("=" * 80 + "\n" + "Evaluating model..." + "\n" + "=" * 80)
|
||||
# logger.info("=" * 80 + "\n" + "Evaluating model..." + "\n" + "=" * 80)
|
||||
|
||||
model.eval()
|
||||
# model.eval()
|
||||
|
||||
eval_trainer = Seq2SeqTrainer(
|
||||
model=model,
|
||||
args=eval_trainer_args,
|
||||
# TODO: use a different collator for test, e.g., more max_len truncation
|
||||
# w/ left padding?
|
||||
# removing label part from input_ids
|
||||
data_collator=generation_collator,
|
||||
)
|
||||
# eval_trainer = Seq2SeqTrainer(
|
||||
# model=model,
|
||||
# args=eval_trainer_args,
|
||||
# # TODO: use a different collator for test, e.g., more max_len truncation
|
||||
# # w/ left padding?
|
||||
# # removing label part from input_ids
|
||||
# data_collator=generation_collator,
|
||||
# )
|
||||
|
||||
for split, ds in zip(["eval", "test"], [val_dataset, test_dataset]):
|
||||
if ds is None:
|
||||
continue
|
||||
eval_generation(eval_trainer, tokenizer, ds, split, gen_kwargs)
|
||||
clear_gpu()
|
||||
|
||||
# if test_dataset is not None:
|
||||
# eval_generation(eval_trainer, tokenizer, test_dataset, "test", gen_kwargs)
|
||||
# for split, ds in zip(["eval", "test"], [val_dataset, test_dataset]):
|
||||
# if ds is None:
|
||||
# continue
|
||||
# eval_generation(eval_trainer, tokenizer, ds, split, gen_kwargs)
|
||||
# clear_gpu()
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue