apunkt/doc-to-lora
1
0
Fork 0
mirror of https://github.com/SakanaAI/doc-to-lora.git synced 2026-07-23 17:01:04 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

add finetuning + gte as ctx encoder + fix ctx input decode

Browse source
This commit is contained in:
51616 2025-01-20 09:19:45 +00:00
parent 930a22b2c0
commit 6dd8f13c28
5 changed files with 83 additions and 44 deletions
Download patch file Download diff file Expand all files Collapse all files

81
intx_sft.py
View file

@ -257,7 +257,7 @@ def main():
run_name = (
get_run_name(seed_str=training_args.logging_dir.strip("runs/"))
if not checkpoint_dir
else checkpoint_dir.split("/")[-2]
else checkpoint_dir.strip("/").split("/")[-2]
)
output_dir = f"train_outputs/runs/{run_name}"
@ -296,40 +296,59 @@ def main():
############ Model setup
model_name = model_args.model_name_or_path
model, tokenizer = get_model_and_tokenizer(
**vars(model_args),
train=True,
requires_grad=ctx_args.exp_setup == ExperimentSetup.FULL_FINETUNE,
peft_config=get_lora_config(model_name, **vars(lora_args)),
)
ctx_name = ctx_encoder_args.ctx_encoder_model_name_or_path
if ctx_name is not None:
ctx_encoder_model_config = AutoConfig.from_pretrained(ctx_name)
if "Llama" in ctx_name and "Vision" in ctx_name:
ctx_encoder_model_config = ctx_encoder_model_config.text_config
ctx_tokenizer = get_tokenizer(ctx_name)
else:
ctx_encoder_model_config = model.config
ctx_tokenizer = tokenizer
if not ctx_args.from_pretrained_checkpoint:
model_name = model_args.model_name_or_path
model, tokenizer = get_model_and_tokenizer(
**vars(model_args),
train=True,
requires_grad=ctx_args.exp_setup == ExperimentSetup.FULL_FINETUNE,
peft_config=get_lora_config(model_name, **vars(lora_args)),
)
ctx_name = ctx_encoder_args.ctx_encoder_model_name_or_path
if ctx_name is not None:
ctx_encoder_model_config = AutoConfig.from_pretrained(
ctx_name, trust_remote_code=True
)
if "Llama" in ctx_name and "Vision" in ctx_name:
ctx_encoder_model_config = ctx_encoder_model_config.text_config
ctx_tokenizer = get_tokenizer(ctx_name)
else:
ctx_encoder_model_config = model.config
ctx_tokenizer = tokenizer
if ctx_args.exp_setup == ExperimentSetup.HYPER_LORA:
logger.info("Using HyperLoRA")
hypernet_config = get_hypernet_config(
model, ctx_encoder_model_config, hypernet_args, aggregator_args
)
if ctx_encoder_args.layer_idx is None:
ctx_encoder_args.layer_idx = ctx_encoder_model_config.num_hidden_layers // 4
logger.info(
f"Using the first {ctx_encoder_args.layer_idx} layers"
" as the context encoder"
if not ctx_args.from_pretrained_checkpoint:
hypernet_config = get_hypernet_config(
model, ctx_encoder_model_config, hypernet_args, aggregator_args
)
model = ModulatedPretrainedModel(
model,
hypernet_config,
ctx_encoder_args,
ctx_args.use_kl_loss,
)
if ctx_encoder_args.layer_idx is None:
ctx_encoder_args.layer_idx = (
ctx_encoder_model_config.num_hidden_layers // 4
)
logger.info(
f"Using the first {ctx_encoder_args.layer_idx} layers"
" as the context encoder"
)
model = ModulatedPretrainedModel(
model,
hypernet_config,
ctx_encoder_args,
ctx_args.use_kl_loss,
)
else:
logger.info(
f"Loading from checkpoint: {ctx_args.from_pretrained_checkpoint}"
)
model = ModulatedPretrainedModel.from_state_dict(
torch.load(open(ctx_args.from_pretrained_checkpoint, "rb")),
train=True,
)
tokenizer = get_tokenizer(model.base_model.config.name_or_path)
ctx_name = model.ctx_encoder_args.ctx_encoder_model_name_or_path
if ctx_name is None:
ctx_name = model.base_model.config.name_or_path
ctx_tokenizer = get_tokenizer(ctx_name)
if len([p for p in model.ctx_encoder.parameters() if p.requires_grad]):
raise ValueError("ctx_encoder contains trainable parameters")
if len([p for p in model.base_model.parameters() if p.requires_grad]):

4
src/ctx_to_lora/configs.py
View file

@ -228,6 +228,10 @@ class ModelArguments:
default=None,
metadata={"help": ("Base model name or path.")},
)
use_flash_attn: bool = field(
default=True,
metadata={"help": "Whether to use flash attention."},
)
# use_peft: bool = field(
# default=False,
# metadata={"help": ("Whether to use PEFT or not for training.")},

16
src/ctx_to_lora/eval.py
View file

@ -143,7 +143,7 @@ def save_generated_text(samples, output_dir, split):
f.write(json.dumps(sample) + "\n")
def decode_test_result(test_dataset, test_result, tokenizer):
def decode_test_result(test_dataset, test_result, tokenizer, ctx_tokenizer):
out = []
for sample, pred_toks in zip(test_dataset, test_result.predictions):
# sample and labels are not padded
@ -168,14 +168,16 @@ def decode_test_result(test_dataset, test_result, tokenizer):
d["input"] = tokenizer.decode(input_toks, skip_special_tokens=True)
d["generated"] = tokenizer.decode(gen_toks, skip_special_tokens=True)
if "ctx_ids" in sample:
d["context"] = tokenizer.decode(sample["ctx_ids"], skip_special_tokens=True)
d["context"] = ctx_tokenizer.decode(
sample["ctx_ids"], skip_special_tokens=True
)
out.append(d)
return out
@torch.no_grad()
def eval_generation(eval_trainer, tokenizer, datasets, split, gen_kwargs):
def eval_generation(eval_trainer, tokenizer, ctx_tokenizer, datasets, split, gen_kwargs):
if not isinstance(datasets, dict):
datasets = {"": datasets}
out = {}
@ -186,7 +188,7 @@ def eval_generation(eval_trainer, tokenizer, datasets, split, gen_kwargs):
metric_key_prefix=split_name,
**gen_kwargs,
)
decoded_txts = decode_test_result(ds, eval_result, tokenizer)
decoded_txts = decode_test_result(ds, eval_result, tokenizer, ctx_tokenizer)
rouge_metrics = compute_rouge(
[txt["generated"] for txt in decoded_txts],
[txt["label"] for txt in decoded_txts],
@ -428,7 +430,9 @@ def evaluate(checkpoint_path, args, split, generative):
out.update(metrics)
else:
eval_trainer = Seq2SeqTrainer(**trainer_kwargs)
metrics = eval_generation(eval_trainer, tokenizer, datasets, split, gen_kwargs)
metrics = eval_generation(
eval_trainer, tokenizer, ctx_tokenizer, datasets, split, gen_kwargs
)
out.update(metrics)
clear_gpu()
@ -466,4 +470,4 @@ if __name__ == "__main__":
shutil.rmtree(args.output_dir)
evaluate(checkpoint_path, args, split="validation", generative=False)
# evaluate(checkpoint_path, args, split="validation", generative=True)
evaluate(checkpoint_path, args, split="validation", generative=True)

13
src/ctx_to_lora/model_loading.py
View file

@ -38,7 +38,8 @@ def get_model_and_tokenizer(
)
tokenizer = get_tokenizer(model_name_or_path, tokenizer_kwargs, peft_config, train)
model.config.pad_token_id = tokenizer.pad_token_id
model.generation_config.pad_token_id = tokenizer.pad_token_id
if getattr(model, "generation_config", None):
model.generation_config.pad_token_id = tokenizer.pad_token_id
return model, tokenizer
@ -65,6 +66,7 @@ def get_tokenizer(
model_name_or_path,
padding_side=padding_side,
truncation_side=truncation_side,
trust_remote_code=True,
**tokenizer_kwargs,
)
@ -128,17 +130,22 @@ def get_model(
if model_kwargs is not None:
model_init_kwargs.update(model_kwargs)
if use_flash_attn:
if "gemma" not in model_name_or_path:
if ("gemma" not in model_name_or_path) and ("gte" not in model_name_or_path):
model_init_kwargs["attn_implementation"] = "flash_attention_2"
if is_vision_model:
model_init_kwargs["attn_implementation"] = "sdpa"
model_init_kwargs.pop("use_cache")
# for training disable cache
elif "gte" in model_name_or_path:
model_init_kwargs["attn_implementation"] = "sdpa"
model_init_kwargs["torch_dtype"] = torch.float32
model_init_kwargs.pop("use_cache")
# if train and not is_vision_model:
# model_init_kwargs["use_cache"] = False
logger.debug(f"Model init kwargs: {model_init_kwargs}")
if not is_vision_model:
model = AutoModelForCausalLM.from_pretrained(**model_init_kwargs)
# model = AutoModelForCausalLM.from_pretrained(**model_init_kwargs)
model = AutoModel.from_pretrained(**model_init_kwargs)
else:
model = MllamaForConditionalGeneration.from_pretrained(**model_init_kwargs)
model = model.language_model

13
src/ctx_to_lora/modeling_utils.py
View file

@ -344,7 +344,10 @@ class EarlyExit(nn.Module):
def __init__(self, base_model: PreTrainedModel, exit_layer: int):
super().__init__()
self.base_model = base_model
self.base_model.layers = base_model.layers[:exit_layer]
if "gte" in base_model.config.name_or_path:
self.base_model.encoder.layer = base_model.encoder.layer[:exit_layer]
else:
self.base_model.layers = base_model.layers[:exit_layer]
# self.exit_layer = exit_layer
@torch.no_grad()
@ -612,7 +615,7 @@ class ModulatedPretrainedModel(nn.Module):
# if any(module_name in name for module_name in trainable_base_modules):
# param.requires_grad = True
self.hypernet = HyperLoRA(self.hypernet_config).to(self.device)
self.hypernet = HyperLoRA(self.hypernet_config).to(self.device).to(torch.float32)
# ctx_encoder_name =
# if self.ctx_encoder_args.ctx_encoder_model_name_or_path is not None:
# encoder_model = get_model(
@ -627,12 +630,14 @@ class ModulatedPretrainedModel(nn.Module):
ctx_model_name = self.base_model.config.name_or_path
# use an explicit copy of the base model
# for using with "modules_to_save"
base_model_attn_impl = self.base_model.config._attn_implementation
logger.debug(f"ctx_model_name: {ctx_model_name}")
logger.debug(f"base_model.config._attn_implementation: {base_model_attn_impl}")
encoder_model = get_model(
ctx_model_name,
train=True,
requires_grad=False,
use_flash_attn=self.base_model.config._attn_implementation
== "flash_attention_2",
use_flash_attn=base_model_attn_impl == "flash_attention_2",
)
self.ctx_encoder = EarlyExit(
get_base_model(encoder_model), self.ctx_encoder_args.layer_idx