mirror of
https://github.com/SakanaAI/doc-to-lora.git
synced 2026-07-23 17:01:04 +02:00
279 lines
7.9 KiB
Python
279 lines
7.9 KiB
Python
import ast
|
|
import gc
|
|
import hashlib
|
|
import logging
|
|
import os
|
|
import random
|
|
import string
|
|
import time
|
|
from collections.abc import Iterable
|
|
from contextlib import contextmanager
|
|
from enum import Enum
|
|
|
|
import torch
|
|
import yaml
|
|
from peft import PeftConfig, PeftModel
|
|
from peft.tuners.tuners_utils import BaseTunerLayer, check_target_module_exists
|
|
from peft.utils import get_peft_model_state_dict
|
|
|
|
TRAINING_TASK = Enum("TRAINING_TASK", ["CAUSAL_LM", "COMPLETION"])
|
|
|
|
|
|
logger = logging.getLogger()
|
|
|
|
|
|
# taken from https://discuss.pytorch.org/t/opinion-eval-should-be-a-context-manager/18998/3
|
|
@contextmanager
|
|
def evaluating(*models):
|
|
"""Temporarily switch to evaluation mode."""
|
|
is_training = [model.training if model is not None else False for model in models]
|
|
try:
|
|
for model in models:
|
|
if model is not None:
|
|
model.eval()
|
|
yield models
|
|
finally:
|
|
for model, training in zip(models, is_training):
|
|
if model is not None:
|
|
model.train(training)
|
|
|
|
|
|
def get_layers(model):
|
|
if hasattr(model, "model"):
|
|
return get_layers(model.model)
|
|
return model.layers
|
|
|
|
|
|
def get_num_layers(model):
|
|
return len(get_layers(model))
|
|
|
|
|
|
def get_base_model(model):
|
|
if hasattr(model, "model"):
|
|
return get_base_model(model.model)
|
|
return model
|
|
|
|
|
|
def get_num_params(model):
|
|
total_params = 0
|
|
trainable_params = 0
|
|
for p in model.parameters():
|
|
total_params += p.numel()
|
|
if p.requires_grad:
|
|
trainable_params += p.numel()
|
|
|
|
return total_params, trainable_params
|
|
|
|
|
|
def log_num_train_params(model):
|
|
logger.debug("Trainable model parameters:")
|
|
for name, p in model.named_parameters():
|
|
if p.requires_grad:
|
|
logger.debug(f"{name}, dtype:{p.dtype}")
|
|
|
|
num_total_params, num_trainable_params = get_num_params(model)
|
|
logger.info(
|
|
f"trainable params: {num_trainable_params:,d} "
|
|
f"|| all params: {num_total_params:,d} "
|
|
f"|| trainable%: {100 * num_trainable_params / num_total_params:.4f}"
|
|
)
|
|
|
|
|
|
def get_run_name(seed_str: str | None = None):
|
|
if not seed_str:
|
|
uuid = "".join(
|
|
[random.choice(string.ascii_letters + string.digits) for _ in range(8)]
|
|
)
|
|
run_name = time.strftime("%Y%m%d-%H%M%S") + f"_{uuid}"
|
|
else:
|
|
# Generate a UUID from the seed string
|
|
hash_object = hashlib.sha256(seed_str.encode())
|
|
uuid = hash_object.hexdigest()[:8] # Take the first 8 characters of the hash
|
|
run_name = seed_str + f"_{uuid}"
|
|
return run_name
|
|
|
|
|
|
def try_convert(s):
|
|
try:
|
|
return ast.literal_eval(s)
|
|
except:
|
|
return s
|
|
|
|
|
|
def extract_cli_args(argv: list[str]):
|
|
out = dict()
|
|
for elem in argv:
|
|
if elem.endswith(".yaml"):
|
|
out["config"] = elem
|
|
|
|
elif elem.startswith("--"):
|
|
k, v = elem.split("=")
|
|
k = k.split("--")[1]
|
|
v = try_convert(v)
|
|
# if k.startswith('env_'):
|
|
# k = k.split('_')[1]
|
|
out[k] = v
|
|
return out
|
|
|
|
|
|
def setup_logging(output_dir, debug=False):
|
|
global logger
|
|
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
log_formatter = logging.Formatter(
|
|
fmt="%(asctime)s %(levelname)s: %(message)s", datefmt="%Y-%m-%d %H:%M:%S"
|
|
)
|
|
stream_level = logging.DEBUG if debug else logging.INFO
|
|
stream_handler = logging.StreamHandler()
|
|
stream_handler.setFormatter(log_formatter)
|
|
stream_handler.setLevel(stream_level)
|
|
logger.addHandler(stream_handler)
|
|
|
|
log_path = f"{output_dir}/debug.log"
|
|
debug_handler = logging.FileHandler(log_path, delay=True)
|
|
debug_handler.setLevel(logging.DEBUG)
|
|
debug_handler.setFormatter(log_formatter)
|
|
logger.addHandler(debug_handler)
|
|
logger.setLevel(logging.DEBUG)
|
|
logger.info(f"Logging to: {log_path}")
|
|
|
|
|
|
def validate_args(args_list):
|
|
# there shouldn't be overlap between args
|
|
keys = set()
|
|
for args in args_list:
|
|
logger.debug(args)
|
|
args_keys = set(vars(args).keys())
|
|
assert len(keys & args_keys) == 0, "Overlap between args"
|
|
keys |= args_keys
|
|
|
|
|
|
def save_yaml(data, path):
|
|
# Filter out non-primitive fields
|
|
data = {
|
|
k: v
|
|
for k, v in data.items()
|
|
if isinstance(v, (int, float, str, bool, list, dict, type(None)))
|
|
}
|
|
|
|
with open(path, "w") as file:
|
|
yaml.dump(data, file)
|
|
|
|
|
|
def get_peft_modules(model: PeftModel, peft_config: PeftConfig) -> list[dict[str, str]]:
|
|
return [
|
|
{"name": name, "module": module}
|
|
for name, module in model.named_modules()
|
|
if name.split(".")[-1] in peft_config.target_modules
|
|
and isinstance(module, BaseTunerLayer)
|
|
and check_target_module_exists(peft_config, name)
|
|
]
|
|
|
|
|
|
def get_peft_in_out_features(
|
|
model: PeftModel,
|
|
peft_config: PeftConfig | None = None,
|
|
) -> tuple[dict[str, int], dict[str, int]]:
|
|
if peft_config is None:
|
|
return None, None
|
|
in_features = dict()
|
|
out_features = dict()
|
|
for module_info in get_peft_modules(model, peft_config):
|
|
module_name = module_info["name"]
|
|
module = module_info["module"]
|
|
# support just Linear layer for now
|
|
# all modules should be a leave module that is Linear layer
|
|
assert isinstance(module.base_layer, torch.nn.Linear), (
|
|
"all modules should be a leave module that is Linear layer"
|
|
)
|
|
|
|
# this should always pass
|
|
name = module_name.split(".")[-1]
|
|
assert name in peft_config.target_modules
|
|
|
|
if name not in in_features:
|
|
in_features[name] = module.in_features
|
|
out_features[name] = module.out_features
|
|
else:
|
|
# assumes each module has the same input and output features
|
|
assert in_features[name] == module.in_features
|
|
assert out_features[name] == module.out_features
|
|
|
|
return in_features, out_features
|
|
|
|
|
|
def generated_lora_to_state_dict(
|
|
lora_dict: dict,
|
|
module_names: dict,
|
|
target_modules: list[str],
|
|
layer_indices: Iterable[int],
|
|
) -> dict:
|
|
lora_state_dict = dict()
|
|
for target_module in target_modules:
|
|
for layer_idx in layer_indices:
|
|
for module_name in module_names[target_module][layer_idx]:
|
|
if "lora_A" in module_name:
|
|
lora_state_dict[module_name] = (
|
|
lora_dict[target_module]["A"][layer_idx].cpu().contiguous()
|
|
)
|
|
elif "lora_B" in module_name:
|
|
lora_state_dict[module_name] = (
|
|
lora_dict[target_module]["B"][layer_idx].cpu().contiguous()
|
|
)
|
|
else:
|
|
raise ValueError(f"Unexpected module name: {module_name}")
|
|
return lora_state_dict
|
|
|
|
|
|
def get_lora_module_names(
|
|
model: PeftModel,
|
|
target_modules: list[str],
|
|
layer_indices: Iterable[int],
|
|
) -> dict[str, list[str]]:
|
|
module_names = {
|
|
target_module: [[] for _ in range(len(layer_indices))]
|
|
for target_module in target_modules
|
|
}
|
|
for k in get_peft_model_state_dict(model):
|
|
if "lora" not in k:
|
|
continue
|
|
layer_idx = int(k.split("layers.")[-1].split(".")[0])
|
|
if layer_idx in layer_indices:
|
|
for target_module in target_modules:
|
|
if target_module in k:
|
|
module_names[target_module][layer_idx].append(k)
|
|
break
|
|
return module_names
|
|
|
|
|
|
def compile_linear(model):
|
|
for name, module in model.named_modules():
|
|
if isinstance(module, torch.nn.Linear):
|
|
module.compile()
|
|
|
|
|
|
def clear_gpu():
|
|
gc.collect()
|
|
if not torch.cuda.is_available():
|
|
return
|
|
torch.cuda.empty_cache()
|
|
try:
|
|
torch.cuda.reset_peak_memory_stats()
|
|
except RuntimeError:
|
|
return
|
|
|
|
|
|
def concat_list(l):
|
|
out = []
|
|
for x in l:
|
|
out += x
|
|
return out
|
|
|
|
|
|
def check_is_iterable(x):
|
|
try:
|
|
iter(x)
|
|
except TypeError:
|
|
return False
|
|
return True
|