mirror of
https://github.com/SakanaAI/doc-to-lora.git
synced 2026-07-23 17:01:04 +02:00
398 lines
13 KiB
Python
Executable file
398 lines
13 KiB
Python
Executable file
import contextlib
|
|
import logging
|
|
import os
|
|
from copy import deepcopy
|
|
from functools import partial
|
|
|
|
import numpy as np
|
|
import torch
|
|
import wandb
|
|
from datasets import disable_caching
|
|
from peft import PeftModel
|
|
from transformers import (
|
|
AutoConfig,
|
|
set_seed,
|
|
)
|
|
|
|
from ctx_to_lora.configs import (
|
|
AggregatorArguments,
|
|
ArgumentParser,
|
|
CtxEncoderArguments,
|
|
CtxTrainingArguments,
|
|
DataArguments,
|
|
ExperimentSetup,
|
|
HypernetArguments,
|
|
LoRAArguments,
|
|
ModelArguments,
|
|
TrainingArguments,
|
|
)
|
|
from ctx_to_lora.data.collator import ( # train_packed_collator,; DefaultDataCollator,
|
|
flatten_if_not_packed,
|
|
)
|
|
from ctx_to_lora.data.processing import get_tokenized_dataset, pack
|
|
from ctx_to_lora.metrics import (
|
|
Evaluator,
|
|
compute_metrics,
|
|
compute_per_token_acc,
|
|
compute_perplexity,
|
|
compute_prefix_matching,
|
|
)
|
|
from ctx_to_lora.model_loading import (
|
|
get_lora_config,
|
|
get_model_and_tokenizer,
|
|
get_tokenizer,
|
|
)
|
|
from ctx_to_lora.modeling.hypernet import (
|
|
ModulatedPretrainedModel,
|
|
get_hypernet_config,
|
|
)
|
|
from ctx_to_lora.trainer import train_model
|
|
from ctx_to_lora.utils import (
|
|
compile_linear,
|
|
extract_cli_args,
|
|
get_run_name,
|
|
log_num_train_params,
|
|
save_yaml,
|
|
setup_logging,
|
|
validate_args,
|
|
)
|
|
|
|
logger = logging.getLogger()
|
|
|
|
LOCAL_RANK = int(os.getenv("LOCAL_RANK", "0"))
|
|
|
|
|
|
def main():
|
|
############ Argument parsing
|
|
parser = ArgumentParser(
|
|
(
|
|
DataArguments,
|
|
CtxTrainingArguments,
|
|
ModelArguments,
|
|
LoRAArguments,
|
|
TrainingArguments,
|
|
HypernetArguments,
|
|
AggregatorArguments,
|
|
CtxEncoderArguments,
|
|
)
|
|
)
|
|
(
|
|
data_args,
|
|
ctx_args,
|
|
model_args,
|
|
lora_args,
|
|
training_args,
|
|
hypernet_args,
|
|
aggregator_args,
|
|
ctx_encoder_args,
|
|
) = parser.parse()
|
|
|
|
# there shouldn't be overlap between args
|
|
validate_args(
|
|
[
|
|
data_args,
|
|
ctx_args,
|
|
model_args,
|
|
lora_args,
|
|
training_args,
|
|
hypernet_args,
|
|
aggregator_args,
|
|
ctx_encoder_args,
|
|
]
|
|
)
|
|
|
|
assert ctx_args.use_sequence_packing, (
|
|
f"Please set use_sequence_packing=True in {ctx_args}. It's faster!"
|
|
)
|
|
|
|
set_seed(training_args.seed)
|
|
checkpoint_dir = training_args.resume_from_checkpoint
|
|
if checkpoint_dir and not os.path.isdir(checkpoint_dir):
|
|
raise NotADirectoryError(f"Checkpoint{checkpoint_dir} is not a directory")
|
|
|
|
# should be the same across processes
|
|
# still possible to have a name crash though
|
|
# logging_dir is just "runs/DATE_TIME_HOSTNAME"
|
|
slurm_job_id = f"_{os.getenv('SLURM_JOB_ID')}" if os.getenv("SLURM_JOB_ID") else ""
|
|
run_name = (
|
|
get_run_name(seed_str=training_args.logging_dir.strip("runs/") + slurm_job_id)
|
|
if not checkpoint_dir
|
|
else checkpoint_dir.strip("/").split("/")[-2]
|
|
)
|
|
|
|
output_dir = f"train_outputs/runs/{run_name}"
|
|
setup_logging(output_dir, debug=os.getenv("DEBUG", False))
|
|
logger.debug(f"CMD: {' '.join(os.sys.argv)}")
|
|
cli_args = extract_cli_args(os.sys.argv)
|
|
save_yaml(cli_args, f"{output_dir}/cli_args.yaml")
|
|
if "config" in cli_args:
|
|
config_name = os.path.basename(cli_args["config"]).split(".yaml")[0]
|
|
os.environ["WANDB_TAGS"] = config_name
|
|
|
|
run_name = os.path.basename(output_dir)
|
|
training_args.run_name = run_name
|
|
training_args.output_dir = output_dir
|
|
training_args.logging_dir = output_dir
|
|
|
|
if (
|
|
training_args.lr_scheduler_type == "cosine_with_min_lr"
|
|
and training_args.lr_scheduler_kwargs is None
|
|
):
|
|
training_args.lr_scheduler_kwargs = {"min_lr": 1e-7}
|
|
args = {
|
|
**vars(deepcopy(data_args)),
|
|
**vars(deepcopy(ctx_args)),
|
|
**vars(deepcopy(model_args)),
|
|
**vars(deepcopy(lora_args)),
|
|
**vars(deepcopy(training_args)),
|
|
**vars(deepcopy(hypernet_args)),
|
|
**vars(deepcopy(aggregator_args)),
|
|
**vars(deepcopy(ctx_encoder_args)),
|
|
}
|
|
args["deepspeed_plugin"] = None
|
|
logger.debug(f"args: {args}")
|
|
save_yaml(args, f"{output_dir}/args.yaml")
|
|
|
|
############ Model setup
|
|
if not ctx_args.from_pretrained_checkpoint:
|
|
model_name = model_args.model_name_or_path
|
|
base_model, tokenizer = get_model_and_tokenizer(
|
|
**vars(model_args),
|
|
train=True,
|
|
requires_grad=False, # ctx_args.exp_setup == ExperimentSetup.FULL_FINETUNE,
|
|
peft_config=get_lora_config(model_name, **vars(lora_args)),
|
|
# use_q_lora=True,
|
|
)
|
|
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, train=True)
|
|
else:
|
|
ctx_name = base_model.base_model.config.name_or_path
|
|
ctx_encoder_model_config = base_model.config
|
|
ctx_tokenizer = tokenizer
|
|
|
|
if ctx_args.exp_setup == ExperimentSetup.HYPERLORA:
|
|
# TODO: handle only extra_modules case (no target_modules)
|
|
logger.info("Using HyperLoRA")
|
|
if not ctx_args.from_pretrained_checkpoint:
|
|
hypernet_config = get_hypernet_config(
|
|
base_model,
|
|
ctx_encoder_model_config,
|
|
hypernet_args,
|
|
aggregator_args,
|
|
ctx_encoder_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"
|
|
)
|
|
model = ModulatedPretrainedModel(
|
|
base_model,
|
|
hypernet_config,
|
|
ctx_encoder_args,
|
|
)
|
|
|
|
else:
|
|
logger.info(
|
|
f"Loading from checkpoint: {ctx_args.from_pretrained_checkpoint}"
|
|
)
|
|
model = ModulatedPretrainedModel.from_state_dict(
|
|
torch.load(ctx_args.from_pretrained_checkpoint, weights_only=False),
|
|
train=True,
|
|
use_flash_attn=model_args.use_flash_attn,
|
|
)
|
|
tokenizer = get_tokenizer(model.base_model.config.name_or_path, train=True)
|
|
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, train=True)
|
|
|
|
training_args.gen_lora_l1_reg_coef = ctx_args.gen_lora_l1_reg_coef
|
|
training_args.use_kl_loss = ctx_args.use_kl_loss
|
|
training_args.use_per_ctx_average_loss = ctx_args.use_per_ctx_average_loss
|
|
|
|
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]):
|
|
raise ValueError("base model contains trainable parameters")
|
|
# we can't compile the base model bc we will be
|
|
# adding/removing forward hooks during training
|
|
model.hypernet.compile(fullgraph=True, mode="max-autotune")
|
|
|
|
else:
|
|
# activate LoRA
|
|
base_model_config = AutoConfig.from_pretrained(
|
|
model_args.model_name_or_path, trust_remote_code=True
|
|
)
|
|
base_model_config.save_pretrained(output_dir)
|
|
logger.info("Using LoRA")
|
|
model.set_adapter("default")
|
|
model = torch.compile(model)
|
|
|
|
model.train()
|
|
logger.debug(model)
|
|
log_num_train_params(model)
|
|
|
|
############ Dataset setup
|
|
logger.info("Loading dataset...")
|
|
|
|
add_ctx_to_chat = not isinstance(model, ModulatedPretrainedModel)
|
|
ctx_model_max_len = model.ctx_encoder.config.max_position_embeddings
|
|
if ctx_args.max_ctx_len > 0:
|
|
ctx_model_max_len = ctx_args.max_ctx_len
|
|
|
|
_get_tokenized_dataset = partial(
|
|
get_tokenized_dataset,
|
|
max_qas_len=ctx_args.max_qas_len,
|
|
max_qas_per_sample=ctx_args.max_qas_per_sample,
|
|
base_model_max_len=model.base_model.config.max_position_embeddings,
|
|
tokenizer=tokenizer,
|
|
ctx_model_max_len=ctx_model_max_len,
|
|
ctx_tokenizer=ctx_tokenizer,
|
|
add_ctx_to_chat=add_ctx_to_chat,
|
|
use_kl_loss=ctx_args.use_kl_loss,
|
|
)
|
|
splits = ["train"]
|
|
if training_args.eval_strategy != "no":
|
|
splits.append("validation")
|
|
tokenized_ds = {split: {} for split in splits}
|
|
for split, ds_names in zip(
|
|
splits,
|
|
[data_args.train_ds_names, data_args.val_ds_names],
|
|
):
|
|
if not ds_names:
|
|
continue
|
|
ctx_mgr = (
|
|
training_args.main_process_first()
|
|
if split == "train"
|
|
else contextlib.nullcontext()
|
|
)
|
|
with ctx_mgr:
|
|
# process and tokenize on the main process
|
|
# then other replicas can just load the cached dataset
|
|
# we dont save cache for validation ds
|
|
for ds_name in ds_names:
|
|
ds = _get_tokenized_dataset(ds_name, split)
|
|
|
|
base_name = os.path.basename(ds_name)
|
|
if ds_name.startswith("self_gen/"):
|
|
ds_name = "self_gen/" + base_name
|
|
else:
|
|
ds_name = base_name
|
|
|
|
tokenized_ds[split][ds_name] = ds
|
|
|
|
train_ds = tokenized_ds["train"]
|
|
logging.info(f"train_ds: {train_ds}")
|
|
|
|
val_ds = dict()
|
|
if "validation" in tokenized_ds:
|
|
n_val_samples = data_args.max_val_samples_per_ds
|
|
for ds_name, ds in tokenized_ds["validation"].items():
|
|
if ds is None:
|
|
# take some samples from train_ds
|
|
ds = train_ds[ds_name].take(n_val_samples)
|
|
train_ds[ds_name] = train_ds[ds_name].skip(n_val_samples)
|
|
|
|
val_ds[ds_name] = ds
|
|
val_indices = np.random.permutation(len(ds))[:n_val_samples]
|
|
val_ds[ds_name] = val_ds[ds_name].select(val_indices)
|
|
|
|
with training_args.main_process_first():
|
|
train_ds = pack(
|
|
train_ds,
|
|
ctx_args.max_packed_inp_len,
|
|
ctx_args.max_packed_ctx_len,
|
|
max_packed_size=-1,
|
|
seed=training_args.seed,
|
|
num_proc=16,
|
|
)
|
|
logger.info("Setting per_device_train_batch_size to 1")
|
|
training_args.per_device_train_batch_size = 1
|
|
|
|
logger.info(f"train_ds: {train_ds}")
|
|
logger.info(f"val_ds: {val_ds}")
|
|
|
|
collator = flatten_if_not_packed
|
|
|
|
# HACK [local patch]: deepspeed model loading problem (for resume training)
|
|
# see https://github.com/microsoft/DeepSpeed/pull/6626/files
|
|
# /home/rujikorn_sakana_ai/.conda/envs/ctx-to-lora/lib/python3.10/site-packages/deepspeed/runtime/engine.py
|
|
# HACK: [local path]: make liger kernel works with gemma3
|
|
# see .venv/lib/python3.10/site-packages/liger_kernel/transformers/monkey_patch.py
|
|
# L810:814 + L781
|
|
|
|
if isinstance(model, ModulatedPretrainedModel):
|
|
if isinstance(model.base_model, PeftModel):
|
|
base_model = model.base_model.base_model
|
|
else:
|
|
base_model = model.base_model
|
|
|
|
if ctx_name is not None:
|
|
logger.info("Compiling ctx_encoder_model")
|
|
ctx_base_model = model.ctx_encoder.base_model
|
|
compile_linear(ctx_base_model)
|
|
# TODO: why we can't compile ctx_base_model?
|
|
# ctx_base_model.compile(fullgraph=True, mode="max-autotune")
|
|
|
|
elif isinstance(model, PeftModel):
|
|
base_model = model.base_model
|
|
|
|
logger.info("Compiling base_model")
|
|
base_model.compile(fullgraph=True, mode="max-autotune")
|
|
|
|
if LOCAL_RANK == 0:
|
|
wandb.init(
|
|
project=os.getenv("WANDB_PROJECT"),
|
|
name=run_name,
|
|
group=run_name,
|
|
config=args,
|
|
tags=os.getenv("WANDB_TAGS").split(","),
|
|
notes=ctx_args.notes,
|
|
resume="allow",
|
|
)
|
|
else:
|
|
wandb.init(mode="disabled")
|
|
|
|
train_model(
|
|
model,
|
|
training_args,
|
|
train_ds,
|
|
val_ds,
|
|
collator,
|
|
compute_metrics=partial(
|
|
compute_metrics,
|
|
evaluator=Evaluator(
|
|
[compute_per_token_acc, compute_prefix_matching, compute_perplexity]
|
|
),
|
|
),
|
|
)
|
|
logger.info(f"Training run finished and saved to {output_dir}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "true"
|
|
os.environ["TOKENIZERS_PARALLELISM"] = "true"
|
|
os.environ["WANDB_DIR"] = ".wandb/"
|
|
os.environ["WANDB_PROJECT"] = "ctx_to_lora"
|
|
os.environ["WANDB_WATCH"] = ""
|
|
os.environ["WANDB_CONSOLE"] = "off"
|
|
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
|
|
os.environ["OMP_NUM_THREADS"] = "23"
|
|
torch._dynamo.config.capture_scalar_outputs = True
|
|
|
|
# DOES THIS CAUSE ACCURACY DEGRADATION?
|
|
torch.backends.cuda.matmul.allow_tf32 = True
|
|
torch.backends.cudnn.allow_tf32 = True
|
|
if os.getenv("DEBUG", False):
|
|
disable_caching()
|
|
main()
|